rag.test.ts · Source

rag.test.tspackages/ai/src/providers/rag.test.tsraw ← providers6,147 B← Back to the app
1/**
2 * TASK-0125 — RAG citation post-processor (the mock's deterministic `rag`
3 * mode).
4 *
5 * The RAG app layer (rag.service) sends the retrieved, numbered context as
6 * `[n] Source: <label>\n<chunk text>` blocks + a `Question: <q>` line and
7 * opts in with `metadata: { mode: "rag" }`. The mock plays the "model" end of
8 * that REAL contract: it parses the first source block (`parseRagSource`) and
9 * streams a deterministic, CITED answer. These tests pin that contract — the
10 * citation post-processor's boundary cases the spec calls out:
11 *   - a well-formed single/multi-source prompt → the right `[n]` + label;
12 *   - a multi-source prompt cites the FIRST block (the deterministic rule);
13 *   - the answer text is truncated to 200 chars of the chunk (no leak);
14 *   - a missing/malformed marker → null → the "no results" answer (never the
15 *     literal `[[…]]`);
16 *   - nested brackets / brackets that are not a source header are not
17 *     mis-parsed as a source.
18 */
19import { describe, expect, it } from "vitest";
20 
21import { MockProvider, parseRagSource } from "./index.js";
22import type { ChatRequest } from "../types.js";
23 
24function ragReq(context: string, question: string): ChatRequest {
25  return {
26    model: "mock",
27    messages: [
28      { role: "system", content: "Cite with [n]. Use only the sources." },
29      { role: "user", content: `${context}\n\nQuestion: ${question}` },
30    ],
31    metadata: { mode: "rag" },
32  };
33}
34 
35async function answerText(req: ChatRequest): Promise<string> {
36  const p = new MockProvider("openai", 8);
37  let text = "";
38  for await (const chunk of p.chat(req as never)) {
39    const c = chunk as { delta?: string; done?: boolean; finishReason?: string };
40    if (c.delta) text += c.delta;
41  }
42  return text;
43}
44 
45describe("parseRagSource (the citation post-processor)", () => {
46  it("parses a single well-formed source block (n, label, text)", () => {
47    const src = parseRagSource(
48      "[1] Source: Launch Plan\nThe launch code word is Zanzibar-7.\n\nQuestion: What is the code?",
49    );
50    expect(src).toEqual({ n: 1, label: "Launch Plan", text: "The launch code word is Zanzibar-7." });
51  });
52 
53  it("parses a multi-source prompt and cites the FIRST block", () => {
54    const src = parseRagSource(
55      "[1] Source: Alpha\nfirst chunk\n[2] Source: Beta\nsecond chunk\n\nQuestion: q?",
56    );
57    expect(src?.n).toBe(1);
58    expect(src?.label).toBe("Alpha");
59    expect(src?.text).toBe("first chunk");
60  });
61 
62  it("trims the label and the chunk text", () => {
63    const src = parseRagSource("[3] Source:   Padded Label   \n   body text   \n\nQuestion: q?");
64    expect(src?.label).toBe("Padded Label");
65    expect(src?.text).toBe("body text");
66  });
67 
68  it("returns null when there is no source header (no results — never a literal marker)", () => {
69    expect(parseRagSource("just a plain question with no context")).toBeNull();
70    expect(parseRagSource("[n] Source: not-a-number\nx")).toBeNull();
71  });
72 
73  it("does not mis-parse a bracketed phrase that is not a `Source:` header", () => {
74    // A `[1]` inside prose (no `Source:` header) is not a source marker.
75    expect(parseRagSource("As noted in [1], the plan is on track.\n\nQuestion: q?")).toBeNull();
76    // Nested brackets are not a header either.
77    expect(parseRagSource("see [[chunk:abc]] for details\n\nQuestion: q?")).toBeNull();
78  });
79 
80  it("a source header must be followed by its text on the next line", () => {
81    // `[2] Source: X` immediately before `Question:` with no body → empty text.
82    const src = parseRagSource("[2] Source: X\n\nQuestion: q?");
83    expect(src).toEqual({ n: 2, label: "X", text: "" });
84  });
85 
86  it("keeps MULTI-LINE chunk text (a whole-page body is not truncated at line 1)", () => {
87    // The RAG service feeds a whole doc's body (title line + content lines) as
88    // the first source when a page doc out-ranks its chunks. The old `/m`
89    // regex made `$` match every line break and cut the text at the first
90    // newline — the TASK-0125 integration's seeded `Launch Runbook\n…
91    // Zanzibar-7.` parent body proved it.
92    const src = parseRagSource(
93      "[1] Source: Launch Runbook\nLaunch Runbook\nThe launch code word is Zanzibar-7.\n\nQuestion: What is the code word?",
94    );
95    expect(src).toEqual({
96      n: 1,
97      label: "Launch Runbook",
98      text: "Launch Runbook\nThe launch code word is Zanzibar-7.",
99    });
100  });
101});
102 
103describe("MockProvider rag mode (the cited answer)", () => {
104  it("streams a deterministic cited answer: `According to [n] (label): <200 chars>`", async () => {
105    const text = await answerText(ragReq("[1] Source: Launch Plan\nThe launch code word is Zanzibar-7.", "What is the launch code word?"));
106    expect(text).toBe("According to [1] (Launch Plan): The launch code word is Zanzibar-7.");
107  });
108 
109  it("truncates the cited chunk text to 200 chars (no leak of the rest)", async () => {
110    const longChunk = "x".repeat(350);
111    const text = await answerText(ragReq("[2] Source: Big Doc\n" + longChunk, "q?"));
112    // 200 chars of the chunk, then the trailing space the per-word streaming adds.
113    expect(text).toBe(`According to [2] (Big Doc): ${"x".repeat(200)}`);
114    expect(text).not.toContain("x".repeat(201));
115  });
116 
117  it("no source context → the deterministic no-results answer (no literal marker leaked)", async () => {
118    const text = await answerText(ragReq("no context here at all", "What is the code?"));
119    expect(text).toBe("I could not find any workspace content that answers this question. Try different wording.");
120    expect(text).not.toContain("[");
121    expect(text).not.toContain("Zanzibar");
122  });
123 
124  it("usage is the ceil(len/4) of the user prompt / answer (deterministic)", async () => {
125    const p = new MockProvider("openai", 8);
126    let usage: { in: number; out: number } | undefined;
127    for await (const chunk of p.chat(ragReq("[1] Source: A\nB", "q?") as never)) {
128      const c = chunk as { usage?: { in: number; out: number } };
129      if (c.usage) usage = { in: c.usage.in, out: c.usage.out };
130    }
131    expect(usage?.in).toBeGreaterThan(0);
132    expect(usage?.out).toBeGreaterThan(0);
133  });
134});
135