feat: RAG — pgvector retrieval and a search_documents tool (ADR) - #41
Open
amal66 wants to merge 1 commit into
Open
feat: RAG — pgvector retrieval and a search_documents tool (ADR)#41amal66 wants to merge 1 commit into
amal66 wants to merge 1 commit into
Conversation
Port of the fork's (amal66/mike) RAG feature: - migrations/20260716_01_document_chunks_pgvector.sql: pgvector-backed document_chunks table (RLS default-deny, anon/authenticated revoked), HNSW cosine index, and the match_document_chunks top-k RPC. The caller-supplied document-id set is the authz boundary; rows are filtered by embedding_model so a model change can't mix widths. - lib/rag/: token-aware markdown chunker (js-tiktoken cl100k_base, page attribution from '## Page N' headers), the idempotent chunk+embed ingestion (delete-then-insert per version; only the CURRENT version is indexed), and the cosine search with an in-process fallback when the RPC is missing. - lib/llm/embeddings/: an embedding-provider registry with OpenAI (text-embedding-3-*, Matryoshka-truncated to the column width) and Gemini (text-embedding-004 / gemini-embedding-001) adapters. One model + one width per deployment. - lib/queue/embeddingQueue.ts + workers/embeddingWorker.ts: the document-embedding BullMQ queue/worker (deterministic per-version jobId, attempts:3 + backoff) wired into the worker registry behind ASYNC_EMBEDDING. maybeEnqueueEmbedding is called from the upload, new-version, copy-version and replace-file paths and is a strict no-op unless the flag is on; enqueue failures never fail the upload. - search_documents assistant tool: semantic top-k search across the chat's documents, scoped to the access-checked doc index, returning passages with doc-label + page for citation. Degrades gracefully (helpful note, no error) when embeddings are unconfigured. - scripts/backfillEmbeddings.ts: idempotent backfill for existing documents (enqueue or --sync inline). ASYNC_EMBEDDING defaults "false": nothing dials Redis or an embedding API, and search_documents tells the model to fall back to read_document. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CEguyEgXa9JjCciXCcVemC
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Design (ADR)
Summarized from the fork's
docs/adr/0001-registry-based-extensibility.md,docs/adr/0003-default-synchronous-env-gated-async.md,docs/adr/0004-database-enforced-idempotency.mdand thedocument_chunksmigration's design notes.Context. On big matters the assistant's document tools are literal:
read_documentpulls whole documents into context (expensive and truncation-prone on long ones) andfind_in_documentis a Ctrl+F that misses paraphrase. The assistant needs meaning-based retrieval over the chat's documents, inside the firm's own database — not a third-party vector service.Decision.
document_chunks: one row per token-aware chunk,vector(768), HNSW cosine index) rather than an external vector DB. The search RPCmatch_document_chunksruns top-k cosine over a caller-supplied, pre-scoped set of document ids — the app's access-checked doc index is the authz boundary; the RPC never reaches beyond it. RLS stays default-deny withanon/authenticatedrevoked, like every other content table.vector(N)column has a single fixed N, soEMBEDDING_MODEL/EMBEDDING_DIMENSIONpin the deployment;embedding_modelis stored per row and filtered in the RPC so a model change is a hard boundary (re-embed viascripts/backfillEmbeddings.ts) rather than silent dimension-mixing.(version_id, chunk_index), deterministic BullMQ jobId per version, and only the document's current version is embedded (superseded jobs are no-ops), so retries and edit-bursts can't duplicate or index dead versions.EmbeddingProviderAdapterwithregister/findon a plain Map — same shape as the LLM-provider registry): OpenAI and Gemini adapters here; adding another is a single-file, zero-core-change operation.ASYNC_EMBEDDING="false"produces no embeddings andsearch_documentsdegrades gracefully (a helpful note telling the model to fall back toread_document— never an error).Consequences. Retrieval quality on long documents without shipping content to a new vendor; the index lives (and is backed up/deleted) with the rest of the firm's data; a pinned model per deployment is a real constraint — switching models requires a backfill; ingestion cost is bounded and repairable (idempotent backfill script).
Alternatives considered. External vector DB (rejected: new infrastructure + data leaves the firm's Postgres); embedding inline on upload (rejected: adds seconds and a new external API call to every upload; as a queue job it retries and never blocks or fails the user's upload); per-user model choice for embeddings (rejected: vector(N) has one width — deployment-pinned by design).
Summary
Better answers on big matters. The assistant gets a
search_documentstool: semantic, meaning-based search across the documents in the chat, returning the most relevant passages with their source document and page so answers stay grounded and citable. On long contracts and large document sets it finds the indemnity/termination/assignment language by concept — even when the wording doesn't match — instead of reading whole documents into context. Everything runs inside the firm's own database (pgvector in the same Postgres/Supabase instance); no document text is sent to any new third-party service beyond the embedding model already configured for the deployment.Changes
backend/migrations/20260716_01_document_chunks_pgvector.sql—create extension if not exists vector;document_chunkstable (RLS default-deny, grants matching the other content tables), HNSW cosine index, and thematch_document_chunkstop-k RPC.backend/src/lib/rag/—chunker.ts(token-aware chunking via js-tiktoken cl100k_base, 512-token windows with 64 overlap, page attribution from the## Page Nmarkers the PDF extractor emits),ingest.ts(idempotent chunk+embed for one version; graceful no-op when no provider is configured),searchDocuments.ts(RPC-first cosine search with an in-process ranking fallback when the RPC is missing).backend/src/lib/llm/embeddings/— provider registry + OpenAI (text-embedding-3-*, asked for exactly the 768 column width) and Gemini (text-embedding-004,gemini-embedding-001) adapters.backend/src/lib/queue/embeddingQueue.ts+backend/src/workers/embeddingWorker.ts— thedocument-embeddingBullMQ queue/worker, registered inWORKER_REGISTRYbehindASYNC_EMBEDDING.maybeEnqueueEmbeddingis called from the upload, new-version, copy-version and replace-file paths; it is a strict no-op unless the flag is on, and an enqueue failure never fails the user's request (the backfill script repairs gaps).search_documentsassistant tool — schema intoolSchemas.ts, dispatch intoolDispatcher.ts: embeds the query with the same model used at ingest, scopes the search to the chat's access-checked document ids, optionaldoc_idnarrowing, and returns each passage under the doc-label citation reminder (page included) likeread_document.backend/scripts/backfillEmbeddings.ts— idempotent backfill (enqueue mode, or--syncto ingest inline with no queue).backend/.env.example,backend/package.json(+ lockfile).lib/rag/__tests__/{chunker,searchDocuments}.test.ts,lib/llm/embeddings/__tests__/registry.test.ts,workers/__tests__/embeddingWorker.test.ts,lib/chat/tools/__tests__/searchDocuments.test.ts.Why
Base branch: this PR is based on
upstream-pr/durable-queues, notupstream-main, because the dependency is real in the fork's code:embeddingQueue.tsbuilds onlib/queue/connection.ts, the embedding worker registers in the queues PR'sWORKER_REGISTRY, and the fork's only production ingestion path is that worker. Retrieval itself (searchDocumentChunks, the tool) has no queue dependency, but without the worker nothing populates the index outside the backfill script.Cost posture for a £20-VPS firm: everything defaults off.
ASYNC_EMBEDDING="false"(the default) produces no embeddings, dials no Redis and no embedding API, andsearch_documentsdegrades to a polite "unavailable, use read_document" note. Turning it on costs: the Redis container already required by the queues PR; the pgvector extension — on Supabase this is the built-invectorextension: enable it in Dashboard → Database → Extensions (or let this migration'screate extension if not exists vectordo it when run with sufficient privileges) — no paid add-on, works on the free tier; and embedding API usage, which at OpenAItext-embedding-3-smallpricing (~$0.02 per 1M tokens) rounds to pennies per thousand document pages. Self-hosted Postgres needs the pgvector package installed (apt install postgresql-16-pgvectoror the Docker image), which is free.New runtime dep (shipped by the fork for this feature):
js-tiktoken ^1.0.21— pure-JS tokenizer for the chunker's token budgeting (no native binary, no network). Embedding calls use the deployment's existing OpenAI/Gemini keys; no new SDK.Testing
npm install && npm run build(tsc) green on the branch as committed.npx vitest run→ 12 test files, 68 tests, all passing (this branch's 5 new suites + the queues branch's 6 + the harness'sdownloadTokens).maybeEnqueueEmbeddingreturns before touching the queue unlessASYNC_EMBEDDING === "true"; the tool test covers the no-provider degradation path.Provenance
All added lines are mechanical ports of amal66/mike@origin/main (b3166dd) — path moves, import rewrites, inlined types; exceptions:
org_idcolumn/index andingest.tsdrops theorg_idrow stamping (tests adjusted accordingly).user_idcolumn istext not null(mirroring upstream'sdocuments.user_id) instead of the fork'suuid references auth.users on delete cascade; deletion cascades through thedocument_idFK instead.backend/migrations/20260716_01_…with a-- Migration date:header, sorting after upstream's latest20260710_01); comment references to fork-only migrations removed. SQL content otherwise verbatim.AIRGAPPEDgating omitted — they belong to the fork's local-models/air-gap feature, whose substrate (Ollama chat provider,OPENAI_BASE_URLSSRF guard) upstream lacks.registerBuiltinEmbeddingProvidersregisters OpenAI + Gemini unconditionally andresolveEmbeddingModeldefaults straight to the cloud model. The registry-test's air-gap cases are dropped with it.https://api.openai.com/v1(the same constant upstream's chat adapter uses) instead of the fork'sresolveOpenAIBaseUrl()guard.spotlight()prompt-injection fencing + turn nonce are fork-only chat infrastructure; chunk bodies are returned under the citation reminder exactly like upstream'sread_document, without the fence. The corresponding test assertions are dropped.documentToolHandlers+ToolExecutionContext) into anelse ifbranch of upstream'srunToolCallsdispatcher — same logic line-for-line where the surfaces align (pushToolResult→toolResults.push,ctx.results.docsFound→docsFound,return→continue); its test is rewritten to driverunToolCalls.logger→console, OpenTelemetry payload carrier removed, and the fork's zodenvmodule reads inlined asprocess.envwith defaults preserved — same conventions as the queues PR.Credits & prior art
🤖 Generated with Claude Code
https://claude.ai/code/session_01CEguyEgXa9JjCciXCcVemC