feat: RAG — pgvector retrieval and a search_documents tool (ADR)

backend/.env.example

@@ -37,3 +37,15 @@ ASYNC_DOCUMENT_CONVERSION=false
 # progress over Redis pub/sub and can be resumed via GET .../generate/stream.
 # Requires REDIS_URL. Default "false" runs extraction inline.
 ASYNC_TABULAR_EXTRACTION=false
+# When "true", new/replaced document versions are chunked + embedded on the
+# BullMQ document-embedding queue so the search_documents semantic-search tool
+# has an index. Requires REDIS_URL and a configured embedding model. Default
+# "false": no embeddings are produced and search_documents degrades gracefully.
+ASYNC_EMBEDDING=false
+# The single embedding model used to BOTH ingest and search (they must match; a
+# vector(N) column has one fixed width). Unset → text-embedding-3-small.
+# Changing it requires re-embedding via scripts/backfillEmbeddings.ts.
+# EMBEDDING_MODEL=text-embedding-3-small
+# Embedding vector width; MUST equal the vector(N) in the document_chunks
+# migration (768). Only change alongside a new migration.
+EMBEDDING_DIMENSION=768

backend/migrations/20260716_01_document_chunks_pgvector.sql

@@ -0,0 +1,107 @@
+-- Migration date: 2026-07-16
+--
+-- RAG / semantic retrieval (pgvector-backed document chunks).
+--
+-- Stores the token-aware chunks + embeddings produced by the async embedding
+-- worker (see backend/src/lib/rag/ingest.ts) so the `search_documents` chat
+-- tool can run top-k cosine search over a user's documents.
+--
+-- Access model — identical to every other content table:
+--   * RLS is ENABLED with NO policy (default-deny).
+--   * anon/authenticated are revoked outright; the API runs as service_role
+--     and enforces authz in the app layer — the search RPC below is only ever
+--     handed the set of document_ids the caller already has access to (scoped
+--     from the chat's doc index), never the whole table.
+--   * user_id mirrors documents.user_id (text, no FK — matching this schema's
+--     documents table); deletion cascades through the document_id FK instead.
+--
+-- Embedding dimension: pinned to 768 to match the default EMBEDDING_DIMENSION.
+-- Every built-in adapter is configured to emit 768-dim vectors (Gemini
+-- text-embedding-004 is natively 768; OpenAI text-embedding-3-* are asked for
+-- `dimensions: 768`). A single vector(N) column has ONE fixed N and the HNSW
+-- index depends on it, so mixing models of different native widths is unsafe —
+-- embedding_model is stored per row and the search RPC filters on it so a model
+-- change is a hard boundary (re-embed via scripts/backfillEmbeddings.ts).
+-- Changing the width requires a new migration.
+--
+-- Does NOT edit any existing migration.
+
+create extension if not exists vector;
+
+create table if not exists public.document_chunks (
+  id uuid primary key default gen_random_uuid(),
+  document_id uuid not null references public.documents(id) on delete cascade,
+  version_id uuid not null references public.document_versions(id) on delete cascade,
+  user_id text not null,
+  chunk_index int not null,
+  content text not null,
+  -- 1-based page number parsed from the "## Page N" markers extractPdfMarkdown
+  -- emits, so citations can carry {page, quote}. Null for DOCX (no pages).
+  page int,
+  token_count int,
+  -- The embedding model that produced this row's vector. The search RPC filters
+  -- on it so vectors of different models/dimensions are never mixed in one query.
+  embedding_model text not null,
+  embedding vector(768) not null,
+  metadata jsonb not null default '{}'::jsonb,
+  created_at timestamptz not null default now(),
+  -- One row per (version, chunk): the ingestion job deletes-by-version then
+  -- re-inserts, so a retry can't create duplicates.
+  unique (version_id, chunk_index)
+);
+
+-- Approximate-nearest-neighbour index for the cosine top-k search. HNSW builds
+-- fine on an empty table and stays correct as rows are added.
+create index if not exists idx_document_chunks_embedding
+  on public.document_chunks using hnsw (embedding vector_cosine_ops);
+
+create index if not exists idx_document_chunks_document on public.document_chunks(document_id);
+create index if not exists idx_document_chunks_version on public.document_chunks(version_id);
+create index if not exists idx_document_chunks_user on public.document_chunks(user_id);
+
+alter table public.document_chunks enable row level security;
+
+-- Explicit service_role DML grant (the default privileges already cover it;
+-- this documents the intent at the table).
+grant select, insert, update, delete on public.document_chunks to service_role;
+revoke all on public.document_chunks from anon, authenticated;
+
+-- ---------------------------------------------------------------------------
+-- match_document_chunks — top-k cosine search
+-- ---------------------------------------------------------------------------
+-- Called by the search_documents tool as service_role. Authz is enforced by the
+-- caller: p_document_ids is the pre-scoped set of documents the user can access
+-- (built from the chat's doc index), so this function never reaches beyond it.
+-- Filters on embedding_model so a model change can't return dimension-mismatched
+-- rows. Ordered by cosine distance (<=> is pgvector's cosine operator).
+
+create or replace function public.match_document_chunks(
+  p_query_embedding vector,
+  p_document_ids uuid[],
+  p_model text,
+  p_match_count int
+)
+returns table (
+  document_id uuid,
+  version_id uuid,
+  chunk_index int,
+  content text,
+  page int,
+  distance double precision
+)
+language sql
+stable
+as $$
+  select
+    c.document_id,
+    c.version_id,
+    c.chunk_index,
+    c.content,
+    c.page,
+    (c.embedding <=> p_query_embedding)::double precision as distance
+  from public.document_chunks c
+  where c.document_id = any (p_document_ids)
+    and c.embedding_model = p_model
+  order by c.embedding <=> p_query_embedding
+  limit greatest(1, p_match_count);
+$$;

backend/package.json

@@ -24,6 +24,7 @@
     "fast-xml-parser": "^5.7.1",
     "helmet": "^8.1.0",
     "ioredis": "5.10.1",
+    "js-tiktoken": "^1.0.21",
     "jszip": "^3.10.1",
     "libreoffice-convert": "^1.6.0",
     "mammoth": "^1.9.0",

backend/scripts/backfillEmbeddings.ts

@@ -0,0 +1,93 @@
+/**
+ * Backfill semantic-search embeddings for documents that don't have any yet.
+ *
+ * Idempotent: it only touches documents whose CURRENT version has no rows in
+ * document_chunks, and the ingestion itself is delete-then-insert, so re-running
+ * is safe. Use it to seed the index after enabling ASYNC_EMBEDDING on an
+ * existing deployment, or after a model change (pass --sync to also wipe/rebuild
+ * a specific set by first truncating; here it simply fills gaps).
+ *
+ * Usage (from backend):
+ *   tsx scripts/backfillEmbeddings.ts            # enqueue jobs (needs a worker)
+ *   tsx scripts/backfillEmbeddings.ts --sync     # run ingestion inline, no queue
+ *
+ * Reuses the SAME ingestion function the worker calls (runEmbeddingIngestion),
+ * so backfilled rows are identical to live-ingested ones.
+ */
+import { createServerSupabase } from "../src/lib/supabase";
+import { runEmbeddingIngestion } from "../src/lib/rag/ingest";
+import { enqueueEmbedding, closeEmbeddingQueue } from "../src/lib/queue/embeddingQueue";
+import { closeRedisConnection } from "../src/lib/queue/connection";
+
+async function main(): Promise<void> {
+    const sync = process.argv.includes("--sync");
+    const db = createServerSupabase();
+
+    const { data: docs, error } = await db
+        .from("documents")
+        .select("id, current_version_id, user_id")
+        .not("current_version_id", "is", null);
+    if (error) {
+        console.error("[backfill-embeddings] failed to list documents", { err: error });
+        process.exitCode = 1;
+        return;
+    }
+
+    // The set of version_ids that already have chunks — one round-trip.
+    const { data: existing } = await db
+        .from("document_chunks")
+        .select("version_id");
+    const indexedVersions = new Set(
+        ((existing ?? []) as { version_id: string }[]).map((r) => r.version_id),
+    );
+
+    const pending = ((docs ?? []) as {
+        id: string;
+        current_version_id: string;
+        user_id: string;
+    }[]).filter((d) => !indexedVersions.has(d.current_version_id));
+
+    console.log(
+        "[backfill-embeddings] starting",
+        { total: docs?.length ?? 0, pending: pending.length, mode: sync ? "sync" : "enqueue" },
+    );
+
+    let embedded = 0;
+    let skipped = 0;
+    for (const d of pending) {
+        const job = {
+            documentId: d.id,
+            versionId: d.current_version_id,
+            userId: d.user_id,
+        };
+        try {
+            if (sync) {
+                const result = await runEmbeddingIngestion(job);
+                if (result.status === "embedded") embedded++;
+                else skipped++;
+            } else {
+                await enqueueEmbedding(job);
+                embedded++;
+            }
+        } catch (err) {
+            skipped++;
+            console.error(
+                "[backfill-embeddings] failed",
+                { err, documentId: d.id, versionId: d.current_version_id },
+            );
+        }
+    }
+
+    console.log(
+        "[backfill-embeddings] done",
+        { processed: embedded, skipped },
+    );
+
+    if (!sync) await closeEmbeddingQueue();
+    await closeRedisConnection();
+}
+
+main().catch((err) => {
+    console.error("[backfill-embeddings] fatal", { err });
+    process.exit(1);
+});

backend/src/lib/chat/tools/__tests__/searchDocuments.test.ts

@@ -0,0 +1,138 @@
+import { describe, it, expect, beforeEach, vi } from "vitest";
+
+// Offline module loading: mock supabase (no client is ever constructed here).
+vi.mock("../../../supabase", () => ({ createServerSupabase: vi.fn() }));
+
+import { runToolCalls } from "../toolDispatcher";
+import type { DocStore, DocIndex, ToolCall } from "../../types";
+import {
+    registerEmbeddingProvider,
+    _resetEmbeddingRegistryForTesting,
+    type EmbeddingProviderAdapter,
+} from "../../../llm/embeddings";
+
+function fakeProvider(): EmbeddingProviderAdapter {
+    return {
+        id: "fake-embed",
+        // Match whatever resolveEmbeddingModel() returns (the cloud default here).
+        matchesModel: () => true,
+        dimensions: 2,
+        models: ["fake"],
+        embed: async (texts) => texts.map(() => [0.1, 0.2]),
+    };
+}
+
+type RpcArgs = Record<string, unknown>;
+
+function makeFixture(opts: {
+    matches: unknown[];
+    onRpc?: (args: RpcArgs) => void;
+}) {
+    const db = {
+        rpc: async (_name: string, args: RpcArgs) => {
+            opts.onRpc?.(args);
+            return { data: opts.matches, error: null };
+        },
+    };
+    const docStore: DocStore = new Map([
+        [
+            "doc-0",
+            { storage_path: "", file_type: "pdf", filename: "Contract.pdf" },
+        ],
+    ]);
+    const docIndex: DocIndex = {
+        "doc-0": { document_id: "DID-0", filename: "Contract.pdf" },
+    };
+    return { db, docStore, docIndex };
+}
+
+function searchCall(args: Record<string, unknown>): ToolCall {
+    return {
+        id: "call-1",
+        function: {
+            name: "search_documents",
+            arguments: JSON.stringify(args),
+        },
+    };
+}
+
+async function runSearch(
+    fixture: ReturnType<typeof makeFixture>,
+    args: Record<string, unknown>,
+) {
+    return runToolCalls(
+        [searchCall(args)],
+        fixture.docStore,
+        "u1",
+        fixture.db as never,
+        () => {},
+        undefined,
+        undefined,
+        fixture.docIndex,
+        undefined,
+        undefined,
+        null,
+        undefined,
+        {},
+    );
+}
+
+beforeEach(() => {
+    _resetEmbeddingRegistryForTesting();
+    registerEmbeddingProvider(fakeProvider());
+});
+
+describe("search_documents tool", () => {
+    it("embeds the query, scopes the search, and emits a citation reminder", async () => {
+        let rpcArgs: RpcArgs | undefined;
+        const fixture = makeFixture({
+            matches: [
+                {
+                    document_id: "DID-0",
+                    version_id: "v1",
+                    chunk_index: 2,
+                    content: "The indemnity clause is unlimited.",
+                    page: 4,
+                    distance: 0.12,
+                },
+            ],
+            onRpc: (args) => (rpcArgs = args),
+        });
+
+        const res = await runSearch(fixture, { query: "indemnity" });
+
+        // Query was embedded via the fake provider and serialized as a literal.
+        expect(rpcArgs?.p_query_embedding).toBe("[0.1,0.2]");
+        // Scoped to the chat's document ids (authz boundary).
+        expect(rpcArgs?.p_document_ids).toEqual(["DID-0"]);
+
+        const content = (res.toolResults[0] as { content: string }).content;
+        expect(content).toContain("The indemnity clause is unlimited.");
+        // Citation reminder maps document_id -> the doc-N label + filename + page.
+        expect(content).toContain('"doc-0"');
+        expect(content).toContain("Contract.pdf");
+        expect(content).toContain("page 4");
+
+        // Records one docsFound entry for the matched document.
+        expect(res.docsFound).toEqual([
+            { filename: "Contract.pdf", query: "indemnity", total_matches: 1 },
+        ]);
+    });
+
+    it("returns an error result for an empty query", async () => {
+        const fixture = makeFixture({ matches: [] });
+        const res = await runSearch(fixture, { query: "   " });
+        const content = (res.toolResults[0] as { content: string }).content;
+        expect(content).toContain("query is required");
+    });
+
+    it("degrades gracefully when no embedding provider is available", async () => {
+        _resetEmbeddingRegistryForTesting(); // leave the registry empty
+        const fixture = makeFixture({ matches: [] });
+        const res = await runSearch(fixture, { query: "anything" });
+        const content = (res.toolResults[0] as { content: string }).content;
+        expect(content).toContain("unavailable");
+        // No matches recorded — the turn is not errored.
+        expect(res.docsFound).toEqual([]);
+    });
+});