diff --git a/src/policy_scout/index.py b/src/policy_scout/index.py index 98fb951..b8efe48 100644 --- a/src/policy_scout/index.py +++ b/src/policy_scout/index.py @@ -12,6 +12,7 @@ critique finding #1. from __future__ import annotations import json +import hashlib from pathlib import Path import numpy as np @@ -22,6 +23,27 @@ from .query import expand RRF_K = 60 # standard reciprocal-rank-fusion constant MMR_LAMBDA = 0.7 # relevance vs diversity trade-off +BACKEND = "hybrid-tfidf-sparsebm25-mmr" + + +def _read_chunks(data: bytes) -> list[dict]: + chunks = [json.loads(line) for line in data.decode("utf-8").splitlines() if line.strip()] + seen = set() + for chunk in chunks: + if not isinstance(chunk, dict): + raise ValueError("Each chunk must be a JSON object") + if any(not isinstance(chunk.get(key), str) or not chunk[key].strip() + for key in ("doc_id", "text", "source_path")): + raise ValueError("Chunks require nonempty doc_id, text, and source_path strings") + if type(chunk.get("chunk_id")) is not int or chunk["chunk_id"] < 0: + raise ValueError("chunk_id must be a nonnegative integer") + key = (chunk["doc_id"], chunk["chunk_id"]) + if key in seen: + raise ValueError("Duplicate chunk citation identifier") + seen.add(key) + if not chunks: + raise ValueError("No chunks found") + return chunks class BM25: @@ -62,10 +84,10 @@ class SearchIndex: self.bm25: BM25 | None = None def build(self, chunks_path: Path) -> SearchIndex: - with chunks_path.open(encoding="utf-8") as f: - self.chunks = [json.loads(line) for line in f] - if not self.chunks: - raise ValueError(f"No chunks found in {chunks_path}") + return self._build_chunks(_read_chunks(chunks_path.read_bytes())) + + def _build_chunks(self, chunks: list[dict]) -> SearchIndex: + self.chunks = chunks texts = [c["text"] for c in self.chunks] self.tfidf = TfidfVectorizer(stop_words="english", ngram_range=(1, 2), sublinear_tf=True) self.tfidf_matrix = self.tfidf.fit_transform(texts) @@ -125,12 +147,15 @@ class SearchIndex: # -- persistence (no pickle, by design) -------------------------------- def save(self, dir_path: Path) -> None: + if self.tfidf is None: + raise RuntimeError("Cannot save an index before building it") dir_path.mkdir(parents=True, exist_ok=True) - with (dir_path / "chunks.jsonl").open("w", encoding="utf-8") as f: - for c in self.chunks: - f.write(json.dumps(c) + "\n") + data = ("\n".join(json.dumps(c, ensure_ascii=False) for c in self.chunks) + "\n").encode("utf-8") + _read_chunks(data) + (dir_path / "chunks.jsonl").write_bytes(data) (dir_path / "manifest.json").write_text( - json.dumps({"version": 3, "n_chunks": len(self.chunks), "backend": "hybrid-tfidf-sparsebm25-mmr"}) + json.dumps({"version": 4, "n_chunks": len(self.chunks), "backend": BACKEND, + "chunks_sha256": hashlib.sha256(data).hexdigest()}), encoding="utf-8" ) @classmethod @@ -138,4 +163,18 @@ class SearchIndex: manifest = dir_path / "manifest.json" if not manifest.exists(): raise FileNotFoundError(f"No index manifest at {dir_path}") - return cls().build(dir_path / "chunks.jsonl") + meta = json.loads(manifest.read_text(encoding="utf-8")) + if not isinstance(meta, dict) or type(meta.get("version")) is not int or meta["version"] not in (3, 4): + raise ValueError("Unsupported index manifest version") + if meta.get("backend") != BACKEND: + raise ValueError("Unsupported index backend") + if type(meta.get("n_chunks")) is not int or meta["n_chunks"] < 1: + raise ValueError("Invalid manifest chunk count") + data = (dir_path / "chunks.jsonl").read_bytes() + if meta["version"] == 4 and meta.get("chunks_sha256") != hashlib.sha256(data).hexdigest(): + raise ValueError("Index chunk checksum mismatch") + chunks = _read_chunks(data) + if len(chunks) != meta["n_chunks"]: + raise ValueError("Index chunk count mismatch") + # Rebuild from exactly the bytes checked, without reopening a mutable file. + return cls()._build_chunks(chunks) diff --git a/tests/test_persistence_integrity.py b/tests/test_persistence_integrity.py new file mode 100644 index 0000000..993ca83 --- /dev/null +++ b/tests/test_persistence_integrity.py @@ -0,0 +1,65 @@ +import json + +import pytest + +from policy_scout.index import SearchIndex + + +@pytest.fixture() +def saved(tmp_path): + records = [{"doc_id": "local-policy", "chunk_id": 0, + "text": "An example policy requires access review and named owners.", + "source_path": "fixtures/local-policy.txt"}] + source = tmp_path / "input.jsonl" + source.write_text(json.dumps(records[0]) + "\n", encoding="utf-8") + SearchIndex().build(source).save(tmp_path) + return tmp_path + + +@pytest.mark.parametrize("key,value", [("version", 99), ("version", True), + ("backend", "pickle"), ("n_chunks", True), ("n_chunks", -1), ("n_chunks", 2)]) +def test_rejects_incompatible_or_inconsistent_manifest(saved, key, value): + path = saved / "manifest.json" + meta = json.loads(path.read_text()) + meta[key] = value + path.write_text(json.dumps(meta)) + with pytest.raises(ValueError): + SearchIndex.load(saved) + + +def test_same_length_text_change_is_detected(saved): + path = saved / "chunks.jsonl" + path.write_bytes(path.read_bytes().replace(b"requires", b"excludes")) + with pytest.raises(ValueError, match="checksum"): + SearchIndex.load(saved) + + +def test_missing_v4_checksum_is_rejected(saved): + path = saved / "manifest.json" + meta = json.loads(path.read_text()) + del meta["chunks_sha256"] + path.write_text(json.dumps(meta)) + with pytest.raises(ValueError, match="checksum"): + SearchIndex.load(saved) + + +def test_v3_remains_readable_with_schema_and_count_validation(saved): + path = saved / "manifest.json" + meta = json.loads(path.read_text()) + meta["version"] = 3 + del meta["chunks_sha256"] + path.write_text(json.dumps(meta)) + assert SearchIndex.load(saved).search("access review") + + +@pytest.mark.parametrize("change", ["duplicate", "empty_text", "boolean_id", "missing_source"]) +def test_invalid_chunk_citations_rejected_before_build(saved, change): + path = saved / "chunks.jsonl" + chunk = json.loads(path.read_text()) + if change == "empty_text": chunk["text"] = " " + if change == "boolean_id": chunk["chunk_id"] = True + if change == "missing_source": del chunk["source_path"] + lines = [chunk, chunk] if change == "duplicate" else [chunk] + path.write_text("\n".join(map(json.dumps, lines)), encoding="utf-8") + with pytest.raises(ValueError): + SearchIndex().build(path)