diff --git a/tests/route_parity/fixture_generator.py b/tests/route_parity/fixture_generator.py index de51f09b8b7..4054a888490 100644 --- a/tests/route_parity/fixture_generator.py +++ b/tests/route_parity/fixture_generator.py @@ -2,7 +2,7 @@ from __future__ import annotations import argparse import logging -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path from typing import Final, Generic, Protocol, TypeVar, cast @@ -15,6 +15,8 @@ from tests.route_parity.fixture_recorder import ProviderSpec, generate_case_inpu LOGGER: Final = logging.getLogger(__name__) InputT = TypeVar("InputT", bound=SdkInputBase) +InputT_contra = TypeVar("InputT_contra", bound=SdkInputBase, contravariant=True) +DependencyT_contra = TypeVar("DependencyT_contra", contravariant=True) CaseT = TypeVar("CaseT", bound=BaseModel) @@ -30,28 +32,29 @@ class FixtureTarget(Generic[InputT]): name: str provider_spec: ProviderSpec strategy: SearchStrategy[InputT] - invoke: Callable[[str, InputT], object] + invocation: FixtureInvocation[InputT] required_inputs: tuple[InputT, ...] = () -class FixtureSdkCall(Protocol): - def __call__(self, **kwargs: object) -> object: ... +class FixtureInvocation(Protocol[InputT_contra]): + def execute(self, provider_url: str, case_input: InputT_contra) -> None: ... -class FixtureProvider(Protocol[InputT]): +class FixtureSource(Protocol[InputT, DependencyT_contra]): def targets( self, environ: Mapping[str, str], - sdk_call: FixtureSdkCall, + dependency: DependencyT_contra, + /, ) -> tuple[FixtureTarget[InputT], ...]: ... def discover_fixture_targets( - providers: tuple[FixtureProvider[InputT], ...], + sources: tuple[FixtureSource[InputT, DependencyT_contra], ...], environ: Mapping[str, str], - sdk_call: FixtureSdkCall, + dependency: DependencyT_contra, ) -> tuple[FixtureTarget[InputT], ...]: - return tuple(target for provider in providers for target in provider.targets(environ, sdk_call)) + return tuple(target for source in sources for target in source.targets(environ, dependency)) def generate_target_fixtures( @@ -67,7 +70,7 @@ def generate_target_fixtures( target.provider_spec, root / target.name, case_inputs, - target.invoke, + target.invocation.execute, case_type, concurrency, ) diff --git a/tests/route_parity/fixture_recorder.py b/tests/route_parity/fixture_recorder.py index 4ee15ac02eb..36a5f41d292 100644 --- a/tests/route_parity/fixture_recorder.py +++ b/tests/route_parity/fixture_recorder.py @@ -163,8 +163,7 @@ class _RecordingHandler(BaseHTTPRequestHandler): self._send_response(502, (), str(error).encode("utf-8")) return - if 200 <= recorded_response.status_code < 300: - provider.responses.put(recorded_response) + provider.responses.put(recorded_response) if isinstance(recorded_response, RecordedHttpResponse): self._send_response( recorded_response.status_code, recorded_response.headers, recorded_response.body_bytes() @@ -191,9 +190,9 @@ class _RecordingHandler(BaseHTTPRequestHandler): headers: tuple[HttpHeader, ...], ) -> RecordedHttpStreamResponse: self.send_response_only(upstream.status_code) + provider: Final = self.server + assert isinstance(provider, _RecordingProvider) for header in headers: - provider: Final = self.server - assert isinstance(provider, _RecordingProvider) self.send_header(header.name, _local_response_header(header.name, header.value, provider.url)) self.send_header("transfer-encoding", "chunked") self.end_headers() @@ -217,9 +216,9 @@ class _RecordingHandler(BaseHTTPRequestHandler): def _send_response(self, status_code: int, headers: tuple[HttpHeader, ...], body: bytes) -> None: self.send_response_only(status_code) + provider: Final = self.server + assert isinstance(provider, _RecordingProvider) for header in headers: - provider: Final = self.server - assert isinstance(provider, _RecordingProvider) self.send_header(header.name, _local_response_header(header.name, header.value, provider.url)) self.send_header("content-length", str(len(body))) self.end_headers() @@ -259,6 +258,21 @@ def fixture_cache_key(case_input: FixtureInput) -> dict[str, object]: return case_input.canonical_input() +def _invoke_and_take_responses( + recorder: _RecordingProvider, + case_input: InputT, + sdk_call: Callable[[str, InputT], object], +) -> tuple[RecordedResponse, ...]: + try: + sdk_call(recorder.url, case_input) + except Exception as invocation_error: + try: + return recorder.take_responses() + except RuntimeError: + raise invocation_error + return recorder.take_responses() + + def _load_fixture(raw_fixture: dict[str, object], path: Path, case_type: type[CaseT]) -> CaseT: schema_version: Final = raw_fixture.get("schema_version") if schema_version != FIXTURE_SCHEMA_VERSION: @@ -287,8 +301,7 @@ def record_case( return RecorderResult(case=_load_fixture(cached, cache.path_for(cache_key), case_type), cache_hit=True) with _recording_provider(spec) as recorder: - sdk_call(recorder.url, case_input) - upstream_responses: Final = recorder.take_responses() + upstream_responses: Final = _invoke_and_take_responses(recorder, case_input, sdk_call) case: Final = case_type.model_validate({"litellm_input": case_input, "provider_responses": upstream_responses}) envelope: Final = FixtureEnvelope( diff --git a/tests/route_parity/models.py b/tests/route_parity/models.py index 1b96cfbd355..898b58d23ee 100644 --- a/tests/route_parity/models.py +++ b/tests/route_parity/models.py @@ -1,8 +1,9 @@ from __future__ import annotations +import base64 from typing import Annotated, Final, Literal, cast -from pydantic import BaseModel, ConfigDict, Field, JsonValue +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter class CapturedRequest(BaseModel): @@ -36,7 +37,60 @@ class SDKError(BaseModel): llm_provider: str | None -SDKReport = Annotated[SDKSuccess | SDKError, Field(discriminator="status")] +class SDKJsonChunk(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["json"] = "json" + value: JsonValue + + +class SDKBytesChunk(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["bytes"] = "bytes" + data_b64: str + + def data_bytes(self) -> bytes: + return base64.b64decode(self.data_b64, validate=True) + + +SDKChunk = Annotated[SDKJsonChunk | SDKBytesChunk, Field(discriminator="kind")] + + +class SDKStreamCompleted(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["completed"] = "completed" + + +class SDKStreamFailed(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["failed"] = "failed" + error: SDKError + + +SDKStreamTerminal = Annotated[SDKStreamCompleted | SDKStreamFailed, Field(discriminator="kind")] + + +class SDKStreamReport(BaseModel): + model_config = ConfigDict(frozen=True) + + status: Literal["stream"] = "stream" + chunks: tuple[SDKChunk, ...] + terminal: SDKStreamTerminal + + +SDKReport = Annotated[SDKSuccess | SDKError | SDKStreamReport, Field(discriminator="status")] +JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + + +def sdk_chunk(value: object) -> SDKChunk: + if isinstance(value, bytes): + return SDKBytesChunk(data_b64=base64.b64encode(value).decode("ascii")) + if isinstance(value, BaseModel): + return SDKJsonChunk(value=JSON_VALUE_ADAPTER.validate_python(value.model_dump(mode="json"))) + return SDKJsonChunk(value=JSON_VALUE_ADAPTER.validate_python(value)) def _string_attribute(error: Exception, name: str) -> str | None: diff --git a/tests/route_parity/stream.py b/tests/route_parity/stream.py index 21997d76003..5f3352dc6d3 100644 --- a/tests/route_parity/stream.py +++ b/tests/route_parity/stream.py @@ -7,6 +7,15 @@ from typing import Final, Literal, TypeAlias from pydantic import BaseModel from tests.route_parity.compare import public_model_copy +from tests.route_parity.models import ( + SDKChunk, + SDKReport, + SDKStreamCompleted, + SDKStreamFailed, + SDKStreamReport, + sdk_chunk, + sdk_error_report, +) @dataclass(frozen=True, slots=True) @@ -40,6 +49,46 @@ class StreamOutcome: ChunkNormalizer: TypeAlias = Callable[[object], object] +def drain_sync_stream(stream: Iterable[object]) -> None: + for _ in stream: + pass + + +async def drain_async_stream(stream: AsyncIterable[object]) -> None: + async for _ in stream: + pass + + +def capture_sync_stream(create: Callable[[], Iterable[object]]) -> SDKReport: + try: + stream: Final = create() + except Exception as error: + return sdk_error_report(error) + + chunks: list[SDKChunk] = [] # mutable-ok: partial chunks must survive an iteration failure + try: + for chunk in stream: + chunks.append(sdk_chunk(chunk)) + except Exception as error: + return SDKStreamReport(chunks=tuple(chunks), terminal=SDKStreamFailed(error=sdk_error_report(error))) + return SDKStreamReport(chunks=tuple(chunks), terminal=SDKStreamCompleted()) + + +async def capture_async_stream(create: Callable[[], Awaitable[AsyncIterable[object]]]) -> SDKReport: + try: + stream: Final = await create() + except Exception as error: + return sdk_error_report(error) + + chunks: list[SDKChunk] = [] # mutable-ok: partial chunks must survive an iteration failure + try: + async for chunk in stream: + chunks.append(sdk_chunk(chunk)) + except Exception as error: + return SDKStreamReport(chunks=tuple(chunks), terminal=SDKStreamFailed(error=sdk_error_report(error))) + return SDKStreamReport(chunks=tuple(chunks), terminal=SDKStreamCompleted()) + + def _failed(phase: Literal["creation", "iteration"], error: BaseException) -> StreamFailed: status_code: Final = getattr(error, "status_code", None) llm_provider: Final = getattr(error, "llm_provider", None) diff --git a/tests/route_parity/test_fixture_generator.py b/tests/route_parity/test_fixture_generator.py index 5b2030d8888..1843729a939 100644 --- a/tests/route_parity/test_fixture_generator.py +++ b/tests/route_parity/test_fixture_generator.py @@ -2,12 +2,12 @@ from __future__ import annotations import queue from collections.abc import Mapping -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Final from hypothesis import strategies as st -from tests.route_parity.fixture_generator import FixtureSdkCall, FixtureTarget, discover_fixture_targets +from tests.route_parity.fixture_generator import FixtureSource, FixtureTarget, discover_fixture_targets from tests.route_parity.fixture_models import SdkInputBase from tests.route_parity.fixture_recorder import ProviderSpec @@ -16,6 +16,21 @@ class ExampleSdkInput(SdkInputBase): model: str +@dataclass(frozen=True, slots=True) +class ExampleInvocation: + calls: queue.SimpleQueue[dict[str, object]] + api_key: str = field(repr=False) + + def execute(self, provider_url: str, case_input: ExampleSdkInput) -> None: + self.calls.put( + { + "api_base": provider_url, + "api_key": self.api_key, + **case_input.as_sdk_kwargs(), + } + ) + + @dataclass(frozen=True, slots=True) class ExampleProvider: name: str @@ -24,22 +39,19 @@ class ExampleProvider: def targets( self, environ: Mapping[str, str], - sdk_call: FixtureSdkCall, + calls: queue.SimpleQueue[dict[str, object]], ) -> tuple[FixtureTarget[ExampleSdkInput], ...]: api_key: Final = environ.get(self.key_name) if not api_key: return () - def invoke(api_base: str, case_input: ExampleSdkInput) -> object: - return sdk_call(api_base=api_base, api_key=api_key, **case_input.as_sdk_kwargs()) - case_input: Final = ExampleSdkInput(model=f"{self.name}/model") return ( FixtureTarget( name=self.name, provider_spec=ProviderSpec(upstream_base=f"https://{self.name}.example"), strategy=st.just(case_input), - invoke=invoke, + invocation=ExampleInvocation(calls=calls, api_key=api_key), required_inputs=(case_input,), ), ) @@ -48,11 +60,7 @@ class ExampleProvider: def test_discover_fixture_targets_flattens_configured_providers_and_injects_sdk_call() -> None: calls: Final[queue.SimpleQueue[dict[str, object]]] = queue.SimpleQueue() - def sdk_call(**kwargs: object) -> object: - calls.put(kwargs) - return "response" - - providers: Final = ( + providers: Final[tuple[FixtureSource[ExampleSdkInput, queue.SimpleQueue[dict[str, object]]], ...]] = ( ExampleProvider(name="first", key_name="FIRST_KEY"), ExampleProvider(name="skipped", key_name="SKIPPED_KEY"), ExampleProvider(name="second", key_name="SECOND_KEY"), @@ -60,11 +68,11 @@ def test_discover_fixture_targets_flattens_configured_providers_and_injects_sdk_ targets: Final = discover_fixture_targets( providers, {"FIRST_KEY": "first-secret", "SECOND_KEY": "second-secret"}, - sdk_call, + calls, ) assert tuple(target.name for target in targets) == ("first", "second") - assert targets[1].invoke("http://127.0.0.1:1234", targets[1].required_inputs[0]) == "response" + targets[1].invocation.execute("http://127.0.0.1:1234", targets[1].required_inputs[0]) assert calls.get_nowait() == { "api_base": "http://127.0.0.1:1234", "api_key": "second-secret", diff --git a/tests/route_parity/test_fixture_recorder.py b/tests/route_parity/test_fixture_recorder.py index 693592841a6..3dde6b5cad7 100644 --- a/tests/route_parity/test_fixture_recorder.py +++ b/tests/route_parity/test_fixture_recorder.py @@ -97,8 +97,6 @@ class _ControlledUpstreamHandler(BaseHTTPRequestHandler): self._send_json(200, b'{"result":{"chunks":[]}}') return if self.path == "/analyze": - upstream: Final = self.server - assert isinstance(upstream, _ControlledUpstream) self.send_response(202) self.send_header("operation-location", f"{upstream.url}/results/1") self.send_header("content-length", "0") @@ -119,6 +117,9 @@ class _ControlledUpstreamHandler(BaseHTTPRequestHandler): self.wfile.write(b"0\r\n\r\n") self.wfile.flush() return + if self.path == "/error": + self._send_json(429, b'{"error":{"message":"rate limited"}}') + return upstream.start_request() try: body: Final = b"{}" @@ -172,6 +173,12 @@ def _stream_sdk_call(api_base: str, case_input: _FixtureInput) -> object: return httpx.post(f"{api_base}/v1/chat/completions", content=b"{}", timeout=5) +def _error_sdk_call(api_base: str, case_input: _FixtureInput) -> object: + response: Final = httpx.post(f"{api_base}/error", content=b"{}", timeout=5) + response.raise_for_status() + return response + + def _multi_sdk_call(api_base: str, case_input: _FixtureInput) -> object: upload: Final = httpx.post(f"{api_base}/upload", json={"document": case_input.identifier}, timeout=5) upload.raise_for_status() @@ -238,6 +245,7 @@ def test_streaming_response_records_and_replays_chunks(tmp_path: Path) -> None: response: Final = result.case.provider_responses[0] assert isinstance(response, RecordedHttpStreamResponse) assert tuple(chunk.data_bytes() for chunk in response.chunks) == _SSE_CHUNKS + assert isinstance(response.model_dump(mode="json")["chunks"], list) with replay_server() as provider: provider.enqueue_response(response) @@ -248,6 +256,20 @@ def test_streaming_response_records_and_replays_chunks(tmp_path: Path) -> None: assert replayed_chunks == _SSE_CHUNKS +def test_non_successful_provider_response_is_recorded(tmp_path: Path) -> None: + with _controlled_upstream() as upstream: + result: Final = record_case( + ProviderSpec(upstream_base=upstream.url), + tmp_path, + _case("provider-error"), + _error_sdk_call, + _ParityCase, + ) + + response: Final = result.case.provider_responses[0] + assert response.status_code == 429 + + def test_stream_response_model_rejects_buffered_body() -> None: with pytest.raises(ValueError, match="Extra inputs are not permitted"): RecordedHttpStreamResponse.model_validate( diff --git a/tests/route_parity/test_stream.py b/tests/route_parity/test_stream.py index 7cb6796c074..503c4b3d93b 100644 --- a/tests/route_parity/test_stream.py +++ b/tests/route_parity/test_stream.py @@ -1,16 +1,22 @@ from __future__ import annotations +import queue from collections.abc import AsyncIterator, Iterator from typing import Final import pytest from pydantic import BaseModel, PrivateAttr +from tests.route_parity.models import SDKBytesChunk, SDKJsonChunk, SDKStreamFailed, SDKStreamReport from tests.route_parity.stream import ( StreamFailed, assert_stream_parity, + capture_async_stream, + capture_sync_stream, consume_async_stream, consume_sync_stream, + drain_async_stream, + drain_sync_stream, ) @@ -23,7 +29,7 @@ class _Chunk(BaseModel): class _SyncStream: - def __init__(self, chunks: tuple[_Chunk, ...], error: BaseException | None = None) -> None: + def __init__(self, chunks: tuple[object, ...], error: BaseException | None = None) -> None: self.chunks: Final = chunks self.error: Final = error @@ -34,7 +40,7 @@ class _SyncStream: class _AsyncStream: - def __init__(self, chunks: tuple[_Chunk, ...], error: BaseException | None = None) -> None: + def __init__(self, chunks: tuple[object, ...], error: BaseException | None = None) -> None: self.chunks: Final = chunks self.error: Final = error @@ -59,7 +65,7 @@ def _creation_error() -> _SyncStream: raise _PublicStreamError(status_code=429, llm_provider="test", model="test-model") -async def _async_stream(chunks: tuple[_Chunk, ...], error: BaseException | None = None) -> _AsyncStream: +async def _async_stream(chunks: tuple[object, ...], error: BaseException | None = None) -> _AsyncStream: return _AsyncStream(chunks, error) @@ -132,3 +138,56 @@ def test_stream_parity_accepts_route_specific_chunk_normalizer() -> None: accelerated: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="rust-generated-id"),))) assert_stream_parity(python, accelerated, normalize=lambda chunk: type(chunk)) + + +def test_drain_sync_stream_exhausts_lazy_iterator() -> None: + consumed: Final[queue.SimpleQueue[str]] = queue.SimpleQueue() + + def chunks() -> Iterator[object]: + yield _Chunk(value="one") + consumed.put("complete") + + drain_sync_stream(chunks()) + + assert consumed.get_nowait() == "complete" + + +@pytest.mark.asyncio +async def test_drain_async_stream_exhausts_lazy_iterator() -> None: + consumed: Final[queue.SimpleQueue[str]] = queue.SimpleQueue() + + async def chunks() -> AsyncIterator[object]: + yield b"one" + consumed.put("complete") + + await drain_async_stream(chunks()) + + assert consumed.get_nowait() == "complete" + + +def test_capture_sync_stream_serializes_model_chunks_and_partial_failure() -> None: + report: Final = capture_sync_stream( + lambda: _SyncStream( + (_Chunk(value="before-error"),), + _PublicStreamError(status_code=429, llm_provider="test", model="test-model"), + ) + ) + + assert isinstance(report, SDKStreamReport) + assert len(report.chunks) == 1 + chunk: Final = report.chunks[0] + assert isinstance(chunk, SDKJsonChunk) + assert chunk.value == {"value": "before-error"} + assert isinstance(report.terminal, SDKStreamFailed) + assert report.terminal.error.status_code == 429 + + +@pytest.mark.asyncio +async def test_capture_async_stream_serializes_message_bytes_in_order() -> None: + report: Final = await capture_async_stream(lambda: _async_stream((b"first", b"second"))) + + assert isinstance(report, SDKStreamReport) + assert tuple(chunk.data_bytes() for chunk in report.chunks if isinstance(chunk, SDKBytesChunk)) == ( + b"first", + b"second", + ) diff --git a/tests/test_litellm/ocr/fixtures/azure.py b/tests/test_litellm/ocr/fixtures/azure.py index c4b23e0d669..119f2a5fca8 100644 --- a/tests/test_litellm/ocr/fixtures/azure.py +++ b/tests/test_litellm/ocr/fixtures/azure.py @@ -6,9 +6,9 @@ from typing import Final, cast from hypothesis import strategies as st from hypothesis.strategies import DrawFn, SearchStrategy -from tests.route_parity.fixture_generator import FixtureSdkCall from tests.route_parity.fixture_recorder import ProviderSpec from tests.test_litellm.ocr.fixtures.common import ( + OcrFixtureClient, OcrFixtureTarget, invoke_with_api_key, pdf_document, @@ -71,8 +71,8 @@ def azure_document_intelligence_input_strategy(draw: DrawFn) -> AzureDocumentInt ) -class AzureMistralFixtureProvider: - def targets(self, environ: Mapping[str, str], sdk_call: FixtureSdkCall) -> tuple[OcrFixtureTarget, ...]: +class AzureMistralFixtureSource: + def targets(self, environ: Mapping[str, str], client: OcrFixtureClient) -> tuple[OcrFixtureTarget, ...]: api_key: Final = environ.get("AZURE_AI_API_KEY") upstream_base: Final = environ.get("AZURE_AI_API_BASE") configured_model: Final = environ.get("AZURE_AI_OCR_MODEL") @@ -87,7 +87,7 @@ class AzureMistralFixtureProvider: SearchStrategy[OcrSdkInputBase], mistral_input_strategy(MISTRAL_MODEL).map(lambda case_input: _as_azure_mistral(case_input, model)), ), - invoke=invoke_with_api_key(sdk_call, api_key), + invocation=invoke_with_api_key(client, api_key), required_inputs=cast( tuple[OcrSdkInputBase, ...], tuple( @@ -98,8 +98,8 @@ class AzureMistralFixtureProvider: ) -class AzureDocumentIntelligenceFixtureProvider: - def targets(self, environ: Mapping[str, str], sdk_call: FixtureSdkCall) -> tuple[OcrFixtureTarget, ...]: +class AzureDocumentIntelligenceFixtureSource: + def targets(self, environ: Mapping[str, str], client: OcrFixtureClient) -> tuple[OcrFixtureTarget, ...]: api_key: Final = environ.get("AZURE_DOCUMENT_INTELLIGENCE_API_KEY") upstream_base: Final = environ.get("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") if not api_key or not upstream_base: @@ -109,7 +109,7 @@ class AzureDocumentIntelligenceFixtureProvider: name="azure-document-intelligence", provider_spec=ProviderSpec(upstream_base=upstream_base.rstrip("/")), strategy=cast(SearchStrategy[OcrSdkInputBase], azure_document_intelligence_input_strategy()), - invoke=invoke_with_api_key(sdk_call, api_key), + invocation=invoke_with_api_key(client, api_key), required_inputs=cast(tuple[OcrSdkInputBase, ...], _required_document_intelligence_inputs()), ), ) diff --git a/tests/test_litellm/ocr/fixtures/common.py b/tests/test_litellm/ocr/fixtures/common.py index 2d21f15dc15..74e4aa915ee 100644 --- a/tests/test_litellm/ocr/fixtures/common.py +++ b/tests/test_litellm/ocr/fixtures/common.py @@ -1,15 +1,15 @@ from __future__ import annotations import base64 -from collections.abc import Callable +from dataclasses import dataclass, field from pathlib import Path -from typing import Final +from typing import Final, Protocol from urllib.parse import quote from hypothesis import strategies as st from hypothesis.strategies import SearchStrategy -from tests.route_parity.fixture_generator import FixtureSdkCall, FixtureTarget +from tests.route_parity.fixture_generator import FixtureTarget from tests.test_litellm.ocr.fixtures.models import ( JsonSchemaDefinition, JsonSchemaResponseFormat, @@ -21,6 +21,23 @@ from tests.test_litellm.ocr.fixtures.models import ( OcrFixtureTarget = FixtureTarget[OcrSdkInputBase] +class OcrFixtureClient(Protocol): + def execute(self, api_base: str, api_key: str, case_input: OcrSdkInputBase) -> None: ... + + +class OcrSdkCall(Protocol): + def __call__(self, **kwargs: object) -> object: ... + + +@dataclass(frozen=True, slots=True) +class ApiKeyOcrInvocation: + client: OcrFixtureClient + api_key: str = field(repr=False) + + def execute(self, provider_url: str, case_input: OcrSdkInputBase) -> None: + self.client.execute(provider_url, self.api_key, case_input) + + def image_document(text: str, font_size: int) -> MistralImageUrlDocument: url: Final = f"https://dummyjson.com/image/800x300/ffffff/000000?text={quote(text)}&fontSize={font_size}" return MistralImageUrlDocument(type="image_url", image_url=url) @@ -57,8 +74,5 @@ def annotation_format(name: str) -> JsonSchemaResponseFormat: ) -def invoke_with_api_key(sdk_call: FixtureSdkCall, api_key: str) -> Callable[[str, OcrSdkInputBase], object]: - def invoke(api_base: str, case_input: OcrSdkInputBase) -> object: - return sdk_call(api_base=api_base, api_key=api_key, **case_input.as_sdk_kwargs()) - - return invoke +def invoke_with_api_key(client: OcrFixtureClient, api_key: str) -> ApiKeyOcrInvocation: + return ApiKeyOcrInvocation(client=client, api_key=api_key) diff --git a/tests/test_litellm/ocr/fixtures/generate.py b/tests/test_litellm/ocr/fixtures/generate.py index 60bbe851b5d..4de560b4064 100644 --- a/tests/test_litellm/ocr/fixtures/generate.py +++ b/tests/test_litellm/ocr/fixtures/generate.py @@ -11,8 +11,7 @@ from dotenv import load_dotenv import litellm from litellm.rust_bridge.ocr import use_litellm_rust from tests.route_parity.fixture_generator import ( - FixtureProvider, - FixtureSdkCall, + FixtureSource, discover_fixture_targets, generate_target_fixtures, parse_generator_args, @@ -20,23 +19,23 @@ from tests.route_parity.fixture_generator import ( from tests.route_parity.fixture_generator import require_targets as require_fixture_targets from tests.route_parity.fixture_recorder import fixture_directory from tests.test_litellm.ocr.fixtures.azure import ( - AzureDocumentIntelligenceFixtureProvider, - AzureMistralFixtureProvider, + AzureDocumentIntelligenceFixtureSource, + AzureMistralFixtureSource, azure_document_intelligence_input_strategy, ) -from tests.test_litellm.ocr.fixtures.common import OcrFixtureTarget +from tests.test_litellm.ocr.fixtures.common import OcrFixtureClient, OcrFixtureTarget, OcrSdkCall from tests.test_litellm.ocr.fixtures.mistral import ( - MistralFixtureProvider, + MistralFixtureSource, mistral_input_strategy, ) from tests.test_litellm.ocr.fixtures.models import OcrParityCase, OcrSdkInputBase from tests.test_litellm.ocr.fixtures.reducto import ( - ReductoFixtureProvider, + ReductoFixtureSource, reducto_legacy_input_strategy, reducto_v3_input_strategy, ) from tests.test_litellm.ocr.fixtures.vertex import ( - VertexFixtureProvider, + VertexFixtureSource, vertex_deepseek_input_strategy, ) @@ -49,20 +48,28 @@ __all__ = ( ) FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" -OCR_FIXTURE_PROVIDERS: Final[tuple[FixtureProvider[OcrSdkInputBase], ...]] = ( - MistralFixtureProvider(), - AzureMistralFixtureProvider(), - AzureDocumentIntelligenceFixtureProvider(), - VertexFixtureProvider(), - ReductoFixtureProvider(), +OCR_FIXTURE_SOURCES: Final[tuple[FixtureSource[OcrSdkInputBase, OcrFixtureClient], ...]] = ( + MistralFixtureSource(), + AzureMistralFixtureSource(), + AzureDocumentIntelligenceFixtureSource(), + VertexFixtureSource(), + ReductoFixtureSource(), ) +class LiteLLMOcrFixtureClient: + def __init__(self, sdk_call: OcrSdkCall) -> None: + self.sdk_call: Final = sdk_call + + def execute(self, api_base: str, api_key: str, case_input: OcrSdkInputBase) -> None: + self.sdk_call(api_base=api_base, api_key=api_key, **case_input.as_sdk_kwargs()) + + def discover_targets( environ: Mapping[str, str], - sdk_call: FixtureSdkCall, + client: OcrFixtureClient, ) -> tuple[OcrFixtureTarget, ...]: - return discover_fixture_targets(OCR_FIXTURE_PROVIDERS, environ, sdk_call) + return discover_fixture_targets(OCR_FIXTURE_SOURCES, environ, client) def require_targets(targets: tuple[OcrFixtureTarget, ...]) -> tuple[OcrFixtureTarget, ...]: @@ -76,8 +83,8 @@ def main() -> None: logging.basicConfig(level=logging.INFO, format="%(message)s") load_dotenv() args: Final = parse_generator_args() - sdk_call: Final = cast(FixtureSdkCall, litellm.ocr) - targets: Final = require_targets(discover_targets(os.environ, sdk_call)) + client: Final = LiteLLMOcrFixtureClient(cast(OcrSdkCall, litellm.ocr)) + targets: Final = require_targets(discover_targets(os.environ, client)) root: Final = fixture_directory( args.fixture_dir, os.environ.get(FIXTURE_DIR_ENV), diff --git a/tests/test_litellm/ocr/fixtures/mistral.py b/tests/test_litellm/ocr/fixtures/mistral.py index fc585b3ec2d..fdf08330ac1 100644 --- a/tests/test_litellm/ocr/fixtures/mistral.py +++ b/tests/test_litellm/ocr/fixtures/mistral.py @@ -6,9 +6,9 @@ from typing import Final, cast from hypothesis import strategies as st from hypothesis.strategies import DrawFn, SearchStrategy -from tests.route_parity.fixture_generator import FixtureSdkCall from tests.route_parity.fixture_recorder import ProviderSpec from tests.test_litellm.ocr.fixtures.common import ( + OcrFixtureClient, OcrFixtureTarget, annotation_format, image_document, @@ -80,8 +80,8 @@ def required_mistral_inputs(model: str) -> tuple[MistralOcrSdkInput, ...]: return tuple(MistralOcrSdkInput.model_validate({"model": model, "document": document, **case}) for case in cases) -class MistralFixtureProvider: - def targets(self, environ: Mapping[str, str], sdk_call: FixtureSdkCall) -> tuple[OcrFixtureTarget, ...]: +class MistralFixtureSource: + def targets(self, environ: Mapping[str, str], client: OcrFixtureClient) -> tuple[OcrFixtureTarget, ...]: api_key: Final = environ.get("MISTRAL_API_KEY") if not api_key: return () @@ -92,7 +92,7 @@ class MistralFixtureProvider: name="mistral-ocr", provider_spec=ProviderSpec(upstream_base=upstream_base), strategy=cast(SearchStrategy[OcrSdkInputBase], mistral_input_strategy(MISTRAL_MODEL)), - invoke=invoke_with_api_key(sdk_call, api_key), + invocation=invoke_with_api_key(client, api_key), required_inputs=cast(tuple[OcrSdkInputBase, ...], required_mistral_inputs(MISTRAL_MODEL)), ), ) diff --git a/tests/test_litellm/ocr/fixtures/reducto.py b/tests/test_litellm/ocr/fixtures/reducto.py index bc28c483f70..ceb44460349 100644 --- a/tests/test_litellm/ocr/fixtures/reducto.py +++ b/tests/test_litellm/ocr/fixtures/reducto.py @@ -6,9 +6,9 @@ from typing import Final, cast from hypothesis import strategies as st from hypothesis.strategies import DrawFn, SearchStrategy -from tests.route_parity.fixture_generator import FixtureSdkCall from tests.route_parity.fixture_recorder import ProviderSpec from tests.test_litellm.ocr.fixtures.common import ( + OcrFixtureClient, OcrFixtureTarget, fixture_pdf_data_uri, invoke_with_api_key, @@ -155,27 +155,27 @@ def _required_v3_inputs(document: ReductoDocumentUrlDocument) -> tuple[ReductoPa ) -class ReductoFixtureProvider: - def targets(self, environ: Mapping[str, str], sdk_call: FixtureSdkCall) -> tuple[OcrFixtureTarget, ...]: +class ReductoFixtureSource: + def targets(self, environ: Mapping[str, str], client: OcrFixtureClient) -> tuple[OcrFixtureTarget, ...]: api_key: Final = environ.get("REDUCTO_API_KEY") if not api_key: return () upstream_base: Final = environ.get("REDUCTO_API_BASE", _REDUCTO_API_BASE).rstrip("/") document: Final = ReductoDocumentUrlDocument(type="document_url", document_url=fixture_pdf_data_uri()) - invoke: Final = invoke_with_api_key(sdk_call, api_key) + invocation: Final = invoke_with_api_key(client, api_key) return ( OcrFixtureTarget( name="reducto-v3", provider_spec=ProviderSpec(upstream_base=upstream_base), strategy=cast(SearchStrategy[OcrSdkInputBase], reducto_v3_input_strategy(document)), - invoke=invoke, + invocation=invocation, required_inputs=cast(tuple[OcrSdkInputBase, ...], _required_v3_inputs(document)), ), OcrFixtureTarget( name="reducto-legacy", provider_spec=ProviderSpec(upstream_base=upstream_base), strategy=cast(SearchStrategy[OcrSdkInputBase], reducto_legacy_input_strategy(document)), - invoke=invoke, + invocation=invocation, required_inputs=cast( tuple[OcrSdkInputBase, ...], ( diff --git a/tests/test_litellm/ocr/fixtures/vertex.py b/tests/test_litellm/ocr/fixtures/vertex.py index aa79b0f9b48..0b71a26a536 100644 --- a/tests/test_litellm/ocr/fixtures/vertex.py +++ b/tests/test_litellm/ocr/fixtures/vertex.py @@ -6,9 +6,9 @@ from typing import Final, cast from hypothesis import strategies as st from hypothesis.strategies import DrawFn, SearchStrategy -from tests.route_parity.fixture_generator import FixtureSdkCall from tests.route_parity.fixture_recorder import ProviderSpec from tests.test_litellm.ocr.fixtures.common import ( + OcrFixtureClient, OcrFixtureTarget, image_document, invoke_with_api_key, @@ -72,15 +72,15 @@ def vertex_deepseek_input_strategy(draw: DrawFn, project: str, location: str) -> ) -class VertexFixtureProvider: - def targets(self, environ: Mapping[str, str], sdk_call: FixtureSdkCall) -> tuple[OcrFixtureTarget, ...]: +class VertexFixtureSource: + def targets(self, environ: Mapping[str, str], client: OcrFixtureClient) -> tuple[OcrFixtureTarget, ...]: api_key: Final = environ.get("VERTEX_AI_API_KEY") project: Final = environ.get("VERTEXAI_PROJECT") or environ.get("VERTEX_PROJECT") location: Final = environ.get("VERTEXAI_LOCATION") or environ.get("VERTEX_LOCATION") or "us-central1" if not api_key or not project: return () upstream_base: Final = environ.get("VERTEX_AI_API_BASE") or f"https://{location}-aiplatform.googleapis.com" - invoke: Final = invoke_with_api_key(sdk_call, api_key) + invocation: Final = invoke_with_api_key(client, api_key) return ( OcrFixtureTarget( name="vertex-mistral", @@ -94,7 +94,7 @@ class VertexFixtureProvider: location=st.just(location), ), ), - invoke=invoke, + invocation=invocation, required_inputs=cast( tuple[OcrSdkInputBase, ...], tuple( @@ -107,7 +107,7 @@ class VertexFixtureProvider: name="vertex-deepseek", provider_spec=ProviderSpec(upstream_base=upstream_base.rstrip("/")), strategy=cast(SearchStrategy[OcrSdkInputBase], vertex_deepseek_input_strategy(project, location)), - invoke=invoke, + invocation=invocation, required_inputs=cast(tuple[OcrSdkInputBase, ...], _required_deepseek_inputs(project, location)), ), ) diff --git a/tests/test_litellm/ocr/test_generate_fixtures.py b/tests/test_litellm/ocr/test_generate_fixtures.py index e2e2360b480..1004a69178c 100644 --- a/tests/test_litellm/ocr/test_generate_fixtures.py +++ b/tests/test_litellm/ocr/test_generate_fixtures.py @@ -1,6 +1,7 @@ from __future__ import annotations import queue +from dataclasses import dataclass from pathlib import Path from typing import Final @@ -12,10 +13,23 @@ from tests.test_litellm.ocr.fixtures.generate import ( parse_generator_args, require_targets, ) +from tests.test_litellm.ocr.fixtures.models import OcrSdkInputBase -def _unused_sdk_call(**kwargs: object) -> object: - raise AssertionError(f"unexpected SDK call with {tuple(kwargs)}") +class _UnusedOcrClient: + def execute(self, api_base: str, api_key: str, case_input: OcrSdkInputBase) -> None: + raise AssertionError(f"unexpected SDK call to {api_base} with {api_key!r} and {case_input!r}") + + +@dataclass(frozen=True, slots=True) +class _RecordingOcrClient: + calls: queue.SimpleQueue[dict[str, object]] + + def execute(self, api_base: str, api_key: str, case_input: OcrSdkInputBase) -> None: + self.calls.put({"api_base": api_base, "api_key": api_key, **case_input.as_sdk_kwargs()}) + + +_UNUSED_OCR_CLIENT: Final = _UnusedOcrClient() def test_parse_args_has_no_model_selection() -> None: @@ -37,7 +51,7 @@ def test_parse_args_has_no_model_selection() -> None: ), ) def test_discovery_requires_provider_specific_key(environ: dict[str, str]) -> None: - assert discover_targets(environ, _unused_sdk_call) == () + assert discover_targets(environ, _UNUSED_OCR_CLIENT) == () def test_no_discovered_targets_has_actionable_error() -> None: @@ -58,7 +72,7 @@ def test_discovery_is_explicit_per_available_provider_boundary() -> None: "VERTEX_AI_API_KEY": "vertex-secret", "VERTEXAI_PROJECT": "project-1", }, - _unused_sdk_call, + _UNUSED_OCR_CLIENT, ) assert tuple(target.name for target in targets) == ( @@ -78,12 +92,12 @@ def test_azure_mistral_discovery_requires_and_normalizes_deployment_model() -> N "AZURE_AI_API_KEY": "azure-secret", "AZURE_AI_API_BASE": "https://azure.example", } - assert discover_targets(incomplete, _unused_sdk_call) == () + assert discover_targets(incomplete, _UNUSED_OCR_CLIENT) == () - target: Final = discover_targets({**incomplete, "AZURE_AI_OCR_MODEL": "mistral-ocr-deployment"}, _unused_sdk_call)[ + target: Final = discover_targets({**incomplete, "AZURE_AI_OCR_MODEL": "mistral-ocr-deployment"}, _UNUSED_OCR_CLIENT)[ 0 ] - assert target.required_inputs[0].model == "azure_ai/mistral-ocr-deployment" + assert target.required_inputs[0].canonical_input()["model"] == "azure_ai/mistral-ocr-deployment" @pytest.mark.parametrize( @@ -102,7 +116,7 @@ def test_mistral_target_uses_canonical_model_and_normalized_base( "MISTRAL_API_KEY": "mistral-secret", **({"MISTRAL_API_BASE": configured} if configured is not None else {}), } - targets: Final = discover_targets(environ, _unused_sdk_call) + targets: Final = discover_targets(environ, _UNUSED_OCR_CLIENT) assert len(targets) == 1 target: Final = targets[0] @@ -139,14 +153,11 @@ def test_mistral_target_uses_canonical_model_and_normalized_base( def test_mistral_target_invocation_forwards_discovered_credentials() -> None: calls: Final[queue.SimpleQueue[dict[str, object]]] = queue.SimpleQueue() - def sdk_call(**kwargs: object) -> object: - calls.put(kwargs) - return object() - - target: Final = discover_targets({"MISTRAL_API_KEY": "mistral-secret"}, sdk_call)[0] + client: Final = _RecordingOcrClient(calls) + target: Final = discover_targets({"MISTRAL_API_KEY": "mistral-secret"}, client)[0] case_input: Final = generate_case_inputs(target.strategy, examples=1)[0] - target.invoke("http://127.0.0.1:1234", case_input) + target.invocation.execute("http://127.0.0.1:1234", case_input) kwargs: Final = calls.get_nowait() assert kwargs["api_base"] == "http://127.0.0.1:1234"