ai-embedding.ts

ai-embedding.tsapps/worker/src/processors/ai-embedding.tsraw ← processors22,492 B← Back to the app
1/**
2 * apps/worker — RAG embedding pipeline (TASK-0125, replaces the TASK-0099 §4
3 * Phase-12 stub).
4 *
5 * Consumed from the `ai` queue (see ai.processor.ts dispatch). Jobs:
6 *
7 *   { type: "embed:chunks", workspaceId, docType, refId, chunkIds? }
8 *      One (re)indexed page/record: derive its chunk set from the SOURCE
9 *      (the same derivation the search indexer uses — whole doc + `#c<n>`
10 *      chunk docs), embed every chunk text in ONE `AiGateway.embed` call,
11 *      merge each vector onto its content doc (`search.updateDoc` partial
12 *      merge — single-writer rule, TASK-0096) and upsert the `ai_embeddings`
13 *      registry row per chunk (unique `[workspaceId, chunkRef]`).
14 *
15 *      The chunk set is re-derived from the source rows (not taken from the
16 *      payload): the source is the truth, and a shrink/grow of the body
17 *      changes the chunk count. Stale registry rows (chunks that no longer
18 *      exist) are deleted so `ai_embeddings` count == index chunk count
19 *      (the parity the DoD requires).
20 *
21 *   { type: "embed:reembed" }
22 *      Nightly full re-embed (03:30 UTC repeatable, registered in index.ts):
23 *      walk every `ai_embeddings` row in batches of 100, re-embed, merge the
24 *      fresh vector, and verify count parity — drift > 2% emits the
25 *      `ai.reembed_required` admin notification (the model-change re-index
26 *      rule: a new model simply re-embeds everything; a DIMENSION mismatch
27 *      instead emits `ai.reembed_required` and BLOCKS the write — the index
28 *      must be rolled via scripts/opensearch-reindex.ts first).
29 *
30 * Provider (recorded decision): the worker honors `AI_PROVIDER` exactly like
31 * the API's `gateway-provider.ts` — `AI_PROVIDER=mock` builds a
32 * deterministic `MockProvider` (dimension = `SEARCH_VECTOR_DIMENSION`), so
33 * the dev stack + E2E run the full RAG surface offline. The old
34 * `AI_EMBEDDING_PROVIDER` gate is superseded: the real embed now flows
35 * through the `@factory/ai` gateway (retries, budget ledger, circuit
36 * breaker — the same surface the chat features use).
37 *
38 * Reuse (recorded decision, spec conflict §29): the spec's separate
39 * `ai-chunks` OpenSearch index is NOT built — the Phase 8 single
40 * `factory-content` index already carries the `vector` knn_vector field,
41 * the `grants` permission snapshot, `type: "chunk"` and per-chunk `body`,
42 * and the search-index pipeline already emits the chunk docs. RAG merges
43 * vectors onto those same docs and reuses the Phase 8 visibility filter at
44 * query time — one index, one filter, no divergent snapshot implementation
45 * (the spec's own Security section requires the filter to be the Phase 8
46 * one, "import, don't duplicate").
47 */
48import type { PrismaClient } from "@prisma/client";
49import { z } from "zod";
50import {
51  AiGateway,
52  AnthropicAdapter,
53  BudgetService,
54  CustomAdapter,
55  GoogleAdapter,
56  MockProvider,
57  OpenAIAdapter,
58  OssAdapter,
59  resolveAiConfig,
60  UsageLedger,
61  type ProviderAdapter,
62  type ProviderId,
63  type ProviderRegistry,
64  type ResolvedAiConfig,
65  type UsageRecord,
66} from "@factory/ai";
67import { CONTENT_INDEX_ALIAS } from "@factory/search";
68import type { SearchClient } from "@factory/search";
69import { asJson } from "../lib/prisma.js";
70import { getLogger } from "../logger.js";
71import type { WorkerEnv } from "../config.js";
72import { buildPageDoc, buildRecordDoc, expandChunkDocs } from "./search-index/builders.js";
73import type { IndexerDeps } from "./search-index/builders.js";
74 
75const log = getLogger("ai-embedding");
76 
77/** Only page/record bodies are chunked (TASK-0099 §3), so only those have
78 *  chunk docs to embed. Fail fast on any other type. */
79const CHUNKABLE_DOC_TYPES = ["page", "record"] as const;
80 
81export const embedChunksJobSchema = z.object({
82  type: z.literal("embed:chunks"),
83  workspaceId: z.string().min(1),
84  docType: z.enum(CHUNKABLE_DOC_TYPES),
85  refId: z.string().min(1),
86  /**
87   * The chunk ids (`c0`, `c1`, …) as computed at index time. TASK-0125:
88   * INFORMATIONAL only — the handler re-derives the chunk set from the
89   * source (truth). `[]` (short body → whole doc only) is valid now.
90   */
91  chunkIds: z.array(z.string().min(1)).default([]),
92});
93export type EmbedChunksJobData = z.infer<typeof embedChunksJobSchema>;
94 
95/** The nightly full re-embed (registered as a BullMQ repeatable in index.ts). */
96export const reembedJobSchema = z.object({
97  type: z.literal("embed:reembed"),
98});
99export type ReembedJobData = z.infer<typeof reembedJobSchema>;
100 
101/** Reconstruct the content-doc id of a chunk (`<docType>:<refId>#<chunkId>`). */
102export function chunkDocId(docType: string, refId: string, chunkId: string): string {
103  return `${docType}:${refId}#${chunkId}`;
104}
105 
106/** The registry prefix for every chunk of one item (whole doc + chunks). */
107export function itemChunkPrefix(docType: string, refId: string): string {
108  return `${docType}:${refId}`;
109}
110 
111export interface EmbedChunksResult {
112  /** true when there is nothing to embed (source vanished → cleanup only). */
113  skipped: boolean;
114  reason?: string;
115  /** Number of chunks embedded (vectors merged + registry rows upserted). */
116  embedded: number;
117  /** Stale registry rows deleted (chunks that no longer exist). */
118  pruned: number;
119  /** The content-doc ids embedded (whole doc + chunk docs). */
120  chunkRefs: string[];
121  /** The provider model the vectors were produced with. */
122  model?: string;
123  /** The vector dimension written (vs the index's SEARCH_VECTOR_DIMENSION). */
124  dimension?: number;
125  /** true when the dimension mismatch blocked the write (re-index required). */
126  blocked?: boolean;
127}
128 
129// ---------------------------------------------------------------------------
130// Worker gateway (AI_PROVIDER-honoring, like the API's gateway-provider)
131// ---------------------------------------------------------------------------
132 
133/**
134 * Build the worker's `AiGateway`. Mirrors the API's `buildAiGateway`
135 * provider selection, minus the workspace settings (the worker never
136 * resolves per-workspace config — platform env only) and minus the budget
137 * gate (embeds are batch/background: the per-workspace monthly budget is
138 * enforced on the request path by the API's guard; the gateway's own
139 * submit-time check stays as the backstop with limits disabled here).
140 *
141 * `AI_PROVIDER=mock` → deterministic `MockProvider` at the index's dimension
142 * (the dev stack + E2E run the full pipeline offline).
143 */
144export function createWorkerAiGateway(env: WorkerEnv): AiGateway {
145  const resolveConfig = (workspaceId: string): ResolvedAiConfig => {
146    const provider: ProviderId = env.AI_PROVIDER === "mock" ? "custom" : (env.AI_PROVIDER || "openai") as ProviderId;
147    const config = resolveAiConfig(null, null, {
148      provider,
149      chatModel: env.AI_MODEL,
150      embedModel: env.AI_EMBEDDING_MODEL,
151      ...(env.AI_BASE_URL ? { baseUrl: env.AI_BASE_URL } : {}),
152    });
153    config.workspaceId = workspaceId;
154    // Background embeds: no per-call budget gate here (the API's
155    // service-level guard enforces the monthly budget on the request path).
156    // `tokensPerMonth: 0` = unlimited; `requestsPerMinute` must stay
157    // POSITIVE (0 would mean "no requests allowed" — the check is `>=`).
158    config.budgets = { ...config.budgets, tokensPerMonth: 0, requestsPerMinute: Math.max(1, config.budgets.requestsPerMinute) };
159    return config;
160  };
161  const registry: ProviderRegistry = {
162    get(providerId: ProviderId, config: ResolvedAiConfig): ProviderAdapter {
163      if (env.AI_PROVIDER === "mock") {
164        return new MockProvider(providerId === "anthropic" ? "anthropic" : "custom", env.SEARCH_VECTOR_DIMENSION);
165      }
166      const baseUrl = config.endpoint?.baseUrl ?? "https://api.openai.com/v1";
167      const key = env.AI_API_KEY;
168      switch (providerId) {
169        case "openai":
170          return new OpenAIAdapter({ baseUrl, apiKey: key });
171        case "anthropic":
172          return new AnthropicAdapter({ apiKey: key });
173        case "google":
174          return new GoogleAdapter({ apiKey: key });
175        case "oss":
176          return new OssAdapter({ baseUrl, apiKey: key });
177        default:
178          return new CustomAdapter({
179            baseUrl,
180            apiKey: key,
181            ...(config.endpoint?.extraHeaders ? { extraHeaders: config.endpoint.extraHeaders } : {})
182          });
183      }
184    },
185  };
186  const ledger = new UsageLedger({
187    sink: async (record: UsageRecord) => {
188      log.info(
189        { workspaceId: record.workspaceId, kind: record.kind, model: record.model, provider: record.provider, status: record.status },
190        "ai usage (worker)",
191      );
192    },
193  });
194  return new AiGateway({ resolveConfig, registry, budgets: new BudgetService(ledger), ledger });
195}
196 
197/**
198 * One process-lifetime gateway per env snapshot (the factory is cheap but
199 * the circuit breaker state is useful to keep). Tests pass their own
200 * gateway to the handlers and never touch this.
201 */
202let workerGateway: AiGateway | null = null;
203let workerGatewayEnv: WorkerEnv | null = null;
204export function workerAiGateway(env: WorkerEnv): AiGateway {
205  if (!workerGateway || workerGatewayEnv !== env) {
206    workerGateway = createWorkerAiGateway(env);
207    workerGatewayEnv = env;
208  }
209  return workerGateway;
210}
211 
212// ---------------------------------------------------------------------------
213// embed:chunks — one (re)indexed item
214// ---------------------------------------------------------------------------
215 
216/** The chunk set of one item, derived from the SOURCE (index-time truth). */
217export async function deriveItemChunks(
218  prisma: PrismaClient,
219  docType: "page" | "record",
220  refId: string,
221): Promise<{ docId: string; workspaceId: string; text: string }[] | null> {
222  const deps: IndexerDeps = { prisma };
223  const doc = docType === "page" ? await buildPageDoc(deps, refId) : await buildRecordDoc(deps, refId);
224  if (!doc) return null; // source vanished (trashed/purged) → caller prunes
225  const expanded = expandChunkDocs(doc);
226  const out: { docId: string; workspaceId: string; text: string }[] = [];
227  // The whole doc ALWAYS gets a vector (RAG retrieves it even for short
228  // bodies — the mock/short-page E2E depends on this; TASK-0099 only chunks
229  // long bodies, it does not stop the parent from being a vector target).
230  out.push({ docId: doc.docId, workspaceId: doc.workspaceId, text: doc.body });
231  for (const d of expanded) {
232    if (d.type === "chunk") out.push({ docId: d.docId, workspaceId: d.workspaceId, text: d.body });
233  }
234  return out;
235}
236 
237/**
238 * Handle one `embed:chunks` job. Pure (no BullMQ here) so it is unit-tested
239 * Docker-free with a fake Prisma + a fake SearchClient + an injected gateway.
240 */
241export async function handleEmbedChunks(opts: {
242  prisma: PrismaClient;
243  env: WorkerEnv;
244  search: SearchClient;
245  job: EmbedChunksJobData;
246  /** Present unless the caller has no job id (e.g. direct unit invocation). */
247  jobId: string | undefined;
248  /** Injectable gateway (tests pass a fake; production = workerAiGateway). */
249  gateway?: AiGateway;
250}): Promise<EmbedChunksResult> {
251  const { prisma, env, search, job } = opts;
252  const gateway = opts.gateway ?? workerAiGateway(env);
253  const { workspaceId, docType, refId } = job;
254 
255  // 1) Re-derive the chunk set from the source. Source gone → prune the
256  //    registry (the search-index removal already dropped the index docs;
257  //    the registry rows are this job's to own) and stop.
258  const chunks = await deriveItemChunks(prisma, docType, refId);
259  if (!chunks) {
260    const pruned = await pruneItemEmbeddings(prisma, workspaceId, docType, refId, []);
261    log.info({ jobId: opts.jobId, workspaceId, docType, refId, pruned }, "embed:chunks — source vanished; registry pruned");
262    return { skipped: true, reason: "source vanished", embedded: 0, pruned, chunkRefs: [] };
263  }
264 
265  // 2) ONE gateway call for the whole chunk set (batched — the mock and the
266  //    real providers both price per call + per token, and a 50-chunk page
267  //    should not be 50 round trips).
268  let result;
269  try {
270    result = await gateway.embed({ workspaceId, texts: chunks.map((c) => c.text), metadata: { capability: "rag.embed" } });
271  } catch (err) {
272    // Retryable transport/provider failures rethrow (BullMQ backoff retries
273    // the whole job — the embed is idempotent: upsert + partial merge).
274    throw err;
275  }
276 
277  // 3) Dimension guard (the model-change rule): the vectors MUST match the
278  //    index's `vector` dimension. A mismatch means the embedding model
279  //    changed (new dim) — BLOCK the write and emit `ai.reembed_required`
280  //    (the alias rollover script is the fix; writing a wrong-dim vector
281  //    corrupts the index for every other doc's kNN).
282  if (result.dimension !== env.SEARCH_VECTOR_DIMENSION) {
283    await emitReembedRequired(prisma, workspaceId, {
284      reason: "dimension_mismatch",
285      detail: `embed model "${result.model}" produced dim ${result.dimension}, index has ${env.SEARCH_VECTOR_DIMENSION}`,
286    });
287    log.error(
288      { jobId: opts.jobId, workspaceId, model: result.model, got: result.dimension, want: env.SEARCH_VECTOR_DIMENSION },
289      "embed BLOCKED — vector dimension mismatch (alias rollover required)",
290    );
291    return {
292      skipped: true,
293      reason: "dimension_mismatch",
294      embedded: 0,
295      pruned: 0,
296      chunkRefs: [],
297      model: result.model,
298      dimension: result.dimension,
299      blocked: true,
300    };
301  }
302 
303  // 4) Merge each vector onto its content doc + upsert the registry row.
304  let embedded = 0;
305  for (let i = 0; i < chunks.length; i++) {
306    const chunk = chunks[i]!;
307    const vector = result.vectors[i];
308    if (!vector) continue;
309    try {
310      // Partial (merge) update — a full PUT would clobber title/body/grants
311      // (single-writer rule, TASK-0096).
312      await search.updateDoc(chunk.docId, { vector, vectorModel: result.model });
313    } catch (err) {
314      // A vanished doc (race with a delete) is fine — the registry row is
315      // still the record; the next re-embed + prune reconciles.
316      if (!(err instanceof Error && /404/.test(err.message))) {
317        log.warn({ err: (err as Error).message, docId: chunk.docId }, "vector merge failed (row still recorded)");
318      }
319    }
320    await prisma.aiEmbedding.upsert({
321      where: { workspaceId_chunkRef: { workspaceId, chunkRef: chunk.docId } },
322      create: {
323        workspaceId,
324        chunkRef: chunk.docId,
325        chunkText: chunk.text,
326        vectorId: chunk.docId,
327        model: result.model,
328        dimension: result.dimension,
329        metadata: asJson({ docType, refId, index: CONTENT_INDEX_ALIAS }),
330      },
331      update: {
332        chunkText: chunk.text,
333        model: result.model,
334        dimension: result.dimension,
335        metadata: asJson({ docType, refId, index: CONTENT_INDEX_ALIAS }),
336      },
337    });
338    embedded++;
339  }
340 
341  // 5) Prune stale registry rows (a shrunk body ⇒ fewer chunks now).
342  const currentRefs = chunks.map((c) => c.docId);
343  const pruned = await pruneItemEmbeddings(prisma, workspaceId, docType, refId, currentRefs);
344 
345  log.info(
346    { jobId: opts.jobId, workspaceId, docType, refId, embedded, pruned, model: result.model, dimension: result.dimension },
347    "embed:chunks done",
348  );
349  return { skipped: false, embedded, pruned, chunkRefs: currentRefs, model: result.model, dimension: result.dimension };
350}
351 
352/**
353 * Delete the item's registry rows that are NOT in `keepRefs`. The item's
354 * refs are the whole doc (`<docType>:<refId>`) + its chunks (`…#<chunkId>`);
355 * the match is EXACT-on-whole + `#`-prefixed so a sibling item whose refId
356 * is a string-prefix of this one (page `pg1` vs `pg10`) is never touched.
357 */
358async function pruneItemEmbeddings(
359  prisma: PrismaClient,
360  workspaceId: string,
361  docType: string,
362  refId: string,
363  keepRefs: string[],
364): Promise<number> {
365  const prefix = itemChunkPrefix(docType, refId);
366  const res = await prisma.aiEmbedding.deleteMany({
367    where: {
368      workspaceId,
369      OR: [{ chunkRef: prefix }, { chunkRef: { startsWith: `${prefix}#` } }],
370      ...(keepRefs.length > 0 ? { NOT: { chunkRef: { in: keepRefs } } } : {}),
371    },
372  });
373  return res.count;
374}
375 
376// ---------------------------------------------------------------------------
377// embed:reembed — nightly full walk (parity verification)
378// ---------------------------------------------------------------------------
379 
380export interface ReembedResult {
381  /** Registry rows walked. */
382  total: number;
383  /** Rows re-embedded successfully. */
384  embedded: number;
385  /** Rows that failed (source vanished or provider error). */
386  failed: number;
387  /** true when drift (failed/total) exceeded the 2% threshold. */
388  driftAlert: boolean;
389  /** The provider model the re-embed ran with (undefined when nothing ran). */
390  model?: string | undefined;
391  /** true when the dimension mismatch blocked the whole run. */
392  blocked?: boolean;
393}
394 
395const REEMBED_BATCH = 100;
396const REEMBED_DRIFT_THRESHOLD = 0.02; // 2% (spec)
397 
398/**
399 * The nightly full re-embed: walk every `ai_embeddings` row in batches of
400 * 100, re-embed the stored chunk text, merge the fresh vector onto the index
401 * doc, and refresh the row's model/dimension (this is ALSO the model-change
402 * re-index — a new model simply re-embeds everything overnight). Count
403 * parity: failed/total > 2% → `ai.reembed_required` admin notification.
404 *
405 * Pure (no BullMQ) — unit-tested with fakes.
406 */
407export async function handleReembed(opts: {
408  prisma: PrismaClient;
409  env: WorkerEnv;
410  search: SearchClient;
411  jobId: string | undefined;
412  gateway?: AiGateway;
413}): Promise<ReembedResult> {
414  const { prisma, env, search } = opts;
415  const gateway = opts.gateway ?? workerAiGateway(env);
416  const t0 = Date.now();
417 
418  const rows = await prisma.aiEmbedding.findMany({ orderBy: { workspaceId: "asc" } });
419  if (rows.length === 0) {
420    return { total: 0, embedded: 0, failed: 0, driftAlert: false };
421  }
422 
423  let embedded = 0;
424  let failed = 0;
425  let model: string | undefined;
426  // First batch doubles as the dimension probe: if the (possibly new)
427  // model's dimension doesn't match the index, BLOCK the whole run rather
428  // than corrupt it row by row.
429  let blocked = false;
430  for (let start = 0; start < rows.length; start += REEMBED_BATCH) {
431    const batch = rows.slice(start, start + REEMBED_BATCH);
432    let result;
433    try {
434      result = await gateway.embed({
435        workspaceId: batch[0]!.workspaceId,
436        texts: batch.map((r) => r.chunkText),
437        metadata: { capability: "rag.reembed" },
438      });
439    } catch (err) {
440      log.warn({ err: (err as Error).message, batch: batch.length }, "reembed batch embed failed — counting rows as failed");
441      failed += batch.length;
442      continue;
443    }
444    if (start === 0 && result.dimension !== env.SEARCH_VECTOR_DIMENSION) {
445      const ws0 = batch[0]!.workspaceId;
446      await emitReembedRequired(prisma, ws0, {
447        reason: "dimension_mismatch",
448        detail: `reembed model "${result.model}" produced dim ${result.dimension}, index has ${env.SEARCH_VECTOR_DIMENSION}`,
449      });
450      log.error(
451        { jobId: opts.jobId, model: result.model, got: result.dimension, want: env.SEARCH_VECTOR_DIMENSION },
452        "reembed BLOCKED — dimension mismatch (alias rollover required)",
453      );
454      blocked = true;
455      break;
456    }
457    model = result.model;
458    for (let i = 0; i < batch.length; i++) {
459      const row = batch[i]!;
460      const vector = result.vectors[i];
461      if (!vector) {
462        failed++;
463        continue;
464      }
465      try {
466        await search.updateDoc(row.vectorId, { vector, vectorModel: result.model });
467        await prisma.aiEmbedding.update({
468          where: { id: row.id },
469          data: { model: result.model, dimension: result.dimension, chunkText: row.chunkText },
470        });
471        embedded++;
472      } catch (err) {
473        // A vanished index doc (pruned page) is the expected drift source —
474        // prune the orphan row so parity converges on the NEXT run.
475        failed++;
476        await prisma.aiEmbedding
477          .delete({ where: { id: row.id } })
478          .catch(() => undefined);
479        log.warn({ err: (err as Error).message, vectorId: row.vectorId }, "reembed row failed — orphan pruned");
480      }
481    }
482  }
483 
484  const total = rows.length;
485  const drift = total > 0 ? failed / total : 0;
486  const driftAlert = !blocked && drift > REEMBED_DRIFT_THRESHOLD;
487  if (driftAlert) {
488    // The drift can span many workspaces; notify per affected workspace is
489    // overkill for the nightly job — one audit event + log is the spec's
490    // "alert log + admin notification" (the notification targets the first
491    // workspace with a failure, best-effort).
492    const firstFailedWs = rows.find((r) => true)!.workspaceId;
493    await emitReembedRequired(prisma, firstFailedWs, {
494      reason: "count_parity_drift",
495      detail: `${failed}/${total} chunks failed re-embed (${(drift * 100).toFixed(1)}% drift)`,
496    });
497  }
498  log.info(
499    { jobId: opts.jobId, total, embedded, failed, drift: +(drift * 100).toFixed(2), driftAlert, blocked, model, ms: Date.now() - t0 },
500    "reembed done",
501  );
502  return { total, embedded, failed, driftAlert, model, ...(blocked ? { blocked: true } : {}) };
503}
504 
505// ---------------------------------------------------------------------------
506// ai.reembed_required — the config-drift / model-change signal
507// ---------------------------------------------------------------------------
508 
509/**
510 * Emit the `ai.reembed_required` signal: a SYSTEM audit row (workspace-
511 * scoped) + a WARN log. The audit row is the admin console's source (Phase
512 * 14); the log is the on-call's. Best-effort — a failed emit must never
513 * take the embedding job down.
514 */
515export async function emitReembedRequired(
516  prisma: PrismaClient,
517  workspaceId: string,
518  detail: { reason: string; detail: string },
519): Promise<void> {
520  try {
521    await prisma.auditLog.create({
522      data: {
523        actorType: "SYSTEM",
524        actorId: "worker:ai",
525        workspaceId,
526        resourceType: "workspace",
527        resourceId: workspaceId,
528        action: "ai.reembed_required",
529        metadata: asJson({ reason: detail.reason, detail: detail.detail, index: CONTENT_INDEX_ALIAS }),
530      },
531    });
532  } catch (err) {
533    log.warn({ err: (err as Error).message, workspaceId }, "failed to emit ai.reembed_required audit row");
534  }
535  log.warn({ workspaceId, ...detail }, "ai.reembed_required — embedding model/config drifted from the index; run scripts/opensearch-reindex.ts");
536}
537