rag.service.ts · Source

rag.service.tsapps/api/src/modules/ai/rag.service.tsraw ← ai17,929 B← Back to the app
1/**
2 * TASK-0125 — RAG workspace search (query path).
3 *
4 * `POST /api/v1/workspaces/:workspaceId/ai/search` (body `{ query, context? }`)
5 * → SSE answer stream. The documented wire contract (docs/api/ai.md §9):
6 *
7 *   sources → delta* → usage → done
8 *
9 * The FIRST frame is `sources` (the server-side citations, built from the
10 * retrieval results — deterministic, independent of the model). A real model
11 * would stream inline `[n]` markers that the UI maps onto this frame; the
12 * deterministic mock (`metadata.mode = "rag"`) instead streams a cited answer
13 * whose `[n]` marker + source label are parsed from the SAME numbered context
14 * it is given, so the integration/E2E path exercises the real
15 * prompt → provider → server-citation contract.
16 *
17 * Retrieval (the permission-isolation proof): the viewer is resolved by the
18 * search module's fail-closed resolver (non-member → 404; zero memberships →
19 * empty) and the query runs the Phase-8 hybrid path (BM25 + kNN, RRF k=60,
20 * both legs filtered by the requester's tenant + visibility snapshot) —
21 * "import, don't duplicate" (the spec's Security requirement). A user can
22 * therefore never retrieve a chunk of a page they cannot read; the kNN leg
23 * carries the visibility filter in its own `filter` clause (queries.ts).
24 *
25 * Budget/cap: the query embed + the answer chat both count against the
26 * workspace budget (the pre-stream `guard.check`) and the shared per-user
27 * `ai.calls` limiter (30 chat+writing+rag/min). The answer is persisted in a
28 * HIDDEN per-call conversation (`metadata.feature="rag"`, the citation list
29 * in `toolResults`) for the usage ledger — it never appears in the panel.
30 *
31 * Top-12: the spec's "top-12 chunks". Chunk hits are de-duplicated to their
32 * parent (a page over the 1600-char threshold can surface several chunks;
33 * the UI cites the parent), so the de-duped set is capped at 12 — a page
34 * with 3 retrieved chunks cites ONCE (its parent), freeing slots for other
35 * docs (recorded decision: the spec counts chunks, the UI cites sources).
36 */
37import { Inject, Injectable } from "@nestjs/common";
38import type { Request, Response } from "express";
39 
40import type { PrismaClient } from "@prisma/client";
41import type { SearchResult, SearchType } from "@factory/search";
42import type { AiGateway } from "@factory/ai";
43 
44import { PrismaService } from "../../common/prisma.service.js";
45import { logger } from "../../common/logger.js";
46import { ApiError, type ErrorCode, type ErrorDetail } from "../../common/errors.js";
47import { assertCan } from "../../lib/authorize.js";
48import { SearchService } from "../search/search.service.js";
49import {
50  AI_GATEWAY,
51  mapAiError,
52  primeWorkspaceCache,
53  type AiBudgetGuard,
54} from "./gateway-provider.js";
55import { AiBudgetGuardService } from "./budget-guard.service.js";
56import { AiRateLimitService } from "./ai-rate-limit.js";
57import { RAG_CONV_PREFIX } from "./prompts.js";
58import type { RagSearchInput } from "./rag.dto.js";
59 
60/** The spec's top-12 (de-duplicated to parent sources). */
61const RAG_TOP_N = 12;
62/** Hybrid candidates retrieved before de-dup + cap (the kNN leg's k). */
63const RAG_CANDIDATES = 50;
64/** The answer context's per-source text cap (keeps the prompt bounded). */
65const RAG_CONTEXT_CHARS = 1200;
66const RAG_TIMEOUT_MS = 60_000;
67const KEEPALIVE_MS = 10_000;
68 
69/** One citation (the `sources` frame element + the persisted toolResults). */
70export interface RagSource {
71  /** 1-based citation number (the inline `[n]` marker). */
72  n: number;
73  /** The content doc id (`page:<id>`, `page:<id>#c<n>`, …). */
74  docId: string;
75  /** The parent resource id (page/record/… id, not the doc id). */
76  id: string;
77  /** The public type (chunk hits surface as their parent's kind). */
78  type: SearchType;
79  title: string;
80  /** The deep link the citation chip navigates to. */
81  url: string;
82  /** A short snippet for the chip hover / sources list. */
83  snippet: string;
84  /** Which retrieval paths produced the hit: bm25 / vec / bm25+vec. */
85  source: SearchResult["source"];
86  /** Deep-link metadata (the chip navigates by the workspace slug + these). */
87  pageId?: string;
88  blockId?: string;
89  databaseId?: string;
90  recordId?: string;
91}
92 
93/** One SSE frame — the documented wire contract (docs/api/ai.md §9). */
94export type RagSseEvent =
95  | { type: "sources"; sources: RagSource[] }
96  | { type: "delta"; text: string }
97  | { type: "usage"; tokensIn: number; tokensOut: number }
98  | { type: "error"; code: string; message: string; retryable?: boolean; retryAfterSeconds?: number }
99  | { type: "done"; assistantMessageId?: string; userMessageId?: string; reason?: "ok" | "aborted" | "error" };
100 
101function sseFrame(event: RagSseEvent): string {
102  return `data: ${JSON.stringify(event)}\n\n`;
103}
104 
105/** A deduplicated source row (the de-duped, RRF-ordered retrieval result). */
106interface RankedSource {
107  hit: SearchResult;
108  /** The full chunk text (from `_source.body`) for the prompt context. */
109  body: string;
110}
111 
112/**
113 * De-duplicate the fused hybrid results to their PARENT source, preserving
114 * RRF order, and cap at `n`. A chunk hit (`page:<id>#c<k>`) collapses onto
115 * its parent (`page:<id>`); the FIRST (highest-ranked) hit per parent wins
116 * its title/url/snippet, but the body kept is the FIRST chunk's (the text
117 * that actually ranked — more specific than the whole doc for long pages).
118 */
119export function dedupeToParentSources(results: SearchResult[], n: number): SearchResult[] {
120  const byParent = new Map<string, SearchResult>();
121  for (const r of results) {
122    const parent = r.docId.includes("#") ? r.docId.slice(0, r.docId.indexOf("#")) : r.docId;
123    const existing = byParent.get(parent);
124    if (!existing) byParent.set(parent, r);
125    // keep the existing (higher-ranked) hit; nothing to merge (title/url are
126    // identical for parent + chunks — they share the parent's facets)
127  }
128  return [...byParent.values()].slice(0, n);
129}
130 
131/** The numbered context block the RAG prompt builder emits for one source. */
132export function renderRagContext(n: number, label: string, body: string): string {
133  const text = body.length > RAG_CONTEXT_CHARS ? `${body.slice(0, RAG_CONTEXT_CHARS)}…` : body;
134  return `[${n}] Source: ${label}\n${text}`;
135}
136 
137/** The RAG system prompt (a real model follows it; the mock ignores it). */
138export const RAG_SYSTEM_PROMPT =
139  "You are a search-augmented assistant. Answer the user's question using ONLY the numbered " +
140  "workspace sources provided. Cite your sources inline with [n] markers that reference the " +
141  "source number. If the sources do not contain the answer, say so. Never invent workspace " +
142  "content.";
143 
144@Injectable()
145export class RagService {
146  constructor(
147    @Inject(PrismaService) private readonly prisma: PrismaClient,
148    @Inject(AI_GATEWAY) private readonly gateway: AiGateway,
149    @Inject(AiBudgetGuardService) private readonly guardService: AiBudgetGuardService,
150    @Inject(AiRateLimitService) private readonly rateLimit: AiRateLimitService,
151    @Inject(SearchService) private readonly search: SearchService,
152  ) {}
153 
154  private get guard(): AiBudgetGuard {
155    return this.guardService.guard;
156  }
157 
158  async ask(workspaceId: string, userId: string, input: RagSearchInput, req: Request, res: Response): Promise<void> {
159    const q = input.query.trim();
160 
161    // ---- fail-closed workspace gate (the search viewer's documented rule) ----
162    // The viewer is the search module's fail-closed resolver: an explicit
163    // `workspaceId` the caller is NOT a member of simply does not exist for
164    // them → 404 NOT_FOUND (existence hidden, NOT 403). This MUST run before
165    // `assertCan` below — assertCan 403s on a NONE level and would otherwise
166    // mask the 404 boundary. Zero memberships + no workspaceId → null → empty
167    // (the query is never sent).
168    const viewer = await this.search.viewerForWorkspace(userId, workspaceId);
169 
170    // ---- pre-stream gates (canonical JSON envelope on failure) ----
171    // Members always have ≥ VIEW (GUEST→VIEW, MEMBER→EDIT), so this passes
172    // for everyone the viewer resolution just admitted; it is the belt-
173    // and-braces level gate + the budget gate's precondition.
174    await assertCan(this.prisma, { workspaceId, userId }, "VIEW");
175    const rl = await this.rateLimit.check(userId);
176    if (!rl.allowed) {
177      throw new ApiError("RATE_LIMITED", 429, "AI call limit reached.", [
178        { reason: "ai_call_rate_limited", retryAfterSeconds: rl.retryAfterSeconds },
179      ]);
180    }
181    await primeWorkspaceCache(this.prisma, workspaceId);
182    try {
183      await this.guard.check(workspaceId, undefined, req);
184    } catch (err) {
185      throw mapAiError(err);
186    }
187 
188    // ---- retrieval (the permission-isolation proof) ----
189    const client = this.search.searchClient();
190    let ranked: RankedSource[] = [];
191    const wantedTypes = input.context?.types;
192    if (viewer) {
193      const queryVector = (await this.gateway.embed({ workspaceId, texts: [q], metadata: { capability: "rag.query" } })).vectors[0];
194      const response = await client.search({ q, mode: "hybrid", pageSize: RAG_CANDIDATES }, viewer, queryVector);
195      // The search page's type chips → a metadata filter. The kNN leg cannot
196      // filter chunk docs by their PARENT's public type (a chunk doc carries
197      // `type:"chunk"`, recorded conflict, task file), so the filter is
198      // applied to the FUSED hits by their resolved display type — a clean,
199      // non-leaky guarantee that only the requested types are cited.
200      const hits =
201        wantedTypes && wantedTypes.length > 0
202          ? response.results.filter((r) => wantedTypes.includes(r.type))
203          : response.results;
204      const top = dedupeToParentSources(hits, RAG_TOP_N);
205      // Full chunk text for the prompt context (the 160-char snippet is too
206      // thin); a doc deleted between retrieval + mget drops out of context
207      // (its source row stays — the citation is still valid).
208      const bodies = await client.getDocs(top.map((r) => r.docId));
209      const bodyById = new Map(bodies.map((b) => [b.docId, String(b.source.body ?? "")]));
210      ranked = top.map((hit) => ({ hit, body: bodyById.get(hit.docId) ?? hit.snippet ?? "" }));
211    }
212 
213    const sources: RagSource[] = ranked.map((r, i) => ({
214      n: i + 1,
215      docId: r.hit.docId,
216      id: r.hit.id,
217      type: r.hit.type,
218      title: r.hit.title,
219      url: r.hit.url,
220      snippet: r.hit.snippet ?? "",
221      source: r.hit.source,
222      // Deep-link metadata (chunk docs inherit the parent's pageId /
223      // databaseId + recordId — the chip navigates by the slug + these).
224      ...(r.hit.pageId ? { pageId: r.hit.pageId } : {}),
225      ...(r.hit.blockId ? { blockId: r.hit.blockId } : {}),
226      ...(r.hit.databaseId ? { databaseId: r.hit.databaseId } : {}),
227      ...(r.hit.recordId ? { recordId: r.hit.recordId } : {}),
228    }));
229 
230    // ---- hidden per-call conversation (never listed in the panel) ----
231    const conversation = await this.prisma.aiConversation.create({
232      data: {
233        workspaceId,
234        userId,
235        title: `${RAG_CONV_PREFIX}${q.slice(0, 60)}`,
236        context: { query: q, context: input.context ?? null } as object,
237      },
238    });
239    const userRow = await this.prisma.aiMessage.create({
240      data: { conversationId: conversation.id, role: "user", content: q },
241    });
242 
243    // ---- SSE header (the first byte — after this, errors are in-band) ----
244    res.status(200);
245    res.setHeader("Content-Type", "text/event-stream; charset=utf-8");
246    res.setHeader("Cache-Control", "no-cache, no-transform");
247    res.setHeader("Connection", "keep-alive");
248    res.setHeader("X-Accel-Buffering", "no");
249    res.flushHeaders?.();
250    if (typeof res.socket?.setNoDelay === "function") res.socket.setNoDelay(true);
251 
252    // The FIRST frame: the server-side citations (deterministic — the UI's
253    // chips + sources list render from it regardless of the model's output).
254    res.write(sseFrame({ type: "sources", sources }));
255 
256    // abort plumbing (same as Ask AI / writing: res close + writableEnded)
257    const ac = new AbortController();
258    let clientClosed = false;
259    const onResClose = (): void => {
260      if (res.writableEnded) return;
261      clientClosed = true;
262      ac.abort();
263    };
264    res.on("close", onResClose);
265    const keepalive = setInterval(() => {
266      if (!res.writableEnded) res.write(": keepalive\n\n");
267    }, KEEPALIVE_MS);
268    let timer: ReturnType<typeof setTimeout> | undefined;
269    const armTimeout = (): void => {
270      if (timer) clearTimeout(timer);
271      timer = setTimeout(() => ac.abort(), RAG_TIMEOUT_MS);
272    };
273    armTimeout();
274 
275    let finished = false;
276    const finish = (event: RagSseEvent): void => {
277      if (finished) return;
278      finished = true;
279      if (timer) clearTimeout(timer);
280      clearInterval(keepalive);
281      res.off("close", onResClose);
282      if (!res.writableEnded) {
283        res.write(sseFrame(event));
284        res.end();
285      }
286    };
287    const sendErrorDelta = (err: unknown): void => {
288      const mapped = mapAiError(err);
289      const detail = (mapped.details?.[0] ?? {}) as Record<string, unknown>;
290      const retryable = detail.retryable === true || detail.retryAfterSeconds != null || mapped.code === "RATE_LIMITED";
291      res.write(
292        sseFrame({
293          type: "error",
294          code: mapped.code,
295          message: mapped.message,
296          ...(retryable ? { retryable: true } : {}),
297          ...(typeof detail.retryAfterSeconds === "number" ? { retryAfterSeconds: detail.retryAfterSeconds } : {}),
298        }),
299      );
300      logger.warn(`rag stream error workspace=${workspaceId} code=${mapped.code}`, {
301        event: "ai.rag.error",
302        workspaceId,
303        code: mapped.code,
304      });
305    };
306 
307    const persistAssistant = async (content: string, truncated: boolean, error?: string): Promise<string | undefined> => {
308      try {
309        const created = await this.prisma.aiMessage.create({
310          data: {
311            conversationId: conversation.id,
312            role: "assistant",
313            content: content || "(empty)",
314            model: this.gateway.resolve(workspaceId).chatModel,
315            ...(usage.in > 0 ? { tokensIn: usage.in } : {}),
316            ...(usage.out > 0 ? { tokensOut: usage.out } : {}),
317            latencyMs: Date.now() - startedAt,
318            toolResults: { query: q, sources } as unknown as object,
319            metadata: {
320              feature: "rag",
321              ...(truncated ? { truncated: true } : {}),
322              ...(error ? { error } : {}),
323            } as object,
324          },
325        });
326        return created.id;
327      } catch (err) {
328        logger.error(`rag assistant persist failed: ${(err as Error).message}`);
329        return undefined;
330      }
331    };
332 
333    const chatModel = this.gateway.resolve(workspaceId).chatModel;
334    const startedAt = Date.now();
335    const usage = { in: 0, out: 0 };
336 
337    // ---- no results: a deterministic empty-state answer (no chat call) ----
338    if (ranked.length === 0) {
339      const text = "I could not find any workspace content that answers this question. Try different wording.";
340      for (const word of text.split(" ")) res.write(sseFrame({ type: "delta", text: word + " " }));
341      const assistantId = await persistAssistant(text, false);
342      finish({ type: "done", assistantMessageId: assistantId, userMessageId: userRow.id, reason: "ok" });
343      return;
344    }
345 
346    // ---- the answer (numbered context + question → the model) ----
347    const contextBlocks = ranked.map((r, i) => renderRagContext(i + 1, r.hit.title || "(untitled)", r.body)).join("\n\n");
348    const messages: Array<{ role: "system" | "user"; content: string }> = [
349      { role: "system", content: RAG_SYSTEM_PROMPT },
350      { role: "user", content: `${contextBlocks}\n\nQuestion: ${q}` },
351    ];
352 
353    let content = "";
354    let doneSeen = false;
355    try {
356      for await (const chunk of this.gateway.chat({
357        workspaceId,
358        userId,
359        model: chatModel,
360        messages,
361        stream: true,
362        capability: "rag",
363        metadata: { mode: "rag" },
364        signal: ac.signal,
365      })) {
366        if (timer) armTimeout();
367        if (chunk.delta) {
368          content += chunk.delta;
369          res.write(sseFrame({ type: "delta", text: chunk.delta }));
370        }
371        if (chunk.done) {
372          doneSeen = true;
373          if (chunk.usage) {
374            usage.in = chunk.usage.in;
375            usage.out = chunk.usage.out;
376            res.write(sseFrame({ type: "usage", tokensIn: chunk.usage.in, tokensOut: chunk.usage.out }));
377          }
378          const assistantId = await persistAssistant(content, false);
379          finish({ type: "done", assistantMessageId: assistantId, userMessageId: userRow.id, reason: "ok" });
380          return;
381        }
382      }
383      // provider stream ended without a done chunk (defensive) — persist.
384      if (!doneSeen) {
385        const assistantId = await persistAssistant(content, false);
386        finish({ type: "done", assistantMessageId: assistantId, userMessageId: userRow.id, reason: "ok" });
387      }
388    } catch (err) {
389      if (clientClosed || ac.signal.aborted) {
390        // the client went away (or the watchdog fired) — persist the partial.
391        const assistantId = await persistAssistant(content, content.length > 0, "aborted");
392        if (!clientClosed) {
393          res.write(sseFrame({ type: "error", code: "TIMEOUT", message: "The AI request took too long.", retryable: true }));
394        }
395        finish({ type: "done", assistantMessageId: assistantId, userMessageId: userRow.id, reason: clientClosed ? "aborted" : "error" });
396        return;
397      }
398      sendErrorDelta(err);
399      const assistantId = await persistAssistant(content, content.length > 0, "error");
400      finish({ type: "done", assistantMessageId: assistantId, userMessageId: userRow.id, reason: "error" });
401    }
402  }
403}
404