digikedai-bot/tests/agent.test.ts

import { describe, it, expect } from "vitest";
import {
  Agent,
  MEMORY_KEY_SUMMARY,
  MEMORY_KEY_LAST_QUERIES,
  isRepeatQuery,
  normalizeQuery,
} from "../src/ai/agent/agent.js";
import type { LlmProvider } from "../src/ai/providers/llm.js";
import type { Db } from "../src/db/db.js";
import type { Retriever, RetrievalResult } from "../src/ai/retrieval/retriever.js";
import type { ConversationMemory } from "../src/ai/memory/memory.js";
import type { Summarizer } from "../src/ai/memory/summarizer.js";

class FakeLlm implements LlmProvider {
  calls: { system: string; messages: { role: string; content: string }[] }[] = [];
  async chat(args: {
    system: string;
    messages: { role: "user" | "assistant"; content: string }[];
  }): Promise<string> {
    this.calls.push(args);
    return "mock reply";
  }
  async chatSmall(): Promise<string> {
    return "small reply";
  }
}

function fakeDb(): Db {
  return {
    botDbName: "bot",
    pool: {} as never,
    saveExchange: async () => ({ conversationId: 1, userId: 1 }),
    recentMessages: async () => [
      { role: "user", content: "earlier question" },
      { role: "assistant", content: "earlier answer" },
    ],
    memorySet: async () => {},
    memoryGet: async () => undefined,
    memoryGetAll: async () => ({}),
    close: async () => {},
  };
}

class FakeMemory implements ConversationMemory {
  store = new Map<string, string>();
  failRecallAll = false;
  async remember(userId: number, key: string, value: string): Promise<void> {
    this.store.set(`${userId}:${key}`, value);
  }
  async recall(userId: number, key: string): Promise<string | undefined> {
    return this.store.get(`${userId}:${key}`);
  }
  async recallAll(userId: number): Promise<Record<string, string>> {
    if (this.failRecallAll) throw new Error("recall boom");
    const out: Record<string, string> = {};
    for (const [k, v] of this.store) {
      if (k.startsWith(`${userId}:`)) out[k.slice(`${userId}:`.length)] = v;
    }
    return out;
  }
  seed(userId: number, key: string, value: string) {
    this.store.set(`${userId}:${key}`, value);
  }
}

class FakeSummarizer implements Summarizer {
  calls: {
    priorSummary?: string;
    priorQueries: string[];
    currentTurn: string;
  }[] = [];
  result = "remembered summary";
  fail = false;
  async extract(args: {
    priorSummary?: string;
    priorQueries: string[];
    currentTurn: string;
  }): Promise<string | undefined> {
    this.calls.push(args);
    if (this.fail) throw new Error("summarizer boom");
    return this.result;
  }
}

function respondArgs(overrides: Partial<{ userText: string; userId: number }> = {}) {
  return {
    conversationId: 1,
    userId: 1,
    userText: "do you have courses?",
    preferredLanguage: "en",
    ...overrides,
  };
}

describe("Agent.respond", () => {
  it("assembles system + history + current message", async () => {
    const llm = new FakeLlm();
    const agent = new Agent(llm, fakeDb());

    const reply = await agent.respond(respondArgs());

    expect(reply).toBe("mock reply");
    expect(llm.calls).toHaveLength(1);
    const call = llm.calls[0];
    expect(call.system).toContain("Digi Kedai");
    expect(call.messages[call.messages.length - 1]).toEqual({
      role: "user",
      content: "do you have courses?",
    });
    // History should be present
    expect(call.messages.some((m) => m.content === "earlier question")).toBe(true);
  });

  it("injects retrieved product facts into the system prompt", async () => {
    const llm = new FakeLlm();
    const retriever: Retriever = {
      retrieve: async (q, topK): Promise<RetrievalResult[]> => [
        {
          text: `得到 全平台（SKU CDD01）→ https://www.digikedai.com/products/cdd01/`,
          source: "https://www.digikedai.com/products/cdd01/",
          score: 1,
        },
      ],
    };
    const agent = new Agent(llm, fakeDb(), retriever);

    await agent.respond(respondArgs({ userText: "do you have CDD01?" }));

    const call = llm.calls[0];
    expect(call.system).toContain("Product facts");
    expect(call.system).toContain("https://www.digikedai.com/products/cdd01/");
  });

  it("keeps the system prompt clean when retrieval returns nothing", async () => {
    const llm = new FakeLlm();
    const agent = new Agent(llm, fakeDb());

    await agent.respond(respondArgs({ userText: "hello" }));

    const call = llm.calls[0];
    expect(call.system).not.toContain("Product facts");
  });

  it("injects remembered user summary into the system prompt", async () => {
    const llm = new FakeLlm();
    const memory = new FakeMemory();
    memory.seed(1, MEMORY_KEY_SUMMARY, "customer speaks Chinese, likes CDD01");
    const agent = new Agent(llm, fakeDb(), undefined, memory);

    await agent.respond(respondArgs());

    const call = llm.calls[0];
    expect(call.system).toContain("User memory");
    expect(call.system).toContain("customer speaks Chinese, likes CDD01");
  });

  it("replying stays safe when memory recall fails", async () => {
    const llm = new FakeLlm();
    const memory = new FakeMemory();
    memory.failRecallAll = true;
    const agent = new Agent(llm, fakeDb(), undefined, memory);

    const reply = await agent.respond(respondArgs());

    expect(reply).toBe("mock reply");
    const call = llm.calls[0];
    expect(call.system).not.toContain("User memory");
  });

  it("injects the repeat-question hint when the same question was asked before", async () => {
    const llm = new FakeLlm();
    const memory = new FakeMemory();
    memory.seed(
      1,
      MEMORY_KEY_LAST_QUERIES,
      JSON.stringify(["do you have courses?"]),
    );
    const agent = new Agent(llm, fakeDb(), undefined, memory);

    await agent.respond(respondArgs({ userText: "Do you have courses?" }));

    expect(llm.calls[0].system).toContain("repeat-question rule");
  });

  it("does not flag a fresh question as a repeat", async () => {
    const llm = new FakeLlm();
    const memory = new FakeMemory();
    memory.seed(
      1,
      MEMORY_KEY_LAST_QUERIES,
      JSON.stringify(["do you have courses?"]),
    );
    const agent = new Agent(llm, fakeDb(), undefined, memory);

    await agent.respond(respondArgs({ userText: "how much is the template?" }));

    expect(llm.calls[0].system).not.toContain("repeat-question rule");
  });

  it("stays silent (empty reply) on WhatsApp when the model emits NO_REPLY", async () => {
    const silentLlm: LlmProvider = {
      chat: async () => "NO_REPLY",
      chatSmall: async () => "x",
    };
    const agent = new Agent(silentLlm, fakeDb());

    const reply = await agent.respond({
      conversationId: 1,
      userId: 1,
      userText: "tell me a joke",
      channel: "whatsapp",
    });

    expect(reply).toBe("");
  });

  it("does not silence NO_REPLY on non-WhatsApp channels", async () => {
    const silentLlm: LlmProvider = {
      chat: async () => "NO_REPLY",
      chatSmall: async () => "x",
    };
    const agent = new Agent(silentLlm, fakeDb());

    const reply = await agent.respond({
      conversationId: 1,
      userId: 1,
      userText: "tell me a joke",
    });

    expect(reply).toBe("NO_REPLY");
  });
});

describe("isRepeatQuery", () => {
  it("normalizes case and whitespace", () => {
    expect(normalizeQuery("  你好 ！  ")).toBe("你好 ！");
    expect(normalizeQuery("Do You Have Courses?")).toBe(
      "do you have courses?",
    );
  });

  it("detects exact repeats", () => {
    expect(
      isRepeatQuery("Do you have courses?", ["do you have courses?"]),
    ).toBe(true);
  });

  it("detects meaningful containment for longer strings", () => {
    expect(
      isRepeatQuery("do you have courses?", [
        "do you have courses? and what is the price",
      ]),
    ).toBe(true);
    expect(
      isRepeatQuery("do you have courses? and what is the price", [
        "do you have courses?",
      ]),
    ).toBe(true);
  });

  it("ignores short fuzzy collisions", () => {
    expect(isRepeatQuery("hi", ["hi there"])).toBe(false);
  });

  it("returns false for different questions and empty input", () => {
    expect(isRepeatQuery("how much?", ["do you have courses?"])).toBe(false);
    expect(isRepeatQuery("", ["anything"])).toBe(false);
    expect(isRepeatQuery("hi", [])).toBe(false);
  });
});

describe("Agent.updateMemoryAsync", () => {
  it("writes lang when it changes", async () => {
    const memory = new FakeMemory();
    memory.seed(1, "lang", "zh");
    const agent = new Agent(new FakeLlm(), fakeDb(), undefined, memory);

    await agent.updateMemoryAsync(respondArgs({ userId: 1 }), memory.recallAll(1));

    expect(await memory.recall(1, "lang")).toBe("en");
  });

  it("keeps last_queries at max 5 items and truncates long entries", async () => {
    const memory = new FakeMemory();
    memory.seed(
      1,
      "last_queries",
      JSON.stringify(["q1", "q2", "q3", "q4", "q5"]),
    );
    const agent = new Agent(new FakeLlm(), fakeDb(), undefined, memory);
    const longText = "x".repeat(300);

    await agent.updateMemoryAsync(
      respondArgs({ userId: 1, userText: longText }),
      await memory.recallAll(1),
    );

    const stored = JSON.parse((await memory.recall(1, "last_queries")) ?? "[]");
    expect(stored).toHaveLength(5);
    expect(stored[4]).toHaveLength(200);
    expect(stored[0]).toBe("q2");
  });

  it("skips bot commands", async () => {
    const memory = new FakeMemory();
    const agent = new Agent(new FakeLlm(), fakeDb(), undefined, memory);

    await agent.updateMemoryAsync(
      respondArgs({ userId: 1, userText: "/start" }),
      memory.recallAll(1),
    );

    expect(await memory.recall(1, "last_queries")).toBeUndefined();
    expect(await memory.recall(1, "lang")).toBeUndefined();
  });

  it("triggers summarizer on a new SKU for a fresh user", async () => {
    const memory = new FakeMemory();
    const summarizer = new FakeSummarizer();
    const agent = new Agent(
      new FakeLlm(),
      fakeDb(),
      undefined,
      memory,
      summarizer,
    );

    await agent.updateMemoryAsync(
      respondArgs({ userId: 1, userText: "how about CDD01?" }),
      memory.recallAll(1),
    );

    expect(summarizer.calls).toHaveLength(1);
    expect(summarizer.calls[0].currentTurn).toBe("how about CDD01?");
    expect(await memory.recall(1, MEMORY_KEY_SUMMARY)).toBe("remembered summary");
  });

  it("does not summarise for plain chit-chat once lang is settled", async () => {
    const memory = new FakeMemory();
    memory.seed(1, "lang", "en");
    const summarizer = new FakeSummarizer();
    const agent = new Agent(
      new FakeLlm(),
      fakeDb(),
      undefined,
      memory,
      summarizer,
    );

    await agent.updateMemoryAsync(
      respondArgs({ userId: 1, userText: "hi" }),
      await memory.recallAll(1),
    );

    expect(summarizer.calls).toHaveLength(0);
  });

  it("summarizer failure keeps the old summary and never throws", async () => {
    const memory = new FakeMemory();
    memory.seed(1, "lang", "en");
    memory.seed(1, MEMORY_KEY_SUMMARY, "old summary");
    const summarizer = new FakeSummarizer();
    summarizer.fail = true;
    const agent = new Agent(
      new FakeLlm(),
      fakeDb(),
      undefined,
      memory,
      summarizer,
    );

    await expect(
      agent.updateMemoryAsync(
        respondArgs({ userId: 1, userText: "我叫小明，想买CDD01" }),
        await memory.recallAll(1),
      ),
    ).resolves.toBeUndefined();

    expect(await memory.recall(1, MEMORY_KEY_SUMMARY)).toBe("old summary");
  });
});