ai-embedding.test.ts

ai-embedding.test.tsapps/worker/src/__tests__/ai-embedding.test.tsraw ← __tests__18,999 B← Back to the app
1/**
2 * TASK-0125 — RAG embedding pipeline unit tests (Docker-free: fake prisma,
3 * fake SearchClient, fake gateway).
4 *
5 * Covers: the real embed path (one gateway call per item, vector merge onto
6 * each content doc + registry upsert), the chunk-id enumeration parity with
7 * the indexer (`page:<id>#c<n>`), the source-vanished prune path, the
8 * dimension-parity BLOCK (model-change → alias rollover required), stale-row
9 * pruning, the nightly re-embed sweep (orphan prune + count-parity drift
10 * alert), and payload validation.
11 */
12import { describe, expect, it, vi } from "vitest";
13import type { PrismaClient } from "@prisma/client";
14import type { AiGateway, EmbedResult } from "@factory/ai";
15import type { SearchClient } from "@factory/search";
16import { aiJobSchema } from "../processors/ai.processor.js";
17import {
18  chunkDocId,
19  deriveItemChunks,
20  handleEmbedChunks,
21  handleReembed,
22  type EmbedChunksJobData,
23} from "../processors/ai-embedding.js";
24import type { WorkerEnv } from "../config.js";
25 
26// ---------------------------------------------------------------------------
27// Fakes
28// ---------------------------------------------------------------------------
29 
30const NOW = new Date("2026-07-01T00:00:00Z");
31 
32/** A page row with one paragraph block of `text` (the search body source). */
33function pageRow(text: string, overrides: Record<string, unknown> = {}) {
34  return {
35    id: "pg1",
36    workspaceId: "ws1",
37    teamId: null,
38    title: "Magic page",
39    icon: null,
40    description: null,
41    isPrivate: false,
42    archivedAt: null,
43    trashedAt: null,
44    deletedAt: null,
45    createdBy: "u1",
46    createdAt: NOW,
47    updatedAt: NOW,
48    blocks: [
49      {
50        id: "b1",
51        type: "paragraph",
52        content: { text },
53        properties: null,
54        position: "a0",
55        deletedAt: null,
56        createdBy: "u1",
57        createdAt: NOW,
58        updatedAt: NOW,
59      },
60    ],
61    ...overrides,
62  };
63}
64 
65function fakePrisma(opts: {
66  page?: unknown | null;
67  embeddings?: Record<string, unknown>[];
68} = {}) {
69  const upsert = vi.fn(async (_args: Record<string, unknown>) => ({}));
70  const deleteMany = vi.fn(async (_args: Record<string, unknown>) => ({ count: 1 }));
71  const update = vi.fn(async (_args: Record<string, unknown>) => ({}));
72  const del = vi.fn(async (_args: Record<string, unknown>) => ({}));
73  const auditCreate = vi.fn(async (_args: Record<string, unknown>) => ({}));
74  const prisma = {
75    page: { findUnique: vi.fn(async () => opts.page) },
76    databaseRecord: { findUnique: vi.fn(async () => null) },
77    share: { findMany: vi.fn(async () => []) },
78    permission: { findMany: vi.fn(async () => []) },
79    aiEmbedding: {
80      upsert,
81      deleteMany,
82      update,
83      delete: del,
84      findMany: vi.fn(async () => opts.embeddings ?? []),
85    },
86    auditLog: { create: auditCreate },
87  } as unknown as PrismaClient;
88  return { prisma, upsert, deleteMany, update, del, auditCreate };
89}
90 
91function fakeSearch(updateDocErr?: (docId: string) => Error | undefined) {
92  const updateDoc = vi.fn(async (docId: string, patch: Record<string, unknown>) => {
93    if (updateDocErr) {
94      const err = updateDocErr(docId);
95      if (err) throw err;
96    }
97    return {};
98  });
99  const search = { updateDoc } as unknown as SearchClient;
100  return { search, updateDoc };
101}
102 
103function fakeGateway(dimension: number, model = "mock-embed", fail?: (texts: string[]) => Error) {
104  const embed = vi.fn(async (params: { texts: string[] }): Promise<EmbedResult> => {
105    if (fail) throw fail(params.texts);
106    const vectors = params.texts.map((t) => {
107      const v = new Array<number>(dimension).fill(0);
108      v[0] = Math.min(1, t.length / 100);
109      if (dimension > 1) v[1] = 0.5;
110      return v;
111    });
112    return {
113      model,
114      dimension,
115      vectors,
116      usage: { tokens: Math.ceil(params.texts.join("").length / 4) },
117      latencyMs: 1,
118      provider: "custom",
119    };
120  });
121  const gateway = { embed } as unknown as AiGateway;
122  return { gateway, embed };
123}
124 
125const ENV: WorkerEnv = { SEARCH_VECTOR_DIMENSION: 8, AI_PROVIDER: "mock" } as unknown as WorkerEnv;
126 
127/** > CHUNK_BODY_THRESHOLD (1600) → multi-chunk page. */
128const LONG_BODY =
129  "The secret launch code is Zanzibar-7. " +
130  "The factory's quarterly target is 4,000 widgets by Q4. ".repeat(40);
131 
132const JOB: EmbedChunksJobData = {
133  type: "embed:chunks",
134  workspaceId: "ws1",
135  docType: "page",
136  refId: "pg1",
137  chunkIds: [],
138};
139 
140// ---------------------------------------------------------------------------
141// chunkDocId
142// ---------------------------------------------------------------------------
143 
144describe("chunkDocId", () => {
145  it("reconstructs <docType>:<refId>#<chunkId>", () => {
146    expect(chunkDocId("page", "pg1", "c0")).toBe("page:pg1#c0");
147    expect(chunkDocId("record", "r9", "c3")).toBe("record:r9#c3");
148  });
149});
150 
151// ---------------------------------------------------------------------------
152// handleEmbedChunks — the real embed path
153// ---------------------------------------------------------------------------
154 
155describe("handleEmbedChunks (TASK-0125 real path)", () => {
156  it("embeds whole doc + every chunk in ONE gateway call, merges vectors, upserts rows", async () => {
157    const { prisma, upsert, deleteMany } = fakePrisma({ page: pageRow(LONG_BODY) });
158    const { search, updateDoc } = fakeSearch();
159    const { gateway, embed } = fakeGateway(8);
160    const res = await handleEmbedChunks({ prisma, env: ENV, search, job: JOB, jobId: "j-1", gateway });
161 
162    expect(res.skipped).toBe(false);
163    expect(res.blocked).toBeUndefined();
164    expect(res.model).toBe("mock-embed");
165    expect(res.dimension).toBe(8);
166    expect(embed).toHaveBeenCalledTimes(1); // batched — ONE call for the whole item
167    expect(res.embedded).toBeGreaterThan(1); // whole doc + ≥ 2 chunks
168 
169    // chunk-id enumeration parity: the refs are page:pg1 then page:pg1#c<n>
170    expect(res.chunkRefs[0]).toBe("page:pg1");
171    expect(res.chunkRefs[1]).toBe("page:pg1#c0");
172    const n = res.embedded;
173    expect(n).toBe(res.chunkRefs.length);
174    for (let i = 0; i < n - 1; i++) expect(res.chunkRefs[i + 1]).toBe(`page:pg1#c${i}`);
175 
176    // vector merged onto each content doc with the model stamp
177    expect(updateDoc).toHaveBeenCalledTimes(n);
178    const firstCall = updateDoc.mock.calls[0]![0] as string;
179    const firstPatch = updateDoc.mock.calls[0]![1] as { vector: number[]; vectorModel: string };
180    expect(firstCall).toBe("page:pg1");
181    expect(firstPatch.vector).toHaveLength(8);
182    expect(firstPatch.vectorModel).toBe("mock-embed");
183    expect(updateDoc.mock.calls[n - 1]![0]).toBe(`page:pg1#c${n - 2}`);
184 
185    // registry row per chunkRef, keyed [workspaceId, chunkRef]
186    expect(upsert).toHaveBeenCalledTimes(n);
187    const firstUpsert = upsert.mock.calls[0]![0] as {
188      where: { workspaceId_chunkRef: { workspaceId: string; chunkRef: string } };
189      create: { workspaceId: string; chunkRef: string; vectorId: string; model: string; dimension: number; chunkText: string };
190    };
191    expect(firstUpsert.where.workspaceId_chunkRef).toEqual({ workspaceId: "ws1", chunkRef: "page:pg1" });
192    expect(firstUpsert.create.vectorId).toBe("page:pg1");
193    expect(firstUpsert.create.model).toBe("mock-embed");
194    expect(firstUpsert.create.dimension).toBe(8);
195    expect(firstUpsert.create.chunkText.length).toBeGreaterThan(0);
196 
197    // the prune pass runs but keeps every ref (NOT-in) — nothing stale to go
198    expect(deleteMany).toHaveBeenCalledTimes(1);
199    const where = (deleteMany.mock.calls[0]![0] as { where: Record<string, unknown> }).where;
200    expect((where.NOT as { chunkRef: { in: string[] } }).chunkRef.in).toEqual(res.chunkRefs);
201  });
202 
203  it("short body (≤ threshold) embeds the whole doc ONLY", async () => {
204    const { prisma } = fakePrisma({ page: pageRow("The secret launch code is Zanzibar-7.") });
205    const { search, updateDoc } = fakeSearch();
206    const { gateway, embed } = fakeGateway(8);
207    const res = await handleEmbedChunks({ prisma, env: ENV, search, job: JOB, jobId: "j-2", gateway });
208    expect(res.embedded).toBe(1);
209    expect(res.chunkRefs).toEqual(["page:pg1"]);
210    expect(embed).toHaveBeenCalledTimes(1);
211    expect(updateDoc).toHaveBeenCalledTimes(1);
212    expect(updateDoc.mock.calls[0]![0]).toBe("page:pg1");
213  });
214 
215  it("source vanished → prunes registry rows, skips cleanly", async () => {
216    const { prisma, deleteMany } = fakePrisma({ page: null });
217    const { search, updateDoc } = fakeSearch();
218    const { gateway, embed } = fakeGateway(8);
219    const res = await handleEmbedChunks({ prisma, env: ENV, search, job: JOB, jobId: "j-3", gateway });
220    expect(res.skipped).toBe(true);
221    expect(res.reason).toBe("source vanished");
222    expect(res.embedded).toBe(0);
223    expect(embed).not.toHaveBeenCalled();
224    expect(updateDoc).not.toHaveBeenCalled();
225    expect(deleteMany).toHaveBeenCalledTimes(1);
226    const where = deleteMany.mock.calls[0]![0] as {
227      where: { workspaceId: string; OR: { chunkRef: string | { startsWith: string } }[] };
228    };
229    expect(where.where.workspaceId).toBe("ws1");
230    // exact whole-doc + #-prefixed chunks (no sibling-prefix collisions)
231    expect(where.where.OR).toEqual([
232      { chunkRef: "page:pg1" },
233      { chunkRef: { startsWith: "page:pg1#" } },
234    ]);
235  });
236 
237  it("prune pass: stale registry rows are deleted, current chunkRefs kept", async () => {
238    const { prisma, deleteMany } = fakePrisma({ page: pageRow(LONG_BODY) });
239    const { search } = fakeSearch();
240    const { gateway } = fakeGateway(8);
241    const res = await handleEmbedChunks({ prisma, env: ENV, search, job: JOB, jobId: "j-4", gateway });
242    expect(deleteMany).toHaveBeenCalledTimes(1);
243    const where = (deleteMany.mock.calls[0]![0] as { where: Record<string, unknown> }).where;
244    // exact whole-doc + #-prefixed chunks (no sibling-prefix collisions)
245    expect(where.OR).toEqual([
246      { chunkRef: "page:pg1" },
247      { chunkRef: { startsWith: "page:pg1#" } },
248    ]);
249    // keep list == exactly the refs just embedded (whole doc + c0..cN)
250    expect((where.NOT as { chunkRef: { in: string[] } }).chunkRef.in).toEqual(res.chunkRefs);
251  });
252 
253  it("source with zero chunks after shrink → prune keeps only the whole doc", async () => {
254    const { prisma, deleteMany } = fakePrisma({ page: pageRow("short body only") });
255    const { search } = fakeSearch();
256    const { gateway } = fakeGateway(8);
257    await handleEmbedChunks({ prisma, env: ENV, search, job: JOB, jobId: "j-6", gateway });
258    const where = (deleteMany.mock.calls[0]![0] as { where: Record<string, unknown> }).where;
259    expect((where.NOT as { chunkRef: { in: string[] } }).chunkRef.in).toEqual(["page:pg1"]);
260  });
261 
262  it("dimension mismatch → BLOCKS the write + emits ai.reembed_required (no upsert, no merge)", async () => {
263    const { prisma, upsert, auditCreate } = fakePrisma({ page: pageRow(LONG_BODY) });
264    const { search, updateDoc } = fakeSearch();
265    const { gateway } = fakeGateway(1536, "new-model-dim1536");
266    const res = await handleEmbedChunks({ prisma, env: ENV, search, job: JOB, jobId: "j-5", gateway });
267    expect(res.skipped).toBe(true);
268    expect(res.blocked).toBe(true);
269    expect(res.reason).toBe("dimension_mismatch");
270    expect(res.dimension).toBe(1536);
271    expect(updateDoc).not.toHaveBeenCalled();
272    expect(upsert).not.toHaveBeenCalled();
273    expect(auditCreate).toHaveBeenCalledTimes(1);
274    const audit = auditCreate.mock.calls[0]![0] as {
275      data: { action: string; metadata: { reason: string } };
276    };
277    expect(audit.data.action).toBe("ai.reembed_required");
278    expect(audit.data.metadata.reason).toBe("dimension_mismatch");
279  });
280});
281 
282// a long body always expands to ≥ 2 chunks (it's 4x the 800-char budget)
283function LongBodyCount(): number {
284  return 4; // c0..c3 (see derivation below — asserted against deriveItemChunks)
285}
286function resChunkRefsFor(n: number): string[] {
287  return Array.from({ length: n }, (_, i) => `page:pg1#c${i}`);
288}
289 
290// ---------------------------------------------------------------------------
291// deriveItemChunks — enumeration parity with the indexer
292// ---------------------------------------------------------------------------
293 
294describe("deriveItemChunks", () => {
295  it("returns whole doc + c0..cN-1 in indexer order (page)", async () => {
296    const { prisma } = fakePrisma({ page: pageRow(LONG_BODY) });
297    const chunks = await deriveItemChunks(prisma, "page", "pg1");
298    expect(chunks).not.toBeNull();
299    const ids = chunks!.map((c) => c.docId);
300    expect(ids[0]).toBe("page:pg1");
301    for (let i = 1; i < ids.length; i++) expect(ids[i]).toBe(`page:pg1#c${i - 1}`);
302    // long body → strictly more chunks than the whole doc
303    expect(ids.length).toBeGreaterThan(2);
304    // the whole-doc text is the search body (title + block text)
305    expect(chunks![0]!.text).toContain("Zanzibar-7");
306    expect(chunks![0]!.text).toContain("Magic page");
307    // chunk texts are the chunked body
308    expect(chunks![1]!.text.length).toBeGreaterThan(0);
309    expect(chunks!.every((c) => c.workspaceId === "ws1")).toBe(true);
310  });
311 
312  it("returns null for a trashed page", async () => {
313    const { prisma } = fakePrisma({ page: pageRow("x", { trashedAt: NOW }) });
314    expect(await deriveItemChunks(prisma, "page", "pg1")).toBeNull();
315  });
316});
317 
318// ---------------------------------------------------------------------------
319// handleReembed — the nightly sweep (parity verification)
320// ---------------------------------------------------------------------------
321 
322describe("handleReembed (nightly sweep)", () => {
323  const row = (i: number, overrides: Record<string, unknown> = {}) => ({
324    id: `row-${i}`,
325    workspaceId: "ws1",
326    chunkRef: `page:pg1#c${i}`,
327    chunkText: `chunk text ${i}`,
328    vectorId: `page:pg1#c${i}`,
329    model: "mock-embed",
330    dimension: 8,
331    metadata: null,
332    createdAt: NOW,
333    ...overrides,
334  });
335 
336  it("re-embeds every row in batches, merges vectors, refreshes model/dimension", async () => {
337    const rows = Array.from({ length: 5 }, (_, i) => row(i));
338    const { prisma, update: embUpdate, del } = fakePrisma({ embeddings: rows });
339    const { search, updateDoc } = fakeSearch();
340    const { gateway, embed } = fakeGateway(8, "mock-embed-v2");
341    const res = await handleReembed({ prisma, env: ENV, search, jobId: "r-1", gateway });
342    expect(res.total).toBe(5);
343    expect(res.embedded).toBe(5);
344    expect(res.failed).toBe(0);
345    expect(res.driftAlert).toBe(false);
346    expect(res.blocked).toBeUndefined();
347    expect(res.model).toBe("mock-embed-v2");
348    expect(embed).toHaveBeenCalledTimes(1); // 5 < batch of 100
349    expect(updateDoc).toHaveBeenCalledTimes(5);
350    expect(updateDoc.mock.calls[0]![1]).toMatchObject({ vectorModel: "mock-embed-v2" });
351    expect(embUpdate).toHaveBeenCalledTimes(5);
352    expect(del).not.toHaveBeenCalled();
353  });
354 
355  it("orphan rows (vanish index doc) are pruned + counted as failed (drift alerts at > 2%)", async () => {
356    const rows = Array.from({ length: 10 }, (_, i) => row(i));
357    const { prisma, del, auditCreate } = fakePrisma({ embeddings: rows });
358    const { search } = fakeSearch((docId) =>
359      docId === "page:pg1#c3" ? new Error("Not Found 404") : undefined,
360    );
361    const { gateway } = fakeGateway(8);
362    const res = await handleReembed({ prisma, env: ENV, search, jobId: "r-2", gateway });
363    expect(res.total).toBe(10);
364    expect(res.embedded).toBe(9);
365    expect(res.failed).toBe(1);
366    expect(del).toHaveBeenCalledTimes(1); // the orphan row is deleted
367    expect(res.driftAlert).toBe(true); // 1/10 = 10% > 2% threshold
368    expect(auditCreate).toHaveBeenCalledTimes(1);
369  });
370 
371  it("drift > 2% emits ai.reembed_required (count parity)", async () => {
372    const rows = Array.from({ length: 10 }, (_, i) => row(i));
373    const { prisma, auditCreate } = fakePrisma({ embeddings: rows });
374    // 3 of 10 fail → 30% drift > 2%
375    const { search } = fakeSearch((docId) => {
376      const n = Number(docId.split("#c")[1]);
377      if (n < 3) return new Error("404 gone");
378      return undefined;
379    });
380    const { gateway } = fakeGateway(8);
381    const res = await handleReembed({ prisma, env: ENV, search, jobId: "r-3", gateway });
382    expect(res.driftAlert).toBe(true);
383    expect(auditCreate).toHaveBeenCalledTimes(1);
384    const audit = auditCreate.mock.calls[0]![0] as {
385      data: { action: string; metadata: { reason: string } };
386    };
387    expect(audit.data.action).toBe("ai.reembed_required");
388    expect(audit.data.metadata.reason).toBe("count_parity_drift");
389  });
390 
391  it("dimension mismatch on the probe batch BLOCKS the whole run", async () => {
392    const rows = Array.from({ length: 5 }, (_, i) => row(i));
393    const { prisma, update: embUpdate, auditCreate } = fakePrisma({ embeddings: rows });
394    const { search, updateDoc } = fakeSearch();
395    const { gateway } = fakeGateway(1536, "new-model-dim1536");
396    const res = await handleReembed({ prisma, env: ENV, search, jobId: "r-4", gateway });
397    expect(res.blocked).toBe(true);
398    expect(res.embedded).toBe(0);
399    expect(updateDoc).not.toHaveBeenCalled();
400    expect(embUpdate).not.toHaveBeenCalled();
401    expect(auditCreate).toHaveBeenCalledTimes(1);
402  });
403 
404  it("empty registry → no-op", async () => {
405    const { prisma } = fakePrisma({ embeddings: [] });
406    const { search } = fakeSearch();
407    const { gateway, embed } = fakeGateway(8);
408    const res = await handleReembed({ prisma, env: ENV, search, jobId: "r-5", gateway });
409    expect(res).toEqual({ total: 0, embedded: 0, failed: 0, driftAlert: false });
410    expect(embed).not.toHaveBeenCalled();
411  });
412});
413 
414// ---------------------------------------------------------------------------
415// payload validation
416// ---------------------------------------------------------------------------
417 
418describe("aiJobSchema — embed:chunks / embed:reembed validation", () => {
419  it("accepts a valid embed:chunks payload (chunkIds optional — re-derived)", () => {
420    expect(aiJobSchema.safeParse(JOB).success).toBe(true);
421    // chunkIds is informational now (the handler re-derives) — [] is valid
422    const parsed = aiJobSchema.parse(JOB) as { type: string; chunkIds: string[] };
423    expect(parsed.type).toBe("embed:chunks");
424    expect(parsed.chunkIds).toEqual([]);
425  });
426 
427  it("accepts embed:reembed (no fields)", () => {
428    expect(aiJobSchema.safeParse({ type: "embed:reembed" }).success).toBe(true);
429  });
430 
431  it("rejects an unknown docType (only page/record are chunkable)", () => {
432    expect(aiJobSchema.safeParse({ ...JOB, docType: "comment" }).success).toBe(false);
433  });
434 
435  it("rejects a missing workspaceId", () => {
436    const { workspaceId: _drop, ...rest } = JOB;
437    expect(aiJobSchema.safeParse(rest).success).toBe(false);
438  });
439 
440  it("coexists with the other variants in the union", () => {
441    expect(aiJobSchema.safeParse({ type: "embed", workspaceId: "ws1", texts: ["x"] }).success).toBe(true);
442    expect(aiJobSchema.safeParse({ type: "usage.rollup" }).success).toBe(true);
443  });
444});
445