"""Tests for content-aware embedding reuse in create_vector_stores. chunk_id (``f"{doc_id}:{index}"``) is stable across content edits, so reuse keyed on ids alone served stale text/embeddings forever when an upstream document changed in place. process_source now compares the stored Chroma document text against the incoming record text and re-embeds on mismatch. Uses a real chromadb PersistentClient on a tmp path; embed_texts is mocked (no Cohere calls). """ from __future__ import annotations import json import os import pickle import tempfile import unittest import uuid from unittest.mock import patch import chromadb import data.scraping_scripts.create_vector_stores as cvs from app.chroma_rag import ChunkRecord from data.scraping_scripts.source_registry import ACTIVE_SOURCE_KEYS class _OpaqueRecord: """A pickled record with no metadata anywhere: source is undeterminable.""" def _fake_embed(client, texts, **kwargs): return [[float(len(text)), 1.0] for text in texts] def _record(doc_id: str, index: int, text: str, heading: str = "") -> ChunkRecord: return ChunkRecord( chunk_id=f"{doc_id}:{index}", doc_id=doc_id, text=text, metadata={ "doc_id": doc_id, "title": "Guide", "url": "https://example.com/guide", "source": "testsource", "retrieve_doc": False, "tokens": 100, "chunk_index": index, "heading_path": heading, }, ) class ProcessSourceContentAwareReuseTestCase(unittest.TestCase): def setUp(self) -> None: self._tmp = tempfile.TemporaryDirectory() self.addCleanup(self._tmp.cleanup) self._old_cwd = os.getcwd() os.chdir(self._tmp.name) self.addCleanup(os.chdir, self._old_cwd) os.makedirs("data", exist_ok=True) # Unique per test: chromadb caches client instances by path, and # process_source uses paths relative to the (per-test) cwd, so a # repeated relative path would resurface another test's client. self.source = f"testsource-{uuid.uuid4().hex[:8]}" self.db_name = f"chroma-db-{self.source}" self.db_path = f"data/{self.db_name}" config = { "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}.json.gz", } with open(config["input_file"], "w", encoding="utf-8") as handle: handle.write( json.dumps( { "doc_id": "doc-1", "content": "# Guide\n\nSome document content.", "name": "Guide", "url": "https://example.com/guide", "source": self.source, "retrieve_doc": False, "tokens": 100, } ) + "\n" ) for patcher in ( patch.dict(cvs.SOURCE_CONFIGS, {self.source: config}), patch.dict(os.environ, {"COHERE_API_KEY": "fake-key"}), patch.object(cvs, "cohere"), ): patcher.start() self.addCleanup(patcher.stop) def _run(self, records: list[ChunkRecord], **kwargs) -> list[list[str]]: """Run process_source with mocked embeddings; return embed call texts.""" calls: list[list[str]] = [] def tracking_embed(client, texts, **embed_kwargs): calls.append(list(texts)) return _fake_embed(client, texts) with ( patch.object(cvs, "load_or_create_chunk_records", return_value=records), patch.object(cvs, "embed_texts", side_effect=tracking_embed), ): cvs.process_source(self.source, **kwargs) return calls def _collection(self): return chromadb.PersistentClient(path=self.db_path).get_collection( name=self.db_name ) def _stored_document(self, chunk_id: str) -> str | None: result = self._collection().get(ids=[chunk_id], include=["documents"]) documents = result.get("documents") or [None] return documents[0] def test_unchanged_chunks_are_reused_without_reembedding(self) -> None: records = [_record("doc-1", 0, "chunk zero"), _record("doc-1", 1, "chunk one")] first_calls = self._run(records) self.assertEqual(first_calls, [["chunk zero", "chunk one"]]) second_calls = self._run(records) self.assertEqual(second_calls, []) self.assertEqual(self._collection().count(), 2) def test_changed_content_with_same_chunk_id_is_reembedded(self) -> None: self._run([_record("doc-1", 0, "old text"), _record("doc-1", 1, "chunk one")]) changed = _record("doc-1", 0, "new text", heading="Guide > Install") calls = self._run([changed, _record("doc-1", 1, "chunk one")]) # Only the changed chunk is re-embedded; the unchanged one is reused. self.assertEqual(calls, [["new text"]]) # Stored document text and metadata are refreshed for the changed id. self.assertEqual(self._stored_document("doc-1:0"), "new text") result = self._collection().get(ids=["doc-1:0"], include=["metadatas"]) self.assertEqual(result["metadatas"][0]["heading_path"], "Guide > Install") self.assertEqual(self._stored_document("doc-1:1"), "chunk one") def test_removed_chunks_are_deleted(self) -> None: self._run([_record("doc-1", 0, "chunk zero"), _record("doc-1", 1, "chunk one")]) self._run([_record("doc-1", 0, "chunk zero")]) collection = self._collection() self.assertEqual(collection.count(), 1) self.assertEqual(collection.get(include=[])["ids"], ["doc-1:0"]) def test_legacy_chunks_with_matching_stored_text_are_reused(self) -> None: # Pre-fix collections always stored documents at upsert time, so an # upgrade run over an unchanged corpus must not re-embed anything. record = _record("doc-1", 0, "legacy text") collection = chromadb.PersistentClient( path=self.db_path ).get_or_create_collection(name=self.db_name) collection.add( ids=[record.chunk_id], embeddings=[[1.0, 2.0]], documents=[record.text], metadatas=[record.metadata], ) calls = self._run([record]) self.assertEqual(calls, []) def test_chunk_missing_stored_document_is_reembedded(self) -> None: # Defensive path: a chunk stored without document text cannot be # compared, so it is treated as changed and rewritten. collection = chromadb.PersistentClient( path=self.db_path ).get_or_create_collection(name=self.db_name) collection.add(ids=["doc-1:0"], embeddings=[[1.0, 2.0]]) calls = self._run([_record("doc-1", 0, "recovered text")]) self.assertEqual(calls, [["recovered text"]]) self.assertEqual(self._stored_document("doc-1:0"), "recovered text") def test_force_rebuild_reembeds_everything(self) -> None: records = [_record("doc-1", 0, "chunk zero")] self._run(records) # force_rebuild rmtree's the Chroma directory; drop chromadb's cached # client for the old path first, as a fresh CLI invocation would. from chromadb.api.shared_system_client import SharedSystemClient SharedSystemClient.clear_system_cache() calls = self._run(records, force_rebuild=True) self.assertEqual(calls, [["chunk zero"]]) class LoadOrCreateChunkRecordsTestCase(unittest.TestCase): def test_unparseable_records_fail_open_with_warning(self) -> None: active_source = next(iter(ACTIVE_SOURCE_KEYS)) active = ChunkRecord( chunk_id="doc-a:0", doc_id="doc-a", text="active", metadata={"doc_id": "doc-a", "source": active_source}, ) retired = ChunkRecord( chunk_id="doc-b:0", doc_id="doc-b", text="retired", metadata={"doc_id": "doc-b", "source": "definitely-retired-source"}, ) opaque = _OpaqueRecord() with tempfile.TemporaryDirectory() as tmp: old_cwd = os.getcwd() os.chdir(tmp) try: os.makedirs("data", exist_ok=True) with open("data/all_sources_contextual_nodes.pkl", "wb") as handle: pickle.dump([opaque, active, retired], handle) with self.assertLogs(cvs.logger, level="WARNING") as logs: kept = cvs.load_or_create_chunk_records("all_sources") finally: os.chdir(old_cwd) self.assertIn(active, kept) self.assertNotIn(retired, kept) self.assertTrue(any(isinstance(item, _OpaqueRecord) for item in kept)) self.assertTrue( any("could not be determined" in message for message in logs.output) ) self.assertTrue(any("_OpaqueRecord" in message for message in logs.output)) if __name__ == "__main__": unittest.main()