diff --git a/py/autoevals/oai.py b/py/autoevals/oai.py index 7a20e29..654bfa3 100644 --- a/py/autoevals/oai.py +++ b/py/autoevals/oai.py @@ -6,7 +6,7 @@ import warnings from collections.abc import Callable from contextvars import ContextVar -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, Optional, Protocol, TypedDict, TypeVar, Union, cast, runtime_checkable GATEWAY_URL = "https://gateway.braintrust.dev" @@ -194,6 +194,7 @@ def complete(self, **kwargs): RateLimitError: type[Exception] = None # type: ignore # Set in __post_init__ is_async: bool = False _is_wrapped: bool = False + _embedding_cache_key: object = field(default_factory=object, init=False, repr=False, compare=False) def __post_init__(self): NamedWrapper, wrap_openai = get_openai_wrappers() diff --git a/py/autoevals/string.py b/py/autoevals/string.py index a14ecde..2acf0db 100644 --- a/py/autoevals/string.py +++ b/py/autoevals/string.py @@ -24,7 +24,7 @@ from autoevals.partial import ScorerWithPartial from autoevals.value import normalize_value -from .oai import LLMClient, arun_cached_request, get_default_embedding_model, run_cached_request +from .oai import LLMClient, _client_var, arun_cached_request, get_default_embedding_model, run_cached_request from .score import Score @@ -101,6 +101,8 @@ async def compare_texts(): expected_min: Minimum similarity threshold (default: 0.7) client: Optional AsyncOpenAI/OpenAI client. If not provided, uses global client from init() + Cached embeddings are isolated by the explicit or currently configured default client. + Returns: Score object with: - score: Normalized similarity (0-1) @@ -133,10 +135,14 @@ 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 + # Include the effective client as well as the request configuration because + # the cache is shared across scorer instances and default client contexts. + client = self.client or _client_var.get(None) + if isinstance(client, LLMClient): + # A token stays unique even after a replaced default client is collected. + client_key = client._embedding_cache_key + else: + client_key = id(client) if 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): diff --git a/py/autoevals/test_embeddings.py b/py/autoevals/test_embeddings.py index de754ba..3d9d70f 100644 --- a/py/autoevals/test_embeddings.py +++ b/py/autoevals/test_embeddings.py @@ -1,9 +1,14 @@ import asyncio +import json import pytest +import respx +from httpx import Response +from openai import AsyncOpenAI, OpenAI import autoevals.string as string_module -from autoevals import EmbeddingSimilarity +from autoevals import EmbeddingSimilarity, init +from autoevals.oai import LLMClient SYNONYMS = [ ("water", ["water", "H2O", "agua"]), @@ -23,6 +28,95 @@ def reset_embedding_cache(): EmbeddingSimilarity._CACHE.clear() +@pytest.fixture +def embedding_providers(): + def handle(request): + body = json.loads(request.content) + vector = ( + [0.0, 1.0] if request.url.host == "provider-b.example" and body["input"] == "reference" else [1.0, 0.0] + ) + return Response( + 200, + json={ + "object": "list", + "model": body["model"], + "data": [{"object": "embedding", "index": 0, "embedding": vector}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + }, + ) + + init() + with respx.mock as router: + routes = [ + router.post(f"https://provider-{provider}.example/v1/embeddings").mock(side_effect=handle) + for provider in ("a", "b") + ] + yield routes + init() + + +@pytest.mark.parametrize("run_async", [False, True]) +@pytest.mark.asyncio +async def test_embedding_cache_isolated_by_default_client(embedding_providers, run_async): + client_type = AsyncOpenAI if run_async else OpenAI + clients = [ + client_type(api_key="test-key", base_url=f"https://provider-{provider}.example/v1") for provider in ("a", "b") + ] + wrappers = [LLMClient(openai=client) for client in clients] + evaluator = EmbeddingSimilarity(expected_min=0) + + async def score(metric): + if run_async: + return (await metric.eval_async("output", "reference")).score + return (await asyncio.to_thread(metric.eval, "output", "reference")).score + + try: + init(client=wrappers[0]) + assert await score(evaluator) == 1.0 + assert await score(EmbeddingSimilarity(expected_min=0)) == 1.0 + assert embedding_providers[0].call_count == 2 + + init(client=wrappers[1]) + assert await score(evaluator) == 0.0 + assert await score(EmbeddingSimilarity(expected_min=0)) == 0.0 + assert embedding_providers[1].call_count == 2 + + assert await score(EmbeddingSimilarity(client=wrappers[0], expected_min=0)) == 1.0 + init(client=wrappers[0]) + assert await score(evaluator) == 1.0 + assert embedding_providers[0].call_count == 2 + finally: + init() + for client in clients: + if run_async: + await client.close() + else: + client.close() + + +@pytest.mark.asyncio +async def test_embedding_cache_isolated_by_async_client_context(embedding_providers): + async with ( + AsyncOpenAI(api_key="test-key", base_url="https://provider-a.example/v1") as client_a, + AsyncOpenAI(api_key="test-key", base_url="https://provider-b.example/v1") as client_b, + ): + ready = [asyncio.Event(), asyncio.Event()] + evaluator = EmbeddingSimilarity(expected_min=0) + + async def evaluate(client, index): + init(client=client) + ready[index].set() + await ready[1 - index].wait() + first = await evaluator.eval_async("output", "reference") + second = await evaluator.eval_async("output", "reference") + return first.score, second.score + + results = await asyncio.gather(evaluate(client_a, 0), evaluate(client_b, 1)) + + assert results == [(1.0, 1.0), (0.0, 0.0)] + assert [route.call_count for route in embedding_providers] == [2, 2] + + @pytest.mark.parametrize("run_async", [False, True]) @pytest.mark.asyncio async def test_embedding_cache_isolated_by_configuration(monkeypatch, run_async):