ai.processor.ts

ai.processor.tsapps/worker/src/processors/ai.processor.tsraw ← processors22,353 B← Back to the app
1/**
2 * apps/worker — AI queue processor (embeddings + generation, stub-real).
3 *
4 * Jobs (queue `ai`):
5 *   {type:"embed", workspaceId, texts: string[], chunkId?: string}
6 *      → POST {AI_BASE_URL}/embeddings {model, input: texts}
7 *      → vector MERGED onto the single content doc (chunkId IS the doc id —
8 *        TASK-0096 removed the 5-index scheme) + an AiEmbedding row
9 *        (workspace/chunkRef upsert). If chunkId is absent a stable auto id
10 *        is used.
11 *   {type:"generate", prompt: string, workspaceId?, conversationId?, userId?}
12 *      → POST {AI_BASE_URL}/chat/completions {model, messages:[{user}]}
13 *      → result returned; when conversationId+userId are present the
14 *        assistant turn is persisted as an AiMessage row.
15 *   {type:"embed:chunks", workspaceId, docType, refId, chunkIds?}
16 *      → TASK-0125 real RAG embedding (ai-embedding.ts): re-derive the item's
17 *        chunk set from the source, ONE AiGateway.embed call (AI_PROVIDER-
18 *        honoring — mock works keyless), merge each vector onto its content
19 *        doc (search.updateDoc partial merge), upsert the ai_embeddings
20 *        registry row, prune stale rows. Dimension mismatch → BLOCK + the
21 *        ai.reembed_required signal (alias rollover required).
22 *   {type:"embed:reembed"}  (TASK-0125)
23 *      → nightly full re-embed (registered as a repeatable in index.ts): walk
24 *        every ai_embeddings row in batches of 100, re-embed + merge, prune
25 *        orphans; count parity drift > 2% → ai.reembed_required.
26 *   {type:"usage.rollup", workspaceId?}  (TASK-0121)
27 *      → nightly (registered as a repeatable in index.ts) + manual: reconcile
28 *        per-workspace month-to-date token totals (sums ai_messages of the
29 *        workspace's own conversations since month start) and enqueue a
30 *        budget.notify for any budget that was crossed but not yet notified
31 *        (the 24 h slot is checked by the notify job — race-safe).
32 *   {type:"budget.notify", workspaceId, used, limit}  (TASK-0121)
33 *      → the budget-exceeded admin notification: IN_APP Notification rows for
34 *        every OWNER/ADMIN + the notification.created stream event (email
35 *        channel via the worker's email decision) + the ai.budget_exceeded
36 *        audit row. Throttled to one per workspace per 24 h (the shared
37 *        ai:budget:notified:<ws> Redis slot, SET NX).
38 *
39 * Dev-safe: without AI_API_KEY every job logs "skipping — no AI key" and
40 * completes (Phase 12 wires the full provider gateway; the request shaping
41 * here is already provider-agnostic and unit-tested with mocked fetch).
42 */
43import { randomUUID } from "node:crypto";
44import { asJson } from "../lib/prisma.js";
45 
46import { Worker, type Job } from "bullmq";
47import { withJobOtel } from "../lib/job-otel.js";
48import type { Redis } from "ioredis";
49import { z } from "zod";
50import type { PrismaClient } from "@prisma/client";
51import { QUEUES, enqueue } from "../queues.js";
52import { getLogger } from "../logger.js";
53import type { WorkerEnv } from "../config.js";
54import type { SearchClient } from "@factory/search";
55import { embedChunksJobSchema, handleEmbedChunks, handleReembed, reembedJobSchema } from "./ai-embedding.js";
56 
57const log = getLogger("ai");
58 
59const embedJobSchema = z.object({
60  type: z.literal("embed"),
61  workspaceId: z.string().min(1),
62  texts: z.array(z.string()).min(1).max(100),
63  /** Content doc id the vector merges into, e.g. "page:abc" (auto if absent). */
64  chunkId: z.string().min(1).optional(),
65  /** DEPRECATED no-op (TASK-0096): the single content index has no per-type
66   * override — accepted for backward compatibility with in-flight jobs. */
67  index: z.string().optional()
68});
69const generateJobSchema = z.object({
70  type: z.literal("generate"),
71  prompt: z.string().min(1),
72  workspaceId: z.string().optional(),
73  conversationId: z.string().optional(),
74  userId: z.string().optional(),
75  system: z.string().optional()
76});
77/**
78 * TASK-0121 — hourly AI usage rollup (worker-owned, registered as a
79 * BullMQ repeatable in `apps/worker/src/index.ts`). Per workspace
80 * (all workspaces when `workspaceId` is absent): sum `ai_messages`
81 * `tokensIn + tokensOut` since month-start (the same range the API's
82 * budget guard uses) and log the metric; when a workspace has crossed
83 * `usageLimits.aiMonthlyTokens` and has NOT been notified this 24 h
84 * (the shared `ai:budget:notified:<ws>` Redis slot), enqueue a
85 * `budget.notify` job — so an admin is alerted within one rollup cycle
86 * even if the workspace never makes another AI call after crossing.
87 */
88const usageRollupJobSchema = z.object({
89  type: z.literal("usage.rollup"),
90  workspaceId: z.string().min(1).optional()
91});
92/**
93 * TASK-0121 — budget-exceeded admin notification (worker-owned): create an
94 * IN_APP Notification row for every workspace OWNER/ADMIN + publish the
95 * `notification.created` domain event (the worker's email-decision consumer
96 * sends the email channel per the user's prefs — the same channel every
97 * notification type uses). The 24 h throttle is the shared Redis slot (the
98 * API's guard hook and this job race on the same key — one notification per
99 * workspace per day).
100 */
101const budgetNotifyJobSchema = z.object({
102  type: z.literal("budget.notify"),
103  workspaceId: z.string().min(1),
104  used: z.number().int().nonnegative(),
105  limit: z.number().int().positive()
106});
107export const aiJobSchema = z.discriminatedUnion("type", [
108  embedJobSchema,
109  generateJobSchema,
110  embedChunksJobSchema,
111  reembedJobSchema,
112  usageRollupJobSchema,
113  budgetNotifyJobSchema
114]);
115export type AiJobData = z.infer<typeof aiJobSchema>;
116 
117// ---------------------------------------------------------------------------
118// Request shaping (exported for unit tests with mocked fetch)
119// ---------------------------------------------------------------------------
120 
121export function embeddingsEndpoint(baseUrl: string): string {
122  return baseUrl.replace(/\/+$/, "") + "/embeddings";
123}
124export function chatEndpoint(baseUrl: string): string {
125  return baseUrl.replace(/\/+$/, "") + "/chat/completions";
126}
127 
128/** Request body for POST /embeddings (shaped independently of the URL). */
129export function buildEmbeddingsRequest(
130  model: string,
131  texts: string[]
132): { model: string; input: string[] } {
133  return { model, input: texts };
134}
135 
136export function buildChatRequest(
137  model: string,
138  prompt: string,
139  system?: string
140): { model: string; messages: { role: string; content: string }[] } {
141  const messages: { role: string; content: string }[] = [];
142  if (system) messages.push({ role: "system", content: system });
143  messages.push({ role: "user", content: prompt });
144  return { model, messages };
145}
146 
147async function aiHeaders(env: WorkerEnv): Promise<Record<string, string>> {
148  return {
149    "content-type": "application/json",
150    authorization: `Bearer ${env.AI_API_KEY}`
151  };
152}
153 
154// ---------------------------------------------------------------------------
155// Chunk → index mapping (vector updates land in the content index)
156// ---------------------------------------------------------------------------
157 
158// ---------------------------------------------------------------------------
159// Chunk → index mapping (TASK-0096: single content index)
160// ---------------------------------------------------------------------------
161 
162// The vector lands on the SAME doc the search indexer wrote: `chunkRef` IS the
163// content doc id (`<type>:<refId>`, user: `user:<ws>:<uid>`). The 5-index
164// scheme (factory_pages/…) was removed in TASK-0096 — there is one
165// `factory-content` index now, so no per-type index map is needed.
166const CONTENT_INDEX = "factory-content";
167 
168// ---------------------------------------------------------------------------
169// Worker
170// ---------------------------------------------------------------------------
171 
172export function createAiWorker(opts: {
173  connection: Redis;
174  prisma: PrismaClient;
175  env: WorkerEnv;
176  search: SearchClient;
177  concurrency?: number;
178  fetchImpl?: typeof fetch;
179}): Worker {
180  const { connection, prisma, env, search } = opts;
181  const doFetch = opts.fetchImpl ?? fetch;
182  const aiBase = env.AI_BASE_URL || "https://api.openai.com/v1";
183 
184  return new Worker(
185    QUEUES.ai,
186    withJobOtel(QUEUES.ai, async (job: Job) => {
187      const started = Date.now();
188      const data = aiJobSchema.parse(job.data);
189 
190      // RAG embedding jobs (TASK-0125) route BEFORE the AI_API_KEY gate:
191      // they flow through the worker's own gateway (AI_PROVIDER-honoring —
192      // `mock` works with no key; a named provider needs AI_API_KEY, which
193      // the gateway resolves from the platform env via `env:AI_API_KEY`).
194      if (data.type === "embed:chunks") {
195        return handleEmbedChunks({ prisma, env, search, job: data, jobId: job.id });
196      }
197      if (data.type === "embed:reembed") {
198        return handleReembed({ prisma, env, search, jobId: job.id });
199      }
200 
201      // TASK-0121 usage/budget jobs are DB+Redis only (no provider call) —
202      // they run regardless of AI_API_KEY. Logic lives in the exported
203      // runUsageRollup / runBudgetNotify (unit-tested without BullMQ).
204      if (data.type === "usage.rollup") {
205        return runUsageRollup(
206          {
207            prisma,
208            redis: connection,
209            enqueueBudgetNotify: (d) => enqueue(connection, QUEUES.ai, d as Record<string, unknown>)
210          },
211          data,
212          job.id ?? null
213        );
214      }
215      if (data.type === "budget.notify") {
216        return runBudgetNotify({ prisma, redis: connection }, data, job.id ?? null);
217      }
218 
219      if (!env.AI_API_KEY) {
220        log.info(
221          { jobId: job.id, type: data.type },
222          "skipping — no AI_API_KEY configured (Phase 12 wires the gateway)"
223        );
224        return { skipped: true, reason: "no ai api key" };
225      }
226 
227      if (data.type === "embed") {
228        return handleEmbed(job, data, started);
229      }
230      return handleGenerate(job, data, started);
231 
232      async function handleEmbed(
233        job: Job,
234        data: z.infer<typeof embedJobSchema>,
235        started: number
236      ): Promise<Record<string, unknown>> {
237        const url = embeddingsEndpoint(aiBase);
238        const body = buildEmbeddingsRequest(env.AI_MODEL, data.texts);
239        const res = await doFetch(url, {
240          method: "POST",
241          headers: await aiHeaders(env),
242          body: JSON.stringify(body)
243        });
244        if (!res.ok) {
245          throw new Error(`embeddings HTTP ${res.status}: ${(await res.text()).slice(0, 300)}`);
246        }
247        const json = (await res.json()) as {
248          data: { index: number; embedding: number[] }[];
249          usage?: { prompt_tokens: number; total_tokens: number };
250        };
251        const vectors = json.data.map((d) => d.embedding);
252        if (vectors.length !== data.texts.length) {
253          throw new Error(`embedding count mismatch: ${vectors.length} != ${data.texts.length}`);
254        }
255 
256        const results: { chunkId: string; dimension: number; indexed: boolean }[] = [];
257        for (let i = 0; i < data.texts.length; i++) {
258          const chunkRef = data.chunkId ?? `auto:${randomUUID()}`;
259          const docId = chunkRef; // chunkRef IS the content doc id
260          const vector = vectors[i]!;
261          let indexed = false;
262          try {
263            // Partial (merge) update of the `vector` field on the single
264            // content doc — a full PUT here would clobber the doc's
265            // title/body/grants (single-writer rule, TASK-0096).
266            await search.updateDoc(docId, { vector });
267            indexed = true;
268          } catch (err) {
269            log.warn(
270              { err: (err as Error).message, docId },
271              "vector merge into content doc failed (row still recorded)"
272            );
273          }
274          await prisma.aiEmbedding.upsert({
275            where: { workspaceId_chunkRef: { workspaceId: data.workspaceId, chunkRef } },
276            create: {
277              workspaceId: data.workspaceId,
278              chunkRef,
279              chunkText: data.texts[i]!,
280              vectorId: docId,
281              model: env.AI_MODEL,
282              dimension: vector.length,
283              metadata: asJson({ index: CONTENT_INDEX })
284            },
285            update: {
286              chunkText: data.texts[i]!,
287              model: env.AI_MODEL,
288              dimension: vector.length,
289              metadata: asJson({ index: CONTENT_INDEX })
290            }
291          });
292          results.push({ chunkId: chunkRef, dimension: vector.length, indexed });
293        }
294        log.info(
295          {
296            jobId: job.id,
297            texts: data.texts.length,
298            usage: json.usage,
299            ms: Date.now() - started
300          },
301          "embeddings stored"
302        );
303        return { embedded: results, usage: json.usage };
304      }
305 
306      async function handleGenerate(
307        job: Job,
308        data: z.infer<typeof generateJobSchema>,
309        started: number
310      ): Promise<Record<string, unknown>> {
311        const url = chatEndpoint(aiBase);
312        const body = buildChatRequest(env.AI_MODEL, data.prompt, data.system);
313        const res = await doFetch(url, {
314          method: "POST",
315          headers: await aiHeaders(env),
316          body: JSON.stringify(body)
317        });
318        if (!res.ok) {
319          throw new Error(`chat HTTP ${res.status}: ${(await res.text()).slice(0, 300)}`);
320        }
321        const json = (await res.json()) as {
322          choices?: { message?: { content?: string } }[];
323          usage?: { prompt_tokens: number; completion_tokens: number; total_tokens: number };
324        };
325        const content = json.choices?.[0]?.message?.content ?? "";
326        if (data.conversationId && data.userId) {
327          await prisma.aiMessage.create({
328            data: {
329              conversationId: data.conversationId,
330              role: "assistant",
331              content,
332              model: env.AI_MODEL,
333              tokensOut: json.usage?.completion_tokens ?? null,
334              latencyMs: Date.now() - started
335            }
336          });
337        }
338        log.info(
339          { jobId: job.id, model: env.AI_MODEL, usage: json.usage, ms: Date.now() - started },
340          "generation complete"
341        );
342        return { content, model: env.AI_MODEL, usage: json.usage };
343      }
344    }),
345    {
346      connection,
347      concurrency: opts.concurrency ?? env.WORKER_CONCURRENCY
348    }
349  );
350}
351 
352// ---------------------------------------------------------------------------
353// TASK-0121 — usage.rollup / budget.notify (exported, unit-testable without
354// BullMQ: the deps are structural, so tests pass plain mocks).
355// ---------------------------------------------------------------------------
356 
357/** The 24 h budget-notification slot key (shared with the API's guard). */
358export function budgetNotifySlot(workspaceId: string): string {
359  return `ai:budget:notified:${workspaceId}`;
360}
361 
362/**
363 * The monthly AI token budget for a workspace (mirrors the API's
364 * `tokensPerMonthFor`): `max(usageLimits.aiMonthlyTokens,
365 * settings.ai.budgets.tokensPerMonth)`; 0 = unlimited.
366 */
367export function aiMonthlyBudget(ws: {
368  usageLimits: unknown;
369  settings: unknown;
370}): number {
371  const usageLimits = (ws.usageLimits ?? {}) as { aiMonthlyTokens?: number };
372  const settings = (ws.settings ?? {}) as { ai?: { budgets?: { tokensPerMonth?: number } } };
373  const fromLimits =
374    typeof usageLimits.aiMonthlyTokens === "number" && Number.isFinite(usageLimits.aiMonthlyTokens)
375      ? Math.max(0, usageLimits.aiMonthlyTokens)
376      : 0;
377  const fromSettings =
378    typeof settings.ai?.budgets?.tokensPerMonth === "number"
379      ? Math.max(0, settings.ai.budgets.tokensPerMonth)
380      : 0;
381  return Math.max(fromLimits, fromSettings);
382}
383 
384export interface RunUsageRollupDeps {
385  /** Structural deps — the worker passes its PrismaClient; unit tests pass a
386   * targeted fake (`as unknown as PrismaClient`, the ai-embedding convention). */
387  prisma: PrismaClient;
388  redis: Redis;
389  /** Enqueue a `budget.notify` job (the worker passes the BullMQ helper;
390   * tests pass a spy). */
391  enqueueBudgetNotify: (data: {
392    type: "budget.notify";
393    workspaceId: string;
394    used: number;
395    limit: number;
396  }) => Promise<unknown>;
397}
398 
399/**
400 * TASK-0121 — `usage.rollup`: reconcile the per-workspace month-to-date
401 * token totals (the AI ledger is `ai_messages`, same range the API budget
402 * guard uses) and sweep for workspaces that CROSSED their budget but have
403 * gone quiet — the API's submit-time guard only notifies on the next call.
404 * A workspace that crossed AND has no 24 h notification slot set gets a
405 * `budget.notify` job (the throttle is checked THERE, race-safe).
406 */
407export async function runUsageRollup(
408  deps: RunUsageRollupDeps,
409  job: z.infer<typeof usageRollupJobSchema>,
410  jobId: string | number | null
411): Promise<Record<string, unknown>> {
412  const t0 = Date.now();
413  const monthStart = new Date();
414  monthStart.setUTCDate(1);
415  monthStart.setUTCHours(0, 0, 0, 0);
416 
417  // Workspaces to reconcile: the single requested one, or — in the
418  // scheduled run — every workspace with at least one AI conversation
419  // (a workspace with none has used 0 and cannot have crossed).
420  const scope = job.workspaceId
421    ? { where: { id: job.workspaceId } }
422    : { where: { aiConversation: { some: {} } } };
423  const workspaces = (await deps.prisma.workspace.findMany({
424    where: scope.where,
425    select: { id: true, settings: true, usageLimits: true }
426  })) as { id: string; settings: unknown; usageLimits: unknown }[];
427 
428  let reconciled = 0;
429  let budgetsExceeded = 0;
430  let enqueued = 0;
431  for (const ws of workspaces) {
432    // Month-to-date tokens for THIS workspace: its conversations' messages,
433    // createdAt >= month start. The aggregate is scoped to the workspace's
434    // own conversation ids — never another workspace's rows.
435    const conversations = await deps.prisma.aiConversation.findMany({
436      where: { workspaceId: ws.id },
437      select: { id: true }
438    });
439    let used = 0;
440    if (conversations.length > 0) {
441      const sums = await deps.prisma.aiMessage.aggregate({
442        where: {
443          conversationId: { in: conversations.map((c) => c.id) },
444          createdAt: { gte: monthStart }
445        },
446        _sum: { tokensIn: true, tokensOut: true }
447      });
448      used = (sums._sum.tokensIn ?? 0) + (sums._sum.tokensOut ?? 0);
449    }
450 
451    const limit = aiMonthlyBudget(ws);
452    reconciled++;
453 
454    if (limit > 0 && used >= limit) {
455      budgetsExceeded++;
456      const notified = await deps.redis.get(budgetNotifySlot(ws.id)).catch(() => null);
457      if (!notified) {
458        await deps.enqueueBudgetNotify({ type: "budget.notify", workspaceId: ws.id, used, limit });
459        enqueued++;
460      }
461    }
462 
463    log.info(
464      { jobId, workspaceId: ws.id, used, limit, ms: Date.now() - t0 },
465      "ai usage rollup (month to date)"
466    );
467  }
468 
469  return { reconciled, budgetsExceeded, enqueued, ms: Date.now() - t0 };
470}
471 
472export interface RunBudgetNotifyDeps {
473  prisma: PrismaClient;
474  redis: Redis;
475}
476 
477/**
478 * TASK-0121 — `budget.notify`: the budget-exceeded admin notification,
479 * worker-owned so a workspace that crossed its budget mid-month and never
480 * calls the API again is still alerted (the API's guard hook only fires on
481 * the NEXT call). Creates an IN_APP Notification for every OWNER/ADMIN +
482 * publishes `notification.created` (the worker's email-decision consumer
483 * owns the email channel — the same path every notification uses).
484 * Throttle: the shared `ai:budget:notified:<ws>` slot (SET NX 24 h) — the
485 * API hook and this job race on it; the loser skips (one per workspace per
486 * 24 h).
487 */
488export async function runBudgetNotify(
489  deps: RunBudgetNotifyDeps,
490  job: { type: "budget.notify"; workspaceId: string; used: number; limit: number },
491  jobId: string | number | null
492): Promise<Record<string, unknown>> {
493  const { workspaceId } = job;
494  const claimed = await deps.redis
495    .set(budgetNotifySlot(workspaceId), "1", "EX", 24 * 3600, "NX")
496    .catch(() => "OK" as const);
497  if (claimed !== "OK") {
498    log.info({ jobId, workspaceId }, "budget notification already sent this 24 h — skipped");
499    return { skipped: true, reason: "already notified" };
500  }
501 
502  const admins = await deps.prisma.workspaceMember.findMany({
503    where: { workspaceId, role: { in: ["OWNER", "ADMIN"] } },
504    select: { userId: true }
505  });
506  if (admins.length === 0) {
507    log.warn({ jobId, workspaceId }, "budget exceeded but no OWNER/ADMIN member — no notification");
508    return { notified: 0 };
509  }
510 
511  const body =
512    `The workspace AI token budget was reached (${job.used} used of ${job.limit} monthly). ` +
513    `New AI calls are blocked until the limit is raised (workspace settings → usageLimits.aiMonthlyTokens).`;
514  const title = "AI budget exceeded";
515  let notified = 0;
516  for (const admin of admins) {
517    const row = await deps.prisma.notification.create({
518      data: {
519        workspaceId,
520        userId: admin.userId,
521        type: "AUTOMATION",
522        channel: "IN_APP",
523        title,
524        body,
525        link: asJson({ type: "workspace", id: workspaceId })
526      },
527      select: { id: true }
528    });
529    // The worker owns the email channel: publish the same
530    // `notification.created` event the API publishes (the email-decision
531    // consumer applies per-user prefs; no actorId → SYSTEM actor).
532    await deps.redis
533      .xadd(
534        "factory:events",
535        "*",
536        "event",
537        "notification.created",
538        "payload",
539        JSON.stringify({
540          workspaceId,
541          notificationId: row.id,
542          userId: admin.userId,
543          type: "AUTOMATION",
544          title,
545          deepLink: { type: "workspace", id: workspaceId }
546        }),
547        "ts",
548        new Date().toISOString()
549      )
550      .catch(() => undefined);
551    notified++;
552  }
553 
554  await deps.prisma.auditLog.create({
555    data: {
556      actorType: "SYSTEM",
557      actorId: "worker:ai",
558      workspaceId,
559      resourceType: "workspace",
560      resourceId: workspaceId,
561      action: "ai.budget_exceeded",
562      metadata: asJson({ used: job.used, limit: job.limit, adminsNotified: notified, source: "usage.rollup" })
563    }
564  });
565 
566  log.info({ jobId, workspaceId, notified }, "ai budget notification sent to admins");
567  return { notified };
568}
569