diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c1ebcb..eb45516 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -201,6 +201,13 @@ ### Changed +- The BM25 index is stored in the database, in `search_indexes`, instead of + `data/corpus/indices/bm25.pkl` (ADR 0021, new migration). A missing index + was a `FileNotFoundError` nothing caught, so an app running anywhere but the + machine that built the corpus answered no questions at all; it now raises an + error naming the command to run. `make build-index` rebuilds it from the + chunks already stored. Dropping the payload's unread `texts` key, a second + copy of the corpus, took it from 3.4 MB to 1.8 MB. - A `robots.txt` the server cannot serve is no longer read as a refusal (ADR 0020, superseding ADR 0013). RFC 9309 puts every 4xx in one "unavailable" class where a crawler may fetch, and probing all seven diff --git a/Makefile b/Makefile index fd6692d..5df78fd 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ -.PHONY: ingest ingest-plans ingest-sbc refresh-sbc sbc-report migrate api \ - ui-dev ui-build ui-lint ui-preview ui-test +.PHONY: ingest build-index ingest-plans ingest-sbc refresh-sbc sbc-report migrate api \ + ui-dev ui-build ui-lint ui-preview ui-test test lint check # The plan year to work on; the commands default to the calendar year, which # is the wrong one once the next year's plans are on sale. e.g. YEAR=2027 @@ -8,6 +8,12 @@ YEAR_ARG = $(if $(YEAR),--year $(YEAR)) ingest: uv run python -m src.ingestion.pipeline +# Rebuild the BM25 index from the chunks already stored, with no fetching, +# chunking or embedding (ADR 0021). `make ingest` does this as its last step; +# this is for a database seeded from a dump. +build-index: + uv run python -m src.ingestion.build_index + # STATES is required: comma-separated codes, or ALL for every HealthCare.gov # state. No default, so a bare `make ingest-plans` cannot start thousands of # requests. e.g. make ingest-plans STATES=TX,FL @@ -55,3 +61,15 @@ ui-preview: ui-test: cd frontend && npm test + +# What CI runs, in one command, so "it passed locally" means the same thing. +test: + uv run pytest -q --cov=src --cov-report=term-missing --cov-fail-under=85 + +lint: + uv run ruff check . + +check: lint test + uv run alembic check + uv run bandit -r src -ll + uv run pip-audit diff --git a/README.md b/README.md index 9dce3c9..4ec97f5 100644 --- a/README.md +++ b/README.md @@ -180,6 +180,11 @@ Runs in order: **fetch** (download each source) → **normalize** (convert to ma **chunk** (per-source splitting + a shared BM25 index) → **embed** (vectorize and upsert to pgvector). Every phase is idempotent, so a re-run only picks up what's missing. +The BM25 index is stored in the database beside the chunks it indexes, not on local disk +(ADR 0021), so the app can run somewhere with no `data/` directory. It is derived data: +`make build-index` rebuilds it from the chunks already stored, with no fetching, chunking +or embedding — which is what makes seeding a fresh database a restore plus one command. + Adding a corpus means writing one `Source` subclass in `src/ingestion/sources/` — `fetch`, `normalize`, `chunk_documents` — and registering it in that package's `__init__.py`. No pipeline, chunking, or embedding code changes. Chunking strategy lives with the source diff --git a/docs/decisions/0021-the-search-index-lives-in-the-database.md b/docs/decisions/0021-the-search-index-lives-in-the-database.md new file mode 100644 index 0000000..e7b7080 --- /dev/null +++ b/docs/decisions/0021-the-search-index-lives-in-the-database.md @@ -0,0 +1,91 @@ +# 0021 — The BM25 index lives in the database, not on local disk + +## Status + +Accepted. Amends ADR 0004's hybrid retrieval, which put the sparse index in +`data/corpus/indices/bm25.pkl`. + +## Context + +Hybrid retrieval (ADR 0004) needs a BM25 index at request time. +`_bm25_index()` read it from a pickle on local disk, cached per process, and +raised `FileNotFoundError` if it was absent. + +Nothing caught that exception. The chat route's catch-all turns it into a 500, +so **a deployment without the file answers no questions at all** — not +degraded, not sparse-only: every question fails. + +That never showed up because the app had only ever run on the machine that +built the corpus. Containerising it makes the gap immediate, and unfixable +where it stands: + +- `data/` is gitignored and stays that way, so CI cannot build an image + containing the file; +- the index is built by `make ingest`, which is a local, long-running job; +- the file is small (3.4 MB), but the constraint is provenance, not size. + +The alternatives were object storage — a storage account, a credential, an SDK +and a new startup failure mode — or building the image locally, which is not +CD at all. + +Meanwhile the index is **derived data**: everything it contains comes from the +`chunks` table, which already holds each chunk's text and id. + +## Decision + +### One row, in the database that holds the chunks it indexes + +`search_indexes` — `name` (primary key), `payload`, `chunks`, `built_at`. The +BM25 index is the only row today; a second index would be a second row, not a +second table. + +- **The request path reads it from Postgres**, still cached per process with + `lru_cache`, so the cost is one query per worker. +- **A missing index raises `MissingSearchIndexError`** naming the command to + run, instead of a `FileNotFoundError` naming a path that will not exist in + the environment that hit it. +- **Nothing writes to local disk.** `bm25_index_path` and `INDEX_DIR` are + gone; there is no path to configure. + +### It is rebuilt from the chunks, not restored with them + +`make build-index` reads `chunks` and rebuilds the index: no fetching, no +chunking, no embedding. `make ingest` still builds it as the last step of a +full run. + +This is what makes seeding a fresh database one restore plus one command, and +it means the index can never be silently out of date with a corpus that was +loaded some other way. + +### The payload no longer carries a second copy of the corpus + +The pickle held `{"bm25", "chunk_ids", "texts"}`. `texts` was read by nothing +outside a test: retrieval uses `bm25` and `chunk_ids`, and the text itself is +a column away in the table being indexed. Dropping it took the payload from +3.4 MB to 1.8 MB for the same 1,568 chunks. + +### Unpickling stays, and why that is not a new risk + +The payload is a pickle because `BM25Okapi` is a Python object. It is written +only by this project's own ingestion, into a table only this application +writes to, and it was already being unpickled from a file. Moving the bytes +from a file the app owns to a row the app owns does not add a source of +untrusted input. Anyone who can write that row can already run code as the +application. + +If that changes — a shared database, or a second writer — this is the thing to +revisit, and the answer is a schema-validated format rather than tighter +permissions. + +## Consequences + +- **The app runs where there is no `data/` directory.** This is the change + that makes a container work, and it is testable: a container with no `data/` + answering a question is the proof. +- **One more thing in the database dump,** and one fewer thing to copy beside + it. +- **A re-chunk must be followed by an index build,** as before — but now the + two are in the same database, so `built_at` and `chunks` say when the index + was made and from how many, which a file's mtime did not. +- **`make ingest` now needs a database.** It always did, at the embed stage; + the chunk stage no longer stands alone. diff --git a/migrations/env.py b/migrations/env.py index e95feec..b7da31a 100644 --- a/migrations/env.py +++ b/migrations/env.py @@ -9,6 +9,7 @@ chunk, plan, sbc, + search_index, usage, user, ) diff --git a/migrations/versions/20260922_1909-f6e50c4aec26_store_the_bm25_index_in_the_database.py b/migrations/versions/20260922_1909-f6e50c4aec26_store_the_bm25_index_in_the_database.py new file mode 100644 index 0000000..37c6b07 --- /dev/null +++ b/migrations/versions/20260922_1909-f6e50c4aec26_store_the_bm25_index_in_the_database.py @@ -0,0 +1,42 @@ +"""Store the BM25 index in the database (ADR 0021) + +One row per built index. It moves off local disk so the request path does +not need a `data/` directory, and so the index travels with the chunks it +indexes. Derived data: `make build-index` rebuilds it from `chunks`. + +Revision ID: f6e50c4aec26 +Revises: c4e91a7d6b28 +Create Date: 2026-09-22 19:09:02.890934+00:00 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "f6e50c4aec26" +down_revision: str | Sequence[str] | None = "c4e91a7d6b28" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.create_table( + "search_indexes", + sa.Column("name", sa.String(length=50), nullable=False), + sa.Column("payload", sa.LargeBinary(), nullable=False), + sa.Column("chunks", sa.Integer(), nullable=False), + sa.Column( + "built_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.PrimaryKeyConstraint("name"), + ) + + +def downgrade() -> None: + op.drop_table("search_indexes") diff --git a/src/core/exceptions.py b/src/core/exceptions.py index c887260..e601756 100644 --- a/src/core/exceptions.py +++ b/src/core/exceptions.py @@ -8,3 +8,7 @@ class InvalidCredentialsError(Exception): class MarketplaceApiKeyMissingError(Exception): """Raised when a Marketplace API call is made without CMS_MARKETPLACE_API_KEY set.""" + + +class MissingSearchIndexError(Exception): + """Raised when the BM25 index has not been built into the database yet.""" diff --git a/src/ingestion/build_index.py b/src/ingestion/build_index.py new file mode 100644 index 0000000..6595183 --- /dev/null +++ b/src/ingestion/build_index.py @@ -0,0 +1,37 @@ +"""Rebuild the BM25 index from the chunks already in the database: `make build-index`. + +The index is derived data (ADR 0021): everything it needs — each chunk's text +and its id — is a column away. So it can be rebuilt without re-fetching, +re-chunking or re-embedding anything, which is what makes seeding a fresh +database a restore plus one command. + +`make ingest` still builds it as the last step of a full run; this is for when +the chunks are already there. +""" +import sys + +from sqlalchemy import select + +from src.core.db import get_session +from src.core.logging import get_logger, setup_logging +from src.models.chunk import Chunk + +from .chunking import build_and_store_index + +logger = get_logger(__name__) + + +def main() -> None: + setup_logging() + with get_session() as session: + rows = session.execute(select(Chunk.chunk_id, Chunk.content).order_by(Chunk.id)).all() + + if not rows: + sys.exit("No chunks in the database. Run `make ingest` first.") + + build_and_store_index([content for _, content in rows], [chunk_id for chunk_id, _ in rows]) + print(f"BM25 index rebuilt from {len(rows)} chunks.") + + +if __name__ == "__main__": + main() diff --git a/src/ingestion/chunking.py b/src/ingestion/chunking.py index a94bac2..7ef6b00 100644 --- a/src/ingestion/chunking.py +++ b/src/ingestion/chunking.py @@ -10,11 +10,13 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter from rank_bm25 import BM25Okapi +from sqlalchemy import func +from sqlalchemy.dialects.postgresql import insert +from src.core.db import get_session from src.core.logging import get_logger from src.core.text import count_tokens - -from .constants import INDEX_DIR +from src.models.search_index import BM25_INDEX, SearchIndex logger = get_logger(__name__) @@ -126,20 +128,23 @@ def make_chunk( def build_and_store_index(texts: list[str], chunk_ids: list[str]) -> None: - INDEX_DIR.mkdir(parents=True, exist_ok=True) - - index_file_path = INDEX_DIR / "bm25.pkl" - - tokenized_texts = [text.lower().split() for text in texts] - - bm25 = BM25Okapi(tokenized_texts) + """Build the BM25 index and store it in the database (ADR 0021). - with index_file_path.open("wb") as file: - pickle.dump({ - "bm25": bm25, - "chunk_ids": chunk_ids, - "texts": texts, - }, - file) - - logger.info(f"BM25 index -> {len(tokenized_texts)} docs saved.") + `texts` is not kept in the payload: retrieval reads only `bm25` and + `chunk_ids`, and the chunk text is a column away in the table this + indexes. + """ + bm25 = BM25Okapi([text.lower().split() for text in texts]) + payload = pickle.dumps({"bm25": bm25, "chunk_ids": chunk_ids}) + + with get_session() as session: + session.execute( + insert(SearchIndex) + .values(name=BM25_INDEX, payload=payload, chunks=len(chunk_ids), built_at=func.now()) + .on_conflict_do_update( + index_elements=[SearchIndex.name], + set_={"payload": payload, "chunks": len(chunk_ids), "built_at": func.now()}, + ) + ) + + logger.info("BM25 index -> %d docs stored (%d KB).", len(chunk_ids), len(payload) // 1024) diff --git a/src/ingestion/constants.py b/src/ingestion/constants.py index 19d4ce0..fc4d9e7 100644 --- a/src/ingestion/constants.py +++ b/src/ingestion/constants.py @@ -65,5 +65,3 @@ MARKDOWN = Path("data/corpus/markdown") CHUNKS_DIR = Path("data/corpus/chunks") - -INDEX_DIR = Path("data/corpus/indices") diff --git a/src/models/search_index.py b/src/models/search_index.py new file mode 100644 index 0000000..0cc3a5c --- /dev/null +++ b/src/models/search_index.py @@ -0,0 +1,28 @@ +from datetime import datetime + +from sqlalchemy import DateTime, Integer, LargeBinary, String, func +from sqlalchemy.orm import Mapped, mapped_column + +from src.core.db import Base + +# The one index there is. A second would be a second row, not a second table. +BM25_INDEX = "bm25" + + +class SearchIndex(Base): + """A built search index, stored beside the chunks it indexes (ADR 0021). + + One row per index; `bm25` is the only one today. It lives here rather than + on disk because the request path needs it and the app has to be able to + run somewhere that has no `data/` directory. It is derived data: rebuilt + from `chunks` whenever they change, never authored. + """ + + __tablename__ = "search_indexes" + + name: Mapped[str] = mapped_column(String(50), primary_key=True) + payload: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + chunks: Mapped[int] = mapped_column(Integer, nullable=False) + built_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), server_default=func.now(), nullable=False + ) diff --git a/src/policypal/config.py b/src/policypal/config.py index 97f558a..22c9055 100644 --- a/src/policypal/config.py +++ b/src/policypal/config.py @@ -48,9 +48,8 @@ class Settings(BaseSettings): # exclusion that follows in the next sentence) at a hard chunk boundary. chunk_overlap: int = 50 - # Hybrid retrieval - bm25_index_path: str = "data/corpus/indices/bm25.pkl" - + # Hybrid retrieval. The BM25 index lives in the database, not on disk + # (ADR 0021), so there is no path to configure. sparse_top_k: int = 20 dense_top_k: int = 30 diff --git a/src/services/retrieval.py b/src/services/retrieval.py index 92404c8..31b6a6b 100644 --- a/src/services/retrieval.py +++ b/src/services/retrieval.py @@ -3,15 +3,16 @@ import re from dataclasses import dataclass from functools import lru_cache -from pathlib import Path from sqlalchemy import select from src.core.db import get_session from src.core.embedding import embed_query +from src.core.exceptions import MissingSearchIndexError from src.core.logging import get_logger from src.core.reranker import rerank from src.models.chunk import Chunk +from src.models.search_index import BM25_INDEX, SearchIndex from src.policypal.config import settings logger = get_logger(__name__) @@ -27,16 +28,23 @@ class RetrievedChunk: @lru_cache def _bm25_index() -> dict: - index_path = Path(settings.bm25_index_path) + """The stored BM25 index, read once per process (ADR 0021). - if not index_path.exists(): - raise FileNotFoundError(f"BM25 index not found at {index_path}. Run the chunk stage first.") - - with index_path.open("rb") as file: - # This file is only ever produced by our own ingestion pipeline - # (src/ingestion/chunk.py), never from user input or an external - # source, so there's no untrusted data to deserialize here. - return pkl.load(file) # nosec B301 + It comes from the database rather than a file so the app can run + somewhere that has no `data/` directory, and so it travels with the + chunks it indexes instead of beside them. + """ + with get_session() as session: + payload = session.scalar(select(SearchIndex.payload).where(SearchIndex.name == BM25_INDEX)) + + if payload is None: + raise MissingSearchIndexError( + "No BM25 index stored. Run `make build-index` (or `make ingest`) first." + ) + # Written only by our own ingestion (src/ingestion/chunking.py) into a + # table only this application writes to, never from user input or an + # external source, so there is no untrusted data to deserialize here. + return pkl.loads(payload) # nosec B301 def _sparse_search(query: str, top_k: int) -> list[str]: diff --git a/tests/test_ingestion_pipeline.py b/tests/test_ingestion_pipeline.py index f64835d..ffd1de7 100644 --- a/tests/test_ingestion_pipeline.py +++ b/tests/test_ingestion_pipeline.py @@ -1,15 +1,14 @@ """Tests for the source-agnostic ingestion orchestration and the Wikipedia source. The orchestrator's failure mode is silence: if it writes a malformed JSONL or -an index whose ids don't line up with its texts, nothing raises — retrieval -just quietly returns the wrong chunks. +hands the indexer ids that don't line up with its texts, nothing raises — +retrieval just quietly returns the wrong chunks. -The only pickle loaded below is the BM25 index these tests just built in a -tmp_path fixture — a build artifact of the code under test, never external -input. Same rationale as the justified load in src/services/retrieval.py. +What reaches the indexer is checked at the call, not by reading back a stored +index: the orchestrator's job is to hand over the right texts and ids, and +storing them is `build_and_store_index`'s (tested separately). """ import json -import pickle from unittest.mock import patch import pytest @@ -49,14 +48,15 @@ def _chunk(chunk_id, text="Some text about coverage.", doc_type="wikipedia"): def _run_execute(tmp_path, sources): - chunks_dir, index_dir = tmp_path / "chunks", tmp_path / "indices" + """Runs the orchestrator, returning where chunks were written and what the indexer got.""" + chunks_dir = tmp_path / "chunks" with patch.object(chunk_module, "SOURCES", sources), \ patch.object(chunk_module, "CHUNKS_DIR", chunks_dir), \ - patch("src.ingestion.chunking.INDEX_DIR", index_dir): + patch.object(chunk_module, "build_and_store_index") as indexer: chunk_module.execute() - return chunks_dir, index_dir + return chunks_dir, indexer # Orchestration @@ -85,27 +85,23 @@ def test_execute_writes_one_valid_json_object_per_line(tmp_path): def test_execute_indexes_the_contextualized_text_not_the_raw_text(tmp_path): """Retrieval depends on this: BM25 searches context, the DB stores raw.""" - _, index_dir = _run_execute(tmp_path, [_FakeSource("a", [_chunk("a_1", text="A deductible.")])]) + _, indexer = _run_execute(tmp_path, [_FakeSource("a", [_chunk("a_1", text="A deductible.")])]) - with (index_dir / "bm25.pkl").open("rb") as f: - index = pickle.load(f) - - assert index["texts"] == ["Context\nA deductible."] - assert index["chunk_ids"] == ["a_1"] + texts, chunk_ids = indexer.call_args.args + assert texts == ["Context\nA deductible."] + assert chunk_ids == ["a_1"] def test_execute_keeps_index_ids_aligned_with_index_texts(tmp_path): """A misalignment here returns the wrong chunk for every BM25 hit.""" chunks = [_chunk(f"a_{i}", text=f"Text {i}") for i in range(5)] - _, index_dir = _run_execute(tmp_path, [_FakeSource("a", chunks)]) - - with (index_dir / "bm25.pkl").open("rb") as f: - index = pickle.load(f) + _, indexer = _run_execute(tmp_path, [_FakeSource("a", chunks)]) - assert len(index["texts"]) == len(index["chunk_ids"]) == 5 - for i, chunk_id in enumerate(index["chunk_ids"]): - assert index["texts"][i].endswith(f"Text {i}") + texts, chunk_ids = indexer.call_args.args + assert len(texts) == len(chunk_ids) == 5 + for i, chunk_id in enumerate(chunk_ids): + assert texts[i].endswith(f"Text {i}") assert chunk_id == f"a_{i}" diff --git a/tests/test_search_index.py b/tests/test_search_index.py new file mode 100644 index 0000000..4f7b37f --- /dev/null +++ b/tests/test_search_index.py @@ -0,0 +1,103 @@ +"""The BM25 index, stored in the database rather than on local disk (ADR 0021). + +Real Postgres through the rolled-back `session` fixture. The only pickle +loaded here is the one these tests just built through the code under test. +""" +import pickle +from contextlib import contextmanager + +import pytest +from sqlalchemy import select + +from src.core.exceptions import MissingSearchIndexError +from src.ingestion import build_index, chunking +from src.models.chunk import Chunk +from src.models.search_index import BM25_INDEX, SearchIndex +from src.services import retrieval + +# BM25Okapi's idf is log((N - df + 0.5) / (df + 0.5)), which is exactly 0 for a +# term in one document of two — and `_sparse_search` keeps only scores above 0. +# Three documents is the smallest corpus where a hit actually scores. +_CORPUS = ["a deductible is what you pay first", "a copay is a flat fee", "a premium is monthly"] +_IDS = ["c1", "c2", "c3"] + + +@pytest.fixture +def stores(monkeypatch, session): + """Point everything that opens its own session at the rolled-back one.""" + @contextmanager + def same_session(): + yield session + session.flush() + + for module in (chunking, build_index, retrieval): + monkeypatch.setattr(module, "get_session", same_session) + # The real database has a built index and a loaded corpus; neither belongs + # in these assertions. Removed inside the transaction that is rolled back. + session.query(SearchIndex).delete() + session.query(Chunk).delete() + session.flush() + retrieval._bm25_index.cache_clear() + yield + retrieval._bm25_index.cache_clear() + + +def _stored(session): + return session.scalar(select(SearchIndex).where(SearchIndex.name == BM25_INDEX)) + + +def test_the_index_is_stored_with_the_chunk_count(stores, session): + chunking.build_and_store_index(["a deductible is what you pay first"], ["c1"]) + + row = _stored(session) + assert (row.name, row.chunks) == (BM25_INDEX, 1) + assert row.built_at is not None + + +def test_the_payload_carries_the_ids_but_not_the_texts(stores, session): + """Retrieval reads `bm25` and `chunk_ids` only; the text is a column away.""" + chunking.build_and_store_index(_CORPUS, _IDS) + + index = pickle.loads(_stored(session).payload) + + assert set(index) == {"bm25", "chunk_ids"} + assert index["chunk_ids"] == _IDS + assert index["bm25"].corpus_size == 3 + + +def test_rebuilding_replaces_the_one_row(stores, session): + chunking.build_and_store_index(["first"], ["c1"]) + chunking.build_and_store_index(_CORPUS, _IDS) + + assert session.scalars(select(SearchIndex.name)).all() == [BM25_INDEX] + assert _stored(session).chunks == 3 + + +def test_retrieval_reads_the_stored_index(stores, session): + chunking.build_and_store_index(_CORPUS, _IDS) + + assert retrieval._sparse_search("deductible", top_k=5) == ["c1"] + + +def test_retrieval_says_what_to_run_when_no_index_is_stored(stores): + with pytest.raises(MissingSearchIndexError, match="make build-index"): + retrieval._bm25_index() + + +def test_build_index_rebuilds_from_the_chunks_already_stored(stores, session): + """A database seeded from a dump has the chunks; the index is derived from them.""" + session.add_all([ + Chunk(content=content, source=f"{chunk_id}.md", chunk_id=chunk_id, embedding=[0.0] * 384) + for content, chunk_id in zip(_CORPUS, _IDS, strict=True) + ]) + session.flush() + + build_index.main() + + assert _stored(session).chunks == 3 + assert retrieval._sparse_search("copay", top_k=5) == ["c2"] + + +def test_build_index_says_what_to_run_when_there_are_no_chunks(stores): + with pytest.raises(SystemExit, match="make ingest"): + build_index.main()