diff --git a/AGENTS.md b/AGENTS.md index 67ca7dd..aa6ebc0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -21,10 +21,11 @@ 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/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) -- `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 @@ -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 7182902..15288a1 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 @@ -247,6 +247,14 @@ 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. + +### 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/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. + 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. @@ -261,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:** @@ -270,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/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..4d6cd89 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 @@ -15,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) @@ -23,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""" @@ -36,10 +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 - """) + # 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) @@ -61,42 +96,104 @@ 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: +# 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() + + +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) + ORDER BY content_hash, id + """, + (list(hashes),), + ) + 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: """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 + +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. + + 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]) + 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: @@ -110,6 +207,23 @@ def setup_meta_table(conn) -> None: last_indexed_at TIMESTAMPTZ ) """) + # 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 ( + 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() @@ -152,3 +266,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..9c51075 --- /dev/null +++ b/rag/voyage_batch.py @@ -0,0 +1,105 @@ +"""Voyage Batch API client: Files and Batches REST endpoints. + +See specs/pipelines/embedding-sync.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 _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.""" + for raw in _stream_lines(f"/files/{output_file_id}/content"): + 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 2b4d862..30cfaf6 100644 --- a/scripts/ingest.py +++ b/scripts/ingest.py @@ -36,7 +36,22 @@ 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 import voyage_batch +from rag.db import ( + add_pending_batch, + content_hash, + delete_chunks_except, + delete_file_chunks, + get_cached_embeddings, + get_conn, + get_embeddings_by_hash, + get_existing_hashes, + 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 +80,8 @@ DOC_CHUNK_OVERLAP = 400 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, @@ -157,63 +174,209 @@ def chunk_file(file_path: Path, repo_root: Path, namespace: str) -> list[dict]: ] -# ── Main pipeline ───────────────────────────────────────────────────────────── +# ── Embedding with reuse ────────────────────────────────────────────────────── -def ingest(repo_path: str, namespace: str, repo_url: str | None = None) -> None: - namespace = sanitise_namespace(namespace) +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] + # 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()))) + reused = sum(1 for h in hashes if h in known) + + # 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: + 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)}) + + deduped = len(hashes) - reused - len(pending) + # 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] + + +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}", flush=True) + 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] - 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): + 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) - if chunks: - all_chunks.extend(chunks) + all_chunks.extend(chunk_file(file_path, repo_root, namespace)) - 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)") + 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: + 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) + conn.commit() + except Exception: + conn.rollback() + raise + print(f" Removed {total_deleted} stale chunks, upserted {len(all_chunks)}") - 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") - - 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.") + +# ── 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: + # 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]: + """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, 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" + tables = [code_table, docs_table] + + conn = get_conn(os.environ["DATABASE_URL"]) + 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() def main() -> None: @@ -234,10 +397,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") @@ -246,7 +418,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/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/reindex.py b/scripts/reindex.py index 942397f..bffd223 100644 --- a/scripts/reindex.py +++ b/scripts/reindex.py @@ -23,63 +23,25 @@ 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, setup_meta_table +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"]) + # 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] 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..59686c9 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/embedding-sync.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,27 +29,78 @@ from dotenv import load_dotenv load_dotenv(Path(__file__).parent.parent / ".env") +from rag import voyage_batch from rag.db import ( - delete_file_chunks, + 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, 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 +CACHE_PAGE = 1000 # batch output rows per cache insert -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"): + # 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) + + # 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, 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 def sync_namespace(conn, ns: dict, token: str | None) -> None: @@ -57,6 +111,27 @@ 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") + # 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 + # 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) + print(f"[{namespace}] Sync complete → {latest_sha[:8]}") + return + if latest_sha == last_sha: print(f"[{namespace}] Up to date ({latest_sha[:8]})") return @@ -64,8 +139,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...") @@ -73,8 +148,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") @@ -84,47 +159,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]}") @@ -141,6 +184,10 @@ def main() -> None: token = os.environ.get("GITHUB_TOKEN") conn = get_conn(os.environ["DATABASE_URL"]) setup_meta_table(conn) + # BL-010: drop week-old batch results + purged = purge_stale_cache(conn) + if purged: + print(f"Purged {purged} stale batch embeddings") namespaces = get_namespaces(conn) @@ -160,6 +207,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/specs/pipelines/embedding-sync.md b/specs/pipelines/embedding-sync.md new file mode 100644 index 0000000..50c5f87 --- /dev/null +++ b/specs/pipelines/embedding-sync.md @@ -0,0 +1,113 @@ +# 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** | `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-12 | + +## 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. +- **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 + +| 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 | +| 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 + +| 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 | +| 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 | +| 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 + +| 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-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/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_batch_embedding.py b/tests/test_batch_embedding.py new file mode 100644 index 0000000..e506f8a --- /dev/null +++ b/tests/test_batch_embedding.py @@ -0,0 +1,264 @@ +"""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 + +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": [], "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) + 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)) + 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") + 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_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_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_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"] == [] + + +def test_ac_016_invalid_threshold_raises_value_error(monkeypatch): + monkeypatch.setenv("VOYAGE_BATCH_THRESHOLD", "lots") + 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": "ā"}) + 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_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) + # 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"] == [] + + +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) + advanced: list[str] = [] + 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, "_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"] == [] + + +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 new file mode 100644 index 0000000..f0360ca --- /dev/null +++ b/tests/test_db_postgres.py @@ -0,0 +1,55 @@ +"""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_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") + 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 diff --git a/tests/test_embedding_reuse.py b/tests/test_embedding_reuse.py new file mode 100644 index 0000000..81e7581 --- /dev/null +++ b/tests/test_embedding_reuse.py @@ -0,0 +1,105 @@ +"""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) + monkeypatch.setattr(ingest, "get_cached_embeddings", lambda conn, hashes: {}) + 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_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 + }) + 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_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 == [] + + +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) + assert [len(c) for c in voyage] == [ingest.EMBED_BATCH, 1] + + +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 + [[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_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): + 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 == [] + + +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") + }) + assert ingest.embed_chunks(FakeConn(), "ns_code", [_chunk("cached"), _chunk("fresh")]) == [[7.0], [5.0]] + assert voyage == [["fresh"]] 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"] == [] 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"