diff --git a/py/autoevals/string.py b/py/autoevals/string.py index 92d9c3a..a14ecde 100644 --- a/py/autoevals/string.py +++ b/py/autoevals/string.py @@ -132,33 +132,42 @@ def __init__( self.client = client + def _cache_key(self, value): + # The cache is shared across scorer instances, so include every setting that + # can select a different embedding provider. A custom client may route to a + # different provider even when the model and base URL match. + client_key = id(self.client) if self.client is not None else None + return (self.extra_args["model"], self.prefix, self.extra_args.get("base_url"), client_key, value) + async def _a_embed(self, value): value = normalize_value(value, maybe_object=False) + cache_key = self._cache_key(value) with self._CACHE_LOCK: - if value in self._CACHE: - return self._CACHE[value] + if cache_key in self._CACHE: + return self._CACHE[cache_key] result = await arun_cached_request( client=self.client, request_type="embed", input=f"{self.prefix}{value}", **self.extra_args ) with self._CACHE_LOCK: - self._CACHE[value] = result + self._CACHE[cache_key] = result return result def _embed(self, value): value = normalize_value(value, maybe_object=False) + cache_key = self._cache_key(value) with self._CACHE_LOCK: - if value in self._CACHE: - return self._CACHE[value] + if cache_key in self._CACHE: + return self._CACHE[cache_key] result = run_cached_request( client=self.client, request_type="embed", input=f"{self.prefix}{value}", **self.extra_args ) with self._CACHE_LOCK: - self._CACHE[value] = result + self._CACHE[cache_key] = result return result diff --git a/py/autoevals/test_embeddings.py b/py/autoevals/test_embeddings.py index 6df632a..de754ba 100644 --- a/py/autoevals/test_embeddings.py +++ b/py/autoevals/test_embeddings.py @@ -1,5 +1,8 @@ import asyncio +import pytest + +import autoevals.string as string_module from autoevals import EmbeddingSimilarity SYNONYMS = [ @@ -11,6 +14,60 @@ UNRELATED = ["water", "The quick brown fox jumps over the lazy dog", "I like to eat apples"] +@pytest.fixture(autouse=True) +def reset_embedding_cache(): + with EmbeddingSimilarity._CACHE_LOCK: + EmbeddingSimilarity._CACHE.clear() + yield + with EmbeddingSimilarity._CACHE_LOCK: + EmbeddingSimilarity._CACHE.clear() + + +@pytest.mark.parametrize("run_async", [False, True]) +@pytest.mark.asyncio +async def test_embedding_cache_isolated_by_configuration(monkeypatch, run_async): + calls = [] + + def fake_run_cached_request(**kwargs): + calls.append((kwargs["model"], kwargs["input"], kwargs.get("base_url"), kwargs.get("client"))) + return {"data": [{"embedding": [1.0, 0.0]}]} + + async def fake_arun_cached_request(**kwargs): + calls.append((kwargs["model"], kwargs["input"], kwargs.get("base_url"), kwargs.get("client"))) + return {"data": [{"embedding": [1.0, 0.0]}]} + + if run_async: + monkeypatch.setattr(string_module, "arun_cached_request", fake_arun_cached_request) + else: + monkeypatch.setattr(string_module, "run_cached_request", fake_run_cached_request) + + value = {"topic": "cache"} + client_a = object() + client_b = object() + evaluators = [ + EmbeddingSimilarity(model="model-a", prefix="query: ", base_url="https://one.example", client=client_a), + EmbeddingSimilarity(model="model-a", prefix="query: ", base_url="https://one.example", client=client_a), + EmbeddingSimilarity(model="model-b", prefix="query: ", base_url="https://one.example", client=client_a), + EmbeddingSimilarity(model="model-a", prefix="document: ", base_url="https://one.example", client=client_a), + EmbeddingSimilarity(model="model-a", prefix="query: ", base_url="https://two.example", client=client_a), + EmbeddingSimilarity(model="model-a", prefix="query: ", base_url="https://one.example", client=client_b), + ] + + for evaluator in evaluators: + if run_async: + await evaluator.eval_async(value, value) + else: + evaluator.eval(value, value) + + assert calls == [ + ("model-a", 'query: {"topic": "cache"}', "https://one.example", client_a), + ("model-b", 'query: {"topic": "cache"}', "https://one.example", client_a), + ("model-a", 'document: {"topic": "cache"}', "https://one.example", client_a), + ("model-a", 'query: {"topic": "cache"}', "https://two.example", client_a), + ("model-a", 'query: {"topic": "cache"}', "https://one.example", client_b), + ] + + def test_embeddings(): evaluator = EmbeddingSimilarity(prefix="resource type: ") for word, synonyms in SYNONYMS: