diff --git a/backend/.env.example b/backend/.env.example index beb62afcb..8e92243ce 100644 --- a/backend/.env.example +++ b/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 diff --git a/backend/migrations/20260716_01_document_chunks_pgvector.sql b/backend/migrations/20260716_01_document_chunks_pgvector.sql new file mode 100644 index 000000000..ea0fdc7a9 --- /dev/null +++ b/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); +$$; diff --git a/backend/package-lock.json b/backend/package-lock.json index 65b7da924..9683d09be 100644 --- a/backend/package-lock.json +++ b/backend/package-lock.json @@ -25,6 +25,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", @@ -4517,6 +4518,15 @@ "url": "https://github.com/sponsors/panva" } }, + "node_modules/js-tiktoken": { + "version": "1.0.21", + "resolved": "https://registry.npmjs.org/js-tiktoken/-/js-tiktoken-1.0.21.tgz", + "integrity": "sha512-biOj/6M5qdgx5TKjDnFT1ymSpM5tbd3ylwDtrQvFQSu0Z7bBYko2dF+W/aUkXUPuk6IVpRxk/3Q2sHOzGlS36g==", + "license": "MIT", + "dependencies": { + "base64-js": "^1.5.1" + } + }, "node_modules/json-bigint": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/json-bigint/-/json-bigint-1.0.0.tgz", diff --git a/backend/package.json b/backend/package.json index 1715fc5ee..4cc3b3b68 100644 --- a/backend/package.json +++ b/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", diff --git a/backend/scripts/backfillEmbeddings.ts b/backend/scripts/backfillEmbeddings.ts new file mode 100644 index 000000000..4f21a7f9c --- /dev/null +++ b/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 { + 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); +}); diff --git a/backend/src/lib/chat/tools/__tests__/searchDocuments.test.ts b/backend/src/lib/chat/tools/__tests__/searchDocuments.test.ts new file mode 100644 index 000000000..662d82393 --- /dev/null +++ b/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; + +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): ToolCall { + return { + id: "call-1", + function: { + name: "search_documents", + arguments: JSON.stringify(args), + }, + }; +} + +async function runSearch( + fixture: ReturnType, + args: Record, +) { + 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([]); + }); +}); diff --git a/backend/src/lib/chat/tools/toolDispatcher.ts b/backend/src/lib/chat/tools/toolDispatcher.ts index 3e6f67447..e76ef0f9f 100644 --- a/backend/src/lib/chat/tools/toolDispatcher.ts +++ b/backend/src/lib/chat/tools/toolDispatcher.ts @@ -55,6 +55,11 @@ import { type DocReplicatedResult, type TextMatch, } from "./documentOps"; +import { + getActiveEmbeddingProvider, + resolveEmbeddingModel, +} from "../../llm/embeddings"; +import { searchDocumentChunks } from "../../rag/searchDocuments"; type CourtlistenerCaseRecord = { @@ -432,6 +437,9 @@ function cachedCaseNotFetchedResult(clusterId: number | null) { }; } +const DEFAULT_SEARCH_TOP_K = 8; +const MAX_SEARCH_TOP_K = 25; + export async function runToolCalls( toolCalls: ToolCall[], docStore: DocStore, @@ -696,6 +704,129 @@ export async function runToolCalls( }); } toolResults.push({ role: "tool", tool_call_id: tc.id, content }); + } else if (tc.function.name === "search_documents") { + // Semantic top-k search across the chat's documents (RAG). Embeds the + // query with the SAME model used at ingest, runs a cosine search scoped + // to the document ids the chat already granted access to, and returns + // each chunk citable. The cosine search is scoped to docIndex's document + // ids (the access-checked set for this turn), so service_role can't + // return another tenant's chunks. + const pushSearchResult = (content: string) => + toolResults.push({ role: "tool", tool_call_id: tc.id, content }); + const query = typeof args.query === "string" ? args.query.trim() : ""; + const rawTopK = + typeof args.top_k === "number" ? args.top_k : DEFAULT_SEARCH_TOP_K; + const topK = Math.max(1, Math.min(MAX_SEARCH_TOP_K, Math.floor(rawTopK))); + + if (!query) { + pushSearchResult(JSON.stringify({ error: "query is required." })); + continue; + } + if (!docIndex) { + pushSearchResult( + JSON.stringify({ matches: [], note: "No documents in context." }), + ); + continue; + } + + // Map document_id -> { label, filename } from the access-checked doc + // index. This is BOTH the citation-label source and the authz scope for + // the search. + const byDocumentId = new Map(); + for (const [label, indexed] of Object.entries(docIndex)) { + if (!byDocumentId.has(indexed.document_id)) { + byDocumentId.set(indexed.document_id, { + label, + filename: indexed.filename, + }); + } + } + + // Optional narrowing to a single doc (accepts a doc-N label or a raw id). + let documentIds = [...byDocumentId.keys()]; + if (typeof args.doc_id === "string" && args.doc_id.trim()) { + const label = + resolveDocLabel(args.doc_id as string, docStore, docIndex) ?? + (args.doc_id as string); + const target = docIndex[label]?.document_id ?? (args.doc_id as string); + documentIds = documentIds.filter((id) => id === target); + } + + if (documentIds.length === 0) { + pushSearchResult( + JSON.stringify({ + matches: [], + note: "No matching documents in context.", + }), + ); + continue; + } + + const provider = getActiveEmbeddingProvider(); + if (!provider) { + // Embeddings unconfigured. Degrade gracefully — the model can still + // fall back to read_document. + pushSearchResult( + JSON.stringify({ + matches: [], + note: "Semantic search is unavailable in this deployment; use read_document or find_in_document instead.", + }), + ); + continue; + } + + const model = resolveEmbeddingModel(); + let matches; + try { + const [queryEmbedding] = await provider.embed([query], apiKeys); + matches = await searchDocumentChunks({ + db, + queryEmbedding: queryEmbedding ?? [], + model, + documentIds, + topK, + }); + } catch (err) { + pushSearchResult( + JSON.stringify({ error: `Semantic search failed: ${String(err)}` }), + ); + continue; + } + + if (matches.length === 0) { + pushSearchResult( + JSON.stringify({ matches: [], note: "No relevant passages found." }), + ); + continue; + } + + // Record one docsFound entry per matched document for the UI chips. + const perDoc = new Map(); + for (const m of matches) + perDoc.set(m.document_id, (perDoc.get(m.document_id) ?? 0) + 1); + for (const [documentId, count] of perDoc) { + const info = byDocumentId.get(documentId); + if (info) { + docsFound.push({ + filename: info.filename, + query, + total_matches: count, + }); + } + } + + // Attach the doc-N citation reminder to each chunk, matching + // read_document/find_in_document. + const parts = matches.map((m) => { + const info = byDocumentId.get(m.document_id); + const label = info?.label ?? m.document_id; + const filename = info?.filename ?? m.document_id; + const pageLine = m.page != null ? ` (page ${m.page})` : ""; + const header = `--- ${label} ("${filename}")${pageLine} ---\n${citationReminder(label, filename)}`; + return `${header}\n\n${m.content}`; + }); + + pushSearchResult(parts.join("\n\n")); } else if (tc.function.name === "list_documents") { const list = Array.from(docStore.entries()).map(([doc_id, info]) => ({ doc_id, diff --git a/backend/src/lib/chat/tools/toolSchemas.ts b/backend/src/lib/chat/tools/toolSchemas.ts index 302ceb9e6..8397d6723 100644 --- a/backend/src/lib/chat/tools/toolSchemas.ts +++ b/backend/src/lib/chat/tools/toolSchemas.ts @@ -468,4 +468,33 @@ export const TOOLS = [ }, }, }, + { + type: "function", + function: { + name: "search_documents", + description: + "Semantic (meaning-based) search across the documents in this chat. Returns the most relevant passages with their source document and page for citation. Use this to locate relevant material by concept when you don't know the exact wording — complementary to find_in_document's literal Ctrl+F match. Prefer this over read_document when the documents are long and you only need the passages relevant to a question.", + parameters: { + type: "object", + properties: { + query: { + type: "string", + description: + "The concept or question to search for. Phrase it as the meaning you want, not exact keywords.", + }, + top_k: { + type: "integer", + description: + "Maximum number of passages to return (default 8, max 25).", + }, + doc_id: { + type: "string", + description: + "Optional. Restrict the search to a single document (e.g. 'doc-0'). Omit to search all documents in the chat.", + }, + }, + required: ["query"], + }, + }, + }, ]; diff --git a/backend/src/lib/llm/embeddings/__tests__/registry.test.ts b/backend/src/lib/llm/embeddings/__tests__/registry.test.ts new file mode 100644 index 000000000..677947a15 --- /dev/null +++ b/backend/src/lib/llm/embeddings/__tests__/registry.test.ts @@ -0,0 +1,74 @@ +import { beforeEach, describe, expect, it } from "vitest"; + +import { + registerBuiltinEmbeddingProviders, + resolveEmbeddingModel, + resolveEmbeddingDimension, + getActiveEmbeddingProvider, +} from "../index"; +import { + registerEmbeddingProvider, + getEmbeddingProvider, + findEmbeddingProviderForModel, + _resetEmbeddingRegistryForTesting, + type EmbeddingProviderAdapter, +} from "../registry"; + +// registerBuiltinEmbeddingProviders ran once at import with the real env; reset +// before each case and re-register against a controlled env for determinism. +beforeEach(() => _resetEmbeddingRegistryForTesting()); + +function fakeAdapter(id: string, models: string[]): EmbeddingProviderAdapter { + const set = new Set(models); + return { + id, + matchesModel: (m) => set.has(m), + dimensions: 768, + models, + embed: async (texts) => texts.map(() => [0, 0, 0]), + }; +} + +describe("embedding registry", () => { + it("registers, finds by model, and resets", async () => { + const adapter = fakeAdapter("fake", ["fake-embed-1"]); + registerEmbeddingProvider(adapter); + expect(getEmbeddingProvider("fake")).toBe(adapter); + expect(findEmbeddingProviderForModel("fake-embed-1")).toBe(adapter); + expect(findEmbeddingProviderForModel("unknown")).toBeUndefined(); + // The fake adapter is deterministic and makes no network call. + expect(await adapter.embed(["a", "b"])).toEqual([[0, 0, 0], [0, 0, 0]]); + + _resetEmbeddingRegistryForTesting(); + expect(getEmbeddingProvider("fake")).toBeUndefined(); + }); + + it("registers the built-in cloud embedding providers", () => { + registerBuiltinEmbeddingProviders({}); + expect(getEmbeddingProvider("openai-embed")).toBeDefined(); + expect(getEmbeddingProvider("gemini-embed")).toBeDefined(); + }); +}); + +describe("resolveEmbeddingModel / resolveEmbeddingDimension", () => { + it("defaults to the cloud model", () => { + expect(resolveEmbeddingModel({})).toBe("text-embedding-3-small"); + }); + + it("honours an explicit EMBEDDING_MODEL override", () => { + expect(resolveEmbeddingModel({ EMBEDDING_MODEL: "text-embedding-3-large" })).toBe( + "text-embedding-3-large", + ); + }); + + it("defaults dimension to 768 and honours EMBEDDING_DIMENSION", () => { + expect(resolveEmbeddingDimension({})).toBe(768); + expect(resolveEmbeddingDimension({ EMBEDDING_DIMENSION: "1536" })).toBe(1536); + expect(resolveEmbeddingDimension({ EMBEDDING_DIMENSION: "bad" })).toBe(768); + }); + + it("getActiveEmbeddingProvider resolves the deployment model to an adapter", () => { + registerBuiltinEmbeddingProviders({}); + expect(getActiveEmbeddingProvider({})?.id).toBe("openai-embed"); + }); +}); diff --git a/backend/src/lib/llm/embeddings/gemini.ts b/backend/src/lib/llm/embeddings/gemini.ts new file mode 100644 index 000000000..eec0004a9 --- /dev/null +++ b/backend/src/lib/llm/embeddings/gemini.ts @@ -0,0 +1,73 @@ +import type { UserApiKeys } from "../types"; +import type { EmbeddingProviderAdapter } from "./registry"; + +// Reuse the same dynamic-ESM import shim the chat adapter uses so @google/genai +// (an ESM-only package) loads from CommonJS without TS rewriting the import. +type GoogleGenAIConstructor = typeof import("@google/genai").GoogleGenAI; +type GoogleGenAIClient = InstanceType; + +const importEsm = new Function("specifier", "return import(specifier)") as ( + specifier: string, +) => Promise<{ GoogleGenAI: GoogleGenAIConstructor }>; + +function apiKey(override?: string | null): string { + const key = override?.trim() || process.env.GEMINI_API_KEY?.trim() || ""; + if (!key) { + throw new Error( + "Gemini API key is not configured. Set GEMINI_API_KEY or add a user Gemini key.", + ); + } + return key; +} + +async function client(override?: string | null): Promise { + const { GoogleGenAI } = await importEsm("@google/genai"); + return new GoogleGenAI({ apiKey: apiKey(override) }); +} + +export const GEMINI_EMBEDDING_MODELS = [ + "text-embedding-004", + "gemini-embedding-001", +] as const; + +type EmbedContentResponse = { + embeddings?: { values?: number[] }[]; + embedding?: { values?: number[] }; +}; + +/** + * Build the cloud Gemini embedding adapter, bound to the deployment's single + * embedding model. text-embedding-004 is natively 768-dim; gemini-embedding-001 + * accepts an outputDimensionality, so we request the column width either way. + */ +export function createGeminiEmbeddingProvider(opts: { + dimensions: number; + model?: string; +}): EmbeddingProviderAdapter { + const modelSet = new Set(GEMINI_EMBEDDING_MODELS); + const model = + opts.model && modelSet.has(opts.model) + ? opts.model + : GEMINI_EMBEDDING_MODELS[0]; + return { + id: "gemini-embed", + matchesModel: (m) => modelSet.has(m), + dimensions: opts.dimensions, + models: GEMINI_EMBEDDING_MODELS, + embed: async (texts: string[], apiKeys?: UserApiKeys) => { + if (texts.length === 0) return []; + const ai = await client(apiKeys?.gemini); + const res = (await ai.models.embedContent({ + model, + contents: texts, + config: { outputDimensionality: opts.dimensions }, + })) as EmbedContentResponse; + // Batch response: one embedding per input, in input order. + if (Array.isArray(res.embeddings)) { + return res.embeddings.map((e) => e.values ?? []); + } + // Defensive: single-content shape. + return res.embedding?.values ? [res.embedding.values] : []; + }, + }; +} diff --git a/backend/src/lib/llm/embeddings/index.ts b/backend/src/lib/llm/embeddings/index.ts new file mode 100644 index 000000000..68b486a1c --- /dev/null +++ b/backend/src/lib/llm/embeddings/index.ts @@ -0,0 +1,65 @@ +import { createOpenAIEmbeddingProvider } from "./openai"; +import { createGeminiEmbeddingProvider } from "./gemini"; +import { + registerEmbeddingProvider, + findEmbeddingProviderForModel, + type EmbeddingProviderAdapter, +} from "./registry"; + +export * from "./registry"; + +// Default embedding model. The whole deployment pins ONE model +// (EMBEDDING_MODEL) + ONE width (EMBEDDING_DIMENSION) because a vector(N) +// column has a single fixed N; see the migration's header for why. +const DEFAULT_CLOUD_EMBEDDING_MODEL = "text-embedding-3-small"; +const DEFAULT_EMBEDDING_DIMENSION = 768; + +/** Column width the adapters emit and the migration pins vector(N) to. */ +export function resolveEmbeddingDimension( + env: NodeJS.ProcessEnv = process.env, +): number { + const raw = env.EMBEDDING_DIMENSION; + const n = raw ? Number.parseInt(raw, 10) : NaN; + return Number.isFinite(n) && n > 0 ? n : DEFAULT_EMBEDDING_DIMENSION; +} + +/** + * The single embedding model this deployment ingests + searches with. Explicit + * EMBEDDING_MODEL wins; otherwise the cloud default. + */ +export function resolveEmbeddingModel( + env: NodeJS.ProcessEnv = process.env, +): string { + if (env.EMBEDDING_MODEL?.trim()) return env.EMBEDDING_MODEL.trim(); + return DEFAULT_CLOUD_EMBEDDING_MODEL; +} + +/** + * Register the built-in embedding providers (OpenAI + Gemini). + * + * Reads process.env directly so importing this module doesn't force full env + * validation — it loads in many unit tests. `env` is injectable so gating can + * be exercised against a controlled environment. + */ +export function registerBuiltinEmbeddingProviders( + env: NodeJS.ProcessEnv = process.env, +): void { + const dimensions = resolveEmbeddingDimension(env); + const model = resolveEmbeddingModel(env); + + registerEmbeddingProvider(createOpenAIEmbeddingProvider({ dimensions, model })); + registerEmbeddingProvider(createGeminiEmbeddingProvider({ dimensions, model })); +} + +registerBuiltinEmbeddingProviders(); + +/** + * Resolve the adapter for the deployment's embedding model, or undefined when + * none is registered. Callers degrade gracefully rather than error the chat + * turn. + */ +export function getActiveEmbeddingProvider( + env: NodeJS.ProcessEnv = process.env, +): EmbeddingProviderAdapter | undefined { + return findEmbeddingProviderForModel(resolveEmbeddingModel(env)); +} diff --git a/backend/src/lib/llm/embeddings/openai.ts b/backend/src/lib/llm/embeddings/openai.ts new file mode 100644 index 000000000..8f0896443 --- /dev/null +++ b/backend/src/lib/llm/embeddings/openai.ts @@ -0,0 +1,94 @@ +import type { UserApiKeys } from "../types"; +import type { EmbeddingProviderAdapter } from "./registry"; + +const OPENAI_BASE_URL = "https://api.openai.com/v1"; + +/** + * OpenAI `/v1/embeddings` client. + */ +export async function embedOpenAICompatible(params: { + model: string; + texts: string[]; + apiKey: string; + baseUrl: string; + /** Requested output width; omitted for backends that don't support it. */ + dimensions?: number; +}): Promise { + if (params.texts.length === 0) return []; + const response = await fetch(`${params.baseUrl}/embeddings`, { + method: "POST", + headers: { + Authorization: `Bearer ${params.apiKey}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + model: params.model, + input: params.texts, + dimensions: params.dimensions, + }), + }); + if (!response.ok) { + const text = await response.text().catch(() => ""); + const err = new Error( + `Embedding request failed (${response.status}): ${text || response.statusText}`, + ); + (err as { status?: number }).status = response.status; + throw err; + } + const json = (await response.json()) as { + data?: { embedding?: number[]; index?: number }[]; + }; + const rows = json.data ?? []; + // The API returns entries with an explicit `index`; sort by it so the output + // order matches the input order even if the server reorders. + const ordered = [...rows].sort((a, b) => (a.index ?? 0) - (b.index ?? 0)); + return ordered.map((r) => r.embedding ?? []); +} + +function apiKey(override?: string | null): string { + const key = override?.trim() || process.env.OPENAI_API_KEY?.trim() || ""; + if (!key) { + throw new Error( + "OpenAI API key is not configured. Set OPENAI_API_KEY or add a user OpenAI key.", + ); + } + return key; +} + +export const OPENAI_EMBEDDING_MODELS = [ + "text-embedding-3-small", + "text-embedding-3-large", +] as const; + +/** + * Build the cloud OpenAI embedding adapter, bound to the deployment's single + * embedding model (see registerBuiltinEmbeddingProviders). `dimensions` is + * passed through to the API — the text-embedding-3-* models support Matryoshka + * truncation, so we ask for exactly the column width (default 768) and every + * stored row stays uniform. Binding one model per deployment is deliberate: a + * vector(N) column has one fixed width, so mixing models is unsafe. + */ +export function createOpenAIEmbeddingProvider(opts: { + dimensions: number; + model?: string; +}): EmbeddingProviderAdapter { + const modelSet = new Set(OPENAI_EMBEDDING_MODELS); + const model = + opts.model && modelSet.has(opts.model) + ? opts.model + : OPENAI_EMBEDDING_MODELS[0]; + return { + id: "openai-embed", + matchesModel: (m) => modelSet.has(m), + dimensions: opts.dimensions, + models: OPENAI_EMBEDDING_MODELS, + embed: (texts: string[], apiKeys?: UserApiKeys) => + embedOpenAICompatible({ + model, + texts, + apiKey: apiKey(apiKeys?.openai), + baseUrl: OPENAI_BASE_URL, + dimensions: opts.dimensions, + }), + }; +} diff --git a/backend/src/lib/llm/embeddings/registry.ts b/backend/src/lib/llm/embeddings/registry.ts new file mode 100644 index 000000000..4e7a5bd27 --- /dev/null +++ b/backend/src/lib/llm/embeddings/registry.ts @@ -0,0 +1,66 @@ +import type { UserApiKeys } from "../types"; + +/** + * Contract every embedding provider adapter must satisfy. + * + * A SIBLING of the chat-provider registry (../registry.ts): the LLM + * `LLMProviderAdapter` only does stream()/complete(); embeddings need their own + * shape (a batch text→vector call and a fixed output width), so they get their + * own registry with the same register/find/reset ergonomics. + * + * Built-ins (OpenAI, Gemini, Ollama) register in ./index.ts on module load, + * behind the SAME air-gap gate as llm/index.ts: cloud adapters are simply not + * registered when AIRGAPPED=true, so a cloud embedding model has no adapter to + * dispatch and semantic search runs against the local (Ollama) adapter only. + */ +export interface EmbeddingProviderAdapter { + /** Stable identifier, e.g. "openai-embed", "gemini-embed", "ollama-embed". */ + readonly id: string; + /** True if this provider handles the given embedding-model string. */ + matchesModel(model: string): boolean; + /** + * Output vector width this adapter is configured to emit. MUST match the + * `vector(N)` column in the migration (pinned to EMBEDDING_DIMENSION) — the + * adapter asks the API for exactly this width where the API supports it. + */ + readonly dimensions: number; + /** Model IDs this provider serves (drives documentation / model lists). */ + readonly models: readonly string[]; + /** Embed a batch of texts, preserving input order in the output. */ + embed(texts: string[], apiKeys?: UserApiKeys): Promise; +} + +const _registry = new Map(); + +/** Register an embedding provider adapter (re-registering an id replaces it). */ +export function registerEmbeddingProvider(adapter: EmbeddingProviderAdapter): void { + _registry.set(adapter.id, adapter); +} + +/** Returns the adapter registered under id, or undefined if none. */ +export function getEmbeddingProvider(id: string): EmbeddingProviderAdapter | undefined { + return _registry.get(id); +} + +/** + * Returns the first registered embedding provider whose matchesModel() returns + * true, or undefined when none match (e.g. a cloud model in air-gapped mode). + */ +export function findEmbeddingProviderForModel( + model: string, +): EmbeddingProviderAdapter | undefined { + for (const p of _registry.values()) { + if (p.matchesModel(model)) return p; + } + return undefined; +} + +/** IDs of all currently registered embedding providers in insertion order. */ +export function registeredEmbeddingProviderIds(): string[] { + return [..._registry.keys()]; +} + +/** Exposed for test isolation only — do not call in production code. */ +export function _resetEmbeddingRegistryForTesting(): void { + _registry.clear(); +} diff --git a/backend/src/lib/queue/embeddingQueue.ts b/backend/src/lib/queue/embeddingQueue.ts new file mode 100644 index 000000000..22db757ef --- /dev/null +++ b/backend/src/lib/queue/embeddingQueue.ts @@ -0,0 +1,77 @@ +import { Queue } from "bullmq"; +import { getRedisConnection } from "./connection"; +import type { EmbeddingJobData } from "../rag/ingest"; + +/** + * BullMQ queue that chunks + embeds a document version off the request thread. + * + * Structurally identical to conversionQueue / extractionQueue: a tiny payload + * (documentId, versionId, userId — NO secrets), a deterministic jobId per + * version so a double-enqueue dedupes into the in-flight job, attempts:3 with + * exponential backoff. The worker re-derives storage path + embedding model + + * API keys at run time (see runEmbeddingIngestion). + */ +export const EMBEDDING_QUEUE = "document-embedding"; + +export type { EmbeddingJobData }; + +let queue: Queue | null = null; + +export function getEmbeddingQueue(): Queue { + if (!queue) { + queue = new Queue(EMBEDDING_QUEUE, { + connection: getRedisConnection(), + }); + } + return queue; +} + +/** Deterministic BullMQ jobId for embedding one version. */ +export function embeddingJobId(versionId: string): string { + return `embed:${versionId}`; +} + +/** + * Enqueue an embedding job. jobId is derived from the (unique-per-version) + * versionId, so re-enqueuing the same version — e.g. a replacement that reuses + * a version row — dedupes instead of racing two ingestions. Durable state lives + * in document_chunks, so removeOnComplete/Fail lets a later re-embed reuse the + * same jobId. + */ +export function enqueueEmbedding(data: EmbeddingJobData) { + return getEmbeddingQueue().add("embed", data, { + jobId: embeddingJobId(data.versionId), + attempts: 3, + backoff: { type: "exponential", delay: 2000 }, + removeOnComplete: true, + removeOnFail: true, + }); +} + +/** + * Enqueue an embedding job iff ASYNC_EMBEDDING is on, swallowing enqueue errors. + * + * Called from the document upload + new-version paths beside enqueueConversion. + * Embedding is a best-effort background index refresh — a Redis hiccup here must + * never fail the user's upload/edit, and the backfill script can repair any gap. + */ +export async function maybeEnqueueEmbedding( + data: EmbeddingJobData, +): Promise { + if (process.env.ASYNC_EMBEDDING !== "true") return; + try { + await enqueueEmbedding(data); + } catch (err) { + console.error( + "[embedding-queue] failed to enqueue embedding job", + { err, documentId: data.documentId, versionId: data.versionId }, + ); + } +} + +export async function closeEmbeddingQueue(): Promise { + if (queue) { + await queue.close(); + queue = null; + } +} diff --git a/backend/src/lib/rag/__tests__/chunker.test.ts b/backend/src/lib/rag/__tests__/chunker.test.ts new file mode 100644 index 000000000..0f7022649 --- /dev/null +++ b/backend/src/lib/rag/__tests__/chunker.test.ts @@ -0,0 +1,76 @@ +import { describe, it, expect } from "vitest"; +import { getEncoding } from "js-tiktoken"; +import { chunkMarkdown, countTokens } from "../chunker"; + +const enc = getEncoding("cl100k_base"); +const tokenLen = (s: string) => enc.encode(s).length; + +describe("chunkMarkdown", () => { + it("returns no chunks for empty / whitespace-only input", () => { + expect(chunkMarkdown("")).toEqual([]); + expect(chunkMarkdown(" \n\n \t ")).toEqual([]); + }); + + it("returns a single chunk for input under the target budget", () => { + const text = "The quick brown fox jumps over the lazy dog."; + const chunks = chunkMarkdown(text, { targetTokens: 512, overlapTokens: 64 }); + expect(chunks).toHaveLength(1); + expect(chunks[0].chunkIndex).toBe(0); + expect(chunks[0].content).toContain("quick brown fox"); + expect(chunks[0].tokenCount).toBe(tokenLen(text)); + // No page headers → null page. + expect(chunks[0].page).toBeNull(); + }); + + it("never exceeds the token budget and indexes chunks sequentially", () => { + // A long single block (no blank lines / page markers) forces windowing. + const text = Array.from({ length: 400 }, (_, i) => `word${i}`).join(" "); + const chunks = chunkMarkdown(text, { targetTokens: 50, overlapTokens: 10 }); + expect(chunks.length).toBeGreaterThan(1); + chunks.forEach((c, i) => { + expect(c.chunkIndex).toBe(i); + expect(c.tokenCount).toBeGreaterThan(0); + expect(c.tokenCount).toBeLessThanOrEqual(50); + }); + }); + + it("produces more (denser) chunks with overlap than without", () => { + const text = Array.from({ length: 400 }, (_, i) => `word${i}`).join(" "); + const noOverlap = chunkMarkdown(text, { targetTokens: 50, overlapTokens: 0 }); + const withOverlap = chunkMarkdown(text, { targetTokens: 50, overlapTokens: 25 }); + expect(withOverlap.length).toBeGreaterThan(noOverlap.length); + }); + + it("attaches the page number from the nearest preceding '## Page N' header", () => { + // Enough tokens per page that a tiny budget splits them apart, so a + // later chunk starts inside page 2. + const page1 = Array.from({ length: 20 }, (_, i) => `alpha${i}`).join(" "); + const page2 = Array.from({ length: 20 }, (_, i) => `beta${i}`).join(" "); + const md = `## Page 1\n\n${page1}\n\n## Page 2\n\n${page2}`; + const chunks = chunkMarkdown(md, { targetTokens: 12, overlapTokens: 0 }); + + const pages = chunks.map((c) => c.page); + expect(pages[0]).toBe(1); + expect(pages).toContain(2); + // The page headers themselves are consumed, not emitted as content. + expect(chunks.every((c) => !/## Page/.test(c.content))).toBe(true); + }); + + it("applies the hard character cap as a safety net", () => { + const text = "x".repeat(5000); + const chunks = chunkMarkdown(text, { + targetTokens: 100000, + overlapTokens: 0, + maxChunkChars: 100, + }); + expect(chunks).toHaveLength(1); + expect(chunks[0].content.length).toBeLessThanOrEqual(100); + }); +}); + +describe("countTokens", () => { + it("counts tokens with the same encoder the chunker uses", () => { + expect(countTokens("hello world")).toBe(tokenLen("hello world")); + expect(countTokens("")).toBe(0); + }); +}); diff --git a/backend/src/lib/rag/__tests__/searchDocuments.test.ts b/backend/src/lib/rag/__tests__/searchDocuments.test.ts new file mode 100644 index 000000000..2c852ea22 --- /dev/null +++ b/backend/src/lib/rag/__tests__/searchDocuments.test.ts @@ -0,0 +1,140 @@ +import { describe, it, expect, vi } from "vitest"; + +// Keep module import side-effects offline: supabase is only touched for its +// types / at construction, never with a real client in this test. +vi.mock("../../supabase", () => ({ createServerSupabase: vi.fn() })); + +import { + searchDocumentChunks, + cosineSimilarity, + parseVectorLiteral, +} from "../searchDocuments"; + +describe("cosineSimilarity / parseVectorLiteral", () => { + it("computes cosine similarity and handles zero vectors", () => { + expect(cosineSimilarity([1, 0], [1, 0])).toBeCloseTo(1); + expect(cosineSimilarity([1, 0], [0, 1])).toBeCloseTo(0); + expect(cosineSimilarity([0, 0], [1, 1])).toBe(0); + }); + + it("parses pgvector text literals back into numbers", () => { + expect(parseVectorLiteral("[1,2,3]")).toEqual([1, 2, 3]); + expect(parseVectorLiteral("[]")).toEqual([]); + expect(parseVectorLiteral([4, 5])).toEqual([4, 5]); + }); +}); + +type RpcResult = { data: unknown; error: unknown }; + +function makeDb(opts: { rpcResult: RpcResult; fallbackRows?: unknown[] }) { + const rpcCalls: { name: string; args: Record }[] = []; + return { + rpcCalls, + rpc: async (name: string, args: Record) => { + rpcCalls.push({ name, args }); + return opts.rpcResult; + }, + from: () => ({ + select: () => ({ + in: () => ({ + eq: async () => ({ data: opts.fallbackRows ?? [] }), + }), + }), + }), + }; +} + +describe("searchDocumentChunks", () => { + it("returns [] without hitting the RPC for empty scope or empty query", async () => { + const db = makeDb({ rpcResult: { data: [], error: null } }); + expect( + await searchDocumentChunks({ + db: db as never, + queryEmbedding: [1, 2], + model: "m", + documentIds: [], + topK: 5, + }), + ).toEqual([]); + expect( + await searchDocumentChunks({ + db: db as never, + queryEmbedding: [], + model: "m", + documentIds: ["d1"], + topK: 5, + }), + ).toEqual([]); + expect(db.rpcCalls).toHaveLength(0); + }); + + it("uses the RPC, passing scoped ids + model, and returns the ranked rows", async () => { + const rows = [ + { document_id: "d1", version_id: "v1", chunk_index: 0, content: "closest", page: 1, distance: 0.1 }, + { document_id: "d2", version_id: "v2", chunk_index: 3, content: "next", page: 2, distance: 0.4 }, + ]; + const db = makeDb({ rpcResult: { data: rows, error: null } }); + + const out = await searchDocumentChunks({ + db: db as never, + queryEmbedding: [0.5, 0.5], + model: "text-embedding-3-small", + documentIds: ["d1", "d2"], + topK: 10, + }); + + expect(out).toEqual(rows); + expect(db.rpcCalls).toHaveLength(1); + const call = db.rpcCalls[0]; + expect(call.name).toBe("match_document_chunks"); + expect(call.args.p_document_ids).toEqual(["d1", "d2"]); + expect(call.args.p_model).toBe("text-embedding-3-small"); + expect(call.args.p_match_count).toBe(10); + // The embedding is serialized as a pgvector literal. + expect(call.args.p_query_embedding).toBe("[0.5,0.5]"); + }); + + it("truncates RPC results to top_k", async () => { + const rows = Array.from({ length: 5 }, (_, i) => ({ + document_id: "d1", + version_id: "v1", + chunk_index: i, + content: `c${i}`, + page: null, + distance: i * 0.1, + })); + const db = makeDb({ rpcResult: { data: rows, error: null } }); + const out = await searchDocumentChunks({ + db: db as never, + queryEmbedding: [1], + model: "m", + documentIds: ["d1"], + topK: 2, + }); + expect(out).toHaveLength(2); + }); + + it("falls back to in-process cosine ranking when the RPC is unavailable", async () => { + const fallbackRows = [ + { document_id: "dA", version_id: "vA", chunk_index: 0, content: "A", page: 1, embedding: "[1,0]" }, + { document_id: "dB", version_id: "vB", chunk_index: 0, content: "B", page: 1, embedding: "[0,1]" }, + { document_id: "dC", version_id: "vC", chunk_index: 0, content: "C", page: 1, embedding: "[0.9,0.1]" }, + ]; + const db = makeDb({ + rpcResult: { data: null, error: { message: "function does not exist" } }, + fallbackRows, + }); + + const out = await searchDocumentChunks({ + db: db as never, + queryEmbedding: [1, 0], + model: "m", + documentIds: ["dA", "dB", "dC"], + topK: 3, + }); + + // Ranked by ascending cosine distance to [1,0]: A (0) < C < B. + expect(out.map((r) => r.document_id)).toEqual(["dA", "dC", "dB"]); + expect(out[0].distance).toBeCloseTo(0); + }); +}); diff --git a/backend/src/lib/rag/chunker.ts b/backend/src/lib/rag/chunker.ts new file mode 100644 index 000000000..1835bd911 --- /dev/null +++ b/backend/src/lib/rag/chunker.ts @@ -0,0 +1,143 @@ +import { getEncoding, type Tiktoken } from "js-tiktoken"; + +/** + * Token-aware markdown chunker for the embedding pipeline. + * + * Consumes the SAME markdown the tabular path produces (extractPdfMarkdown / + * extractDocxMarkdown): extractPdfMarkdown emits "## Page N" headers, which we + * parse so every chunk carries the page it started on — that page flows into + * the citation metadata the search_documents tool returns ({page, quote}). + * + * Tokenisation uses js-tiktoken's cl100k_base (pure JS, bundles its ranks — no + * native binary, no network, sandbox-safe). It is only an APPROXIMATION of the + * embedding model's own tokenizer (esp. Gemini/Ollama), so budgets are best- + * effort and we also apply a hard character cap per chunk as a safety valve. + */ + +export interface ChunkOptions { + /** Target tokens per chunk. */ + targetTokens?: number; + /** Tokens of overlap between consecutive chunks (context continuity). */ + overlapTokens?: number; + /** Hard per-chunk character cap (safety net against tokenizer drift). */ + maxChunkChars?: number; +} + +export interface DocumentChunk { + chunkIndex: number; + content: string; + /** 1-based page from the nearest preceding "## Page N", or null (e.g. DOCX). */ + page: number | null; + tokenCount: number; +} + +const DEFAULTS = { + targetTokens: 512, + overlapTokens: 64, + maxChunkChars: 8000, +}; + +const PAGE_HEADER_RE = /^##\s+Page\s+(\d+)\s*$/i; + +let _encoder: Tiktoken | null = null; +function encoder(): Tiktoken { + if (!_encoder) _encoder = getEncoding("cl100k_base"); + return _encoder; +} + +/** A run of body text tagged with the page it belongs to. */ +interface PagedUnit { + text: string; + page: number | null; +} + +/** + * Split markdown into page-tagged units. "## Page N" headers are consumed (not + * emitted as content) but advance the current page for everything after them. + */ +function parsePagedUnits(markdown: string): PagedUnit[] { + const units: PagedUnit[] = []; + let page: number | null = null; + let buffer: string[] = []; + + const flush = () => { + const text = buffer.join("\n").trim(); + if (text) units.push({ text, page }); + buffer = []; + }; + + for (const line of markdown.split(/\r?\n/)) { + const m = line.match(PAGE_HEADER_RE); + if (m) { + flush(); + page = Number.parseInt(m[1], 10); + continue; + } + buffer.push(line); + } + flush(); + return units; +} + +/** + * Chunk markdown into overlapping, token-budgeted windows with page attribution. + * + * Empty / whitespace-only input yields no chunks. Input smaller than the target + * yields a single chunk. A single oversized unit is split across windows too — + * windowing runs over one flat token stream, so nothing exceeds the budget. + */ +export function chunkMarkdown( + markdown: string, + options: ChunkOptions = {}, +): DocumentChunk[] { + const targetTokens = options.targetTokens ?? DEFAULTS.targetTokens; + const overlapTokens = Math.min( + options.overlapTokens ?? DEFAULTS.overlapTokens, + Math.max(0, targetTokens - 1), + ); + const maxChunkChars = options.maxChunkChars ?? DEFAULTS.maxChunkChars; + const step = Math.max(1, targetTokens - overlapTokens); + + const enc = encoder(); + const units = parsePagedUnits(markdown); + + // Flatten to one token stream while remembering the page of each token, so a + // window that spans a page break is attributed to the page it starts on. + const tokens: number[] = []; + const tokenPages: (number | null)[] = []; + const separator = enc.encode("\n\n"); + units.forEach((unit, i) => { + if (i > 0) { + for (const t of separator) { + tokens.push(t); + tokenPages.push(unit.page); + } + } + for (const t of enc.encode(unit.text)) { + tokens.push(t); + tokenPages.push(unit.page); + } + }); + + const chunks: DocumentChunk[] = []; + for (let start = 0; start < tokens.length; start += step) { + const slice = tokens.slice(start, start + targetTokens); + let content = enc.decode(slice).trim(); + if (content.length > maxChunkChars) content = content.slice(0, maxChunkChars); + if (content) { + chunks.push({ + chunkIndex: chunks.length, + content, + page: tokenPages[start] ?? null, + tokenCount: slice.length, + }); + } + if (start + targetTokens >= tokens.length) break; + } + return chunks; +} + +/** Token count for a string under the same encoder the chunker uses. */ +export function countTokens(text: string): number { + return encoder().encode(text).length; +} diff --git a/backend/src/lib/rag/ingest.ts b/backend/src/lib/rag/ingest.ts new file mode 100644 index 000000000..53f9a6ab4 --- /dev/null +++ b/backend/src/lib/rag/ingest.ts @@ -0,0 +1,186 @@ +import { createServerSupabase } from "../supabase"; +import { downloadFile as defaultDownloadFile } from "../storage"; +import { getUserApiKeys } from "../userSettings"; +import type { UserApiKeys } from "../llm"; +import { + getActiveEmbeddingProvider, + resolveEmbeddingModel, + type EmbeddingProviderAdapter, +} from "../llm/embeddings"; +import { + extractPdfMarkdown, + extractDocxMarkdown, +} from "../tabular/tabular.extract"; +import { chunkMarkdown } from "./chunker"; + +type Db = ReturnType; + +/** Tiny job payload — carries NO secrets; everything else is re-derived here. */ +export interface EmbeddingJobData { + documentId: string; + versionId: string; + /** Enqueuer — only used for logging/attribution; the row owner is the doc's. */ + userId: string; +} + +export interface EmbeddingIngestDeps { + db: Db; + downloadFile: (key: string) => Promise; + getApiKeys: (userId: string, db: Db) => Promise; + /** Resolve the deployment's embedding adapter, or undefined (unconfigured). */ + resolveProvider: () => EmbeddingProviderAdapter | undefined; + /** The single embedding model this deployment ingests + searches with. */ + resolveModel: () => string; + extractMarkdown: (buf: ArrayBuffer, fileType: string) => Promise; +} + +async function defaultExtractMarkdown( + buf: ArrayBuffer, + fileType: string, +): Promise { + const t = fileType.toLowerCase(); + if (t === "pdf") return extractPdfMarkdown(buf); + if (t === "docx" || t === "doc") return extractDocxMarkdown(buf); + return ""; +} + +function defaultDeps(): EmbeddingIngestDeps { + return { + db: createServerSupabase(), + downloadFile: defaultDownloadFile, + getApiKeys: (userId, db) => getUserApiKeys(userId, db), + resolveProvider: () => getActiveEmbeddingProvider(), + resolveModel: () => resolveEmbeddingModel(), + extractMarkdown: defaultExtractMarkdown, + }; +} + +/** Serialize a float vector to a pgvector text literal ('[1,2,3]'). */ +export function toVectorLiteral(vec: number[]): string { + return `[${vec.join(",")}]`; +} + +const EMBED_BATCH_SIZE = 64; + +export type IngestResult = + | { status: "embedded"; chunks: number } + | { status: "skipped"; reason: string } + | { status: "cleared"; chunks: 0 }; + +/** + * Chunk + embed one document version and upsert its rows into document_chunks. + * + * Shared by the BullMQ worker and the backfill script. Idempotent + retry-safe: + * it deletes every chunk for this version before inserting the fresh set, so a + * retry can't leave duplicates or stale rows. Only the document's CURRENT + * version is embedded — an enqueued version that a newer edit has already + * superseded is skipped, so a burst of rapid edits doesn't index dead versions. + * + * Throws (so BullMQ retries) only on transient failures: a missing download or + * an embedding provider returning the wrong count. A missing provider + * (embeddings unconfigured) is a graceful no-op, never a thrown error. + */ +export async function runEmbeddingIngestion( + data: EmbeddingJobData, + deps: EmbeddingIngestDeps = defaultDeps(), +): Promise { + const { documentId, versionId } = data; + const { db } = deps; + + const { data: doc } = await db + .from("documents") + .select("id, current_version_id, user_id") + .eq("id", documentId) + .single(); + if (!doc) return { status: "skipped", reason: "document_not_found" }; + + // Only index the current version. A superseded version's enqueued job is a + // no-op — the newer version has (or will have) its own job. + if ((doc.current_version_id as string | null) !== versionId) { + return { status: "skipped", reason: "superseded" }; + } + + const { data: version } = await db + .from("document_versions") + .select("id, storage_path, file_type") + .eq("id", versionId) + .single(); + const storagePath = + version && typeof version.storage_path === "string" + ? version.storage_path + : null; + if (!storagePath) return { status: "skipped", reason: "no_storage_path" }; + + const provider = deps.resolveProvider(); + if (!provider) { + // No embedding provider configured / available. Degrade gracefully — + // never fail the job. + console.warn( + "[embedding-ingest] no embedding provider registered; skipping", + { documentId, versionId }, + ); + return { status: "skipped", reason: "no_provider" }; + } + const model = deps.resolveModel(); + + const bytes = await deps.downloadFile(storagePath); + if (!bytes) { + // Transient (storage hiccup / eventual consistency) — throw to retry. + throw new Error( + `[embedding-ingest] document bytes not found at ${storagePath}`, + ); + } + + const fileType = + version && typeof version.file_type === "string" ? version.file_type : ""; + const markdown = await deps.extractMarkdown(bytes, fileType); + const chunks = chunkMarkdown(markdown); + + if (chunks.length === 0) { + // Nothing to index (empty/unsupported doc) — clear any stale rows so the + // index reflects the current version, then finish. + await db.from("document_chunks").delete().eq("version_id", versionId); + return { status: "cleared", chunks: 0 }; + } + + // Embed in batches; a batch returning the wrong count is a hard error so the + // job retries rather than silently dropping chunks. + const apiKeys = await deps.getApiKeys(data.userId, db); + const vectors: number[][] = []; + for (let i = 0; i < chunks.length; i += EMBED_BATCH_SIZE) { + const batch = chunks.slice(i, i + EMBED_BATCH_SIZE); + const embedded = await provider.embed( + batch.map((c) => c.content), + apiKeys, + ); + if (embedded.length !== batch.length) { + throw new Error( + `[embedding-ingest] provider returned ${embedded.length} vectors for ${batch.length} inputs`, + ); + } + vectors.push(...embedded); + } + + const rows = chunks.map((chunk, i) => ({ + document_id: documentId, + version_id: versionId, + user_id: doc.user_id as string, + chunk_index: chunk.chunkIndex, + content: chunk.content, + page: chunk.page, + token_count: chunk.tokenCount, + embedding_model: model, + embedding: toVectorLiteral(vectors[i]), + })); + + // Delete-then-insert keeps the operation idempotent under retry. + await db.from("document_chunks").delete().eq("version_id", versionId); + const { error } = await db.from("document_chunks").insert(rows); + if (error) { + throw new Error( + `[embedding-ingest] failed to insert chunks: ${error.message}`, + ); + } + + return { status: "embedded", chunks: rows.length }; +} diff --git a/backend/src/lib/rag/searchDocuments.ts b/backend/src/lib/rag/searchDocuments.ts new file mode 100644 index 000000000..5dd2dfe91 --- /dev/null +++ b/backend/src/lib/rag/searchDocuments.ts @@ -0,0 +1,116 @@ +import { createServerSupabase } from "../supabase"; +import { toVectorLiteral } from "./ingest"; + +type Db = ReturnType; + +export interface ChunkMatch { + document_id: string; + version_id: string; + chunk_index: number; + content: string; + page: number | null; + /** Cosine distance (0 = identical); lower is more relevant. */ + distance: number; +} + +/** Cosine similarity of two equal-length vectors (0 for a zero vector). */ +export function cosineSimilarity(a: number[], b: number[]): number { + const n = Math.min(a.length, b.length); + let dot = 0; + let na = 0; + let nb = 0; + for (let i = 0; i < n; i++) { + dot += a[i] * b[i]; + na += a[i] * a[i]; + nb += b[i] * b[i]; + } + if (na === 0 || nb === 0) return 0; + return dot / (Math.sqrt(na) * Math.sqrt(nb)); +} + +/** Parse a pgvector text literal ('[1,2,3]') back into a number[]. */ +export function parseVectorLiteral(raw: unknown): number[] { + if (Array.isArray(raw)) return raw as number[]; + if (typeof raw !== "string") return []; + const inner = raw.trim().replace(/^\[/, "").replace(/\]$/, ""); + if (!inner) return []; + return inner.split(",").map((s) => Number.parseFloat(s)); +} + +/** + * Top-k cosine search over the chunks of a PRE-SCOPED set of documents. + * + * `documentIds` MUST already be the set the caller has access to — this runs as + * service_role (RLS bypassed), so the document-id filter is the authz boundary + * and passing unscoped ids would leak cross-tenant chunks. It is filtered on + * `model` so a model change can't return dimension-mismatched vectors. + * + * Prefers the SQL RPC (HNSW index); falls back to fetching this document set's + * chunks and ranking in-process only if the RPC is unavailable. + */ +export async function searchDocumentChunks(params: { + db: Db; + queryEmbedding: number[]; + model: string; + documentIds: string[]; + topK: number; +}): Promise { + const { db, queryEmbedding, model, documentIds, topK } = params; + if (documentIds.length === 0 || queryEmbedding.length === 0) return []; + + const { data, error } = await db.rpc("match_document_chunks", { + p_query_embedding: toVectorLiteral(queryEmbedding), + p_document_ids: documentIds, + p_model: model, + p_match_count: topK, + }); + if (!error && Array.isArray(data)) { + return (data as ChunkMatch[]).slice(0, topK); + } + if (error) { + console.warn( + "[search-documents] RPC unavailable; falling back to in-process ranking", + { err: error }, + ); + } + return fallbackRank({ db, queryEmbedding, model, documentIds, topK }); +} + +/** + * In-process ranking fallback: fetch the scoped documents' chunks and sort by + * cosine distance. Correct but O(rows) — the RPC's HNSW path is preferred; this + * only runs where the migration's function is missing (e.g. an old DB). + */ +async function fallbackRank(params: { + db: Db; + queryEmbedding: number[]; + model: string; + documentIds: string[]; + topK: number; +}): Promise { + const { db, queryEmbedding, model, documentIds, topK } = params; + const { data } = await db + .from("document_chunks") + .select("document_id, version_id, chunk_index, content, page, embedding") + .in("document_id", documentIds) + .eq("embedding_model", model); + const rows = (data ?? []) as { + document_id: string; + version_id: string; + chunk_index: number; + content: string; + page: number | null; + embedding: unknown; + }[]; + return rows + .map((r) => ({ + document_id: r.document_id, + version_id: r.version_id, + chunk_index: r.chunk_index, + content: r.content, + page: r.page, + distance: 1 - cosineSimilarity(queryEmbedding, parseVectorLiteral(r.embedding)), + })) + .sort((a, b) => a.distance - b.distance) + .slice(0, topK); +} diff --git a/backend/src/routes/documents.ts b/backend/src/routes/documents.ts index b28f2e49d..1c00caf1e 100644 --- a/backend/src/routes/documents.ts +++ b/backend/src/routes/documents.ts @@ -12,6 +12,7 @@ import { } from "../lib/storage"; import { docxToPdf, convertedPdfKey } from "../lib/convert"; import { enqueueConversion } from "../lib/queue/conversionQueue"; +import { maybeEnqueueEmbedding } from "../lib/queue/embeddingQueue"; import { extractTrackedChangeIds, resolveTrackedChange, @@ -557,6 +558,14 @@ documentsRouter.post( .json({ detail: "Failed to update document current version." }); } + // Re-index the new current version for semantic search (no-op unless + // ASYNC_EMBEDDING); mirrors the conversion enqueue on the upload path. + await maybeEnqueueEmbedding({ + documentId, + versionId: versionRow.id as string, + userId, + }); + if (willDeleteSource) { const { error: deleteErr } = await deleteDocumentAndVersionFiles( db, @@ -729,6 +738,12 @@ documentsRouter.post( .json({ detail: "Failed to update document current version." }); } + await maybeEnqueueEmbedding({ + documentId, + versionId: versionRow.id as string, + userId, + }); + res.status(201).json(versionRow); }, ); @@ -923,6 +938,10 @@ documentsRouter.put( .map((path) => deleteFile(path).catch(() => {})), ); + // The version's bytes changed in place — re-index it if it is the current + // version (the ingestion job skips it otherwise). + await maybeEnqueueEmbedding({ documentId, versionId, userId }); + res.json(updated); }, ); @@ -1437,6 +1456,13 @@ export async function handleDocumentUpload( }); } + // Index the new version for semantic search (no-op unless ASYNC_EMBEDDING). + await maybeEnqueueEmbedding({ + documentId: docId, + versionId: versionRow.id, + userId, + }); + const { data: updated } = await db .from("documents") .select("*") diff --git a/backend/src/workers/__tests__/embeddingWorker.test.ts b/backend/src/workers/__tests__/embeddingWorker.test.ts new file mode 100644 index 000000000..4437fd6f8 --- /dev/null +++ b/backend/src/workers/__tests__/embeddingWorker.test.ts @@ -0,0 +1,187 @@ +import { describe, it, expect, vi } from "vitest"; + +// Keep transitive imports offline (mirrors conversionWorker.test). +vi.mock("../../lib/supabase", () => ({ createServerSupabase: vi.fn() })); +vi.mock("../../lib/storage", () => ({ downloadFile: vi.fn() })); + +import { + runEmbeddingIngestion, + toVectorLiteral, + type EmbeddingIngestDeps, + type EmbeddingJobData, +} from "../../lib/rag/ingest"; +import { isPermanentFailure } from "../embeddingWorker"; +import type { Job } from "bullmq"; + +const JOB: EmbeddingJobData = { + documentId: "doc-1", + versionId: "ver-1", + userId: "user-1", +}; + +type Row = Record; + +function makeDb(opts: { doc: Row | null; version: Row | null }) { + const seq: string[] = []; + const deletes: { col: string; val: unknown }[] = []; + const inserts: Row[][] = []; + const db = { + seq, + deletes, + inserts, + from(table: string) { + if (table === "documents" || table === "document_versions") { + const data = table === "documents" ? opts.doc : opts.version; + return { + select: () => ({ + eq: () => ({ single: async () => ({ data }) }), + }), + }; + } + if (table === "document_chunks") { + return { + delete: () => ({ + eq: async (col: string, val: unknown) => { + seq.push("delete"); + deletes.push({ col, val }); + return {}; + }, + }), + insert: async (rows: Row[]) => { + seq.push("insert"); + inserts.push(rows); + return { error: null }; + }, + }; + } + throw new Error(`unexpected table ${table}`); + }, + }; + return db; +} + +function makeDeps( + db: ReturnType, + overrides: Partial = {}, +): EmbeddingIngestDeps { + return { + db: db as never, + downloadFile: vi.fn(async () => new ArrayBuffer(8)), + getApiKeys: vi.fn(async () => ({})), + resolveProvider: () => ({ + id: "fake-embed", + matchesModel: () => true, + dimensions: 3, + models: ["fake-model"], + embed: async (texts: string[]) => texts.map(() => [0, 0, 0]), + }), + resolveModel: () => "fake-model", + extractMarkdown: async () => "Some document body text to embed.", + ...overrides, + }; +} + +describe("runEmbeddingIngestion", () => { + it("chunks, embeds, and delete-then-inserts chunk rows for the current version", async () => { + const db = makeDb({ + doc: { id: "doc-1", current_version_id: "ver-1", user_id: "owner-1" }, + version: { id: "ver-1", storage_path: "uploads/doc-1.pdf", file_type: "pdf" }, + }); + const result = await runEmbeddingIngestion(JOB, makeDeps(db)); + + expect(result).toEqual({ status: "embedded", chunks: 1 }); + // Idempotent: delete for this version happens BEFORE the insert. + expect(db.seq).toEqual(["delete", "insert"]); + expect(db.deletes[0]).toEqual({ col: "version_id", val: "ver-1" }); + + const rows = db.inserts[0]; + expect(rows).toHaveLength(1); + expect(rows[0]).toMatchObject({ + document_id: "doc-1", + version_id: "ver-1", + user_id: "owner-1", + chunk_index: 0, + embedding_model: "fake-model", + embedding: toVectorLiteral([0, 0, 0]), + }); + expect(typeof rows[0].content).toBe("string"); + expect(typeof rows[0].token_count).toBe("number"); + }); + + it("skips a superseded version without downloading or writing", async () => { + const download = vi.fn(async () => new ArrayBuffer(8)); + const db = makeDb({ + doc: { id: "doc-1", current_version_id: "ver-2", user_id: "owner-1" }, + version: { id: "ver-1", storage_path: "x", file_type: "pdf" }, + }); + const result = await runEmbeddingIngestion(JOB, makeDeps(db, { downloadFile: download })); + expect(result).toEqual({ status: "skipped", reason: "superseded" }); + expect(download).not.toHaveBeenCalled(); + expect(db.inserts).toHaveLength(0); + }); + + it("skips gracefully (no throw) when no embedding provider is registered", async () => { + const db = makeDb({ + doc: { id: "doc-1", current_version_id: "ver-1", user_id: "owner-1" }, + version: { id: "ver-1", storage_path: "x", file_type: "pdf" }, + }); + const result = await runEmbeddingIngestion( + JOB, + makeDeps(db, { resolveProvider: () => undefined }), + ); + expect(result).toEqual({ status: "skipped", reason: "no_provider" }); + expect(db.inserts).toHaveLength(0); + }); + + it("throws so BullMQ retries when the document bytes are missing", async () => { + const db = makeDb({ + doc: { id: "doc-1", current_version_id: "ver-1", user_id: "owner-1" }, + version: { id: "ver-1", storage_path: "x", file_type: "pdf" }, + }); + await expect( + runEmbeddingIngestion(JOB, makeDeps(db, { downloadFile: async () => null })), + ).rejects.toThrow(/bytes not found/); + }); + + it("throws when the provider returns the wrong number of vectors", async () => { + const db = makeDb({ + doc: { id: "doc-1", current_version_id: "ver-1", user_id: "owner-1" }, + version: { id: "ver-1", storage_path: "x", file_type: "pdf" }, + }); + const deps = makeDeps(db, { + resolveProvider: () => ({ + id: "fake", + matchesModel: () => true, + dimensions: 3, + models: [], + embed: async () => [], // 0 vectors for 1 input + }), + }); + await expect(runEmbeddingIngestion(JOB, deps)).rejects.toThrow(/returned 0 vectors/); + }); + + it("clears stale rows and writes nothing when the document has no extractable text", async () => { + const db = makeDb({ + doc: { id: "doc-1", current_version_id: "ver-1", user_id: "owner-1" }, + version: { id: "ver-1", storage_path: "x", file_type: "pdf" }, + }); + const result = await runEmbeddingIngestion( + JOB, + makeDeps(db, { extractMarkdown: async () => " " }), + ); + expect(result).toEqual({ status: "cleared", chunks: 0 }); + expect(db.seq).toEqual(["delete"]); + expect(db.inserts).toHaveLength(0); + }); +}); + +describe("isPermanentFailure", () => { + const job = (attemptsMade: number, attempts?: number) => + ({ attemptsMade, opts: { attempts } }) as unknown as Job; + + it("is false while retries remain, true once exhausted", () => { + expect(isPermanentFailure(job(1, 3))).toBe(false); + expect(isPermanentFailure(job(3, 3))).toBe(true); + expect(isPermanentFailure(job(1))).toBe(true); + }); +}); diff --git a/backend/src/workers/embeddingWorker.ts b/backend/src/workers/embeddingWorker.ts new file mode 100644 index 000000000..94e7645ca --- /dev/null +++ b/backend/src/workers/embeddingWorker.ts @@ -0,0 +1,80 @@ +import { Worker, type Job } from "bullmq"; +import { getRedisConnection } from "../lib/queue/connection"; +import { + EMBEDDING_QUEUE, + type EmbeddingJobData, +} from "../lib/queue/embeddingQueue"; +import { runEmbeddingIngestion } from "../lib/rag/ingest"; + +/** + * In-process BullMQ worker that runs the chunk+embed ingestion for one document + * version. Mirrors conversionWorker / extractionWorker: the job body just calls + * the dependency-injected core (runEmbeddingIngestion), which is what the unit + * tests exercise directly without a live queue. + */ + +/** True once a job has exhausted its retries (BullMQ 'failed', no attempts left). */ +export function isPermanentFailure(job: Job): boolean { + const maxAttempts = job.opts.attempts ?? 1; + return job.attemptsMade >= maxAttempts; +} + +let worker: Worker | null = null; + +export function createEmbeddingWorker(): Worker { + if (worker) return worker; + worker = new Worker( + EMBEDDING_QUEUE, + async (job: Job) => { + const result = await runEmbeddingIngestion(job.data); + console.log("[embedding-worker] ingestion finished", { + jobId: job.id, + documentId: job.data.documentId, + versionId: job.data.versionId, + result, + }); + }, + { + connection: getRedisConnection(), + concurrency: 2, + // Recover jobs orphaned by a worker crash mid-run. + stalledInterval: 30_000, + maxStalledCount: 2, + }, + ); + worker.on("stalled", (jobId) => { + console.warn("[embedding-worker] job stalled; will be re-queued", { jobId }); + }); + worker.on("failed", (job, err) => { + if (!job) { + console.error("[embedding-worker] job failed (no job)", { err }); + return; + } + if (!isPermanentFailure(job)) { + console.error( + "[embedding-worker] job failed (will retry, attempts remain)", + { jobId: job.id, err }, + ); + return; + } + // No terminal DB flip needed: the document stays usable, semantic search + // just misses this version until the next edit or a backfill re-enqueues. + console.error( + "[embedding-worker] job permanently failed; version left unindexed", + { + jobId: job.id, + documentId: job.data.documentId, + versionId: job.data.versionId, + err, + }, + ); + }); + return worker; +} + +export async function stopEmbeddingWorker(): Promise { + if (worker) { + await worker.close(); + worker = null; + } +} diff --git a/backend/src/workers/registry.ts b/backend/src/workers/registry.ts index ecde7a645..e780669b2 100644 --- a/backend/src/workers/registry.ts +++ b/backend/src/workers/registry.ts @@ -6,8 +6,13 @@ import { createExtractionWorker, stopExtractionWorker, } from "./extractionWorker"; +import { + createEmbeddingWorker, + stopEmbeddingWorker, +} from "./embeddingWorker"; import { closeConversionQueue } from "../lib/queue/conversionQueue"; import { closeExtractionQueue } from "../lib/queue/extractionQueue"; +import { closeEmbeddingQueue } from "../lib/queue/embeddingQueue"; /** * One background queue's lifecycle, described declaratively. `startWorkers()` / @@ -48,4 +53,11 @@ export const WORKER_REGISTRY: WorkerDescriptor[] = [ stop: stopExtractionWorker, closeQueue: closeExtractionQueue, }, + { + name: "document-embedding", + enabled: () => process.env.ASYNC_EMBEDDING === "true", + create: createEmbeddingWorker, + stop: stopEmbeddingWorker, + closeQueue: closeEmbeddingQueue, + }, ];