Skip to content

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

Open
amal66 wants to merge 1 commit into
upstream-pr/durable-queuesfrom
upstream-pr/rag-pgvector
Open

feat: RAG — pgvector retrieval and a search_documents tool (ADR)#41
amal66 wants to merge 1 commit into
upstream-pr/durable-queuesfrom
upstream-pr/rag-pgvector

Conversation

@amal66

@amal66 amal66 commented Jul 17, 2026

Copy link
Copy Markdown
Owner

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.md and the document_chunks migration's design notes.

Context. On big matters the assistant's document tools are literal: read_document pulls whole documents into context (expensive and truncation-prone on long ones) and find_in_document is 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.

  • pgvector in the existing Postgres (document_chunks: one row per token-aware chunk, vector(768), HNSW cosine index) rather than an external vector DB. The search RPC match_document_chunks runs 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 with anon/authenticated revoked, like every other content table.
  • One embedding model + one width per deployment. A vector(N) column has a single fixed N, so EMBEDDING_MODEL/EMBEDDING_DIMENSION pin the deployment; embedding_model is stored per row and filtered in the RPC so a model change is a hard boundary (re-embed via scripts/backfillEmbeddings.ts) rather than silent dimension-mixing.
  • Ingestion is DB-enforced idempotent: delete-by-version then insert, unique (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.
  • Embedding providers go through the fork's registry pattern (EmbeddingProviderAdapter with register/find on 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.
  • Env-gated async, off by default (same lever as the queues PR): ASYNC_EMBEDDING="false" produces no embeddings and search_documents degrades gracefully (a helpful note telling the model to fall back to read_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_documents tool: 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.sqlcreate extension if not exists vector; document_chunks table (RLS default-deny, grants matching the other content tables), HNSW cosine index, and the match_document_chunks top-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 N markers 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 — the document-embedding BullMQ queue/worker, registered in WORKER_REGISTRY behind ASYNC_EMBEDDING. maybeEnqueueEmbedding is 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_documents assistant tool — schema in toolSchemas.ts, dispatch in toolDispatcher.ts: embeds the query with the same model used at ingest, scopes the search to the chat's access-checked document ids, optional doc_id narrowing, and returns each passage under the doc-label citation reminder (page included) like read_document.
  • backend/scripts/backfillEmbeddings.ts — idempotent backfill (enqueue mode, or --sync to ingest inline with no queue).
  • backend/.env.example, backend/package.json (+ lockfile).
  • Tests: 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, not upstream-main, because the dependency is real in the fork's code: embeddingQueue.ts builds on lib/queue/connection.ts, the embedding worker registers in the queues PR's WORKER_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, and search_documents degrades 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-in vector extension: enable it in Dashboard → Database → Extensions (or let this migration's create extension if not exists vector do it when run with sufficient privileges) — no paid add-on, works on the free tier; and embedding API usage, which at OpenAI text-embedding-3-small pricing (~$0.02 per 1M tokens) rounds to pennies per thousand document pages. Self-hosted Postgres needs the pgvector package installed (apt install postgresql-16-pgvector or 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.
  • With the vitest harness merged locally (not part of this branch): npx vitest run12 test files, 68 tests, all passing (this branch's 5 new suites + the queues branch's 6 + the harness's downloadTokens).
  • Flag-off path: maybeEnqueueEmbedding returns before touching the queue unless ASYNC_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:

  • Multi-tenancy (org) stripped: upstream has no organizations, so the migration drops the fork's org_id column/index and ingest.ts drops the org_id row stamping (tests adjusted accordingly).
  • user_id column is text not null (mirroring upstream's documents.user_id) instead of the fork's uuid references auth.users on delete cascade; deletion cascades through the document_id FK instead.
  • Migration renamed/dated to upstream's convention (backend/migrations/20260716_01_… with a -- Migration date: header, sorting after upstream's latest 20260710_01); comment references to fork-only migrations removed. SQL content otherwise verbatim.
  • Ollama/local embedding provider and AIRGAPPED gating omitted — they belong to the fork's local-models/air-gap feature, whose substrate (Ollama chat provider, OPENAI_BASE_URL SSRF guard) upstream lacks. registerBuiltinEmbeddingProviders registers OpenAI + Gemini unconditionally and resolveEmbeddingModel defaults straight to the cloud model. The registry-test's air-gap cases are dropped with it.
  • OpenAI embed adapter uses the constant https://api.openai.com/v1 (the same constant upstream's chat adapter uses) instead of the fork's resolveOpenAIBaseUrl() guard.
  • spotlight() prompt-injection fencing + turn nonce are fork-only chat infrastructure; chunk bodies are returned under the citation reminder exactly like upstream's read_document, without the fence. The corresponding test assertions are dropped.
  • The tool handler is translated from the fork's tool-registry architecture (documentToolHandlers + ToolExecutionContext) into an else if branch of upstream's runToolCalls dispatcher — same logic line-for-line where the surfaces align (pushToolResulttoolResults.push, ctx.results.docsFounddocsFound, returncontinue); its test is rewritten to drive runToolCalls.
  • pino loggerconsole, OpenTelemetry payload carrier removed, and the fork's zod env module reads inlined as process.env with defaults preserved — same conventions as the queues PR.

Credits & prior art

  • @nwhitehouse (nwhitehouse/mike) — independently parallels this work: their fork built RAG over tabular documents using pgvector with an HNSW index — the same in-database retrieval choice this PR makes. Independent implementations that converged on pgvector/HNSW inside the app's own Postgres rather than an external vector service.

🤖 Generated with Claude Code

https://claude.ai/code/session_01CEguyEgXa9JjCciXCcVemC

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant