From 8cbc82021afa83484625dfc51a42dd095e366957 Mon Sep 17 00:00:00 2001 From: julianam-w <77252335+julianam-w@users.noreply.github.com> Date: Sun, 11 Oct 2026 17:29:43 +1100 Subject: [PATCH 1/8] feat(indexing): reuse stored embeddings for unchanged chunk text Store a SHA-256 content_hash per chunk and look up existing embeddings by hash before calling Voyage, so ingest, sync and reindex only pay for text the namespace has not seen. Duplicate texts within a run are embedded once. Incremental sync/reindex now embed before deleting and swap old chunks for new in one transaction, so a Voyage failure no longer leaves gaps. Existing tables are backfilled with content_hash on next setup. Co-Authored-By: Claude Opus 5.5 (1M context) --- AGENTS.md | 2 +- README.md | 2 + pyproject.toml | 3 ++ rag/db.py | 64 +++++++++++++++++++---- scripts/ingest.py | 97 ++++++++++++++++++++++++++++++----- scripts/reindex.py | 54 +++---------------- scripts/sync.py | 44 ++-------------- tests/__init__.py | 0 tests/test_embedding_reuse.py | 95 ++++++++++++++++++++++++++++++++++ uv.lock | 53 ++++++++++++++++++- 10 files changed, 305 insertions(+), 109 deletions(-) create mode 100644 tests/__init__.py create mode 100644 tests/test_embedding_reuse.py diff --git a/AGENTS.md b/AGENTS.md index 67ca7dd..9200db5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -24,7 +24,7 @@ The MCP server does **retrieval only** — it never calls an LLM. The calling ag - `rag/auth.py` — Google OAuth token verifier for HTTP transport - `scripts/ingest.py` — ingest a repo into pgvector - `scripts/ask.py` — local CLI for full retrieve + answer pipeline (uses Anthropic API; dev convenience only) -- `scripts/reindex.py` — drop and re-ingest a namespace +- `scripts/reindex.py` — re-embed an explicit list of changed/deleted files (the primitive `scripts/sync.py` builds on) ## Transport modes diff --git a/README.md b/README.md index 7182902..17d1582 100644 --- a/README.md +++ b/README.md @@ -247,6 +247,8 @@ uv run python scripts/sync.py --namespace tupaia # sync one namespace In production this runs on a schedule via a Railway cron service — see [Deployment](#deployment). +Each chunk stores a SHA-256 `content_hash` of its text. Before calling Voyage, every ingest, sync and reindex looks up existing embeddings by hash and only embeds text the namespace hasn't seen; identical chunks within a run are embedded once. An incremental sync embeds everything first, then swaps the old chunks for the new in a single transaction, so a Voyage failure leaves the index as it was. Tables created before this column existed are backfilled the first time they are set up again. + Two lower-level scripts back this: - `scripts/register.py` — register (or re-register) a namespace's baseline commit SHA without re-ingesting. Useful after a manual `ingest.py` run, or to clear the SHA and force a full reindex on the next sync. - `scripts/reindex.py` — re-embed an explicit set of changed/deleted files, given as `CHANGED_FILES` / `DELETED_FILES` (space-separated repo-relative paths). This is the primitive `sync.py` builds on; call it directly only if you already know exactly which files changed. diff --git a/pyproject.toml b/pyproject.toml index ef2080a..a686ac7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,9 @@ dependencies = [ [project.optional-dependencies] ask = ["anthropic>=0.40.0"] +[dependency-groups] +dev = ["pytest>=8.0.0"] + [project.scripts] github-repo-rag-mcp = "mcp_server:main" diff --git a/rag/db.py b/rag/db.py index d3528b0..3f70b5f 100644 --- a/rag/db.py +++ b/rag/db.py @@ -1,6 +1,8 @@ """Database helpers: setup, upsert.""" +import hashlib import re + import psycopg2 from pgvector.psycopg2 import register_vector from psycopg2.extras import execute_values @@ -40,6 +42,17 @@ def setup_table(conn, table: str) -> None: ALTER TABLE {table} ADD COLUMN IF NOT EXISTS chunk_index INT NOT NULL DEFAULT 0 """) + # content_hash lets re-indexing reuse an existing embedding for unchanged text. + # Backfilled in SQL; must match content_hash() below. + cur.execute(f""" + ALTER TABLE {table} + ADD COLUMN IF NOT EXISTS content_hash TEXT + """) + cur.execute(f""" + UPDATE {table} + SET content_hash = encode(sha256(convert_to(text, 'UTF8')), 'hex') + WHERE content_hash IS NULL + """) cur.execute(f""" CREATE INDEX IF NOT EXISTS {table}_embedding_idx ON {table} USING hnsw (embedding vector_cosine_ops) @@ -52,6 +65,10 @@ def setup_table(conn, table: str) -> None: CREATE INDEX IF NOT EXISTS {table}_file_path_idx ON {table} (file_path) """) + cur.execute(f""" + CREATE INDEX IF NOT EXISTS {table}_content_hash_idx + ON {table} (content_hash) + """) conn.commit() @@ -61,39 +78,68 @@ def setup_db(conn, tables: list[str]) -> None: setup_table(conn, table) -def upsert_chunks(conn, table: str, chunks: list[dict], embeddings: list[list[float]]) -> None: +def content_hash(text: str) -> str: + """SHA-256 hex digest of a chunk's text, matching the SQL backfill in setup_table().""" + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def get_embeddings_by_hash(conn, table: str, hashes: list[str]) -> dict[str, list[float]]: + """Return {content_hash: embedding} for any hashes already stored in the table.""" + if not hashes: + return {} + with conn.cursor() as cur: + cur.execute( + f""" + SELECT DISTINCT ON (content_hash) content_hash, embedding + FROM {table} + WHERE content_hash = ANY(%s) + """, + (list(hashes),), + ) + return {h: emb for h, emb in cur.fetchall()} + + +def upsert_chunks( + conn, table: str, chunks: list[dict], embeddings: list[list[float]], commit: bool = True +) -> None: """Upsert a batch of chunks into a RAG table.""" rows = [ - (chunk["id"], emb, chunk["text"], chunk["file_path"], chunk["package"], chunk.get("chunk_index", 0)) - for chunk, emb in zip(chunks, embeddings) + ( + chunk["id"], emb, chunk["text"], chunk["file_path"], chunk["package"], + chunk.get("chunk_index", 0), content_hash(chunk["text"]), + ) + for chunk, emb in zip(chunks, embeddings, strict=True) ] with conn.cursor() as cur: execute_values( cur, f""" - INSERT INTO {table} (id, embedding, text, file_path, package, chunk_index) + INSERT INTO {table} (id, embedding, text, file_path, package, chunk_index, content_hash) VALUES %s ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding, text = EXCLUDED.text, file_path = EXCLUDED.file_path, package = EXCLUDED.package, - chunk_index = EXCLUDED.chunk_index + chunk_index = EXCLUDED.chunk_index, + content_hash = EXCLUDED.content_hash """, rows, - template="(%s, %s::vector, %s, %s, %s, %s)", + template="(%s, %s::vector, %s, %s, %s, %s, %s)", ) - conn.commit() + if commit: + conn.commit() -def delete_file_chunks(conn, tables: list[str], file_path: str) -> int: +def delete_file_chunks(conn, tables: list[str], file_path: str, commit: bool = True) -> int: """Delete all chunks for a file across the given tables. Returns count deleted.""" total = 0 for table in tables: with conn.cursor() as cur: cur.execute(f"DELETE FROM {table} WHERE file_path = %s", (file_path,)) total += cur.rowcount - conn.commit() + if commit: + conn.commit() return total diff --git a/scripts/ingest.py b/scripts/ingest.py index 2b4d862..f213fc5 100644 --- a/scripts/ingest.py +++ b/scripts/ingest.py @@ -36,7 +36,17 @@ from dotenv import load_dotenv load_dotenv(Path(__file__).parent.parent / ".env") -from rag.db import get_conn, register_namespace, sanitise_namespace, setup_db, setup_meta_table, upsert_chunks +from rag.db import ( + content_hash, + delete_file_chunks, + get_conn, + get_embeddings_by_hash, + register_namespace, + sanitise_namespace, + setup_db, + setup_meta_table, + upsert_chunks, +) from rag.github import shallow_clone from rag.query import embed_batch @@ -65,6 +75,7 @@ DOC_CHUNK_OVERLAP = 400 EMBED_BATCH = 64 # texts per voyage API call +INDEX_GROUP = 1000 # chunks per hash lookup + upsert during a full ingest _LANG_MAP = { ".ts": Language.TS, ".tsx": Language.TS, @@ -157,6 +168,77 @@ def chunk_file(file_path: Path, repo_root: Path, namespace: str) -> list[dict]: ] +# ── Embedding with reuse ────────────────────────────────────────────────────── + +def embed_chunks(conn, table: str, chunks: list[dict]) -> list[list[float]]: + """Return an embedding per chunk, calling Voyage only for text not already in the table. + + Identical texts within `chunks` are embedded once. + """ + hashes = [content_hash(c["text"]) for c in chunks] + known = get_embeddings_by_hash(conn, table, sorted(set(hashes))) + + pending: dict[str, str] = {} + for h, c in zip(hashes, chunks, strict=True): + if h not in known: + pending.setdefault(h, c["text"]) + + items = list(pending.items()) + for i in range(0, len(items), EMBED_BATCH): + batch = items[i:i + EMBED_BATCH] + embeddings = embed_batch([text for _, text in batch], input_type="document") + known.update({h: emb for (h, _), emb in zip(batch, embeddings, strict=True)}) + + print(f" {len(chunks)} chunks: {len(chunks) - len(pending)} reused, {len(pending)} embedded") + return [known[h] for h in hashes] + + +def index_chunks(conn, table: str, chunks: list[dict]) -> None: + """Embed (with reuse) and upsert chunks in groups, committing per group.""" + total_groups = (len(chunks) + INDEX_GROUP - 1) // INDEX_GROUP + for i in range(0, len(chunks), INDEX_GROUP): + group = chunks[i:i + INDEX_GROUP] + print(f" Group {i // INDEX_GROUP + 1}/{total_groups}") + upsert_chunks(conn, table, group, embed_chunks(conn, table, group)) + + +def replace_file_chunks( + conn, namespace: str, repo_root: Path, changed: list[str], deleted: list[str] +) -> None: + """Re-index changed files and drop deleted ones in a single transaction. + + Embeddings are fetched before anything is deleted, so stored embeddings for + unchanged text can be reused and the index never has a gap if embedding fails. + """ + code_table = f"{namespace}_code" + docs_table = f"{namespace}_docs" + tables = [code_table, docs_table] + + all_chunks: list[dict] = [] + for rel_path in changed: + file_path = repo_root / rel_path + if not file_path.exists() or should_skip(Path(rel_path)): + continue + all_chunks.extend(chunk_file(file_path, repo_root, namespace)) + + by_table = {t: [c for c in all_chunks if c["table"] == t] for t in tables} + print(f" New chunks: {len(all_chunks)} ({len(by_table[code_table])} code, {len(by_table[docs_table])} docs)") + embeddings = {t: embed_chunks(conn, t, chunks) for t, chunks in by_table.items() if chunks} + + try: + total_deleted = sum( + delete_file_chunks(conn, tables, rel_path, commit=False) for rel_path in changed + deleted + ) + for table, chunks in by_table.items(): + if chunks: + upsert_chunks(conn, table, chunks, embeddings[table], commit=False) + conn.commit() + except Exception: + conn.rollback() + raise + print(f" Removed {total_deleted} stale chunks, upserted {len(all_chunks)}") + + # ── Main pipeline ───────────────────────────────────────────────────────────── def ingest(repo_path: str, namespace: str, repo_url: str | None = None) -> None: @@ -190,17 +272,8 @@ def ingest(repo_path: str, namespace: str, repo_url: str | None = None) -> None: for table, chunks in [(code_table, code_chunks), (docs_table, doc_chunks)]: if not chunks: continue - total_batches = (len(chunks) + EMBED_BATCH - 1) // EMBED_BATCH - print(f"\nUpserting {len(chunks)} chunks to {table} ({total_batches} batches)...") - - for i in range(0, len(chunks), EMBED_BATCH): - batch = chunks[i:i + EMBED_BATCH] - batch_num = i // EMBED_BATCH + 1 - print(f" Batch {batch_num}/{total_batches} — embedding {len(batch)} chunks...", end=" ", flush=True) - embeddings = embed_batch([c["text"] for c in batch], input_type="document") - print("upserting...", end=" ", flush=True) - upsert_chunks(conn, table, batch, embeddings) - print("done") + print(f"\nUpserting {len(chunks)} chunks to {table}...") + index_chunks(conn, table, chunks) if repo_url: try: diff --git a/scripts/reindex.py b/scripts/reindex.py index 942397f..850d2ac 100644 --- a/scripts/reindex.py +++ b/scripts/reindex.py @@ -23,63 +23,23 @@ from dotenv import load_dotenv load_dotenv(Path(__file__).parent.parent / ".env") -from scripts.ingest import EMBED_BATCH, chunk_file, should_skip -from rag.db import delete_file_chunks, get_conn, upsert_chunks -from rag.query import embed_batch +from rag.db import get_conn, sanitise_namespace, setup_db +from scripts.ingest import replace_file_chunks def reindex(repo_path: str, namespace: str) -> None: - code_table = f"{namespace}_code" - docs_table = f"{namespace}_docs" - tables = [code_table, docs_table] - + namespace = sanitise_namespace(namespace) conn = get_conn(os.environ["DATABASE_URL"]) - repo_root = Path(repo_path).resolve() + setup_db(conn, [f"{namespace}_code", f"{namespace}_docs"]) changed = [f for f in os.environ.get("CHANGED_FILES", "").split() if f] deleted = [f for f in os.environ.get("DELETED_FILES", "").split() if f] print(f"Changed: {len(changed)} files, Deleted: {len(deleted)} files") - # Delete stale chunks for changed and deleted files - total_deleted = 0 - for rel_path in changed + deleted: - n = delete_file_chunks(conn, tables, rel_path) - if n: - print(f" Deleted {n} chunks for {rel_path}") - total_deleted += n - print(f"Deleted {total_deleted} stale chunks total") - - # Re-chunk and re-embed changed files - all_chunks: list[dict] = [] - for rel_path in changed: - file_path = repo_root / rel_path - if not file_path.exists(): - continue - if should_skip(Path(rel_path)): - continue - chunks = chunk_file(file_path, repo_root, namespace) - if chunks: - all_chunks.extend(chunks) - - if not all_chunks: - print("No new chunks to index.") + try: + replace_file_chunks(conn, namespace, Path(repo_path).resolve(), changed, deleted) + finally: conn.close() - return - - code_chunks = [c for c in all_chunks if c["table"] == code_table] - doc_chunks = [c for c in all_chunks if c["table"] == docs_table] - print(f"\nNew chunks: {len(all_chunks)} ({len(code_chunks)} code, {len(doc_chunks)} docs)") - - for table, chunks in [(code_table, code_chunks), (docs_table, doc_chunks)]: - if not chunks: - continue - print(f"Upserting {len(chunks)} chunks to {table}...") - for i in range(0, len(chunks), EMBED_BATCH): - batch = chunks[i:i + EMBED_BATCH] - embeddings = embed_batch([c["text"] for c in batch], input_type="document") - upsert_chunks(conn, table, batch, embeddings) - - conn.close() print("Reindex complete.") diff --git a/scripts/sync.py b/scripts/sync.py index 9c08ff2..4aeceb6 100644 --- a/scripts/sync.py +++ b/scripts/sync.py @@ -27,17 +27,15 @@ load_dotenv(Path(__file__).parent.parent / ".env") from rag.db import ( - delete_file_chunks, get_conn, get_namespaces, register_namespace, + setup_db, setup_meta_table, update_namespace_sha, - upsert_chunks, ) from rag.github import get_changed_files, get_latest_sha, parse_repo_url, shallow_clone -from rag.query import embed_batch -from scripts.ingest import EMBED_BATCH, chunk_file, ingest, should_skip +from scripts.ingest import ingest, replace_file_chunks def _full_reindex(repo_url: str, namespace: str, token: str | None) -> None: @@ -84,47 +82,15 @@ def sync_namespace(conn, ns: dict, token: str | None) -> None: print(f"[{namespace}] Up to date (no indexed files changed)") return - code_table = f"{namespace}_code" - docs_table = f"{namespace}_docs" - tables = [code_table, docs_table] - - # Remove stale chunks for every changed or deleted file - total_deleted = 0 - for rel_path in changed + deleted: - n = delete_file_chunks(conn, tables, rel_path) - if n: - print(f" Deleted {n} stale chunks: {rel_path}") - total_deleted += n - print(f" {total_deleted} stale chunks removed") + # Adds content_hash to tables created before embedding reuse existed + setup_db(conn, [f"{namespace}_code", f"{namespace}_docs"]) # Re-embed changed files — requires a checkout for file content - all_chunks: list[dict] = [] with tempfile.TemporaryDirectory() as tmpdir: repo_path = os.path.join(tmpdir, "repo") print(f" Cloning {repo_url} (shallow)...") shallow_clone(repo_url, repo_path, token) - repo_root = Path(repo_path) - - for rel_path in changed: - file_path = repo_root / rel_path - if not file_path.exists() or should_skip(Path(rel_path)): - continue - chunks = chunk_file(file_path, repo_root, namespace) - all_chunks.extend(chunks) - - if all_chunks: - code_chunks = [c for c in all_chunks if c["table"] == code_table] - doc_chunks = [c for c in all_chunks if c["table"] == docs_table] - print(f" New chunks: {len(all_chunks)} ({len(code_chunks)} code, {len(doc_chunks)} docs)") - - for table, chunks in [(code_table, code_chunks), (docs_table, doc_chunks)]: - if not chunks: - continue - print(f" Upserting {len(chunks)} chunks to {table}...") - for i in range(0, len(chunks), EMBED_BATCH): - batch = chunks[i:i + EMBED_BATCH] - embeddings = embed_batch([c["text"] for c in batch], input_type="document") - upsert_chunks(conn, table, batch, embeddings) + replace_file_chunks(conn, namespace, Path(repo_path), changed, deleted) update_namespace_sha(conn, namespace, latest_sha) print(f"[{namespace}] Sync complete → {latest_sha[:8]}") diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_embedding_reuse.py b/tests/test_embedding_reuse.py new file mode 100644 index 0000000..38e8b8f --- /dev/null +++ b/tests/test_embedding_reuse.py @@ -0,0 +1,95 @@ +"""Embedding reuse in scripts/ingest.py — DB and Voyage calls are faked.""" + +from pathlib import Path + +import pytest + +from rag.db import content_hash +from scripts import ingest + + +class FakeConn: + def __init__(self) -> None: + self.commits = 0 + self.rollbacks = 0 + + def commit(self) -> None: + self.commits += 1 + + def rollback(self) -> None: + self.rollbacks += 1 + + +@pytest.fixture +def voyage(monkeypatch): + calls: list[list[str]] = [] + + def fake_embed_batch(texts, input_type="document"): + calls.append(list(texts)) + return [[float(len(t))] for t in texts] + + monkeypatch.setattr(ingest, "embed_batch", fake_embed_batch) + return calls + + +def _chunk(text: str, file_path: str = "a.py", i: int = 0) -> dict: + return {"id": f"{file_path}:{i}", "text": text, "file_path": file_path, "package": "", "chunk_index": i, + "table": "ns_code"} + + +def test_reuses_stored_embeddings_and_dedupes_within_run(monkeypatch, voyage): + stored = {content_hash("old"): [99.0]} + monkeypatch.setattr(ingest, "get_embeddings_by_hash", lambda conn, table, hashes: { + h: stored[h] for h in hashes if h in stored + }) + chunks = [_chunk("old"), _chunk("new"), _chunk("new", "b.py"), _chunk("newer")] + + result = ingest.embed_chunks(FakeConn(), "ns_code", chunks) + + assert voyage == [["new", "newer"]] + assert result == [[99.0], [3.0], [3.0], [5.0]] + + +def test_no_voyage_call_when_everything_is_stored(monkeypatch, voyage): + monkeypatch.setattr(ingest, "get_embeddings_by_hash", lambda conn, table, hashes: {h: [1.0] for h in hashes}) + assert ingest.embed_chunks(FakeConn(), "ns_code", [_chunk("a"), _chunk("b")]) == [[1.0], [1.0]] + assert voyage == [] + + +def test_splits_new_texts_into_voyage_sized_batches(monkeypatch, voyage): + monkeypatch.setattr(ingest, "get_embeddings_by_hash", lambda conn, table, hashes: {}) + chunks = [_chunk(f"t{i}", i=i) for i in range(ingest.EMBED_BATCH + 1)] + ingest.embed_chunks(FakeConn(), "ns_code", chunks) + assert [len(c) for c in voyage] == [ingest.EMBED_BATCH, 1] + + +def test_replace_embeds_before_deleting_and_commits_once(monkeypatch, tmp_path: Path): + (tmp_path / "a.py").write_text("print('hi')\n", encoding="utf-8") + events: list[str] = [] + monkeypatch.setattr(ingest, "embed_chunks", lambda conn, table, chunks: events.append("embed") or + [[0.0]] * len(chunks)) + monkeypatch.setattr(ingest, "delete_file_chunks", lambda conn, tables, path, commit: events.append( + f"delete:{path}:{commit}") or 1) + monkeypatch.setattr(ingest, "upsert_chunks", lambda conn, table, chunks, embs, commit: events.append( + f"upsert:{table}:{commit}")) + conn = FakeConn() + + ingest.replace_file_chunks(conn, "ns", tmp_path, ["a.py"], ["gone.py"]) + + assert events == ["embed", "delete:a.py:False", "delete:gone.py:False", "upsert:ns_code:False"] + assert conn.commits == 1 + + +def test_replace_leaves_index_untouched_when_embedding_fails(monkeypatch, tmp_path: Path): + (tmp_path / "a.py").write_text("print('hi')\n", encoding="utf-8") + + def boom(conn, table, chunks): + raise RuntimeError("voyage down") + + deletes: list[str] = [] + monkeypatch.setattr(ingest, "embed_chunks", boom) + monkeypatch.setattr(ingest, "delete_file_chunks", lambda *a, **k: deletes.append(a[2]) or 0) + + with pytest.raises(RuntimeError): + ingest.replace_file_chunks(FakeConn(), "ns", tmp_path, ["a.py"], []) + assert deletes == [] diff --git a/uv.lock b/uv.lock index a189697..7f9812f 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.11" resolution-markers = [ "python_full_version >= '3.14'", @@ -634,6 +634,11 @@ ask = [ { name = "anthropic" }, ] +[package.dev-dependencies] +dev = [ + { name = "pytest" }, +] + [package.metadata] requires-dist = [ { name = "anthropic", marker = "extra == 'ask'", specifier = ">=0.40.0" }, @@ -647,6 +652,9 @@ requires-dist = [ ] provides-extras = ["ask"] +[package.metadata.requires-dev] +dev = [{ name = "pytest", specifier = ">=8.0.0" }] + [[package]] name = "h11" version = "0.16.0" @@ -746,6 +754,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/01/e1/2069291243c926a2ff1cd706c7f3eeb9b62144bf60f77c9fb9ff2fb26bd3/iniconfig-2.3.1.tar.gz", hash = "sha256:67f4b9c50da0dedf52af349e7749a80a9057a5031199791b906c3bb3ae878960", size = 21209, upload-time = "2026-10-06T22:48:38.076Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/43/4ca9e49d27a1fcf6bece6f6aec0ea46bb9112489b93d4b688fb415457bdb/iniconfig-2.3.1-py3-none-any.whl", hash = "sha256:9121e2c1fdb355232495be3194c8dfe87ccc2d5dee45947b78e68f499790d7a7", size = 7552, upload-time = "2026-10-06T22:48:36.959Z" }, +] + [[package]] name = "jiter" version = "0.16.0" @@ -1400,6 +1417,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/36/54/0169bc772ec491108b62f644f8ecf1fe5d8ae5ebafde2ee2142210166903/pillow-12.3.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:04f01d28a6aaff387bf842a13be313df23ba0597a44f1a976c9feb3c6ff4711a", size = 7231786, upload-time = "2026-07-01T11:56:35.046Z" }, ] +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + [[package]] name = "propcache" version = "0.5.2" @@ -1703,6 +1729,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/77/c1/6e422f34e569cf8e18df68d1939c81c099d2b61e4f7d9621c8a77560799c/pydantic_settings-2.14.2-py3-none-any.whl", hash = "sha256:a20c97b37910b6550d5ea50fbcc2d4187defe58cd57070b73863d069419c9440", size = 61715, upload-time = "2026-06-19T13:44:55.02Z" }, ] +[[package]] +name = "pygments" +version = "2.21.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/49/2e/ced460408999b33da6b31b0021b0f37d329e202d4169aeb164493778f25b/pygments-2.21.0.tar.gz", hash = "sha256:610ca751c9bc2492b38eb9a38a7fbc93edbbb2d7182edaf34e66ae493dee5c8c", size = 5005329, upload-time = "2026-08-17T08:02:48.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/46/17f022dd3e953bf20a04a028a21ec746d942f8d2af30fa0f124fa0e6a684/pygments-2.21.0-py3-none-any.whl", hash = "sha256:2363c69b61c4a97c838da3b130dcd6468f4848992b21a82f2a63ec34377137d9", size = 1250147, upload-time = "2026-08-17T08:02:44.912Z" }, +] + [[package]] name = "pyjwt" version = "2.13.0" @@ -1717,6 +1752,22 @@ crypto = [ { name = "cryptography" }, ] +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + [[package]] name = "python-dotenv" version = "1.2.2" From 1c8db14a19d833bbd6120ee7c3ce2157eb299099 Mon Sep 17 00:00:00 2001 From: julianam-w <77252335+julianam-w@users.noreply.github.com> Date: Sun, 11 Oct 2026 17:34:26 +1100 Subject: [PATCH 2/8] feat(sync): embed large full reindexes via the Voyage Batch API Full reindexes with at least VOYAGE_BATCH_THRESHOLD (default 2000) new texts are submitted to the Voyage Batch API (33% cheaper, <= 12 h). Sync records the batch in rag_pending_batches, skips the namespace while it is in flight, then loads the output into rag_embedding_cache and runs the full reindex from the cache. Incremental syncs stay real time. ingest.py gains --batch for first-time indexing. voyageai 0.5.0 has no batch support, so rag/voyage_batch.py calls the Files and Batches REST endpoints directly. Spec: specs/pipelines/batch-embedding.md Co-Authored-By: Claude Opus 5.5 (1M context) --- AGENTS.md | 2 + README.md | 6 ++ rag/db.py | 105 ++++++++++++++++++++++ rag/voyage_batch.py | 95 ++++++++++++++++++++ scripts/ingest.py | 74 +++++++++++++++- scripts/sync.py | 77 ++++++++++++++-- specs/pipelines/batch-embedding.md | 82 +++++++++++++++++ tests/test_batch_embedding.py | 136 +++++++++++++++++++++++++++++ tests/test_embedding_reuse.py | 10 +++ 9 files changed, 576 insertions(+), 11 deletions(-) create mode 100644 rag/voyage_batch.py create mode 100644 specs/pipelines/batch-embedding.md create mode 100644 tests/test_batch_embedding.py diff --git a/AGENTS.md b/AGENTS.md index 9200db5..9164723 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -21,6 +21,7 @@ The MCP server does **retrieval only** — it never calls an LLM. The calling ag - `mcp_server.py` — FastMCP server; exposes `search_codebase`, `get_file`, `list_files`, `get_repo_structure`, `list_namespaces`, `namespace_info` tools - `rag/query.py` — `embed()` and `retrieve()` (hybrid pgvector search) - `rag/db.py` — DB setup and upsert helpers +- `rag/voyage_batch.py` — Voyage Batch API REST client used for large full reindexes (`specs/pipelines/batch-embedding.md`) - `rag/auth.py` — Google OAuth token verifier for HTTP transport - `scripts/ingest.py` — ingest a repo into pgvector - `scripts/ask.py` — local CLI for full retrieve + answer pipeline (uses Anthropic API; dev convenience only) @@ -41,6 +42,7 @@ Each indexed repo gets a namespace (e.g. `tupaia`), creating `{namespace}_code` |---|---| | `DATABASE_URL` | All DB operations | | `VOYAGE_API_KEY` | Embedding (ingest + query) | +| `VOYAGE_BATCH_THRESHOLD` | Sync / `ingest.py --batch` — new texts before a full reindex uses the Batch API (default 2000) | | `ANTHROPIC_API_KEY` | `scripts/ask.py` only — not required for MCP | | `GOOGLE_ALLOWED_DOMAIN` / `GOOGLE_ALLOWED_EMAILS` | HTTP transport auth (optional) | diff --git a/README.md b/README.md index 17d1582..f0c5ad2 100644 --- a/README.md +++ b/README.md @@ -249,6 +249,12 @@ In production this runs on a schedule via a Railway cron service — see [Deploy Each chunk stores a SHA-256 `content_hash` of its text. Before calling Voyage, every ingest, sync and reindex looks up existing embeddings by hash and only embeds text the namespace hasn't seen; identical chunks within a run are embedded once. An incremental sync embeds everything first, then swaps the old chunks for the new in a single transaction, so a Voyage failure leaves the index as it was. Tables created before this column existed are backfilled the first time they are set up again. +### Voyage Batch API for full reindexes + +A full reindex (no baseline SHA, or a diff too large for GitHub's compare API) with at least `VOYAGE_BATCH_THRESHOLD` (default 2000) new texts goes through the [Voyage Batch API](https://docs.voyageai.com/docs/batch-inference), which costs 33% less and finishes within 12 hours. That sync run submits the batch and leaves the namespace as it is. Later runs skip the namespace while the batch is in flight, then load the results into `rag_embedding_cache` and do the full reindex from the cache. Incremental syncs always embed in real time. Spec: [specs/pipelines/batch-embedding.md](specs/pipelines/batch-embedding.md). + +To index a new repo the same way, run `ingest.py --repo --batch` (it can't be combined with `--ref`). The next scheduled sync after Voyage finishes completes the index. Set `VOYAGE_BATCH_THRESHOLD` to a very large number to turn batching off. + Two lower-level scripts back this: - `scripts/register.py` — register (or re-register) a namespace's baseline commit SHA without re-ingesting. Useful after a manual `ingest.py` run, or to clear the SHA and force a full reindex on the next sync. - `scripts/reindex.py` — re-embed an explicit set of changed/deleted files, given as `CHANGED_FILES` / `DELETED_FILES` (space-separated repo-relative paths). This is the primitive `sync.py` builds on; call it directly only if you already know exactly which files changed. diff --git a/rag/db.py b/rag/db.py index 3f70b5f..27dcb5a 100644 --- a/rag/db.py +++ b/rag/db.py @@ -99,6 +99,19 @@ def get_embeddings_by_hash(conn, table: str, hashes: list[str]) -> dict[str, lis return {h: emb for h, emb in cur.fetchall()} +def get_existing_hashes(conn, tables: list[str], hashes: list[str]) -> set[str]: + """Return the hashes that already have an embedding in any of the tables or the batch cache.""" + found: set[str] = set() + with conn.cursor() as cur: + for table in [*tables, "rag_embedding_cache"]: + cur.execute( + f"SELECT DISTINCT content_hash FROM {table} WHERE content_hash = ANY(%s)", + (list(hashes),), + ) + found.update(r[0] for r in cur.fetchall()) + return found + + def upsert_chunks( conn, table: str, chunks: list[dict], embeddings: list[list[float]], commit: bool = True ) -> None: @@ -156,6 +169,23 @@ def setup_meta_table(conn) -> None: last_indexed_at TIMESTAMPTZ ) """) + # Voyage Batch API state — see specs/pipelines/batch-embedding.md + cur.execute("CREATE EXTENSION IF NOT EXISTS vector") + cur.execute(""" + CREATE TABLE IF NOT EXISTS rag_pending_batches ( + batch_id TEXT PRIMARY KEY, + namespace TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """) + cur.execute(""" + CREATE TABLE IF NOT EXISTS rag_embedding_cache ( + content_hash TEXT PRIMARY KEY, + embedding vector(1024) NOT NULL, + batch_id TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """) conn.commit() @@ -198,3 +228,78 @@ def update_namespace_sha(conn, namespace: str, commit_sha: str) -> None: (commit_sha, namespace), ) conn.commit() + + +# ── Batch embedding state ───────────────────────────────────────────────────── + +def add_pending_batch(conn, batch_id: str, namespace: str) -> None: + with conn.cursor() as cur: + cur.execute( + "INSERT INTO rag_pending_batches (batch_id, namespace) VALUES (%s, %s)", + (batch_id, namespace), + ) + conn.commit() + + +def get_pending_batches(conn, namespace: str) -> list[str]: + with conn.cursor() as cur: + cur.execute( + "SELECT batch_id FROM rag_pending_batches WHERE namespace = %s ORDER BY created_at", + (namespace,), + ) + return [r[0] for r in cur.fetchall()] + + +def delete_pending_batches(conn, batch_ids: list[str]) -> None: + with conn.cursor() as cur: + cur.execute("DELETE FROM rag_pending_batches WHERE batch_id = ANY(%s)", (batch_ids,)) + conn.commit() + + +def cache_embeddings(conn, batch_id: str, rows: list[tuple[str, list[float]]]) -> None: + """Store (content_hash, embedding) pairs from a batch output.""" + if not rows: + return + with conn.cursor() as cur: + execute_values( + cur, + """ + INSERT INTO rag_embedding_cache (content_hash, embedding, batch_id) + VALUES %s + ON CONFLICT (content_hash) DO UPDATE SET + embedding = EXCLUDED.embedding, + batch_id = EXCLUDED.batch_id, + created_at = NOW() + """, + [(h, emb, batch_id) for h, emb in rows], + template="(%s, %s::vector, %s)", + ) + conn.commit() + + +def get_cached_embeddings(conn, hashes: list[str]) -> dict[str, list[float]]: + if not hashes: + return {} + with conn.cursor() as cur: + cur.execute( + "SELECT content_hash, embedding FROM rag_embedding_cache WHERE content_hash = ANY(%s)", + (list(hashes),), + ) + return {h: emb for h, emb in cur.fetchall()} + + +def delete_cached_embeddings(conn, batch_ids: list[str]) -> None: + with conn.cursor() as cur: + cur.execute("DELETE FROM rag_embedding_cache WHERE batch_id = ANY(%s)", (batch_ids,)) + conn.commit() + + +def purge_stale_cache(conn, max_age_days: int = 7) -> int: + with conn.cursor() as cur: + cur.execute( + "DELETE FROM rag_embedding_cache WHERE created_at < NOW() - make_interval(days => %s)", + (max_age_days,), + ) + n = cur.rowcount + conn.commit() + return n diff --git a/rag/voyage_batch.py b/rag/voyage_batch.py new file mode 100644 index 0000000..49f25b7 --- /dev/null +++ b/rag/voyage_batch.py @@ -0,0 +1,95 @@ +"""Voyage Batch API client (Files + Batches REST endpoints). + +The voyageai Python package (0.5.0) has no batch support, so this talks to the +REST API directly. See specs/pipelines/batch-embedding.md. +""" + +import json +import os +import urllib.error +import urllib.request +import uuid +from collections.abc import Iterator + +_API = "https://api.voyageai.com/v1" +MODEL = "voyage-code-3" + +# BL-004: Voyage caps a batch at 100K inputs; the file size cap is undocumented (OQ-001) +MAX_LINES = 100_000 +MAX_BYTES = 100 * 1024 * 1024 + +IN_FLIGHT = {"validating", "in_progress", "finalizing", "cancelling"} +FAILED = {"failed", "cancelled"} + + +def _request(method: str, path: str, body: bytes | None = None, content_type: str | None = None) -> bytes: + headers = {"Authorization": f"Bearer {os.environ['VOYAGE_API_KEY']}"} + if content_type: + headers["Content-Type"] = content_type + req = urllib.request.Request(f"{_API}{path}", data=body, headers=headers, method=method) + try: + with urllib.request.urlopen(req) as resp: + return resp.read() + except urllib.error.HTTPError as e: + detail = e.read().decode(errors="replace") + raise RuntimeError(f"Voyage API {e.code} for {method} {path}: {detail}") from e + + +def build_input_files(texts: dict[str, str]) -> list[bytes]: + """Encode {content_hash: text} as JSONL batch input files, split at the line/byte caps.""" + files: list[bytes] = [] + lines: list[bytes] = [] + size = 0 + for h, text in texts.items(): + # BL-003: one text per line, keyed by its content hash + line = json.dumps({"custom_id": h, "body": {"input": [text]}}, ensure_ascii=False).encode("utf-8") + b"\n" + if lines and (len(lines) >= MAX_LINES or size + len(line) > MAX_BYTES): + files.append(b"".join(lines)) + lines, size = [], 0 + lines.append(line) + size += len(line) + if lines: + files.append(b"".join(lines)) + return files + + +def upload_file(content: bytes, filename: str = "batch.jsonl") -> str: + """Upload a JSONL batch input file and return its file id.""" + boundary = uuid.uuid4().hex + body = ( + f'--{boundary}\r\nContent-Disposition: form-data; name="purpose"\r\n\r\nbatch\r\n' + f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="{filename}"\r\n' + "Content-Type: application/jsonl\r\n\r\n" + ).encode() + content + f"\r\n--{boundary}--\r\n".encode() + resp = _request("POST", "/files", body, f"multipart/form-data; boundary={boundary}") + return json.loads(resp)["id"] + + +def create_batch(input_file_id: str, metadata: dict[str, str] | None = None) -> dict: + payload = { + "endpoint": "/v1/embeddings", + "completion_window": "12h", + "request_params": {"model": MODEL, "input_type": "document"}, + "input_file_id": input_file_id, + } + if metadata: + payload["metadata"] = metadata + return json.loads(_request("POST", "/batches", json.dumps(payload).encode(), "application/json")) + + +def get_batch(batch_id: str) -> dict: + return json.loads(_request("GET", f"/batches/{batch_id}")) + + +def iter_output_embeddings(output_file_id: str) -> Iterator[tuple[str, list[float]]]: + """Yield (content_hash, embedding) for each successful line of a batch output file.""" + content = _request("GET", f"/files/{output_file_id}/content") + for raw in content.splitlines(): + if not raw.strip(): + continue + line = json.loads(raw) + response = line.get("response") or {} + if line.get("error") or response.get("status_code") != 200: + continue + data = response["body"]["data"] + yield line["custom_id"], data[0]["embedding"] diff --git a/scripts/ingest.py b/scripts/ingest.py index f213fc5..f2ced5b 100644 --- a/scripts/ingest.py +++ b/scripts/ingest.py @@ -36,11 +36,15 @@ from dotenv import load_dotenv load_dotenv(Path(__file__).parent.parent / ".env") +from rag import voyage_batch from rag.db import ( + add_pending_batch, content_hash, delete_file_chunks, + get_cached_embeddings, get_conn, get_embeddings_by_hash, + get_existing_hashes, register_namespace, sanitise_namespace, setup_db, @@ -76,6 +80,7 @@ EMBED_BATCH = 64 # texts per voyage API call INDEX_GROUP = 1000 # chunks per hash lookup + upsert during a full ingest +DEFAULT_BATCH_THRESHOLD = 2000 # new texts before a full reindex goes via the Voyage Batch API _LANG_MAP = { ".ts": Language.TS, ".tsx": Language.TS, @@ -177,6 +182,8 @@ def embed_chunks(conn, table: str, chunks: list[dict]) -> list[list[float]]: """ hashes = [content_hash(c["text"]) for c in chunks] known = get_embeddings_by_hash(conn, table, sorted(set(hashes))) + # BL-007: results of a collected Voyage batch + known.update(get_cached_embeddings(conn, sorted(set(hashes) - known.keys()))) pending: dict[str, str] = {} for h, c in zip(hashes, chunks, strict=True): @@ -239,9 +246,46 @@ def replace_file_chunks( print(f" Removed {total_deleted} stale chunks, upserted {len(all_chunks)}") +# ── Voyage Batch API ────────────────────────────────────────────────────────── + +def batch_threshold() -> int: + raw = os.environ.get("VOYAGE_BATCH_THRESHOLD", "") + if not raw: + return DEFAULT_BATCH_THRESHOLD + try: + return int(raw) + except ValueError: + raise SystemExit(f"VOYAGE_BATCH_THRESHOLD must be an integer, got {raw!r}") from None + + +def pending_texts(conn, tables: list[str], chunks: list[dict]) -> dict[str, str]: + """Return {content_hash: text} for chunk texts with no stored or cached embedding.""" + by_hash = {content_hash(c["text"]): c["text"] for c in chunks} + hashes = list(by_hash) + existing: set[str] = set() + for i in range(0, len(hashes), INDEX_GROUP * 10): + existing |= get_existing_hashes(conn, tables, hashes[i:i + INDEX_GROUP * 10]) + return {h: t for h, t in by_hash.items() if h not in existing} + + +def submit_batches(conn, namespace: str, texts: dict[str, str]) -> list[str]: + """Submit texts to the Voyage Batch API and record each batch against the namespace.""" + batch_ids = [] + files = voyage_batch.build_input_files(texts) + for n, content in enumerate(files, 1): + file_id = voyage_batch.upload_file(content, f"{namespace}-{n}.jsonl") + batch = voyage_batch.create_batch(file_id, metadata={"namespace": namespace}) + # BL-005: record the batch; chunks and last_commit_sha are left as they are + add_pending_batch(conn, batch["id"], namespace) + batch_ids.append(batch["id"]) + print(f" Submitted batch {n}/{len(files)}: {batch['id']}") + return batch_ids + + # ── Main pipeline ───────────────────────────────────────────────────────────── -def ingest(repo_path: str, namespace: str, repo_url: str | None = None) -> None: +def ingest(repo_path: str, namespace: str, repo_url: str | None = None, use_batch: bool = False) -> bool: + """Index every file in the repo. Returns False if embedding was handed to the Voyage Batch API.""" namespace = sanitise_namespace(namespace) code_table = f"{namespace}_code" docs_table = f"{namespace}_docs" @@ -269,6 +313,20 @@ def ingest(repo_path: str, namespace: str, repo_url: str | None = None) -> None: doc_chunks = [c for c in all_chunks if c["table"] == docs_table] print(f"Total: {len(all_chunks)} chunks ({len(code_chunks)} code, {len(doc_chunks)} docs)") + # BL-001 / BL-002: large full reindexes go via the Batch API + if use_batch: + texts = pending_texts(conn, tables, all_chunks) + threshold = batch_threshold() + print(f"{len(texts)} distinct texts need embedding (batch threshold {threshold})") + if len(texts) >= threshold: + submit_batches(conn, namespace, texts) + if repo_url: + # No SHA, so the next sync collects the batch and does the full reindex + register_namespace(conn, namespace, repo_url, None) + conn.close() + print("\nBatch submitted — scripts/sync.py will finish indexing once Voyage completes it.") + return False + for table, chunks in [(code_table, code_chunks), (docs_table, doc_chunks)]: if not chunks: continue @@ -287,6 +345,7 @@ def ingest(repo_path: str, namespace: str, repo_url: str | None = None) -> None: conn.close() print("\nIngestion complete.") + return True def main() -> None: @@ -307,10 +366,19 @@ def main() -> None: "--ref", default=None, help="Branch name or release tag to index (e.g. 2.50.5, v2.51.3, main). Defaults to the repo's default branch.", ) + parser.add_argument( + "--batch", action="store_true", + help="Embed via the Voyage Batch API (33%% cheaper, up to 12 h) when enough text is new; " + "scripts/sync.py finishes the index. Requires --repo, not --ref.", + ) args = parser.parse_args() + if args.batch and (not args.repo or args.ref): + # sync.py collects the batch by re-cloning the registered repo's default branch + parser.error("--batch requires --repo and cannot be combined with --ref") + if args.repo_path: - ingest(args.repo_path, args.namespace, repo_url=args.repo) + ingest(args.repo_path, args.namespace, repo_url=args.repo, use_batch=args.batch) else: if not args.repo: parser.error("--repo URL is required when no local repo_path is given") @@ -319,7 +387,7 @@ def main() -> None: ref_label = args.ref or "default branch" print(f"Cloning {args.repo} @ {ref_label} (shallow)...") shallow_clone(args.repo, repo_path, os.environ.get("GITHUB_TOKEN"), ref=args.ref) - ingest(repo_path, args.namespace, repo_url=args.repo) + ingest(repo_path, args.namespace, repo_url=args.repo, use_batch=args.batch) if __name__ == "__main__": diff --git a/scripts/sync.py b/scripts/sync.py index 4aeceb6..f17e093 100644 --- a/scripts/sync.py +++ b/scripts/sync.py @@ -4,7 +4,9 @@ landed since the last index and re-embed only the changed files. Falls back to a full reindex when the diff exceeds GitHub's compare API limits -(> 250 commits or >= 300 files changed). +(> 250 commits or >= 300 files changed). A large full reindex is submitted to the +Voyage Batch API and finished by a later run once the batch completes +(specs/pipelines/batch-embedding.md). Usage: python scripts/sync.py # sync all registered namespaces @@ -13,6 +15,7 @@ Environment variables: DATABASE_URL — PostgreSQL connection string VOYAGE_API_KEY — Voyage AI API key + VOYAGE_BATCH_THRESHOLD — new texts before a full reindex uses the Voyage Batch API (default 2000) GITHUB_TOKEN — GitHub API token (optional; raises rate limit from 60 to 5000 req/hr) """ @@ -26,9 +29,15 @@ from dotenv import load_dotenv load_dotenv(Path(__file__).parent.parent / ".env") +from rag import voyage_batch from rag.db import ( + cache_embeddings, + delete_cached_embeddings, + delete_pending_batches, get_conn, get_namespaces, + get_pending_batches, + purge_stale_cache, register_namespace, setup_db, setup_meta_table, @@ -38,13 +47,47 @@ from scripts.ingest import ingest, replace_file_chunks -def _full_reindex(repo_url: str, namespace: str, token: str | None) -> None: - """Clone and fully reindex. Caller updates the commit SHA afterwards.""" +def _full_reindex(repo_url: str, namespace: str, token: str | None, use_batch: bool = True) -> bool: + """Clone and fully reindex. Returns False if a Voyage batch was submitted instead. + + Caller updates the commit SHA afterwards, only when this returns True. + """ with tempfile.TemporaryDirectory() as tmpdir: repo_path = os.path.join(tmpdir, "repo") print(f" Cloning {repo_url} (shallow)...") shallow_clone(repo_url, repo_path, token) - ingest(repo_path, namespace) # repo_url omitted — caller handles SHA registration + # repo_url omitted — caller handles SHA registration + return ingest(repo_path, namespace, use_batch=use_batch) + + +def _collect_batches(conn, namespace: str, batch_ids: list[str]) -> bool: + """Load finished batch output into the embedding cache. Returns False while any batch is in flight.""" + batches = [voyage_batch.get_batch(b) for b in batch_ids] + + # BL-006: wait for every batch in the namespace to finish + in_flight = [b for b in batches if b["status"] in voyage_batch.IN_FLIGHT] + if in_flight: + for b in in_flight: + counts = b.get("request_counts") or {} + progress = f"{counts.get('completed', 0)}/{counts.get('total', '?')}" + print(f"[{namespace}] Batch {b['id']} {b['status']} ({progress})") + return False + + failed = [] + for b in batches: + if b.get("output_file_id"): + rows = list(voyage_batch.iter_output_embeddings(b["output_file_id"])) + cache_embeddings(conn, b["id"], rows) + print(f"[{namespace}] Batch {b['id']} {b['status']}: cached {len(rows)} embeddings") + if b["status"] in voyage_batch.FAILED: + failed.append(b) + + # BL-009: drop failed batches and fail the sync; the next run resubmits what isn't cached + if failed: + delete_pending_batches(conn, [b["id"] for b in failed]) + detail = ", ".join(f"{b['id']} ({b['status']}: {b.get('errors')})" for b in failed) + raise RuntimeError(f"Voyage batch failed: {detail}") + return True def sync_namespace(conn, ns: dict, token: str | None) -> None: @@ -55,6 +98,20 @@ def sync_namespace(conn, ns: dict, token: str | None) -> None: owner, repo = parse_repo_url(repo_url) latest_sha = get_latest_sha(owner, repo, token) + batch_ids = get_pending_batches(conn, namespace) + if batch_ids: + if not _collect_batches(conn, namespace, batch_ids): + return + # BL-007: real-time full reindex, reading batch results from the cache + print(f"[{namespace}] Batches complete — running full reindex from cache") + _full_reindex(repo_url, namespace, token, use_batch=False) + # BL-008 + register_namespace(conn, namespace, repo_url, latest_sha) + delete_pending_batches(conn, batch_ids) + delete_cached_embeddings(conn, batch_ids) + print(f"[{namespace}] Sync complete → {latest_sha[:8]}") + return + if latest_sha == last_sha: print(f"[{namespace}] Up to date ({latest_sha[:8]})") return @@ -62,8 +119,8 @@ def sync_namespace(conn, ns: dict, token: str | None) -> None: # No previous SHA recorded — do a full reindex to establish a baseline if last_sha is None: print(f"[{namespace}] No baseline commit — running full reindex") - _full_reindex(repo_url, namespace, token) - register_namespace(conn, namespace, repo_url, latest_sha) + if _full_reindex(repo_url, namespace, token): + register_namespace(conn, namespace, repo_url, latest_sha) return print(f"[{namespace}] New commits {last_sha[:8]}...{latest_sha[:8]}, checking diff...") @@ -71,8 +128,8 @@ def sync_namespace(conn, ns: dict, token: str | None) -> None: if too_large: print(f"[{namespace}] Diff too large — falling back to full reindex") - _full_reindex(repo_url, namespace, token) - update_namespace_sha(conn, namespace, latest_sha) + if _full_reindex(repo_url, namespace, token): + update_namespace_sha(conn, namespace, latest_sha) return print(f"[{namespace}] {len(changed)} changed, {len(deleted)} deleted files") @@ -107,6 +164,10 @@ def main() -> None: token = os.environ.get("GITHUB_TOKEN") conn = get_conn(os.environ["DATABASE_URL"]) setup_meta_table(conn) + # BL-010 + purged = purge_stale_cache(conn) + if purged: + print(f"Purged {purged} stale batch embeddings") namespaces = get_namespaces(conn) diff --git a/specs/pipelines/batch-embedding.md b/specs/pipelines/batch-embedding.md new file mode 100644 index 0000000..e620575 --- /dev/null +++ b/specs/pipelines/batch-embedding.md @@ -0,0 +1,82 @@ +# Pipeline Spec: `batch-embedding` + +## Identity + +| Field | Value | +|---|---| +| **Name** | Voyage Batch API embedding for full reindexes | +| **Type** | Scheduled script (`scripts/sync.py`) + CLI (`scripts/ingest.py --batch`) | +| **Status** | `draft` | +| **Owner** | `@julianam-w` | +| **Linear issue** | N/A | +| **Repo** | `github-repo-rag` | +| **Trigger** | Railway cron (sync) / manual (ingest) | +| **Schedule** | Railway cron service schedule | +| **Created** | 2026-10-11 | +| **Last updated** | 2026-10-11 | + +## Purpose + +Cut Voyage spend on large embedding runs by sending them through the Voyage Batch API, which is priced 33% below real-time `/v1/embeddings` in exchange for an up-to-12-hour completion window. Small incremental syncs stay real time so the index stays fresh. + +## Inputs + +- **External sources:** Voyage Files API (`/v1/files`) and Batch API (`/v1/batches`) over HTTPS; GitHub clone as today. +- **Required configuration:** `VOYAGE_API_KEY`, `DATABASE_URL`; optional `VOYAGE_BATCH_THRESHOLD`. + +## Outputs + +- **Tables:** `rag_pending_batches` (one row per in-flight batch), `rag_embedding_cache` (batch results awaiting upsert); `{namespace}_code` / `{namespace}_docs` as today. +- **Downstream consumers:** MCP server search (reads the namespace tables only). + +## Processing logic + +- **BL-001:** Only full reindexes (no baseline, diff too large, or `ingest.py --batch`) are eligible for the Batch API; incremental syncs always embed in real time. +- **BL-002:** A full reindex uses the Batch API when the number of distinct texts still needing an embedding, after reuse from the namespace tables and the cache, is at least `VOYAGE_BATCH_THRESHOLD` (default 2000). +- **BL-003:** Each batch input line embeds one text with `custom_id` set to its `content_hash`, using model `voyage-code-3` and `input_type` `document`. +- **BL-004:** A batch holds at most 100,000 lines and 100 MB of input; more pending texts are split across several batches. +- **BL-005:** Submitting a batch records it in `rag_pending_batches` and leaves the namespace's `last_commit_sha` and indexed chunks unchanged. +- **BL-006:** While a namespace has an in-flight batch, sync skips that namespace. +- **BL-007:** When every batch for a namespace has reached a terminal status, sync loads each successful output line into `rag_embedding_cache`, then runs a real-time full reindex that reads embeddings from the cache before calling Voyage. +- **BL-008:** After a successful collecting reindex, sync advances `last_commit_sha`, deletes the namespace's `rag_pending_batches` rows and deletes the cache rows loaded from those batches. +- **BL-009:** A batch that ends `failed` or `cancelled` has its row deleted and makes the sync for that namespace fail; the next run resubmits only the texts not already cached. +- **BL-010:** Cache rows older than 7 days are deleted at the start of each sync. + +## SLAs + +| SLA | Target | +|---|---| +| Freshness (incremental) | Unchanged — one cron interval | +| Freshness (full reindex via batch) | Up to one cron interval after the batch completes (≤ 12 h + interval) | + +## Error handling + +| Failure mode | Behaviour | Alerting | +|---|---|---| +| Voyage upload / create error | Raise; namespace sync fails; nothing recorded | Cron job exit code | +| Batch `failed` / `cancelled` | BL-009 | Cron job exit code | +| Batch completed with some failed requests | Missing texts embed in real time during collection | Logged count | + +## Acceptance criteria + +| ID | Criterion | Implements | +|---|---|---| +| AC-001 | A full reindex below the threshold embeds in real time and submits no batch | BL-002 | +| AC-002 | A full reindex at or above the threshold submits batches and upserts nothing | BL-002, BL-005 | +| AC-003 | Batch input lines carry `custom_id` = content hash and one input each, split at the line/byte cap | BL-003, BL-004 | +| AC-004 | Sync skips a namespace whose batch is in progress | BL-006 | +| AC-005 | Sync loads completed output into the cache, reindexes with batching disabled, then clears pending rows and cache | BL-007, BL-008 | +| AC-006 | A failed batch row is deleted and the sync raises | BL-009 | + +## Open questions + +| ID | Question | Owner | Due | +|---|---|---|---| +| OQ-001 | Voyage does not document a max upload file size; 100 MB per batch is a conservative assumption | @julianam-w | before merge | +| OQ-002 | Should the Railway cron run more often than daily so completed batches are collected sooner? Up-to-date namespaces cost one GitHub API call each | @julianam-w | before merge | + +## Change log + +| Date | Author | Change | +|---|---|---| +| 2026-10-11 | @julianam-w | Initial draft | diff --git a/tests/test_batch_embedding.py b/tests/test_batch_embedding.py new file mode 100644 index 0000000..a8f1e42 --- /dev/null +++ b/tests/test_batch_embedding.py @@ -0,0 +1,136 @@ +"""Voyage Batch API flow — specs/pipelines/batch-embedding.md. DB, Voyage and git are faked.""" + +import json +from pathlib import Path + +import pytest + +from rag import voyage_batch +from scripts import ingest, sync + + +class FakeConn: + def commit(self) -> None: ... + def rollback(self) -> None: ... + def close(self) -> None: ... + + +@pytest.fixture +def repo(tmp_path: Path) -> Path: + for i in range(5): + (tmp_path / f"f{i}.py").write_text(f"def f{i}():\n return {i}\n", encoding="utf-8") + return tmp_path + + +@pytest.fixture +def db(monkeypatch): + """Stub every DB call ingest() makes and record what it does.""" + calls: dict[str, list] = {"upserts": [], "batches": [], "registered": []} + monkeypatch.setattr(ingest, "get_conn", lambda url: FakeConn()) + monkeypatch.setattr(ingest, "setup_db", lambda conn, tables: None) + monkeypatch.setattr(ingest, "setup_meta_table", lambda conn: None) + monkeypatch.setattr(ingest, "get_existing_hashes", lambda conn, tables, hashes: set()) + monkeypatch.setattr(ingest, "get_embeddings_by_hash", lambda conn, table, hashes: {}) + monkeypatch.setattr(ingest, "get_cached_embeddings", lambda conn, hashes: {}) + monkeypatch.setattr(ingest, "embed_batch", lambda texts, input_type: [[0.0]] * len(texts)) + monkeypatch.setattr(ingest, "upsert_chunks", lambda conn, table, chunks, embs: calls["upserts"].append(table)) + monkeypatch.setattr(ingest, "add_pending_batch", lambda conn, b, ns: calls["batches"].append((b, ns))) + monkeypatch.setattr(ingest, "register_namespace", lambda conn, ns, url, sha: calls["registered"].append(sha)) + monkeypatch.setattr(voyage_batch, "upload_file", lambda content, name: "file-1") + monkeypatch.setattr(voyage_batch, "create_batch", lambda file_id, metadata: {"id": "batch-1"}) + monkeypatch.setenv("DATABASE_URL", "postgresql://fake") + return calls + + +def test_ac_001_below_threshold_embeds_in_real_time(monkeypatch, repo, db): + monkeypatch.setenv("VOYAGE_BATCH_THRESHOLD", "6") + assert ingest.ingest(str(repo), "ns", use_batch=True) is True + assert db["batches"] == [] + assert db["upserts"] == ["ns_code"] + + +def test_ac_002_at_threshold_submits_batch_and_upserts_nothing(monkeypatch, repo, db): + monkeypatch.setenv("VOYAGE_BATCH_THRESHOLD", "5") + assert ingest.ingest(str(repo), "ns", repo_url="https://github.com/o/r", use_batch=True) is False + assert db["batches"] == [("batch-1", "ns")] + assert db["upserts"] == [] + assert db["registered"] == [None] + + +def test_ac_002_batching_off_by_default(monkeypatch, repo, db): + monkeypatch.setenv("VOYAGE_BATCH_THRESHOLD", "0") + assert ingest.ingest(str(repo), "ns") is True + assert db["batches"] == [] + + +def test_invalid_threshold_fails_loudly(monkeypatch): + monkeypatch.setenv("VOYAGE_BATCH_THRESHOLD", "lots") + with pytest.raises(SystemExit): + ingest.batch_threshold() + + +def test_ac_003_input_lines_keyed_by_hash_and_split_at_caps(monkeypatch): + monkeypatch.setattr(voyage_batch, "MAX_LINES", 2) + files = voyage_batch.build_input_files({"h1": "a", "h2": "b", "h3": "ā"}) + lines = [json.loads(line) for f in files for line in f.splitlines()] + assert [f.count(b"\n") for f in files] == [2, 1] + assert lines[2] == {"custom_id": "h3", "body": {"input": ["ā"]}} + + +def test_ac_003_split_at_byte_cap(monkeypatch): + monkeypatch.setattr(voyage_batch, "MAX_BYTES", 60) + files = voyage_batch.build_input_files({"h1": "x" * 20, "h2": "y" * 20}) + assert len(files) == 2 + + +# ── sync collection ─────────────────────────────────────────────────────────── + +@pytest.fixture +def sync_env(monkeypatch): + calls: dict[str, list] = {"reindex": [], "registered": [], "deleted": [], "cached": [], "cleared": []} + monkeypatch.setattr(sync, "get_latest_sha", lambda owner, repo, token: "b" * 40) + monkeypatch.setattr(sync, "get_pending_batches", lambda conn, ns: ["batch-1"]) + monkeypatch.setattr(sync, "_full_reindex", lambda url, ns, token, use_batch=True: + calls["reindex"].append(use_batch) or True) + monkeypatch.setattr(sync, "register_namespace", lambda conn, ns, url, sha: calls["registered"].append(sha)) + monkeypatch.setattr(sync, "delete_pending_batches", lambda conn, ids: calls["deleted"].append(ids)) + monkeypatch.setattr(sync, "cache_embeddings", lambda conn, b, rows: calls["cached"].append((b, rows))) + monkeypatch.setattr(sync, "delete_cached_embeddings", lambda conn, ids: calls["cleared"].append(ids)) + monkeypatch.setattr(voyage_batch, "iter_output_embeddings", lambda file_id: iter([("h1", [1.0])])) + return calls + + +NS = {"namespace": "ns", "repo_url": "https://github.com/o/r", "last_commit_sha": "a" * 40} + + +def test_ac_004_skips_namespace_while_batch_in_progress(monkeypatch, sync_env): + monkeypatch.setattr(voyage_batch, "get_batch", lambda b: {"id": b, "status": "in_progress"}) + sync.sync_namespace(FakeConn(), NS, None) + assert sync_env["reindex"] == [] and sync_env["registered"] == [] and sync_env["cached"] == [] + + +def test_ac_005_collects_completed_batch_then_reindexes_from_cache(monkeypatch, sync_env): + monkeypatch.setattr(voyage_batch, "get_batch", lambda b: {"id": b, "status": "completed", "output_file_id": "out"}) + sync.sync_namespace(FakeConn(), NS, None) + assert sync_env["cached"] == [("batch-1", [("h1", [1.0])])] + assert sync_env["reindex"] == [False] + assert sync_env["registered"] == ["b" * 40] + assert sync_env["deleted"] == [["batch-1"]] and sync_env["cleared"] == [["batch-1"]] + + +def test_ac_006_failed_batch_is_dropped_and_sync_raises(monkeypatch, sync_env): + monkeypatch.setattr(voyage_batch, "get_batch", lambda b: {"id": b, "status": "failed", "errors": "bad"}) + with pytest.raises(RuntimeError, match="batch-1"): + sync.sync_namespace(FakeConn(), NS, None) + assert sync_env["deleted"] == [["batch-1"]] + assert sync_env["reindex"] == [] and sync_env["registered"] == [] + + +def test_submitted_full_reindex_does_not_advance_sha(monkeypatch, sync_env): + monkeypatch.setattr(sync, "get_pending_batches", lambda conn, ns: []) + monkeypatch.setattr(sync, "get_changed_files", lambda *a: ([], [], True)) + monkeypatch.setattr(sync, "_full_reindex", lambda url, ns, token, use_batch=True: False) + advanced: list[str] = [] + monkeypatch.setattr(sync, "update_namespace_sha", lambda conn, ns, sha: advanced.append(sha)) + sync.sync_namespace(FakeConn(), NS, None) + assert advanced == [] diff --git a/tests/test_embedding_reuse.py b/tests/test_embedding_reuse.py index 38e8b8f..f6c8834 100644 --- a/tests/test_embedding_reuse.py +++ b/tests/test_embedding_reuse.py @@ -29,6 +29,7 @@ def fake_embed_batch(texts, input_type="document"): return [[float(len(t))] for t in texts] monkeypatch.setattr(ingest, "embed_batch", fake_embed_batch) + monkeypatch.setattr(ingest, "get_cached_embeddings", lambda conn, hashes: {}) return calls @@ -93,3 +94,12 @@ def boom(conn, table, chunks): with pytest.raises(RuntimeError): ingest.replace_file_chunks(FakeConn(), "ns", tmp_path, ["a.py"], []) assert deletes == [] + + +def test_uses_batch_cache_for_texts_not_in_table(monkeypatch, voyage): + monkeypatch.setattr(ingest, "get_embeddings_by_hash", lambda conn, table, hashes: {}) + monkeypatch.setattr(ingest, "get_cached_embeddings", lambda conn, hashes: { + h: [7.0] for h in hashes if h == content_hash("cached") + }) + assert ingest.embed_chunks(FakeConn(), "ns_code", [_chunk("cached"), _chunk("fresh")]) == [[7.0], [5.0]] + assert voyage == [["fresh"]] From 09bccb37708c8ad099b49795531c4d574e5cd8a0 Mon Sep 17 00:00:00 2001 From: julianam-w <77252335+julianam-w@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:21:48 +1100 Subject: [PATCH 3/8] feat(indexing): prune stale chunks on full reindex; document hourly sync A full reindex now deletes chunks the checkout no longer produces (removed or newly skipped files, and chunk indexes past a file's new length) once every upsert has succeeded. It refuses to run when the checkout yields no indexable files, so a bad clone cannot wipe the index. Document the sync cron as hourly: up-to-date namespaces cost one GitHub API call, only changed text is embedded, and Railway skips a run while the previous one is still going. Co-Authored-By: Claude Opus 5.5 (1M context) --- README.md | 8 ++++---- rag/db.py | 11 +++++++++++ scripts/ingest.py | 12 ++++++++++++ specs/pipelines/batch-embedding.md | 7 +++---- tests/test_batch_embedding.py | 31 ++++++++++++++++++++++++++++-- 5 files changed, 59 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index f0c5ad2..8bed2cc 100644 --- a/README.md +++ b/README.md @@ -238,7 +238,7 @@ uv run python scripts/ask.py --namespace tamanu "How are encounters structured?" ## Incremental reindex -`scripts/ingest.py --repo` registers its namespace automatically, so anything indexed via [Indexing a repository](#indexing-a-repository) is already eligible for sync. `scripts/sync.py` then keeps every registered namespace up to date: for each one it compares the last-indexed commit SHA against the latest on the default branch, re-embeds only the changed files, and deletes chunks for removed files. It falls back to a full reindex when the diff is too large for GitHub's compare API (> 250 commits or ≥ 300 files changed). +`scripts/ingest.py --repo` registers its namespace automatically, so anything indexed via [Indexing a repository](#indexing-a-repository) is already eligible for sync. `scripts/sync.py` then keeps every registered namespace up to date: for each one it compares the last-indexed commit SHA against the latest on the default branch, re-embeds only the changed files, and deletes chunks for removed files. It falls back to a full reindex when the diff is too large for GitHub's compare API (> 250 commits or ≥ 300 files changed). A full reindex (here or via `ingest.py`) removes chunks the checkout no longer produces — deleted or newly skipped files, and chunks past a file's new length — after every upsert succeeds, and refuses to run if the checkout yields no indexable files. ```bash uv run python scripts/sync.py # sync all registered namespaces @@ -269,7 +269,7 @@ Deploys from [`railway.toml`](railway.toml) (Dockerfile build). Needs the base ` ### Scheduled sync (cron service) -The daily incremental sync runs as a separate Railway service (not defined in a repo file — configured directly in the Railway dashboard): +The hourly incremental sync runs as a separate Railway service (not defined in a repo file — configured directly in the Railway dashboard): 1. In the same Railway project, **+ New → Empty Service**, connected to this repo so it builds from the same `Dockerfile` (which copies `scripts/` into the image). 2. **Settings → Deploy → Custom Start Command:** @@ -278,9 +278,9 @@ The daily incremental sync runs as a separate Railway service (not defined in a ``` 3. **Settings → Deploy → Cron Schedule:** ``` - 0 15 * * * + 0 * * * * ``` - Setting a cron schedule makes Railway run the container to completion on that schedule instead of keeping it up as a long-lived service. + Setting a cron schedule makes Railway run the container to completion on that schedule instead of keeping it up as a long-lived service. Hourly costs little: a namespace with no new commits is one GitHub API call, and only changed text is embedded. Railway skips a run while the previous one is still going, so a long full reindex never overlaps the next run. 4. **Variables:** - `DATABASE_URL` — reference the **internal** URL from the Postgres service (e.g. `${{Postgres.DATABASE_URL}}`), not `DATABASE_PUBLIC_URL`. This is why the cron service lives inside the Railway project: the sync job runs over the private network, so Postgres never needs a public port. - `VOYAGE_API_KEY` — Voyage AI key. diff --git a/rag/db.py b/rag/db.py index 27dcb5a..5f01455 100644 --- a/rag/db.py +++ b/rag/db.py @@ -156,6 +156,17 @@ def delete_file_chunks(conn, tables: list[str], file_path: str, commit: bool = T return total + +def delete_chunks_except(conn, table: str, keep_ids: list[str]) -> int: + """Delete every chunk in the table whose id is not in keep_ids. Returns count deleted.""" + with conn.cursor() as cur: + cur.execute("CREATE TEMP TABLE keep_ids (id TEXT PRIMARY KEY) ON COMMIT DROP") + execute_values(cur, "INSERT INTO keep_ids (id) VALUES %s ON CONFLICT DO NOTHING", [(i,) for i in keep_ids]) + cur.execute(f"DELETE FROM {table} t WHERE NOT EXISTS (SELECT 1 FROM keep_ids k WHERE k.id = t.id)") + n = cur.rowcount + conn.commit() + return n + # ── Namespace registry ──────────────────────────────────────────────────────── def setup_meta_table(conn) -> None: diff --git a/scripts/ingest.py b/scripts/ingest.py index f2ced5b..fc7d471 100644 --- a/scripts/ingest.py +++ b/scripts/ingest.py @@ -40,6 +40,7 @@ from rag.db import ( add_pending_batch, content_hash, + delete_chunks_except, delete_file_chunks, get_cached_embeddings, get_conn, @@ -312,6 +313,10 @@ def ingest(repo_path: str, namespace: str, repo_url: str | None = None, use_batc code_chunks = [c for c in all_chunks if c["table"] == code_table] doc_chunks = [c for c in all_chunks if c["table"] == docs_table] print(f"Total: {len(all_chunks)} chunks ({len(code_chunks)} code, {len(doc_chunks)} docs)") + if not all_chunks: + # Guards the stale-chunk delete below from wiping the index after a bad checkout + conn.close() + raise RuntimeError(f"No indexable files found in {repo_root} — refusing to reindex '{namespace}'") # BL-001 / BL-002: large full reindexes go via the Batch API if use_batch: @@ -333,6 +338,13 @@ def ingest(repo_path: str, namespace: str, repo_url: str | None = None, use_batc print(f"\nUpserting {len(chunks)} chunks to {table}...") index_chunks(conn, table, chunks) + # Only after every upsert succeeded: drop chunks for removed or newly skipped files + # and chunk indexes past a file's new chunk count + for table, chunks in [(code_table, code_chunks), (docs_table, doc_chunks)]: + n = delete_chunks_except(conn, table, [c["id"] for c in chunks]) + if n: + print(f"Removed {n} stale chunks from {table}") + if repo_url: try: sha = subprocess.check_output( diff --git a/specs/pipelines/batch-embedding.md b/specs/pipelines/batch-embedding.md index e620575..d702f58 100644 --- a/specs/pipelines/batch-embedding.md +++ b/specs/pipelines/batch-embedding.md @@ -11,7 +11,7 @@ | **Linear issue** | N/A | | **Repo** | `github-repo-rag` | | **Trigger** | Railway cron (sync) / manual (ingest) | -| **Schedule** | Railway cron service schedule | +| **Schedule** | `0 * * * *` (hourly, Railway cron service) | | **Created** | 2026-10-11 | | **Last updated** | 2026-10-11 | @@ -46,8 +46,8 @@ Cut Voyage spend on large embedding runs by sending them through the Voyage Batc | SLA | Target | |---|---| -| Freshness (incremental) | Unchanged — one cron interval | -| Freshness (full reindex via batch) | Up to one cron interval after the batch completes (≤ 12 h + interval) | +| Freshness (incremental) | Within one hour | +| Freshness (full reindex via batch) | Within one hour of the batch completing (≤ 13 h after submission) | ## Error handling @@ -73,7 +73,6 @@ Cut Voyage spend on large embedding runs by sending them through the Voyage Batc | ID | Question | Owner | Due | |---|---|---|---| | OQ-001 | Voyage does not document a max upload file size; 100 MB per batch is a conservative assumption | @julianam-w | before merge | -| OQ-002 | Should the Railway cron run more often than daily so completed batches are collected sooner? Up-to-date namespaces cost one GitHub API call each | @julianam-w | before merge | ## Change log diff --git a/tests/test_batch_embedding.py b/tests/test_batch_embedding.py index a8f1e42..9a78cf4 100644 --- a/tests/test_batch_embedding.py +++ b/tests/test_batch_embedding.py @@ -25,7 +25,7 @@ def repo(tmp_path: Path) -> Path: @pytest.fixture def db(monkeypatch): """Stub every DB call ingest() makes and record what it does.""" - calls: dict[str, list] = {"upserts": [], "batches": [], "registered": []} + calls: dict[str, list] = {"upserts": [], "batches": [], "registered": [], "events": []} monkeypatch.setattr(ingest, "get_conn", lambda url: FakeConn()) monkeypatch.setattr(ingest, "setup_db", lambda conn, tables: None) monkeypatch.setattr(ingest, "setup_meta_table", lambda conn: None) @@ -33,7 +33,16 @@ def db(monkeypatch): monkeypatch.setattr(ingest, "get_embeddings_by_hash", lambda conn, table, hashes: {}) monkeypatch.setattr(ingest, "get_cached_embeddings", lambda conn, hashes: {}) monkeypatch.setattr(ingest, "embed_batch", lambda texts, input_type: [[0.0]] * len(texts)) - monkeypatch.setattr(ingest, "upsert_chunks", lambda conn, table, chunks, embs: calls["upserts"].append(table)) + def upsert(conn, table, chunks, embs): + calls["upserts"].append(table) + calls["events"].append(f"upsert:{table}") + + def delete_except(conn, table, keep_ids): + calls["events"].append(f"prune:{table}:{len(keep_ids)}") + return 0 + + monkeypatch.setattr(ingest, "upsert_chunks", upsert) + monkeypatch.setattr(ingest, "delete_chunks_except", delete_except) monkeypatch.setattr(ingest, "add_pending_batch", lambda conn, b, ns: calls["batches"].append((b, ns))) monkeypatch.setattr(ingest, "register_namespace", lambda conn, ns, url, sha: calls["registered"].append(sha)) monkeypatch.setattr(voyage_batch, "upload_file", lambda content, name: "file-1") @@ -63,6 +72,24 @@ def test_ac_002_batching_off_by_default(monkeypatch, repo, db): assert db["batches"] == [] +def test_full_reindex_prunes_stale_chunks_after_all_upserts(monkeypatch, repo, db): + (repo / "README.md").write_text("# Docs\n", encoding="utf-8") + assert ingest.ingest(str(repo), "ns") is True + assert db["events"] == ["upsert:ns_code", "upsert:ns_docs", "prune:ns_code:5", "prune:ns_docs:1"] + + +def test_submitted_batch_prunes_nothing(monkeypatch, repo, db): + monkeypatch.setenv("VOYAGE_BATCH_THRESHOLD", "1") + assert ingest.ingest(str(repo), "ns", repo_url="https://github.com/o/r", use_batch=True) is False + assert db["events"] == [] + + +def test_empty_checkout_refuses_to_reindex(tmp_path, db): + with pytest.raises(RuntimeError, match="No indexable files"): + ingest.ingest(str(tmp_path), "ns") + assert db["events"] == [] + + def test_invalid_threshold_fails_loudly(monkeypatch): monkeypatch.setenv("VOYAGE_BATCH_THRESHOLD", "lots") with pytest.raises(SystemExit): From 6d6d70f334e68d4b504dc060f490a8db9592ccc0 Mon Sep 17 00:00:00 2001 From: julianam-w <77252335+julianam-w@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:48:53 +1100 Subject: [PATCH 4/8] fix(sync): resubmit failed batches; align code and tests with spec Spec-conformance pass against the (renamed) embedding-sync spec: - BL-009: a failed batch now drops every pending row for the namespace, not just the failed one, so the next run resubmits uncached texts instead of embedding them in real time - BL-005: document that ingest.py --batch registers no SHA - Add BL-011..015 for hash reuse, in-run dedup, the transactional incremental swap, full-reindex pruning and the empty-checkout guard - Tests named by AC; add AC-006 (mixed batches), AC-007 (model and input_type), AC-008 (output parsing) and AC-010 (SQL backfill hash, runs when TEST_DATABASE_URL is set) Co-Authored-By: Claude Opus 5.5 (1M context) --- AGENTS.md | 2 +- README.md | 2 +- rag/db.py | 17 ++++-- rag/voyage_batch.py | 2 +- scripts/ingest.py | 10 +++- scripts/sync.py | 11 ++-- specs/pipelines/batch-embedding.md | 81 ------------------------- specs/pipelines/embedding-sync.md | 95 ++++++++++++++++++++++++++++++ tests/test_batch_embedding.py | 51 +++++++++++++--- tests/test_db_postgres.py | 45 ++++++++++++++ tests/test_embedding_reuse.py | 10 ++-- 11 files changed, 215 insertions(+), 111 deletions(-) delete mode 100644 specs/pipelines/batch-embedding.md create mode 100644 specs/pipelines/embedding-sync.md create mode 100644 tests/test_db_postgres.py diff --git a/AGENTS.md b/AGENTS.md index 9164723..aa6ebc0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -21,7 +21,7 @@ The MCP server does **retrieval only** — it never calls an LLM. The calling ag - `mcp_server.py` — FastMCP server; exposes `search_codebase`, `get_file`, `list_files`, `get_repo_structure`, `list_namespaces`, `namespace_info` tools - `rag/query.py` — `embed()` and `retrieve()` (hybrid pgvector search) - `rag/db.py` — DB setup and upsert helpers -- `rag/voyage_batch.py` — Voyage Batch API REST client used for large full reindexes (`specs/pipelines/batch-embedding.md`) +- `rag/voyage_batch.py` — Voyage Batch API REST client used for large full reindexes (`specs/pipelines/embedding-sync.md`) - `rag/auth.py` — Google OAuth token verifier for HTTP transport - `scripts/ingest.py` — ingest a repo into pgvector - `scripts/ask.py` — local CLI for full retrieve + answer pipeline (uses Anthropic API; dev convenience only) diff --git a/README.md b/README.md index 8bed2cc..15288a1 100644 --- a/README.md +++ b/README.md @@ -251,7 +251,7 @@ Each chunk stores a SHA-256 `content_hash` of its text. Before calling Voyage, e ### Voyage Batch API for full reindexes -A full reindex (no baseline SHA, or a diff too large for GitHub's compare API) with at least `VOYAGE_BATCH_THRESHOLD` (default 2000) new texts goes through the [Voyage Batch API](https://docs.voyageai.com/docs/batch-inference), which costs 33% less and finishes within 12 hours. That sync run submits the batch and leaves the namespace as it is. Later runs skip the namespace while the batch is in flight, then load the results into `rag_embedding_cache` and do the full reindex from the cache. Incremental syncs always embed in real time. Spec: [specs/pipelines/batch-embedding.md](specs/pipelines/batch-embedding.md). +A full reindex (no baseline SHA, or a diff too large for GitHub's compare API) with at least `VOYAGE_BATCH_THRESHOLD` (default 2000) new texts goes through the [Voyage Batch API](https://docs.voyageai.com/docs/batch-inference), which costs 33% less and finishes within 12 hours. That sync run submits the batch and leaves the namespace as it is. Later runs skip the namespace while the batch is in flight, then load the results into `rag_embedding_cache` and do the full reindex from the cache. Incremental syncs always embed in real time. Spec: [specs/pipelines/embedding-sync.md](specs/pipelines/embedding-sync.md). To index a new repo the same way, run `ingest.py --repo --batch` (it can't be combined with `--ref`). The next scheduled sync after Voyage finishes completes the index. Set `VOYAGE_BATCH_THRESHOLD` to a very large number to turn batching off. diff --git a/rag/db.py b/rag/db.py index 5f01455..303414e 100644 --- a/rag/db.py +++ b/rag/db.py @@ -17,6 +17,14 @@ def sanitise_namespace(namespace: str) -> str: return re.sub(r"[^a-z0-9_]", "_", namespace.lower()) +# BL-011: must produce the same digest as content_hash() +BACKFILL_HASH_SQL = """ + UPDATE {table} + SET content_hash = encode(sha256(convert_to(text, 'UTF8')), 'hex') + WHERE content_hash IS NULL +""" + + def get_conn(database_url: str): """Open a psycopg2 connection with pgvector registered.""" conn = psycopg2.connect(database_url) @@ -48,11 +56,7 @@ def setup_table(conn, table: str) -> None: ALTER TABLE {table} ADD COLUMN IF NOT EXISTS content_hash TEXT """) - cur.execute(f""" - UPDATE {table} - SET content_hash = encode(sha256(convert_to(text, 'UTF8')), 'hex') - WHERE content_hash IS NULL - """) + cur.execute(BACKFILL_HASH_SQL.format(table=table)) cur.execute(f""" CREATE INDEX IF NOT EXISTS {table}_embedding_idx ON {table} USING hnsw (embedding vector_cosine_ops) @@ -78,6 +82,7 @@ def setup_db(conn, tables: list[str]) -> None: setup_table(conn, table) +# BL-011 def content_hash(text: str) -> str: """SHA-256 hex digest of a chunk's text, matching the SQL backfill in setup_table().""" return hashlib.sha256(text.encode("utf-8")).hexdigest() @@ -180,7 +185,7 @@ def setup_meta_table(conn) -> None: last_indexed_at TIMESTAMPTZ ) """) - # Voyage Batch API state — see specs/pipelines/batch-embedding.md + # Voyage Batch API state — see specs/pipelines/embedding-sync.md cur.execute("CREATE EXTENSION IF NOT EXISTS vector") cur.execute(""" CREATE TABLE IF NOT EXISTS rag_pending_batches ( diff --git a/rag/voyage_batch.py b/rag/voyage_batch.py index 49f25b7..84e93ff 100644 --- a/rag/voyage_batch.py +++ b/rag/voyage_batch.py @@ -1,7 +1,7 @@ """Voyage Batch API client (Files + Batches REST endpoints). The voyageai Python package (0.5.0) has no batch support, so this talks to the -REST API directly. See specs/pipelines/batch-embedding.md. +REST API directly. See specs/pipelines/embedding-sync.md. """ import json diff --git a/scripts/ingest.py b/scripts/ingest.py index fc7d471..af737e5 100644 --- a/scripts/ingest.py +++ b/scripts/ingest.py @@ -182,10 +182,12 @@ def embed_chunks(conn, table: str, chunks: list[dict]) -> list[list[float]]: Identical texts within `chunks` are embedded once. """ hashes = [content_hash(c["text"]) for c in chunks] + # BL-011: reuse stored embeddings known = get_embeddings_by_hash(conn, table, sorted(set(hashes))) # BL-007: results of a collected Voyage batch known.update(get_cached_embeddings(conn, sorted(set(hashes) - known.keys()))) + # BL-012: one Voyage input per distinct text pending: dict[str, str] = {} for h, c in zip(hashes, chunks, strict=True): if h not in known: @@ -231,6 +233,7 @@ def replace_file_chunks( by_table = {t: [c for c in all_chunks if c["table"] == t] for t in tables} print(f" New chunks: {len(all_chunks)} ({len(by_table[code_table])} code, {len(by_table[docs_table])} docs)") + # BL-013: embed first, then swap in one transaction embeddings = {t: embed_chunks(conn, t, chunks) for t, chunks in by_table.items() if chunks} try: @@ -314,11 +317,12 @@ def ingest(repo_path: str, namespace: str, repo_url: str | None = None, use_batc doc_chunks = [c for c in all_chunks if c["table"] == docs_table] print(f"Total: {len(all_chunks)} chunks ({len(code_chunks)} code, {len(doc_chunks)} docs)") if not all_chunks: - # Guards the stale-chunk delete below from wiping the index after a bad checkout + # BL-015: guards the stale-chunk delete below from wiping the index after a bad checkout conn.close() raise RuntimeError(f"No indexable files found in {repo_root} — refusing to reindex '{namespace}'") - # BL-001 / BL-002: large full reindexes go via the Batch API + # BL-001: only full reindexes reach here with use_batch + # BL-002: large ones go via the Batch API if use_batch: texts = pending_texts(conn, tables, all_chunks) threshold = batch_threshold() @@ -338,7 +342,7 @@ def ingest(repo_path: str, namespace: str, repo_url: str | None = None, use_batc print(f"\nUpserting {len(chunks)} chunks to {table}...") index_chunks(conn, table, chunks) - # Only after every upsert succeeded: drop chunks for removed or newly skipped files + # BL-014: only after every upsert succeeded, drop chunks for removed or newly skipped files # and chunk indexes past a file's new chunk count for table, chunks in [(code_table, code_chunks), (docs_table, doc_chunks)]: n = delete_chunks_except(conn, table, [c["id"] for c in chunks]) diff --git a/scripts/sync.py b/scripts/sync.py index f17e093..d3a1b91 100644 --- a/scripts/sync.py +++ b/scripts/sync.py @@ -6,7 +6,7 @@ Falls back to a full reindex when the diff exceeds GitHub's compare API limits (> 250 commits or >= 300 files changed). A large full reindex is submitted to the Voyage Batch API and finished by a later run once the batch completes -(specs/pipelines/batch-embedding.md). +(specs/pipelines/embedding-sync.md). Usage: python scripts/sync.py # sync all registered namespaces @@ -82,9 +82,10 @@ def _collect_batches(conn, namespace: str, batch_ids: list[str]) -> bool: if b["status"] in voyage_batch.FAILED: failed.append(b) - # BL-009: drop failed batches and fail the sync; the next run resubmits what isn't cached + # BL-009: drop every pending row (completed output stays cached) so the next run's full + # reindex resubmits only the uncached texts instead of embedding them in real time if failed: - delete_pending_batches(conn, [b["id"] for b in failed]) + delete_pending_batches(conn, batch_ids) detail = ", ".join(f"{b['id']} ({b['status']}: {b.get('errors')})" for b in failed) raise RuntimeError(f"Voyage batch failed: {detail}") return True @@ -105,7 +106,7 @@ def sync_namespace(conn, ns: dict, token: str | None) -> None: # BL-007: real-time full reindex, reading batch results from the cache print(f"[{namespace}] Batches complete — running full reindex from cache") _full_reindex(repo_url, namespace, token, use_batch=False) - # BL-008 + # BL-008: advance the SHA, then clear batch state register_namespace(conn, namespace, repo_url, latest_sha) delete_pending_batches(conn, batch_ids) delete_cached_embeddings(conn, batch_ids) @@ -164,7 +165,7 @@ def main() -> None: token = os.environ.get("GITHUB_TOKEN") conn = get_conn(os.environ["DATABASE_URL"]) setup_meta_table(conn) - # BL-010 + # BL-010: drop week-old batch results purged = purge_stale_cache(conn) if purged: print(f"Purged {purged} stale batch embeddings") diff --git a/specs/pipelines/batch-embedding.md b/specs/pipelines/batch-embedding.md deleted file mode 100644 index d702f58..0000000 --- a/specs/pipelines/batch-embedding.md +++ /dev/null @@ -1,81 +0,0 @@ -# Pipeline Spec: `batch-embedding` - -## Identity - -| Field | Value | -|---|---| -| **Name** | Voyage Batch API embedding for full reindexes | -| **Type** | Scheduled script (`scripts/sync.py`) + CLI (`scripts/ingest.py --batch`) | -| **Status** | `draft` | -| **Owner** | `@julianam-w` | -| **Linear issue** | N/A | -| **Repo** | `github-repo-rag` | -| **Trigger** | Railway cron (sync) / manual (ingest) | -| **Schedule** | `0 * * * *` (hourly, Railway cron service) | -| **Created** | 2026-10-11 | -| **Last updated** | 2026-10-11 | - -## Purpose - -Cut Voyage spend on large embedding runs by sending them through the Voyage Batch API, which is priced 33% below real-time `/v1/embeddings` in exchange for an up-to-12-hour completion window. Small incremental syncs stay real time so the index stays fresh. - -## Inputs - -- **External sources:** Voyage Files API (`/v1/files`) and Batch API (`/v1/batches`) over HTTPS; GitHub clone as today. -- **Required configuration:** `VOYAGE_API_KEY`, `DATABASE_URL`; optional `VOYAGE_BATCH_THRESHOLD`. - -## Outputs - -- **Tables:** `rag_pending_batches` (one row per in-flight batch), `rag_embedding_cache` (batch results awaiting upsert); `{namespace}_code` / `{namespace}_docs` as today. -- **Downstream consumers:** MCP server search (reads the namespace tables only). - -## Processing logic - -- **BL-001:** Only full reindexes (no baseline, diff too large, or `ingest.py --batch`) are eligible for the Batch API; incremental syncs always embed in real time. -- **BL-002:** A full reindex uses the Batch API when the number of distinct texts still needing an embedding, after reuse from the namespace tables and the cache, is at least `VOYAGE_BATCH_THRESHOLD` (default 2000). -- **BL-003:** Each batch input line embeds one text with `custom_id` set to its `content_hash`, using model `voyage-code-3` and `input_type` `document`. -- **BL-004:** A batch holds at most 100,000 lines and 100 MB of input; more pending texts are split across several batches. -- **BL-005:** Submitting a batch records it in `rag_pending_batches` and leaves the namespace's `last_commit_sha` and indexed chunks unchanged. -- **BL-006:** While a namespace has an in-flight batch, sync skips that namespace. -- **BL-007:** When every batch for a namespace has reached a terminal status, sync loads each successful output line into `rag_embedding_cache`, then runs a real-time full reindex that reads embeddings from the cache before calling Voyage. -- **BL-008:** After a successful collecting reindex, sync advances `last_commit_sha`, deletes the namespace's `rag_pending_batches` rows and deletes the cache rows loaded from those batches. -- **BL-009:** A batch that ends `failed` or `cancelled` has its row deleted and makes the sync for that namespace fail; the next run resubmits only the texts not already cached. -- **BL-010:** Cache rows older than 7 days are deleted at the start of each sync. - -## SLAs - -| SLA | Target | -|---|---| -| Freshness (incremental) | Within one hour | -| Freshness (full reindex via batch) | Within one hour of the batch completing (≤ 13 h after submission) | - -## Error handling - -| Failure mode | Behaviour | Alerting | -|---|---|---| -| Voyage upload / create error | Raise; namespace sync fails; nothing recorded | Cron job exit code | -| Batch `failed` / `cancelled` | BL-009 | Cron job exit code | -| Batch completed with some failed requests | Missing texts embed in real time during collection | Logged count | - -## Acceptance criteria - -| ID | Criterion | Implements | -|---|---|---| -| AC-001 | A full reindex below the threshold embeds in real time and submits no batch | BL-002 | -| AC-002 | A full reindex at or above the threshold submits batches and upserts nothing | BL-002, BL-005 | -| AC-003 | Batch input lines carry `custom_id` = content hash and one input each, split at the line/byte cap | BL-003, BL-004 | -| AC-004 | Sync skips a namespace whose batch is in progress | BL-006 | -| AC-005 | Sync loads completed output into the cache, reindexes with batching disabled, then clears pending rows and cache | BL-007, BL-008 | -| AC-006 | A failed batch row is deleted and the sync raises | BL-009 | - -## Open questions - -| ID | Question | Owner | Due | -|---|---|---|---| -| OQ-001 | Voyage does not document a max upload file size; 100 MB per batch is a conservative assumption | @julianam-w | before merge | - -## Change log - -| Date | Author | Change | -|---|---|---| -| 2026-10-11 | @julianam-w | Initial draft | diff --git a/specs/pipelines/embedding-sync.md b/specs/pipelines/embedding-sync.md new file mode 100644 index 0000000..a952b31 --- /dev/null +++ b/specs/pipelines/embedding-sync.md @@ -0,0 +1,95 @@ +# Pipeline Spec: `embedding-sync` + +## Identity + +| Field | Value | +|---|---| +| **Name** | Embedding reuse, Voyage Batch API full reindexes, and index pruning | +| **Type** | Scheduled script (`scripts/sync.py`) + CLIs (`scripts/ingest.py`, `scripts/reindex.py`) | +| **Status** | `draft` | +| **Owner** | `@julianam-w` | +| **Linear issue** | N/A | +| **Repo** | `github-repo-rag` | +| **Trigger** | Railway cron (sync) / manual (ingest, reindex) | +| **Schedule** | `0 * * * *` (hourly, Railway cron service) | +| **Created** | 2026-10-11 | +| **Last updated** | 2026-10-11 | + +## Purpose + +Keep each namespace's index in step with its repo while minimising Voyage spend. Text that already has an embedding is never re-embedded. Large full reindexes go through the Voyage Batch API, which costs 33% less than real-time `/v1/embeddings` but can take up to 12 hours. Incremental syncs stay real time so the index stays fresh. Full reindexes also remove chunks the repo no longer produces. + +## Inputs + +- **External sources:** Voyage `/v1/embeddings`, Files API (`/v1/files`) and Batch API (`/v1/batches`) over HTTPS; GitHub clone and compare API. +- **Required configuration:** `VOYAGE_API_KEY`, `DATABASE_URL`; optional `VOYAGE_BATCH_THRESHOLD`, `GITHUB_TOKEN`. + +## Outputs + +- **Tables:** `{namespace}_code` / `{namespace}_docs` (now with `content_hash`), `rag_pending_batches` (one row per in-flight batch), `rag_embedding_cache` (batch results awaiting upsert). +- **Downstream consumers:** MCP server search (reads the namespace tables only). + +## Processing logic + +- **BL-001:** Only full reindexes (no baseline, diff too large, or `ingest.py --batch`) are eligible for the Batch API; incremental syncs always embed in real time. +- **BL-002:** A full reindex uses the Batch API when the number of distinct texts still needing an embedding, after reuse from the namespace tables and the cache, is at least `VOYAGE_BATCH_THRESHOLD` (default 2000). +- **BL-003:** Each batch input line embeds one text with `custom_id` set to its `content_hash`, using model `voyage-code-3` and `input_type` `document`. +- **BL-004:** A batch holds at most 100,000 lines and 100 MB of input; more pending texts are split across several batches. +- **BL-005:** Submitting a batch records it in `rag_pending_batches` and leaves indexed chunks unchanged; a submitting sync leaves `last_commit_sha` unchanged, while `ingest.py --batch` registers the namespace with no SHA so sync collects it. +- **BL-006:** While a namespace has an in-flight batch, sync skips that namespace. +- **BL-007:** When every batch for a namespace has reached a terminal status, sync loads each successful output line into `rag_embedding_cache`, then runs a real-time full reindex that reads embeddings from the cache before calling Voyage. +- **BL-008:** After a successful collecting reindex, sync advances `last_commit_sha`, deletes the namespace's `rag_pending_batches` rows and deletes the cache rows loaded from those batches. +- **BL-009:** A batch that ends `failed` or `cancelled` makes the namespace's sync fail and deletes all of that namespace's pending batch rows while keeping cached results, so the next run resubmits only uncached texts. +- **BL-010:** Cache rows older than 7 days are deleted at the start of each sync. +- **BL-011:** Every stored chunk carries a SHA-256 `content_hash` of its text, and embedding reuses a stored embedding with the same hash in the target table instead of calling Voyage. +- **BL-012:** Identical texts within one run are embedded once. +- **BL-013:** An incremental sync or reindex embeds every changed chunk before deleting anything, then deletes stale chunks and upserts new ones in a single transaction. +- **BL-014:** Once every upsert of a full reindex has succeeded, chunks whose ids the checkout no longer produces are deleted. +- **BL-015:** A full reindex of a checkout with no indexable files fails without changing the index. + +## SLAs + +| SLA | Target | +|---|---| +| Freshness (incremental) | Within one hour | +| Freshness (full reindex via batch) | Within one hour of the batch completing (≤ 13 h after submission) | + +## Error handling + +| Failure mode | Behaviour | Alerting | +|---|---|---| +| Voyage real-time error during incremental sync | Transaction never opened; index unchanged; SHA not advanced (BL-013) | Cron job exit code | +| Voyage upload / create error | Raise; namespace sync fails | Cron job exit code | +| Batch `failed` / `cancelled` | BL-009 | Cron job exit code | +| Batch completed with some failed requests | Missing texts embed in real time during collection | Logged count | +| Empty or broken checkout | BL-015 | Cron job exit code | + +## Acceptance criteria + +| ID | Criterion | Implements | +|---|---|---| +| AC-001 | A full reindex below the threshold embeds in real time and submits no batch | BL-002 | +| AC-002 | A full reindex at or above the threshold submits batches, upserts nothing, and `ingest.py --batch` registers no SHA | BL-002, BL-005 | +| AC-003 | Batch input lines carry `custom_id` = content hash and one input each, split at the line/byte cap | BL-003, BL-004 | +| AC-004 | Sync skips a namespace whose batch is in progress | BL-006 | +| AC-005 | Sync loads completed output into the cache, reindexes with batching disabled, then clears pending rows and cache | BL-007, BL-008 | +| AC-006 | A failed batch deletes every pending row for the namespace and the sync raises | BL-009 | +| AC-007 | Batch creation requests model `voyage-code-3` and `input_type` `document` | BL-003 | +| AC-008 | Only successful output lines are loaded into the cache | BL-007 | +| AC-009 | Embedding reuses stored or cached embeddings by hash and sends each new distinct text to Voyage once | BL-007, BL-011, BL-012 | +| AC-010 | The SQL backfill hash equals the Python `content_hash` | BL-011 | +| AC-011 | Incremental sync embeds before deleting, commits once, and leaves the index untouched if embedding fails | BL-013 | +| AC-012 | A full reindex prunes stale chunks only after every upsert, and a batch submission prunes nothing | BL-014 | +| AC-013 | A full reindex of an empty checkout raises before touching the index | BL-015 | + +## Open questions + +| ID | Question | Owner | Due | +|---|---|---|---| +| OQ-001 | Voyage does not document a max upload file size; 100 MB per batch is a conservative assumption | @julianam-w | before merge | + +## Change log + +| Date | Author | Change | +|---|---|---| +| 2026-10-11 | @julianam-w | Initial draft | diff --git a/tests/test_batch_embedding.py b/tests/test_batch_embedding.py index 9a78cf4..8f07c28 100644 --- a/tests/test_batch_embedding.py +++ b/tests/test_batch_embedding.py @@ -1,4 +1,4 @@ -"""Voyage Batch API flow — specs/pipelines/batch-embedding.md. DB, Voyage and git are faked.""" +"""Voyage Batch API flow — specs/pipelines/embedding-sync.md. DB, Voyage and git are faked.""" import json from pathlib import Path @@ -72,19 +72,19 @@ def test_ac_002_batching_off_by_default(monkeypatch, repo, db): assert db["batches"] == [] -def test_full_reindex_prunes_stale_chunks_after_all_upserts(monkeypatch, repo, db): +def test_ac_012_full_reindex_prunes_stale_chunks_after_all_upserts(monkeypatch, repo, db): (repo / "README.md").write_text("# Docs\n", encoding="utf-8") assert ingest.ingest(str(repo), "ns") is True assert db["events"] == ["upsert:ns_code", "upsert:ns_docs", "prune:ns_code:5", "prune:ns_docs:1"] -def test_submitted_batch_prunes_nothing(monkeypatch, repo, db): +def test_ac_012_submitted_batch_prunes_nothing(monkeypatch, repo, db): monkeypatch.setenv("VOYAGE_BATCH_THRESHOLD", "1") assert ingest.ingest(str(repo), "ns", repo_url="https://github.com/o/r", use_batch=True) is False assert db["events"] == [] -def test_empty_checkout_refuses_to_reindex(tmp_path, db): +def test_ac_013_empty_checkout_refuses_to_reindex(tmp_path, db): with pytest.raises(RuntimeError, match="No indexable files"): ingest.ingest(str(tmp_path), "ns") assert db["events"] == [] @@ -145,11 +145,18 @@ def test_ac_005_collects_completed_batch_then_reindexes_from_cache(monkeypatch, assert sync_env["deleted"] == [["batch-1"]] and sync_env["cleared"] == [["batch-1"]] -def test_ac_006_failed_batch_is_dropped_and_sync_raises(monkeypatch, sync_env): - monkeypatch.setattr(voyage_batch, "get_batch", lambda b: {"id": b, "status": "failed", "errors": "bad"}) - with pytest.raises(RuntimeError, match="batch-1"): +def test_ac_006_failed_batch_drops_all_pending_rows_and_sync_raises(monkeypatch, sync_env): + monkeypatch.setattr(sync, "get_pending_batches", lambda conn, ns: ["batch-1", "batch-2"]) + statuses = { + "batch-1": {"id": "batch-1", "status": "completed", "output_file_id": "out"}, + "batch-2": {"id": "batch-2", "status": "failed", "errors": "bad"}, + } + monkeypatch.setattr(voyage_batch, "get_batch", lambda b: statuses[b]) + with pytest.raises(RuntimeError, match="batch-2"): sync.sync_namespace(FakeConn(), NS, None) - assert sync_env["deleted"] == [["batch-1"]] + # Completed output stays cached; every row goes so the next run resubmits the rest + assert sync_env["cached"] == [("batch-1", [("h1", [1.0])])] + assert sync_env["deleted"] == [["batch-1", "batch-2"]] and sync_env["cleared"] == [] assert sync_env["reindex"] == [] and sync_env["registered"] == [] @@ -161,3 +168,31 @@ def test_submitted_full_reindex_does_not_advance_sha(monkeypatch, sync_env): monkeypatch.setattr(sync, "update_namespace_sha", lambda conn, ns, sha: advanced.append(sha)) sync.sync_namespace(FakeConn(), NS, None) assert advanced == [] + + +# ── Voyage REST payloads ────────────────────────────────────────────────────── + +def test_ac_007_batch_requests_voyage_code_3_documents(monkeypatch): + sent = {} + + def fake_request(method, path, body=None, content_type=None): + sent.update(method=method, path=path, payload=json.loads(body)) + return b'{"id": "batch-1"}' + + monkeypatch.setattr(voyage_batch, "_request", fake_request) + voyage_batch.create_batch("file-1", metadata={"namespace": "ns"}) + assert (sent["method"], sent["path"]) == ("POST", "/batches") + assert sent["payload"]["endpoint"] == "/v1/embeddings" + assert sent["payload"]["request_params"] == {"model": "voyage-code-3", "input_type": "document"} + assert sent["payload"]["input_file_id"] == "file-1" + + +def test_ac_008_only_successful_output_lines_are_loaded(monkeypatch): + def line(custom_id, status=200, error=None): + body = {"data": [{"embedding": [0.5], "index": 0}]} if status == 200 else {"detail": "x"} + return json.dumps({"custom_id": custom_id, "response": {"status_code": status, "body": body}, + "error": error}) + + output = "\n".join([line("ok"), line("bad-status", 400), line("errored", 200, {"message": "x"}), ""]) + monkeypatch.setattr(voyage_batch, "_request", lambda method, path: output.encode()) + assert list(voyage_batch.iter_output_embeddings("out")) == [("ok", [0.5])] diff --git a/tests/test_db_postgres.py b/tests/test_db_postgres.py new file mode 100644 index 0000000..cd76d7a --- /dev/null +++ b/tests/test_db_postgres.py @@ -0,0 +1,45 @@ +"""SQL checks against a real Postgres. Set TEST_DATABASE_URL to run them (pgvector not required). + +Each test works in its own temporary tables, so a scratch database is enough. +""" + +import os + +import pytest + +from rag import db + +DATABASE_URL = os.environ.get("TEST_DATABASE_URL") +pytestmark = pytest.mark.skipif(not DATABASE_URL, reason="TEST_DATABASE_URL not set") + + +@pytest.fixture +def conn(): + import psycopg2 + + c = psycopg2.connect(DATABASE_URL) + yield c + c.rollback() + c.close() + + +def test_ac_010_sql_backfill_hash_matches_python(conn): + texts = ["def foo():\n return 1\n", "Māori: tēnā koe 🌏\r\n\ttabs", "x" * 5000] + with conn.cursor() as cur: + cur.execute("CREATE TEMP TABLE t_hash (text TEXT, content_hash TEXT)") + for text in texts: + cur.execute("INSERT INTO t_hash (text) VALUES (%s)", (text,)) + cur.execute(db.BACKFILL_HASH_SQL.format(table="t_hash")) + cur.execute("SELECT text, content_hash FROM t_hash") + rows = cur.fetchall() + assert len(rows) == len(texts) + assert all(h == db.content_hash(text) for text, h in rows) + + +def test_delete_chunks_except_keeps_only_given_ids(conn): + with conn.cursor() as cur: + cur.execute("CREATE TEMP TABLE t_prune (id TEXT PRIMARY KEY)") + cur.execute("INSERT INTO t_prune SELECT 'id' || g FROM generate_series(1, 100) g") + assert db.delete_chunks_except(conn, "t_prune", [f"id{i}" for i in range(1, 91)]) == 10 + # Second call in the same session: the keep_ids temp table was dropped on commit + assert db.delete_chunks_except(conn, "t_prune", ["id1"]) == 89 diff --git a/tests/test_embedding_reuse.py b/tests/test_embedding_reuse.py index f6c8834..5ee42fc 100644 --- a/tests/test_embedding_reuse.py +++ b/tests/test_embedding_reuse.py @@ -38,7 +38,7 @@ def _chunk(text: str, file_path: str = "a.py", i: int = 0) -> dict: "table": "ns_code"} -def test_reuses_stored_embeddings_and_dedupes_within_run(monkeypatch, voyage): +def test_ac_009_reuses_stored_embeddings_and_dedupes_within_run(monkeypatch, voyage): stored = {content_hash("old"): [99.0]} monkeypatch.setattr(ingest, "get_embeddings_by_hash", lambda conn, table, hashes: { h: stored[h] for h in hashes if h in stored @@ -51,7 +51,7 @@ def test_reuses_stored_embeddings_and_dedupes_within_run(monkeypatch, voyage): assert result == [[99.0], [3.0], [3.0], [5.0]] -def test_no_voyage_call_when_everything_is_stored(monkeypatch, voyage): +def test_ac_009_no_voyage_call_when_everything_is_stored(monkeypatch, voyage): monkeypatch.setattr(ingest, "get_embeddings_by_hash", lambda conn, table, hashes: {h: [1.0] for h in hashes}) assert ingest.embed_chunks(FakeConn(), "ns_code", [_chunk("a"), _chunk("b")]) == [[1.0], [1.0]] assert voyage == [] @@ -64,7 +64,7 @@ def test_splits_new_texts_into_voyage_sized_batches(monkeypatch, voyage): assert [len(c) for c in voyage] == [ingest.EMBED_BATCH, 1] -def test_replace_embeds_before_deleting_and_commits_once(monkeypatch, tmp_path: Path): +def test_ac_011_replace_embeds_before_deleting_and_commits_once(monkeypatch, tmp_path: Path): (tmp_path / "a.py").write_text("print('hi')\n", encoding="utf-8") events: list[str] = [] monkeypatch.setattr(ingest, "embed_chunks", lambda conn, table, chunks: events.append("embed") or @@ -81,7 +81,7 @@ def test_replace_embeds_before_deleting_and_commits_once(monkeypatch, tmp_path: assert conn.commits == 1 -def test_replace_leaves_index_untouched_when_embedding_fails(monkeypatch, tmp_path: Path): +def test_ac_011_replace_leaves_index_untouched_when_embedding_fails(monkeypatch, tmp_path: Path): (tmp_path / "a.py").write_text("print('hi')\n", encoding="utf-8") def boom(conn, table, chunks): @@ -96,7 +96,7 @@ def boom(conn, table, chunks): assert deletes == [] -def test_uses_batch_cache_for_texts_not_in_table(monkeypatch, voyage): +def test_ac_009_uses_batch_cache_for_texts_not_in_table(monkeypatch, voyage): monkeypatch.setattr(ingest, "get_embeddings_by_hash", lambda conn, table, hashes: {}) monkeypatch.setattr(ingest, "get_cached_embeddings", lambda conn, hashes: { h: [7.0] for h in hashes if h == content_hash("cached") From e6b83fce980c0332bf43a19ab912a15c3c851392 Mon Sep 17 00:00:00 2001 From: julianam-w <77252335+julianam-w@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:52:30 +1100 Subject: [PATCH 5/8] fix(sync): stream batch output, roll back failed namespaces Findings from /code-review --fix: - reindex.py now creates the batch tables it reads (rag_embedding_cache) so it cannot fail with UndefinedTable on a fresh deploy - Batch output is streamed line by line and cached in pages of 1000, so a 100K-line output no longer has to fit in memory - sync rolls back the shared connection after a namespace fails, so one failure no longer aborts every later namespace in the run Co-Authored-By: Claude Opus 5.5 (1M context) --- rag/voyage_batch.py | 15 +++++++++++++-- scripts/reindex.py | 4 +++- scripts/sync.py | 20 +++++++++++++++++--- tests/test_batch_embedding.py | 2 +- 4 files changed, 34 insertions(+), 7 deletions(-) diff --git a/rag/voyage_batch.py b/rag/voyage_batch.py index 84e93ff..d41a338 100644 --- a/rag/voyage_batch.py +++ b/rag/voyage_batch.py @@ -81,10 +81,21 @@ def get_batch(batch_id: str) -> dict: return json.loads(_request("GET", f"/batches/{batch_id}")) +def _stream_lines(path: str) -> Iterator[bytes]: + """GET a file and yield it line by line without holding the whole body in memory.""" + headers = {"Authorization": f"Bearer {os.environ['VOYAGE_API_KEY']}"} + req = urllib.request.Request(f"{_API}{path}", headers=headers, method="GET") + try: + with urllib.request.urlopen(req) as resp: + yield from resp + except urllib.error.HTTPError as e: + detail = e.read().decode(errors="replace") + raise RuntimeError(f"Voyage API {e.code} for GET {path}: {detail}") from e + + def iter_output_embeddings(output_file_id: str) -> Iterator[tuple[str, list[float]]]: """Yield (content_hash, embedding) for each successful line of a batch output file.""" - content = _request("GET", f"/files/{output_file_id}/content") - for raw in content.splitlines(): + for raw in _stream_lines(f"/files/{output_file_id}/content"): if not raw.strip(): continue line = json.loads(raw) diff --git a/scripts/reindex.py b/scripts/reindex.py index 850d2ac..bffd223 100644 --- a/scripts/reindex.py +++ b/scripts/reindex.py @@ -23,7 +23,7 @@ from dotenv import load_dotenv load_dotenv(Path(__file__).parent.parent / ".env") -from rag.db import get_conn, sanitise_namespace, setup_db +from rag.db import get_conn, sanitise_namespace, setup_db, setup_meta_table from scripts.ingest import replace_file_chunks @@ -31,6 +31,8 @@ def reindex(repo_path: str, namespace: str) -> None: namespace = sanitise_namespace(namespace) conn = get_conn(os.environ["DATABASE_URL"]) setup_db(conn, [f"{namespace}_code", f"{namespace}_docs"]) + # embed_chunks reads rag_embedding_cache + setup_meta_table(conn) changed = [f for f in os.environ.get("CHANGED_FILES", "").split() if f] deleted = [f for f in os.environ.get("DELETED_FILES", "").split() if f] diff --git a/scripts/sync.py b/scripts/sync.py index d3a1b91..d5bf557 100644 --- a/scripts/sync.py +++ b/scripts/sync.py @@ -46,6 +46,8 @@ from rag.github import get_changed_files, get_latest_sha, parse_repo_url, shallow_clone from scripts.ingest import ingest, replace_file_chunks +CACHE_PAGE = 1000 # batch output rows per cache insert + def _full_reindex(repo_url: str, namespace: str, token: str | None, use_batch: bool = True) -> bool: """Clone and fully reindex. Returns False if a Voyage batch was submitted instead. @@ -76,9 +78,19 @@ def _collect_batches(conn, namespace: str, batch_ids: list[str]) -> bool: failed = [] for b in batches: if b.get("output_file_id"): - rows = list(voyage_batch.iter_output_embeddings(b["output_file_id"])) - cache_embeddings(conn, b["id"], rows) - print(f"[{namespace}] Batch {b['id']} {b['status']}: cached {len(rows)} embeddings") + # Cache in pages so a 100K-line output never sits in memory at once + cached = 0 + page: list[tuple[str, list[float]]] = [] + for row in voyage_batch.iter_output_embeddings(b["output_file_id"]): + page.append(row) + if len(page) >= CACHE_PAGE: + cache_embeddings(conn, b["id"], page) + cached += len(page) + page = [] + if page: + cache_embeddings(conn, b["id"], page) + cached += len(page) + print(f"[{namespace}] Batch {b['id']} {b['status']}: cached {cached} embeddings") if b["status"] in voyage_batch.FAILED: failed.append(b) @@ -188,6 +200,8 @@ def main() -> None: try: sync_namespace(conn, ns, token) except Exception as e: + # Clear any aborted transaction so later namespaces can still use the connection + conn.rollback() print(f"[{ns['namespace']}] ERROR: {e}") errors.append(ns["namespace"]) diff --git a/tests/test_batch_embedding.py b/tests/test_batch_embedding.py index 8f07c28..722e495 100644 --- a/tests/test_batch_embedding.py +++ b/tests/test_batch_embedding.py @@ -194,5 +194,5 @@ def line(custom_id, status=200, error=None): "error": error}) output = "\n".join([line("ok"), line("bad-status", 400), line("errored", 200, {"message": "x"}), ""]) - monkeypatch.setattr(voyage_batch, "_request", lambda method, path: output.encode()) + monkeypatch.setattr(voyage_batch, "_stream_lines", lambda path: iter(output.encode().splitlines(keepends=True))) assert list(voyage_batch.iter_output_embeddings("out")) == [("ok", [0.5])] From 69db8d8295983f1aed9cad0437b2adc0d0062303 Mon Sep 17 00:00:00 2001 From: julianam-w <77252335+julianam-w@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:36:06 +1100 Subject: [PATCH 6/8] fix(indexing): stop emptying a table with zero chunks; isolate a bad threshold /code-review --fix (8-angle review) surfaced and this fixes three bugs: - delete_chunks_except(conn, table, []) wiped the whole table instead of deleting nothing, because execute_values is a silent no-op on an empty list, leaving the keep_ids temp table empty and the anti-join DELETE matching every row. A full reindex of a repo whose docs (or code) table produces zero chunks this run silently wiped that table. Confirmed with a direct Postgres repro before and after the fix. Fixed at both layers: the function now no-ops on an empty keep set, and ingest()'s prune loop skips a zero-chunk table rather than calling it with one. - batch_threshold() raised SystemExit, which is a BaseException, not an Exception, so sync.py's per-namespace `except Exception` never caught it. A malformed VOYAGE_BATCH_THRESHOLD would abort the whole cron run instead of failing just the namespace attempting a full reindex. Now raises ValueError. - The collecting-reindex branch of sync_namespace advanced the SHA and cleared batch state without checking whether the reindex it just ran actually succeeded, unlike the other two full-reindex call sites. Hardened to match them; unreachable as a live bug today but closes the gap for any future ingest() change. Adds BL-016..018 / AC-014..017 to specs/pipelines/embedding-sync.md and tests for all three (one against real Postgres). Co-Authored-By: Claude Sonnet 5 --- rag/db.py | 9 +++++- scripts/ingest.py | 13 ++++++-- scripts/sync.py | 5 ++- specs/pipelines/embedding-sync.md | 11 +++++++ tests/test_batch_embedding.py | 51 +++++++++++++++++++++++++++++-- tests/test_db_postgres.py | 10 ++++++ 6 files changed, 93 insertions(+), 6 deletions(-) diff --git a/rag/db.py b/rag/db.py index 303414e..eb16bda 100644 --- a/rag/db.py +++ b/rag/db.py @@ -163,7 +163,14 @@ def delete_file_chunks(conn, tables: list[str], file_path: str, commit: bool = T def delete_chunks_except(conn, table: str, keep_ids: list[str]) -> int: - """Delete every chunk in the table whose id is not in keep_ids. Returns count deleted.""" + """Delete every chunk in the table whose id is not in keep_ids. Returns count deleted. + + BL-016: a no-op when keep_ids is empty. A run that found zero chunks for this + table (a genuinely empty repo subtree, or a partial/broken checkout) must never + be read as "keep nothing" — that would silently empty an already-indexed table. + """ + if not keep_ids: + return 0 with conn.cursor() as cur: cur.execute("CREATE TEMP TABLE keep_ids (id TEXT PRIMARY KEY) ON COMMIT DROP") execute_values(cur, "INSERT INTO keep_ids (id) VALUES %s ON CONFLICT DO NOTHING", [(i,) for i in keep_ids]) diff --git a/scripts/ingest.py b/scripts/ingest.py index af737e5..b2566df 100644 --- a/scripts/ingest.py +++ b/scripts/ingest.py @@ -186,6 +186,7 @@ def embed_chunks(conn, table: str, chunks: list[dict]) -> list[list[float]]: known = get_embeddings_by_hash(conn, table, sorted(set(hashes))) # BL-007: results of a collected Voyage batch known.update(get_cached_embeddings(conn, sorted(set(hashes) - known.keys()))) + reused = sum(1 for h in hashes if h in known) # BL-012: one Voyage input per distinct text pending: dict[str, str] = {} @@ -199,7 +200,8 @@ def embed_chunks(conn, table: str, chunks: list[dict]) -> list[list[float]]: embeddings = embed_batch([text for _, text in batch], input_type="document") known.update({h: emb for (h, _), emb in zip(batch, embeddings, strict=True)}) - print(f" {len(chunks)} chunks: {len(chunks) - len(pending)} reused, {len(pending)} embedded") + deduped = len(hashes) - reused - len(pending) + print(f" {len(chunks)} chunks: {reused} reused, {deduped} deduped, {len(pending)} embedded") return [known[h] for h in hashes] @@ -253,13 +255,16 @@ def replace_file_chunks( # ── Voyage Batch API ────────────────────────────────────────────────────────── def batch_threshold() -> int: + """Read VOYAGE_BATCH_THRESHOLD, raising ValueError (not SystemExit) on a bad value.""" raw = os.environ.get("VOYAGE_BATCH_THRESHOLD", "") if not raw: return DEFAULT_BATCH_THRESHOLD try: return int(raw) except ValueError: - raise SystemExit(f"VOYAGE_BATCH_THRESHOLD must be an integer, got {raw!r}") from None + # BL-017: a plain ValueError, not SystemExit (a BaseException), so sync.py's + # per-namespace handling catches it and fails only this namespace + raise ValueError(f"VOYAGE_BATCH_THRESHOLD must be an integer, got {raw!r}") from None def pending_texts(conn, tables: list[str], chunks: list[dict]) -> dict[str, str]: @@ -344,7 +349,11 @@ def ingest(repo_path: str, namespace: str, repo_url: str | None = None, use_batc # BL-014: only after every upsert succeeded, drop chunks for removed or newly skipped files # and chunk indexes past a file's new chunk count + # BL-016: a table with zero chunks this run is left untouched, not wiped for table, chunks in [(code_table, code_chunks), (docs_table, doc_chunks)]: + if not chunks: + print(f"No chunks for {table} this run — leaving its existing chunks untouched") + continue n = delete_chunks_except(conn, table, [c["id"] for c in chunks]) if n: print(f"Removed {n} stale chunks from {table}") diff --git a/scripts/sync.py b/scripts/sync.py index d5bf557..23f3bfb 100644 --- a/scripts/sync.py +++ b/scripts/sync.py @@ -117,7 +117,10 @@ def sync_namespace(conn, ns: dict, token: str | None) -> None: return # BL-007: real-time full reindex, reading batch results from the cache print(f"[{namespace}] Batches complete — running full reindex from cache") - _full_reindex(repo_url, namespace, token, use_batch=False) + # BL-018: only advance state if the reindex actually completed + if not _full_reindex(repo_url, namespace, token, use_batch=False): + print(f"[{namespace}] Full reindex from cache did not complete — retrying next run") + return # BL-008: advance the SHA, then clear batch state register_namespace(conn, namespace, repo_url, latest_sha) delete_pending_batches(conn, batch_ids) diff --git a/specs/pipelines/embedding-sync.md b/specs/pipelines/embedding-sync.md index a952b31..0607656 100644 --- a/specs/pipelines/embedding-sync.md +++ b/specs/pipelines/embedding-sync.md @@ -46,6 +46,9 @@ Keep each namespace's index in step with its repo while minimising Voyage spend. - **BL-013:** An incremental sync or reindex embeds every changed chunk before deleting anything, then deletes stale chunks and upserts new ones in a single transaction. - **BL-014:** Once every upsert of a full reindex has succeeded, chunks whose ids the checkout no longer produces are deleted. - **BL-015:** A full reindex of a checkout with no indexable files fails without changing the index. +- **BL-016:** A table that produces zero chunks in a run is left untouched rather than pruned; an empty keep set is never read as "keep nothing". +- **BL-017:** A malformed `VOYAGE_BATCH_THRESHOLD` fails only the namespace attempting a full reindex, not the whole sync run. +- **BL-018:** The full reindex that collects a completed batch advances `last_commit_sha` and clears batch state only once that reindex has itself completed. ## SLAs @@ -63,6 +66,9 @@ Keep each namespace's index in step with its repo while minimising Voyage spend. | Batch `failed` / `cancelled` | BL-009 | Cron job exit code | | Batch completed with some failed requests | Missing texts embed in real time during collection | Logged count | | Empty or broken checkout | BL-015 | Cron job exit code | +| One table produces zero chunks this run | BL-016 | Logged, not an error | +| Malformed `VOYAGE_BATCH_THRESHOLD` | Raises `ValueError`; only that namespace's sync fails (BL-017) | Logged per namespace, cron job continues | +| Collecting reindex does not complete | Batch state and `last_commit_sha` left unchanged for retry (BL-018) | Logged, cron job continues | ## Acceptance criteria @@ -81,6 +87,10 @@ Keep each namespace's index in step with its repo while minimising Voyage spend. | AC-011 | Incremental sync embeds before deleting, commits once, and leaves the index untouched if embedding fails | BL-013 | | AC-012 | A full reindex prunes stale chunks only after every upsert, and a batch submission prunes nothing | BL-014 | | AC-013 | A full reindex of an empty checkout raises before touching the index | BL-015 | +| AC-014 | `delete_chunks_except` returns 0 and deletes nothing when `keep_ids` is empty | BL-016 | +| AC-015 | A full reindex with one table at zero chunks this run skips pruning that table | BL-016 | +| AC-016 | A malformed `VOYAGE_BATCH_THRESHOLD` raises `ValueError`, caught per namespace by sync's loop | BL-017 | +| AC-017 | The collecting reindex only advances `last_commit_sha` and clears batch state on its own success | BL-018 | ## Open questions @@ -93,3 +103,4 @@ Keep each namespace's index in step with its repo while minimising Voyage spend. | Date | Author | Change | |---|---|---| | 2026-10-11 | @julianam-w | Initial draft | +| 2026-10-12 | @julianam-w | `/code-review --fix`: add BL-016..018 / AC-014..017 for a data-loss bug in table pruning, a `SystemExit` that broke per-namespace isolation, and an ungated state advance in the collecting reindex | diff --git a/tests/test_batch_embedding.py b/tests/test_batch_embedding.py index 722e495..9f99391 100644 --- a/tests/test_batch_embedding.py +++ b/tests/test_batch_embedding.py @@ -1,6 +1,7 @@ """Voyage Batch API flow — specs/pipelines/embedding-sync.md. DB, Voyage and git are faked.""" import json +import sys from pathlib import Path import pytest @@ -90,12 +91,42 @@ def test_ac_013_empty_checkout_refuses_to_reindex(tmp_path, db): assert db["events"] == [] -def test_invalid_threshold_fails_loudly(monkeypatch): +def test_ac_016_invalid_threshold_raises_value_error(monkeypatch): monkeypatch.setenv("VOYAGE_BATCH_THRESHOLD", "lots") - with pytest.raises(SystemExit): + with pytest.raises(ValueError, match="VOYAGE_BATCH_THRESHOLD"): ingest.batch_threshold() +def test_ac_016_sync_main_isolates_a_namespace_that_raises_value_error(monkeypatch): + """A bad VOYAGE_BATCH_THRESHOLD must fail only the namespace that hits it, not the whole run.""" + monkeypatch.setenv("DATABASE_URL", "postgresql://fake") + monkeypatch.setenv("VOYAGE_BATCH_THRESHOLD", "lots") + monkeypatch.setattr(sync, "get_conn", lambda url: FakeConn()) + monkeypatch.setattr(sync, "setup_meta_table", lambda conn: None) + monkeypatch.setattr(sync, "purge_stale_cache", lambda conn: 0) + namespaces = [ + {"namespace": "bad", "repo_url": "https://github.com/o/bad", "last_commit_sha": None}, + {"namespace": "good", "repo_url": "https://github.com/o/good", "last_commit_sha": "a" * 40}, + ] + monkeypatch.setattr(sync, "get_namespaces", lambda conn: namespaces) + monkeypatch.setattr(sync, "get_latest_sha", lambda owner, repo, token: "a" * 40) + monkeypatch.setattr(sync, "get_pending_batches", lambda conn, ns: []) + + def fake_full_reindex(url, ns, token, use_batch=True): + # Exercises the real batch_threshold() call for the "bad" namespace's full reindex + if ns == "bad": + ingest.batch_threshold() + return True + + monkeypatch.setattr(sync, "_full_reindex", fake_full_reindex) + monkeypatch.setattr(sys, "argv", ["sync.py"]) + + with pytest.raises(SystemExit) as exc_info: + sync.main() + # main() exits 1 (recorded error) rather than crashing with an uncaught ValueError + assert exc_info.value.code == 1 + + def test_ac_003_input_lines_keyed_by_hash_and_split_at_caps(monkeypatch): monkeypatch.setattr(voyage_batch, "MAX_LINES", 2) files = voyage_batch.build_input_files({"h1": "a", "h2": "b", "h3": "ā"}) @@ -196,3 +227,19 @@ def line(custom_id, status=200, error=None): output = "\n".join([line("ok"), line("bad-status", 400), line("errored", 200, {"message": "x"}), ""]) monkeypatch.setattr(voyage_batch, "_stream_lines", lambda path: iter(output.encode().splitlines(keepends=True))) assert list(voyage_batch.iter_output_embeddings("out")) == [("ok", [0.5])] + + +def test_ac_015_ingest_leaves_a_zero_chunk_table_untouched(monkeypatch, repo, db): + """repo fixture has only .py files, so ns_docs gets zero chunks this run — it must be + skipped (no delete_chunks_except call at all), not pruned with an empty keep list.""" + assert ingest.ingest(str(repo), "ns") is True + assert db["events"] == ["upsert:ns_code", "prune:ns_code:5"] + + +def test_ac_017_collecting_reindex_does_not_advance_state_when_reindex_fails(monkeypatch, sync_env): + monkeypatch.setattr(voyage_batch, "get_batch", lambda b: {"id": b, "status": "completed", "output_file_id": "out"}) + monkeypatch.setattr(sync, "_full_reindex", lambda url, ns, token, use_batch=True: + sync_env["reindex"].append(use_batch) or False) + sync.sync_namespace(FakeConn(), NS, None) + assert sync_env["reindex"] == [False] + assert sync_env["registered"] == [] and sync_env["deleted"] == [] and sync_env["cleared"] == [] diff --git a/tests/test_db_postgres.py b/tests/test_db_postgres.py index cd76d7a..f81f432 100644 --- a/tests/test_db_postgres.py +++ b/tests/test_db_postgres.py @@ -43,3 +43,13 @@ def test_delete_chunks_except_keeps_only_given_ids(conn): assert db.delete_chunks_except(conn, "t_prune", [f"id{i}" for i in range(1, 91)]) == 10 # Second call in the same session: the keep_ids temp table was dropped on commit assert db.delete_chunks_except(conn, "t_prune", ["id1"]) == 89 + + +def test_ac_014_delete_chunks_except_is_a_noop_on_empty_keep_ids(conn): + with conn.cursor() as cur: + cur.execute("CREATE TEMP TABLE t_empty_keep (id TEXT PRIMARY KEY)") + cur.execute("INSERT INTO t_empty_keep SELECT 'id' || g FROM generate_series(1, 20) g") + assert db.delete_chunks_except(conn, "t_empty_keep", []) == 0 + with conn.cursor() as cur: + cur.execute("SELECT count(*) FROM t_empty_keep") + assert cur.fetchone()[0] == 20 From ea4dc93f4d7eaa7cdf68a81a62301c6099b90f96 Mon Sep 17 00:00:00 2001 From: julianam-w <77252335+julianam-w@users.noreply.github.com> Date: Mon, 12 Oct 2026 22:21:17 +1100 Subject: [PATCH 7/8] docs(spec): set embedding-sync to implemented, squash change log MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All 17 acceptance criteria now pass: by automated test, and AC-001, AC-002 and AC-005 additionally confirmed live — a real-time ingest below threshold, a live Voyage Batch API submission, and its collection into a full reindex that reused 100% of the batch's cached embeddings with zero real-time Voyage calls. Co-Authored-By: Claude Sonnet 5 --- specs/pipelines/embedding-sync.md | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/specs/pipelines/embedding-sync.md b/specs/pipelines/embedding-sync.md index 0607656..24b3bde 100644 --- a/specs/pipelines/embedding-sync.md +++ b/specs/pipelines/embedding-sync.md @@ -6,14 +6,14 @@ |---|---| | **Name** | Embedding reuse, Voyage Batch API full reindexes, and index pruning | | **Type** | Scheduled script (`scripts/sync.py`) + CLIs (`scripts/ingest.py`, `scripts/reindex.py`) | -| **Status** | `draft` | +| **Status** | `implemented` | | **Owner** | `@julianam-w` | | **Linear issue** | N/A | | **Repo** | `github-repo-rag` | | **Trigger** | Railway cron (sync) / manual (ingest, reindex) | | **Schedule** | `0 * * * *` (hourly, Railway cron service) | | **Created** | 2026-10-11 | -| **Last updated** | 2026-10-11 | +| **Last updated** | 2026-10-12 | ## Purpose @@ -102,5 +102,4 @@ Keep each namespace's index in step with its repo while minimising Voyage spend. | Date | Author | Change | |---|---|---| -| 2026-10-11 | @julianam-w | Initial draft | -| 2026-10-12 | @julianam-w | `/code-review --fix`: add BL-016..018 / AC-014..017 for a data-loss bug in table pruning, a `SystemExit` that broke per-namespace isolation, and an ungated state advance in the collecting reindex | +| 2026-10-12 | @julianam-w | Implemented: embedding reuse, Voyage Batch API full reindexes, index pruning, hourly sync. All ACs pass by automated test; AC-001, AC-002 and AC-005 also confirmed live against the real Voyage API and database | From f2ff506e611a132a8707399e6365805fabb3d6c2 Mon Sep 17 00:00:00 2001 From: julianam-w <77252335+julianam-w@users.noreply.github.com> Date: Mon, 12 Oct 2026 23:55:16 +1100 Subject: [PATCH 8/8] fix(indexing): stop repeated DDL locks, close conn on error, recover stuck batches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second /code-review --fix pass (8-angle review) on the pushed PR, with all findings consolidated and verified before fixing: - setup_table() ran its column/backfill/index migration unconditionally on every call. ALTER TABLE takes its lock before checking IF NOT EXISTS, so with sync now hourly this repeated a lock acquisition against a live table the MCP server concurrently reads, on every tick, forever, for no effect after the first run. Gated behind an information_schema existence check; confirmed with a mocked test that the ALTER/backfill/index calls are skipped once content_hash exists. - ingest() never closed its database connection if embedding, batch submission or pruning raised — a real connection leak, worse now that sync.py calls it in a per-namespace loop. Wrapped the function body in try/finally. - get_embeddings_by_hash's DISTINCT ON had no ORDER BY, so its tie-break between duplicate rows for a hash was plan-dependent. Added a deterministic ORDER BY. - register.py --sha (documented as forcing a full reindex) didn't clear a namespace's pending Voyage batch, so a stuck or in-flight batch would silently override it — sync.py checks pending batches before ever looking at last_commit_sha. Now also clears pending batch rows and cached embeddings, giving operators a supported recovery path for a batch that can never resolve. - Restored per-file delete logging and flush=True on progress prints, both dropped by the reindex/sync refactor and useful for diagnosing a stuck job against Railway's piped logs. - Three test names and one module docstring brought in line with the team's conventions (test_ac_XXX prefix; don't justify an implementation by the alternative it replaces). Adds BL-019..021 / AC-018..020 to specs/pipelines/embedding-sync.md. Not fixed, considered and documented as known limitations: a claimed Voyage 'expired' batch status (not in Voyage's documented set, unlike OpenAI's Batch API), and several reuse/duplication suggestions (shared SQL/HTTP helpers) deferred as broad refactors. Co-Authored-By: Claude Sonnet 5 --- rag/db.py | 47 +++++---- rag/voyage_batch.py | 5 +- scripts/ingest.py | 154 ++++++++++++++++-------------- scripts/register.py | 23 ++++- scripts/sync.py | 6 +- specs/pipelines/embedding-sync.md | 8 ++ tests/test_batch_embedding.py | 21 +++- tests/test_db.py | 50 ++++++++++ tests/test_db_postgres.py | 2 +- tests/test_embedding_reuse.py | 2 +- tests/test_register.py | 54 +++++++++++ 11 files changed, 273 insertions(+), 99 deletions(-) create mode 100644 tests/test_db.py create mode 100644 tests/test_register.py diff --git a/rag/db.py b/rag/db.py index eb16bda..4d6cd89 100644 --- a/rag/db.py +++ b/rag/db.py @@ -33,7 +33,14 @@ def get_conn(database_url: str): def setup_table(conn, table: str) -> None: - """Create a single RAG table with HNSW, FTS, and file_path indexes.""" + """Create a single RAG table with HNSW, FTS, file_path and content_hash indexes. + + The column/backfill migration below runs only while content_hash is missing — + this is called on every ingest and sync tick (hourly, per namespace), and an + ALTER TABLE takes its lock before checking IF NOT EXISTS, so re-running it + unconditionally would repeat that lock acquisition against a live table the + MCP server concurrently reads, forever, for no effect after the first run. + """ with conn.cursor() as cur: cur.execute("CREATE EXTENSION IF NOT EXISTS vector") cur.execute(f""" @@ -46,17 +53,28 @@ def setup_table(conn, table: str) -> None: chunk_index INT NOT NULL DEFAULT 0 ) """) - cur.execute(f""" - ALTER TABLE {table} - ADD COLUMN IF NOT EXISTS chunk_index INT NOT NULL DEFAULT 0 - """) - # content_hash lets re-indexing reuse an existing embedding for unchanged text. - # Backfilled in SQL; must match content_hash() below. - cur.execute(f""" - ALTER TABLE {table} - ADD COLUMN IF NOT EXISTS content_hash TEXT - """) - cur.execute(BACKFILL_HASH_SQL.format(table=table)) + # BL-019: only migrate while content_hash is missing + cur.execute( + "SELECT 1 FROM information_schema.columns " + "WHERE table_name = %s AND column_name = 'content_hash'", + (table,), + ) + if cur.fetchone() is None: + cur.execute(f""" + ALTER TABLE {table} + ADD COLUMN IF NOT EXISTS chunk_index INT NOT NULL DEFAULT 0 + """) + # content_hash lets re-indexing reuse an existing embedding for unchanged text. + # Backfilled in SQL; must match content_hash() below. + cur.execute(f""" + ALTER TABLE {table} + ADD COLUMN IF NOT EXISTS content_hash TEXT + """) + cur.execute(BACKFILL_HASH_SQL.format(table=table)) + cur.execute(f""" + CREATE INDEX IF NOT EXISTS {table}_content_hash_idx + ON {table} (content_hash) + """) cur.execute(f""" CREATE INDEX IF NOT EXISTS {table}_embedding_idx ON {table} USING hnsw (embedding vector_cosine_ops) @@ -69,10 +87,6 @@ def setup_table(conn, table: str) -> None: CREATE INDEX IF NOT EXISTS {table}_file_path_idx ON {table} (file_path) """) - cur.execute(f""" - CREATE INDEX IF NOT EXISTS {table}_content_hash_idx - ON {table} (content_hash) - """) conn.commit() @@ -98,6 +112,7 @@ def get_embeddings_by_hash(conn, table: str, hashes: list[str]) -> dict[str, lis SELECT DISTINCT ON (content_hash) content_hash, embedding FROM {table} WHERE content_hash = ANY(%s) + ORDER BY content_hash, id """, (list(hashes),), ) diff --git a/rag/voyage_batch.py b/rag/voyage_batch.py index d41a338..9c51075 100644 --- a/rag/voyage_batch.py +++ b/rag/voyage_batch.py @@ -1,7 +1,6 @@ -"""Voyage Batch API client (Files + Batches REST endpoints). +"""Voyage Batch API client: Files and Batches REST endpoints. -The voyageai Python package (0.5.0) has no batch support, so this talks to the -REST API directly. See specs/pipelines/embedding-sync.md. +See specs/pipelines/embedding-sync.md. """ import json diff --git a/scripts/ingest.py b/scripts/ingest.py index b2566df..30cfaf6 100644 --- a/scripts/ingest.py +++ b/scripts/ingest.py @@ -201,7 +201,8 @@ def embed_chunks(conn, table: str, chunks: list[dict]) -> list[list[float]]: known.update({h: emb for (h, _), emb in zip(batch, embeddings, strict=True)}) deduped = len(hashes) - reused - len(pending) - print(f" {len(chunks)} chunks: {reused} reused, {deduped} deduped, {len(pending)} embedded") + # flush=True: a long reindex's only progress signal on a piped/cron log stream + print(f" {len(chunks)} chunks: {reused} reused, {deduped} deduped, {len(pending)} embedded", flush=True) return [known[h] for h in hashes] @@ -210,7 +211,7 @@ def index_chunks(conn, table: str, chunks: list[dict]) -> None: total_groups = (len(chunks) + INDEX_GROUP - 1) // INDEX_GROUP for i in range(0, len(chunks), INDEX_GROUP): group = chunks[i:i + INDEX_GROUP] - print(f" Group {i // INDEX_GROUP + 1}/{total_groups}") + print(f" Group {i // INDEX_GROUP + 1}/{total_groups}", flush=True) upsert_chunks(conn, table, group, embed_chunks(conn, table, group)) @@ -239,9 +240,12 @@ def replace_file_chunks( embeddings = {t: embed_chunks(conn, t, chunks) for t, chunks in by_table.items() if chunks} try: - total_deleted = sum( - delete_file_chunks(conn, tables, rel_path, commit=False) for rel_path in changed + deleted - ) + total_deleted = 0 + for rel_path in changed + deleted: + n = delete_file_chunks(conn, tables, rel_path, commit=False) + if n: + print(f" Deleted {n} chunks for {rel_path}") + total_deleted += n for table, chunks in by_table.items(): if chunks: upsert_chunks(conn, table, chunks, embeddings[table], commit=False) @@ -301,76 +305,78 @@ def ingest(repo_path: str, namespace: str, repo_url: str | None = None, use_batc tables = [code_table, docs_table] conn = get_conn(os.environ["DATABASE_URL"]) - setup_db(conn, tables) - setup_meta_table(conn) - - repo_root = Path(repo_path).resolve() - print(f"Scanning {repo_root} (namespace: {namespace})...") - all_chunks: list[dict] = [] - - for file_path in repo_root.rglob("*"): - if not file_path.is_file(): - continue - rel = file_path.relative_to(repo_root) - if should_skip(rel): - continue - chunks = chunk_file(file_path, repo_root, namespace) - if chunks: - all_chunks.extend(chunks) - - code_chunks = [c for c in all_chunks if c["table"] == code_table] - doc_chunks = [c for c in all_chunks if c["table"] == docs_table] - print(f"Total: {len(all_chunks)} chunks ({len(code_chunks)} code, {len(doc_chunks)} docs)") - if not all_chunks: - # BL-015: guards the stale-chunk delete below from wiping the index after a bad checkout + try: + setup_db(conn, tables) + setup_meta_table(conn) + + repo_root = Path(repo_path).resolve() + print(f"Scanning {repo_root} (namespace: {namespace})...") + all_chunks: list[dict] = [] + + for file_path in repo_root.rglob("*"): + if not file_path.is_file(): + continue + rel = file_path.relative_to(repo_root) + if should_skip(rel): + continue + chunks = chunk_file(file_path, repo_root, namespace) + if chunks: + all_chunks.extend(chunks) + + code_chunks = [c for c in all_chunks if c["table"] == code_table] + doc_chunks = [c for c in all_chunks if c["table"] == docs_table] + print(f"Total: {len(all_chunks)} chunks ({len(code_chunks)} code, {len(doc_chunks)} docs)") + if not all_chunks: + # BL-015: guards the stale-chunk delete below from wiping the index after a bad checkout + raise RuntimeError(f"No indexable files found in {repo_root} — refusing to reindex '{namespace}'") + + # BL-001: only full reindexes reach here with use_batch + # BL-002: large ones go via the Batch API + if use_batch: + texts = pending_texts(conn, tables, all_chunks) + threshold = batch_threshold() + print(f"{len(texts)} distinct texts need embedding (batch threshold {threshold})") + if len(texts) >= threshold: + submit_batches(conn, namespace, texts) + if repo_url: + # No SHA, so the next sync collects the batch and does the full reindex + register_namespace(conn, namespace, repo_url, None) + print("\nBatch submitted — scripts/sync.py will finish indexing once Voyage completes it.") + return False + + for table, chunks in [(code_table, code_chunks), (docs_table, doc_chunks)]: + if not chunks: + continue + print(f"\nUpserting {len(chunks)} chunks to {table}...") + index_chunks(conn, table, chunks) + + # BL-014: only after every upsert succeeded, drop chunks for removed or newly skipped files + # and chunk indexes past a file's new chunk count + # BL-016: a table with zero chunks this run is left untouched, not wiped + for table, chunks in [(code_table, code_chunks), (docs_table, doc_chunks)]: + if not chunks: + print(f"No chunks for {table} this run — leaving its existing chunks untouched") + continue + n = delete_chunks_except(conn, table, [c["id"] for c in chunks]) + if n: + print(f"Removed {n} stale chunks from {table}") + + if repo_url: + try: + sha = subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=str(repo_root) + ).decode().strip() + register_namespace(conn, namespace, repo_url, sha) + print(f"Registered namespace '{namespace}' at {sha[:8]}") + except Exception as e: + print(f"Warning: could not register namespace: {e}") + + print("\nIngestion complete.") + return True + finally: + # BL-020: guarantees the connection closes on every path, including an + # unhandled error from embedding, batch submission or pruning conn.close() - raise RuntimeError(f"No indexable files found in {repo_root} — refusing to reindex '{namespace}'") - - # BL-001: only full reindexes reach here with use_batch - # BL-002: large ones go via the Batch API - if use_batch: - texts = pending_texts(conn, tables, all_chunks) - threshold = batch_threshold() - print(f"{len(texts)} distinct texts need embedding (batch threshold {threshold})") - if len(texts) >= threshold: - submit_batches(conn, namespace, texts) - if repo_url: - # No SHA, so the next sync collects the batch and does the full reindex - register_namespace(conn, namespace, repo_url, None) - conn.close() - print("\nBatch submitted — scripts/sync.py will finish indexing once Voyage completes it.") - return False - - for table, chunks in [(code_table, code_chunks), (docs_table, doc_chunks)]: - if not chunks: - continue - print(f"\nUpserting {len(chunks)} chunks to {table}...") - index_chunks(conn, table, chunks) - - # BL-014: only after every upsert succeeded, drop chunks for removed or newly skipped files - # and chunk indexes past a file's new chunk count - # BL-016: a table with zero chunks this run is left untouched, not wiped - for table, chunks in [(code_table, code_chunks), (docs_table, doc_chunks)]: - if not chunks: - print(f"No chunks for {table} this run — leaving its existing chunks untouched") - continue - n = delete_chunks_except(conn, table, [c["id"] for c in chunks]) - if n: - print(f"Removed {n} stale chunks from {table}") - - if repo_url: - try: - sha = subprocess.check_output( - ["git", "rev-parse", "HEAD"], cwd=str(repo_root) - ).decode().strip() - register_namespace(conn, namespace, repo_url, sha) - print(f"Registered namespace '{namespace}' at {sha[:8]}") - except Exception as e: - print(f"Warning: could not register namespace: {e}") - - conn.close() - print("\nIngestion complete.") - return True def main() -> None: diff --git a/scripts/register.py b/scripts/register.py index 2c7cefc..d334768 100644 --- a/scripts/register.py +++ b/scripts/register.py @@ -12,7 +12,8 @@ # Register at a specific commit SHA (e.g. the one you actually indexed): python scripts/register.py --repo https://github.com/beyondessential/tupaia --namespace tupaia --sha abc1234 - # Clear the SHA so the next sync.py run triggers a full reindex: + # Clear the SHA so the next sync.py run triggers a full reindex + # (also clears any pending Voyage batch recorded for this namespace): python scripts/register.py --repo https://github.com/beyondessential/tupaia --namespace tupaia --sha "" Environment variables: @@ -29,7 +30,14 @@ from dotenv import load_dotenv load_dotenv(Path(__file__).parent.parent / ".env") -from rag.db import get_conn, register_namespace, setup_meta_table +from rag.db import ( + delete_cached_embeddings, + delete_pending_batches, + get_conn, + get_pending_batches, + register_namespace, + setup_meta_table, +) from rag.github import get_latest_sha, parse_repo_url @@ -64,6 +72,17 @@ def main() -> None: conn = get_conn(os.environ["DATABASE_URL"]) setup_meta_table(conn) register_namespace(conn, args.namespace, args.repo, sha) + + if args.sha == "": + # BL-021: clearing the SHA is meant to force a full reindex. sync.py checks + # for a pending batch before it ever looks at last_commit_sha, so a stuck or + # in-flight batch would otherwise silently override this — clear it too. + batch_ids = get_pending_batches(conn, args.namespace) + if batch_ids: + delete_pending_batches(conn, batch_ids) + delete_cached_embeddings(conn, batch_ids) + print(f" Cleared {len(batch_ids)} pending batch(es) for '{args.namespace}'") + conn.close() status = sha[:8] if sha else "None (will full-reindex on next sync)" diff --git a/scripts/sync.py b/scripts/sync.py index 23f3bfb..59686c9 100644 --- a/scripts/sync.py +++ b/scripts/sync.py @@ -117,7 +117,11 @@ def sync_namespace(conn, ns: dict, token: str | None) -> None: return # BL-007: real-time full reindex, reading batch results from the cache print(f"[{namespace}] Batches complete — running full reindex from cache") - # BL-018: only advance state if the reindex actually completed + # BL-018: only advance state if the reindex actually completed. ingest(use_batch=False) + # can only return True or raise — it never returns False — so this mirrors the gating + # at the other two full-reindex call sites rather than guarding a reachable path; a + # genuine failure here raises and is caught by main()'s per-namespace handler instead, + # leaving batch state untouched the same way. if not _full_reindex(repo_url, namespace, token, use_batch=False): print(f"[{namespace}] Full reindex from cache did not complete — retrying next run") return diff --git a/specs/pipelines/embedding-sync.md b/specs/pipelines/embedding-sync.md index 24b3bde..50c5f87 100644 --- a/specs/pipelines/embedding-sync.md +++ b/specs/pipelines/embedding-sync.md @@ -49,6 +49,9 @@ Keep each namespace's index in step with its repo while minimising Voyage spend. - **BL-016:** A table that produces zero chunks in a run is left untouched rather than pruned; an empty keep set is never read as "keep nothing". - **BL-017:** A malformed `VOYAGE_BATCH_THRESHOLD` fails only the namespace attempting a full reindex, not the whole sync run. - **BL-018:** The full reindex that collects a completed batch advances `last_commit_sha` and clears batch state only once that reindex has itself completed. +- **BL-019:** A table's `content_hash` column, backfill and index are created only while the column is missing; once present, a later `setup_table` call does not repeat the migration. +- **BL-020:** `ingest()` closes its database connection on every exit path, including an unhandled error partway through embedding, batch submission or pruning. +- **BL-021:** Clearing a namespace's SHA to force a full reindex also clears any pending Voyage batch and cached embeddings recorded for it. ## SLAs @@ -69,6 +72,7 @@ Keep each namespace's index in step with its repo while minimising Voyage spend. | One table produces zero chunks this run | BL-016 | Logged, not an error | | Malformed `VOYAGE_BATCH_THRESHOLD` | Raises `ValueError`; only that namespace's sync fails (BL-017) | Logged per namespace, cron job continues | | Collecting reindex does not complete | Batch state and `last_commit_sha` left unchanged for retry (BL-018) | Logged, cron job continues | +| Namespace stuck on an unresolvable pending batch | `register.py --sha ""` clears it (BL-021) | Manual operator action | ## Acceptance criteria @@ -91,6 +95,9 @@ Keep each namespace's index in step with its repo while minimising Voyage spend. | AC-015 | A full reindex with one table at zero chunks this run skips pruning that table | BL-016 | | AC-016 | A malformed `VOYAGE_BATCH_THRESHOLD` raises `ValueError`, caught per namespace by sync's loop | BL-017 | | AC-017 | The collecting reindex only advances `last_commit_sha` and clears batch state on its own success | BL-018 | +| AC-018 | `setup_table`'s migration block runs only while `content_hash` is missing from the table | BL-019 | +| AC-019 | `ingest()` closes its connection even when an embedding call raises | BL-020 | +| AC-020 | `register.py --sha ""` clears the namespace's pending batch rows and cached embeddings; a real SHA leaves them untouched | BL-021 | ## Open questions @@ -103,3 +110,4 @@ Keep each namespace's index in step with its repo while minimising Voyage spend. | Date | Author | Change | |---|---|---| | 2026-10-12 | @julianam-w | Implemented: embedding reuse, Voyage Batch API full reindexes, index pruning, hourly sync. All ACs pass by automated test; AC-001, AC-002 and AC-005 also confirmed live against the real Voyage API and database | +| 2026-10-12 | @julianam-w | Second `/code-review --fix` pass (8-angle): add BL-019..021 / AC-018..020 for a repeated-lock migration guard, a connection leak on error, and a `register.py` recovery path for a stuck pending batch | diff --git a/tests/test_batch_embedding.py b/tests/test_batch_embedding.py index 9f99391..e506f8a 100644 --- a/tests/test_batch_embedding.py +++ b/tests/test_batch_embedding.py @@ -191,7 +191,7 @@ def test_ac_006_failed_batch_drops_all_pending_rows_and_sync_raises(monkeypatch, assert sync_env["reindex"] == [] and sync_env["registered"] == [] -def test_submitted_full_reindex_does_not_advance_sha(monkeypatch, sync_env): +def test_ac_002_too_large_diff_batch_submission_does_not_advance_sha(monkeypatch, sync_env): monkeypatch.setattr(sync, "get_pending_batches", lambda conn, ns: []) monkeypatch.setattr(sync, "get_changed_files", lambda *a: ([], [], True)) monkeypatch.setattr(sync, "_full_reindex", lambda url, ns, token, use_batch=True: False) @@ -243,3 +243,22 @@ def test_ac_017_collecting_reindex_does_not_advance_state_when_reindex_fails(mon sync.sync_namespace(FakeConn(), NS, None) assert sync_env["reindex"] == [False] assert sync_env["registered"] == [] and sync_env["deleted"] == [] and sync_env["cleared"] == [] + + +def test_ac_019_ingest_closes_connection_even_when_embedding_raises(monkeypatch, repo, db): + closed = [] + + class ClosingConn(FakeConn): + def close(self): + closed.append(True) + + monkeypatch.setattr(ingest, "get_conn", lambda url: ClosingConn()) + + def boom(texts, input_type): + raise RuntimeError("voyage down") + + monkeypatch.setattr(ingest, "embed_batch", boom) + + with pytest.raises(RuntimeError, match="voyage down"): + ingest.ingest(str(repo), "ns") + assert closed == [True] diff --git a/tests/test_db.py b/tests/test_db.py new file mode 100644 index 0000000..57c6743 --- /dev/null +++ b/tests/test_db.py @@ -0,0 +1,50 @@ +"""Mocked unit tests for rag/db.py logic that doesn't need a real Postgres connection.""" + +from rag import db + + +class _FakeCursor: + def __init__(self, log: list[str], content_hash_exists: bool) -> None: + self.log = log + self.content_hash_exists = content_hash_exists + + def execute(self, sql, params=None) -> None: + self.log.append(" ".join(sql.split())) + + def fetchone(self): + return (1,) if self.content_hash_exists else None + + def __enter__(self): + return self + + def __exit__(self, *exc) -> bool: + return False + + +class _FakeConn: + def __init__(self, content_hash_exists: bool) -> None: + self.log: list[str] = [] + self._content_hash_exists = content_hash_exists + + def cursor(self): + return _FakeCursor(self.log, self._content_hash_exists) + + def commit(self) -> None: + pass + + +def test_ac_018_setup_table_skips_migration_when_content_hash_exists(): + conn = _FakeConn(content_hash_exists=True) + db.setup_table(conn, "ns_code") + assert not any("ALTER TABLE" in q for q in conn.log) + assert not any(q.startswith("UPDATE ns_code") for q in conn.log) + # The non-migration indexes still get their (idempotent) IF NOT EXISTS check + assert any("ns_code_embedding_idx" in q for q in conn.log) + + +def test_ac_018_setup_table_runs_migration_when_content_hash_missing(): + conn = _FakeConn(content_hash_exists=False) + db.setup_table(conn, "ns_code") + assert any("ALTER TABLE" in q and "content_hash" in q for q in conn.log) + assert any(q.startswith("UPDATE ns_code") for q in conn.log) + assert any("ns_code_content_hash_idx" in q for q in conn.log) diff --git a/tests/test_db_postgres.py b/tests/test_db_postgres.py index f81f432..f0360ca 100644 --- a/tests/test_db_postgres.py +++ b/tests/test_db_postgres.py @@ -36,7 +36,7 @@ def test_ac_010_sql_backfill_hash_matches_python(conn): assert all(h == db.content_hash(text) for text, h in rows) -def test_delete_chunks_except_keeps_only_given_ids(conn): +def test_ac_012_delete_chunks_except_keeps_only_given_ids(conn): with conn.cursor() as cur: cur.execute("CREATE TEMP TABLE t_prune (id TEXT PRIMARY KEY)") cur.execute("INSERT INTO t_prune SELECT 'id' || g FROM generate_series(1, 100) g") diff --git a/tests/test_embedding_reuse.py b/tests/test_embedding_reuse.py index 5ee42fc..81e7581 100644 --- a/tests/test_embedding_reuse.py +++ b/tests/test_embedding_reuse.py @@ -57,7 +57,7 @@ def test_ac_009_no_voyage_call_when_everything_is_stored(monkeypatch, voyage): assert voyage == [] -def test_splits_new_texts_into_voyage_sized_batches(monkeypatch, voyage): +def test_ac_009_splits_new_texts_into_voyage_sized_batches(monkeypatch, voyage): monkeypatch.setattr(ingest, "get_embeddings_by_hash", lambda conn, table, hashes: {}) chunks = [_chunk(f"t{i}", i=i) for i in range(ingest.EMBED_BATCH + 1)] ingest.embed_chunks(FakeConn(), "ns_code", chunks) diff --git a/tests/test_register.py b/tests/test_register.py new file mode 100644 index 0000000..49b062e --- /dev/null +++ b/tests/test_register.py @@ -0,0 +1,54 @@ +"""scripts/register.py — forcing a reindex also clears any pending Voyage batch state.""" + +import sys + +import pytest + +from scripts import register + + +class FakeConn: + def commit(self) -> None: ... + def close(self) -> None: ... + + +@pytest.fixture +def env(monkeypatch): + calls: dict[str, list] = {"registered": [], "cleared_pending": [], "cleared_cache": []} + monkeypatch.setenv("DATABASE_URL", "postgresql://fake") + monkeypatch.setattr(register, "get_conn", lambda url: FakeConn()) + monkeypatch.setattr(register, "setup_meta_table", lambda conn: None) + monkeypatch.setattr(register, "register_namespace", lambda conn, ns, url, sha: calls["registered"].append(sha)) + monkeypatch.setattr(register, "delete_pending_batches", lambda conn, ids: calls["cleared_pending"].append(ids)) + monkeypatch.setattr(register, "delete_cached_embeddings", lambda conn, ids: calls["cleared_cache"].append(ids)) + return calls + + +def test_ac_021_clearing_sha_also_clears_a_pending_batch(monkeypatch, env): + monkeypatch.setattr(register, "get_pending_batches", lambda conn, ns: ["batch-1"]) + monkeypatch.setattr( + sys, "argv", ["register.py", "--repo", "https://github.com/o/r", "--namespace", "ns", "--sha", ""] + ) + register.main() + assert env["registered"] == [None] + assert env["cleared_pending"] == [["batch-1"]] + assert env["cleared_cache"] == [["batch-1"]] + + +def test_ac_021_clearing_sha_with_no_pending_batch_clears_nothing(monkeypatch, env): + monkeypatch.setattr(register, "get_pending_batches", lambda conn, ns: []) + monkeypatch.setattr( + sys, "argv", ["register.py", "--repo", "https://github.com/o/r", "--namespace", "ns", "--sha", ""] + ) + register.main() + assert env["cleared_pending"] == [] and env["cleared_cache"] == [] + + +def test_ac_021_registering_a_real_sha_does_not_touch_pending_batches(monkeypatch, env): + monkeypatch.setattr(register, "get_pending_batches", lambda conn, ns: ["batch-1"]) + monkeypatch.setattr( + sys, "argv", ["register.py", "--repo", "https://github.com/o/r", "--namespace", "ns", "--sha", "abc123"] + ) + register.main() + assert env["registered"] == ["abc123"] + assert env["cleared_pending"] == [] and env["cleared_cache"] == []