234 lines
7.8 KiB
Python
234 lines
7.8 KiB
Python
"""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"
|