fix(retrieval): store the BM25 index as gzipped JSON instead of pickle (#13)

HF's malware scanner flagged the public bundle's pkl as unsafe: BM25
vocabulary tokens from Python docs ("subprocess", "urllib.request", ...)
are byte-identical to the STACK_GLOBAL import operands its signatures
hunt for. A pickle also asks public downloaders to trust arbitrary code
execution on load.

The index is now a versioned gzipped-JSON payload (68 MB -> 21 MB) with
full parity verified against the old pkl (equal postings, frequencies,
lengths; identical search scores). Both HF bundles already ship the
.json.gz and the .pkl was removed; stale local bundles self-heal because
_bundle_complete() checks the new path.

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

app/chroma_rag.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import gzip
 import json
 import logging
 import math
@@ -1068,18 +1069,50 @@ def tokenize_for_bm25(text: str) -> list[str]:
     return tokens
 
 
+# Gzipped JSON, not pickle: the index ships in public HF bundles, and a pickle
+# full of Python-doc vocabulary tokens ("subprocess", "urllib.request", ...)
+# both trips Hub malware scanners and asks downloaders to trust arbitrary code
+# execution on load. Bump the version on any incompatible payload change.
+BM25_INDEX_FORMAT_VERSION = 1
+
+
 def save_bm25_index(index: BM25Index, output_file: str) -> None:
     ensure_parent_dir(output_file)
-    with open(output_file, "wb") as handle:
-        pickle.dump(index, handle)
+    payload = {
+        "format_version": BM25_INDEX_FORMAT_VERSION,
+        "k1": index.k1,
+        "b": index.b,
+        "average_document_length": index.average_document_length,
+        "document_lengths": index.document_lengths,
+        "document_frequencies": index.document_frequencies,
+        "postings": index.postings,
+        "records": [asdict(record) for record in index.records],
+    }
+    with gzip.open(output_file, "wt", encoding="utf-8") as handle:
+        json.dump(payload, handle, ensure_ascii=False, separators=(",", ":"))
 
 
 def load_bm25_index(path: str) -> BM25Index | None:
     if not os.path.exists(path):
         return None
     try:
-        with open(path, "rb") as handle:
-            index = pickle.load(handle)
+        with gzip.open(path, "rt", encoding="utf-8") as handle:
+            payload = json.load(handle)
+        version = payload.get("format_version")
+        if version != BM25_INDEX_FORMAT_VERSION:
+            raise ValueError(f"unsupported format_version={version!r}")
+        index = BM25Index(
+            records=[ChunkRecord(**record) for record in payload["records"]],
+            postings={
+                term: [(entry[0], entry[1]) for entry in entries]
+                for term, entries in payload["postings"].items()
+            },
+            document_frequencies=payload["document_frequencies"],
+            document_lengths=payload["document_lengths"],
+            average_document_length=payload["average_document_length"],
+            k1=payload["k1"],
+            b=payload["b"],
+        )
     except Exception as exc:
         logger.warning(
             "Failed to load BM25 index; falling back to dense-only retrieval. "
@@ -1089,14 +1122,7 @@ def load_bm25_index(path: str) -> BM25Index | None:
             exc,
         )
         return None
-    if isinstance(index, BM25Index):
-        return index
-    logger.warning(
-        "BM25 index has unexpected type %s; falling back to dense-only retrieval. path=%s",
-        type(index).__name__,
-        path,
-    )
-    return None
+    return index
 
 
 def default_bm25_index_path(document_dict_path: str) -> str:
@@ -1104,8 +1130,8 @@ def default_bm25_index_path(document_dict_path: str) -> str:
     name = path.name
     if name.startswith("document_dict_") and name.endswith(".pkl"):
         source = name.removeprefix("document_dict_").removesuffix(".pkl")
-        return str(path.with_name(f"bm25_index_{source}.pkl"))
-    return str(path.with_name("bm25_index.pkl"))
+        return str(path.with_name(f"bm25_index_{source}.json.gz"))
+    return str(path.with_name("bm25_index.json.gz"))
 
 
 def result_dedupe_key(result: SearchResult) -> str:

app/config.py

@@ -45,7 +45,7 @@ if langsmith_tracing_enabled():
 VECTOR_DB_DIR = "data/chroma-db-all_sources"
 VECTOR_COLLECTION_NAME = "chroma-db-all_sources"
 DOCUMENT_DICT_PATH = f"{VECTOR_DB_DIR}/document_dict_all_sources.pkl"
-BM25_INDEX_PATH = f"{VECTOR_DB_DIR}/bm25_index_all_sources.pkl"
+BM25_INDEX_PATH = f"{VECTOR_DB_DIR}/bm25_index_all_sources.json.gz"
 CHROMA_SQLITE_PATH = f"{VECTOR_DB_DIR}/chroma.sqlite3"
 # The full bundle (all sources, incl. gated course content) lives in the
 # private repo; cold-starting from it needs an HF_TOKEN with read access. When

data/scraping_scripts/build_public_docs_bundle.py

@@ -140,7 +140,7 @@ logger = logging.getLogger(__name__)
 VECTOR_DB_DIR = "chroma-db-all_sources"
 VECTOR_COLLECTION_NAME = "chroma-db-all_sources"
 DOCUMENT_DICT_FILE = "document_dict_all_sources.pkl"
-BM25_INDEX_FILE = "bm25_index_all_sources.pkl"
+BM25_INDEX_FILE = "bm25_index_all_sources.json.gz"
 KB_DIR_NAME = "kb"
 KB_ARCHIVE_NAME = "kb.tar.gz"
 
@@ -244,11 +244,11 @@ def stage_chroma(source_dir: Path, stage_dir: Path, *, vacuum: bool = True) -> d
     }
 
 
-def rebuild_retrieval_pkls(source_dir: Path, stage_dir: Path) -> dict:
+def rebuild_retrieval_artifacts(source_dir: Path, stage_dir: Path) -> dict:
     """Rebuild the BM25 index and document dict from a docs-only JSONL.
 
     Mirrors ``create_vector_stores.write_retrieval_artifacts`` so the public
-    pkls are byte-for-byte what prod would write without the course rows.
+    artifacts are byte-for-byte what prod would write without the course rows.
     """
     jsonl_path = source_dir / Path(ALL_SOURCES_JSONL).name
     if not jsonl_path.exists():
@@ -707,7 +707,7 @@ def build_bundle(
 
     summary: dict = {"sources": sorted(DOC_SOURCE_KEYS)}
     summary.update(stage_chroma(source_dir, stage_dir, vacuum=vacuum))
-    summary.update(rebuild_retrieval_pkls(source_dir, stage_dir))
+    summary.update(rebuild_retrieval_artifacts(source_dir, stage_dir))
     summary.update(stage_kb(source_dir, stage_dir))
     summary.update(publicize_wiki(stage_dir))
     write_dataset_card(stage_dir)

data/scraping_scripts/source_registry.py

@@ -485,7 +485,7 @@ def vector_store_source_configs() -> dict[str, dict[str, str]]:
             "input_file": str(config["output_file"]),
             "db_name": f"chroma-db-{source}",
             "document_dict_file": f"document_dict_{source}.pkl",
-            "bm25_index_file": f"bm25_index_{source}.pkl",
+            "bm25_index_file": f"bm25_index_{source}.json.gz",
         }
         for source, config in SOURCE_CONFIGS.items()
     }
@@ -493,6 +493,6 @@ def vector_store_source_configs() -> dict[str, dict[str, str]]:
         "input_file": ALL_SOURCES_JSONL,
         "db_name": "chroma-db-all_sources",
         "document_dict_file": "document_dict_all_sources.pkl",
-        "bm25_index_file": "bm25_index_all_sources.pkl",
+        "bm25_index_file": "bm25_index_all_sources.json.gz",
     }
     return configs

tests/test_build_public_docs_bundle.py

@@ -227,7 +227,7 @@ def test_stage_kb_prunes_non_allowlisted_wiki_framework_pages(tmp_path: Path) ->
     builder.audit_staged_kb(stage_dir)
 
 
-def test_rebuild_retrieval_pkls_drops_course_documents(tmp_path: Path) -> None:
+def test_rebuild_retrieval_artifacts_drops_course_documents(tmp_path: Path) -> None:
     source_dir = tmp_path / "data"
     rows = [
         {
@@ -252,7 +252,7 @@ def test_rebuild_retrieval_pkls_drops_course_documents(tmp_path: Path) -> None:
     _write_jsonl(source_dir / "all_sources_data.jsonl", rows)
     (tmp_path / "stage" / builder.VECTOR_DB_DIR).mkdir(parents=True)
 
-    summary = builder.rebuild_retrieval_pkls(source_dir, tmp_path / "stage")
+    summary = builder.rebuild_retrieval_artifacts(source_dir, tmp_path / "stage")
     assert summary["documents"] == 1
 
     dict_path = tmp_path / "stage" / builder.VECTOR_DB_DIR / builder.DOCUMENT_DICT_FILE

tests/test_chroma_rag.py

@@ -115,14 +115,14 @@ After the example.
             count = write_retrieval_artifacts(
                 config={
                     "document_dict_file": "document_dict_test.pkl",
-                    "bm25_index_file": "bm25_index_test.pkl",
+                    "bm25_index_file": "bm25_index_test.json.gz",
                 },
                 document_rows=document_rows,
                 db_path=str(db_path),
             )
 
             document_dict_path = db_path / "document_dict_test.pkl"
-            bm25_path = db_path / "bm25_index_test.pkl"
+            bm25_path = db_path / "bm25_index_test.json.gz"
 
             self.assertGreaterEqual(count, 1)
             self.assertTrue(document_dict_path.exists())

tests/test_config.py

@@ -30,7 +30,7 @@ def _patched_bundle(tmp_path: Path) -> ExitStack:
         ("VECTOR_DB_DIR", db_dir),
         ("CHROMA_SQLITE_PATH", db_dir / "chroma.sqlite3"),
         ("DOCUMENT_DICT_PATH", db_dir / "document_dict.pkl"),
-        ("BM25_INDEX_PATH", db_dir / "bm25.pkl"),
+        ("BM25_INDEX_PATH", db_dir / "bm25.json.gz"),
         ("KB_MANIFEST_PATH", kb_dir / "generated" / "corpus_manifest.jsonl"),
         ("KB_INDEX_PATH", kb_dir / "wiki" / "index.md"),
         ("KB_AGENTS_PATH", kb_dir / "AGENTS.md"),
@@ -47,7 +47,7 @@ def _write_bundle_files(tmp_path: Path, *, include_bm25: bool = True) -> None:
     (db_dir / "chroma.sqlite3").write_bytes(b"")
     (db_dir / "document_dict.pkl").write_bytes(b"")
     if include_bm25:
-        (db_dir / "bm25.pkl").write_bytes(b"")
+        (db_dir / "bm25.json.gz").write_bytes(b"")
     (tmp_path / "kb" / "generated").mkdir(parents=True, exist_ok=True)
     (tmp_path / "kb" / "generated" / "corpus_manifest.jsonl").write_text(
         "", encoding="utf-8"
@@ -120,7 +120,7 @@ def test_complete_bundle_skips_download_and_caches_readiness(tmp_path: Path) ->
 
             # Once verified, later calls are flag-checks: no re-stat, no
             # download attempt even if files vanish mid-process.
-            (tmp_path / "chroma" / "bm25.pkl").unlink()
+            (tmp_path / "chroma" / "bm25.json.gz").unlink()
             config.ensure_local_vector_db()
             assert snapshot_download.call_count == 0
 

tests/test_create_vector_stores.py

@@ -72,7 +72,7 @@ class ProcessSourceContentAwareReuseTestCase(unittest.TestCase):
             "input_file": f"data/{self.source}_data.jsonl",
             "db_name": self.db_name,
             "document_dict_file": f"document_dict_{self.source}.pkl",
-            "bm25_index_file": f"bm25_index_{self.source}.pkl",
+            "bm25_index_file": f"bm25_index_{self.source}.json.gz",
         }
         with open(config["input_file"], "w", encoding="utf-8") as handle:
             handle.write(