MockLanguageModelV4: a readable stand-in

// A stand-in language model for the AI SDK: keyword rules instead of a neural network.
import { simulateReadableStream } from "ai";
import { MockLanguageModelV4 } from "ai/test";

const KEYWORDS: Record<string, string[]> = {
  billing: ["charged", "refund", "invoice", "payment"],
  shipping: ["parcel", "delivery", "courier", "arrived"],
  account: ["password", "login", "email", "account"],
};

function sortTicket(text: string) {
  const lower = text.toLowerCase();
  for (const [category, words] of Object.entries(KEYWORDS)) {
    if (words.some((word) => lower.includes(word))) return category;
  }
  return "other";
}

function decide(options: any): any[] {
  const messages = options.prompt;
  const last = messages[messages.length - 1];
  const userText = messages.filter((m: any) => m.role === "user").at(-1)?.content.map((p: any) => p.text ?? "").join(" ") ?? "";
  const category = sortTicket(userText);

  // A tool ran: answer with what it returned.
  if (last.role === "tool") {
    const result = last.content.find((p: any) => p.type === "tool-result");
    if (result) {
      const label = result.toolName.startsWith("lookup") ? "Order status" : "Done";
      return [{ type: "text", text: `${label}: ${result.output.value}.` }];
    }
    return [{ type: "text", text: "I could not do that without approval." }];
  }
  // An order id and a tool: refund tickets call a refund tool, others a lookup tool.
  const order = userText.match(/A-\d{4}/);
  if (order && options.tools?.length) {
    const names: string[] = options.tools.map((t: any) => t.name);
    const amount = userText.match(/(\d+) euros/);
    const refund = /refund/i.test(userText) && amount ? names.find((n) => /refund/i.test(n)) : undefined;
    const toolName = refund ?? names.find((n) => /lookup/i.test(n)) ?? names[0];
    const input = refund ? { orderId: order[0], amount: Number(amount![1]) } : { orderId: order[0] };
    return [{ type: "tool-call", toolCallId: "call-1", toolName, input: JSON.stringify(input) }];
  }
  // A JSON schema was asked for: fill in the fields it knows.
  if (options.responseFormat?.type === "json") {
    const fields = options.responseFormat.schema?.properties ?? {};
    const answer: Record<string, unknown> = {};
    if ("category" in fields) answer.category = category;
    if ("priority" in fields) answer.priority = category === "billing" ? 4 : 2;
    if ("summary" in fields) answer.summary = userText.slice(0, 40);
    return [{ type: "text", text: JSON.stringify(answer) }];
  }
  const team = category === "other" ? "support" : category;
  return [{ type: "text", text: `Thanks for your message. Our ${team} team will help.` }];
}

function usage(options: any, content: any[]) {
  const words = (text: string) => text.split(/\s+/).filter(Boolean).length;
  const input = words(JSON.stringify(options.prompt));
  const output = words(content.map((p) => p.text ?? p.input ?? "").join(" "));
  return { inputTokens: { total: input, noCache: input, cacheRead: undefined, cacheWrite: undefined }, outputTokens: { total: output, text: output, reasoning: undefined } };
}

export const shopModel = new MockLanguageModelV4({
  provider: "shop",
  modelId: "keywords",
  doGenerate: async (options) => {
    const content = decide(options);
    const finish = content[0].type === "tool-call" ? "tool-calls" : "stop";
    return { content, finishReason: { unified: finish, raw: undefined }, usage: usage(options, content), warnings: [] };
  },
  doStream: async (options) => {
    const content = decide(options);
    const chunks: any[] = [];
    for (const part of content) {
      if (part.type === "text") {
        chunks.push({ type: "text-start", id: "t" });
        for (const word of part.text.split(/(?<= )/)) chunks.push({ type: "text-delta", id: "t", delta: word });
        chunks.push({ type: "text-end", id: "t" });
      } else chunks.push(part);
    }
    const finish = content[0].type === "tool-call" ? "tool-calls" : "stop";
    chunks.push({ type: "finish", finishReason: { unified: finish, raw: undefined }, usage: usage(options, content) });
    return { stream: simulateReadableStream({ chunks, chunkDelayInMs: 0 }) };
  },
});

import { generateText } from "ai";
import { shopModel } from "./shop-model.ts";

for (const ticket of ["I was charged twice", "My parcel never arrived", "Do you sell gift cards?"]) {
  const { text } = await generateText({ model: shopModel, prompt: ticket });
  console.log(`${ticket} -> ${text}`);
}

npx tsx sort.ts