"""The embedding cache, exercised with a counting stub and no network. The claim these tests exist to make checkable is narrow and financial: running the corpus a second time must cost nothing. That is asserted by counting calls the *inner* provider received, not by trusting a hit counter. The other half is the inverse — the cases where a hit would be wrong. Serving a vector after its text was edited, across two models, or across Cohere's document/query subspaces would each be silent: no error, just worse recall or a corpus of mixed vectors. There is a test per direction. """ import json import pytest from ingestion.embed import INPUT_DOCUMENT, EmbeddingVector from ingestion.embed.cache import ( CachingEmbeddingProvider, EmbeddingCache, cache_key, ) from ingestion.embed.ports import EmbeddingProvider DIMENSIONS = 8 class CountingProvider(EmbeddingProvider): """Deterministic vectors, and a record of every text it was asked for.""" def __init__(self, model_id="stub-model-v1", dimensions=DIMENSIONS, batch=96): self._model_id = model_id self._dimensions = dimensions self._batch = batch self.embedded_texts = [] self.batch_calls = 0 @property def name(self): return "stub" @property def model_id(self): return self._model_id @property def dimensions(self): return self._dimensions @property def max_batch_size(self): return self._batch def _embed_batch(self, texts, input_kind): self.batch_calls += 1 self.embedded_texts.extend(texts) return [self._vector(text, input_kind) for text in texts] def _vector(self, text, input_kind): from ingestion.embed import text_digest seed = len(text) + (0 if input_kind == INPUT_DOCUMENT else 1000) return EmbeddingVector( values=[float(seed + i) for i in range(self._dimensions)], text_sha256=text_digest(text), provider="stub", model_id=self._model_id, dimensions=self._dimensions, input_kind=input_kind, normalized=True, input_token_count=len(text.split()), ) @pytest.fixture() def cache_path(tmp_path): return tmp_path / "embeddings.jsonl" def test_second_run_over_the_same_texts_makes_zero_provider_requests(cache_path): texts = ["paracetamol", "chống chỉ định", "liều dùng cho trẻ em"] inner = CountingProvider() first = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path)) cold = first.embed_documents(texts) assert cold.request_count == 1 assert inner.embedded_texts == texts reopened = EmbeddingCache(cache_path) warm = CachingEmbeddingProvider(inner, reopened).embed_documents(texts) assert warm.request_count == 0 assert inner.embedded_texts == texts, "no text reached the provider twice" assert reopened.stats.hits == len(texts) assert reopened.stats.misses == 0 assert reopened.stats.hit_rate == 1.0 assert [v.values for v in warm.vectors] == [v.values for v in cold.vectors] def test_a_repeated_text_in_one_call_is_embedded_once(cache_path): inner = CountingProvider() provider = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path)) batch = provider.embed_documents(["Abacavir.", "Abacavir.", "Abacavir."]) assert inner.embedded_texts == ["Abacavir."] assert len(batch.vectors) == 3 assert batch.vectors[0].values == batch.vectors[2].values def test_editing_the_text_is_a_miss_not_a_stale_hit(cache_path): inner = CountingProvider() cache = EmbeddingCache(cache_path) CachingEmbeddingProvider(inner, cache).embed_documents(["liều 500 mg"]) CachingEmbeddingProvider(inner, cache).embed_documents(["liều 250 mg"]) assert inner.embedded_texts == ["liều 500 mg", "liều 250 mg"] def test_a_second_model_never_reuses_the_first_models_vectors(cache_path): cache = EmbeddingCache(cache_path) titan = CountingProvider(model_id="amazon.titan-embed-text-v2:0") cohere = CountingProvider(model_id="cohere.embed-v4:0") CachingEmbeddingProvider(titan, cache).embed_documents(["metformin"]) CachingEmbeddingProvider(cohere, cache).embed_documents(["metformin"]) assert titan.embedded_texts == ["metformin"] assert cohere.embedded_texts == ["metformin"] assert len(cache) == 2 def test_query_and_document_kinds_are_cached_separately(cache_path): inner = CountingProvider() cache = EmbeddingCache(cache_path) provider = CachingEmbeddingProvider(inner, cache) as_document = provider.embed_documents(["aspirin"]) as_query = provider.embed_queries(["aspirin"]) assert inner.embedded_texts == ["aspirin", "aspirin"] assert as_document.vectors[0].values != as_query.vectors[0].values assert len(cache) == 2 def test_index_and_values_survive_reopening_the_file(cache_path): inner = CountingProvider() original = CachingEmbeddingProvider( inner, EmbeddingCache(cache_path) ).embed_documents(["ACETAZOLAMID", "ADENOSIN"]) reopened = EmbeddingCache(cache_path) assert len(reopened) == 2 restored = reopened.get(cache_key(inner.model_id, INPUT_DOCUMENT, "ADENOSIN")) assert restored is not None assert restored.values == original.vectors[1].values assert restored.input_kind == INPUT_DOCUMENT assert restored.normalized is True assert restored.input_token_count == 1 def test_putting_a_key_twice_does_not_append_a_second_record(cache_path): cache = EmbeddingCache(cache_path) inner = CountingProvider() vector = inner.embed_documents(["digoxin"]).vectors[0] assert cache.put(vector) is True assert cache.put(vector) is False lines = cache_path.read_text(encoding="utf-8").strip().splitlines() assert len(lines) == 1 assert len(cache) == 1 def test_a_cached_record_whose_length_contradicts_its_dimensions_is_rejected( cache_path, ): record = { "model_id": "stub-model-v1", "input_kind": INPUT_DOCUMENT, "text_sha256": "a" * 64, "provider": "stub", "dimensions": 1024, "normalized": True, "input_token_count": 3, "values": [0.1, 0.2], } cache_path.write_text(json.dumps(record) + "\n", encoding="utf-8") cache = EmbeddingCache(cache_path) with pytest.raises(ValueError, match="declares 1024 dimensions"): cache.get(("stub-model-v1", INPUT_DOCUMENT, "a" * 64)) def test_a_record_missing_key_fields_is_rejected_at_index_time(cache_path): cache_path.write_text( json.dumps({"model_id": "stub", "values": []}) + "\n", encoding="utf-8" ) with pytest.raises(ValueError, match="missing key fields"): EmbeddingCache(cache_path) def test_the_cache_wrapper_still_refuses_empty_text_and_bad_input_kind(cache_path): provider = CachingEmbeddingProvider(CountingProvider(), EmbeddingCache(cache_path)) with pytest.raises(ValueError, match="empty or whitespace-only"): provider.embed_documents(["paracetamol", " "]) with pytest.raises(ValueError, match="input_kind must be one of"): provider.embed(["paracetamol"], "search_document") def test_a_missing_cache_file_starts_empty_and_is_created_on_first_put(cache_path): cache = EmbeddingCache(cache_path) assert len(cache) == 0 assert not cache_path.exists() CachingEmbeddingProvider(CountingProvider(), cache).embed_documents(["insulin"]) assert cache_path.exists() assert len(cache) == 1 def test_wrapper_reports_the_inner_models_identity_not_its_own(cache_path): inner = CountingProvider(model_id="cohere.embed-v4:0", dimensions=DIMENSIONS) provider = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path)) assert provider.model_id == "cohere.embed-v4:0" assert provider.dimensions == DIMENSIONS assert provider.max_batch_size == inner.max_batch_size assert provider.name == "cached:stub"