From 4246fafcb9b504353a3196c0e5730ff2cb630034 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 29 Aug 2026 16:49:49 -0700 Subject: [PATCH] refactor(tests): generalize route parity harness --- tests/test_litellm/_fixture_generator.py | 76 +++++++++++++++++++++ tests/test_litellm/_fixture_models.py | 41 +++++++++++ tests/test_litellm/_fixture_recorder.py | 44 ++++++++++++ tests/test_litellm/ocr/conftest.py | 54 ++++----------- tests/test_litellm/ocr/fixture_models.py | 70 ++++++++----------- tests/test_litellm/ocr/generate_fixtures.py | 64 ++++------------- tests/test_litellm/ocr/test_sdk_parity.py | 30 +++----- tests/test_litellm/parity/compare.py | 16 ++--- tests/test_litellm/parity/models.py | 5 +- tests/test_litellm/parity/replay.py | 2 +- tests/test_litellm/parity/runner.py | 42 ++++++++++-- tests/test_litellm/parity/test_parity.py | 13 ++-- 12 files changed, 276 insertions(+), 181 deletions(-) create mode 100644 tests/test_litellm/_fixture_generator.py create mode 100644 tests/test_litellm/_fixture_models.py diff --git a/tests/test_litellm/_fixture_generator.py b/tests/test_litellm/_fixture_generator.py new file mode 100644 index 00000000000..40b9f0f9777 --- /dev/null +++ b/tests/test_litellm/_fixture_generator.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import argparse +import logging +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Generic, TypeVar, cast + +from hypothesis.strategies import SearchStrategy +from pydantic import BaseModel + +from tests.test_litellm._fixture_models import SdkInputBase +from tests.test_litellm._fixture_recorder import ProviderSpec, generate_case_inputs, record_cases + +LOGGER: Final = logging.getLogger(__name__) +InputT = TypeVar("InputT", bound=SdkInputBase) +CaseT = TypeVar("CaseT", bound=BaseModel) + + +@dataclass(frozen=True, slots=True) +class GeneratorArgs: + concurrency: int + examples: int + fixture_dir: Path | None + + +@dataclass(frozen=True, slots=True) +class FixtureTarget(Generic[InputT]): + name: str + provider_spec: ProviderSpec + strategy: SearchStrategy[InputT] + invoke: Callable[[str, InputT], object] + + +def generate_target_fixtures( + target: FixtureTarget[InputT], + root: Path, + examples: int, + concurrency: int, + case_type: type[CaseT], +) -> None: + case_inputs: Final = generate_case_inputs(target.strategy, examples) + results: Final = record_cases( + target.provider_spec, + root, + case_inputs, + target.invoke, + case_type, + concurrency, + ) + for result in results: + LOGGER.info( + "%s %s", + "cached" if result.cache_hit else "recorded", + target.name, + ) + + +def require_targets(targets: tuple[FixtureTarget[InputT], ...], error_message: str) -> tuple[FixtureTarget[InputT], ...]: + if targets: + return targets + raise SystemExit(error_message) + + +def parse_generator_args(argv: Sequence[str] | None = None) -> GeneratorArgs: + parser: Final = argparse.ArgumentParser() + parser.add_argument("--concurrency", type=int, default=4) + parser.add_argument("--examples", type=int, default=4) + parser.add_argument("--fixture-dir", type=Path) + namespace: Final = parser.parse_args(argv) + return GeneratorArgs( + concurrency=cast(int, namespace.concurrency), + examples=cast(int, namespace.examples), + fixture_dir=cast(Path | None, namespace.fixture_dir), + ) diff --git a/tests/test_litellm/_fixture_models.py b/tests/test_litellm/_fixture_models.py new file mode 100644 index 00000000000..a9aecf0730a --- /dev/null +++ b/tests/test_litellm/_fixture_models.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from typing import Generic, Literal, TypeVar, cast + +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +from tests.test_litellm._recorded_http import RecordedResponse + +JsonObject = dict[str, JsonValue] + + +class FixtureModel(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True, serialize_by_alias=True) + + +class SdkInputBase(FixtureModel): + def as_sdk_kwargs(self) -> dict[str, object]: + return cast(dict[str, object], self.model_dump(mode="python", exclude_unset=True)) + + def canonical_input(self) -> dict[str, object]: + return cast(dict[str, object], self.model_dump(mode="json", exclude_unset=True)) + + +class JsonSchemaDefinition(FixtureModel): + name: str + description: str | None = None + schema_definition: JsonObject = Field(alias="schema") + strict: bool = False + + +class JsonSchemaResponseFormat(FixtureModel): + type: Literal["json_schema"] + json_schema: JsonSchemaDefinition + + +InputT = TypeVar("InputT", bound=SdkInputBase) + + +class ParityCase(FixtureModel, Generic[InputT]): + litellm_input: InputT + provider_response: RecordedResponse diff --git a/tests/test_litellm/_fixture_recorder.py b/tests/test_litellm/_fixture_recorder.py index f791b08d13f..18939f46ef5 100644 --- a/tests/test_litellm/_fixture_recorder.py +++ b/tests/test_litellm/_fixture_recorder.py @@ -1,6 +1,7 @@ from __future__ import annotations import hashlib +import os import queue import threading from collections.abc import Callable, Generator, Iterable @@ -13,6 +14,7 @@ from pathlib import Path from typing import Final, Generic, Protocol, TypeVar, cast import httpx +import pytest from hypothesis import given, settings from hypothesis.strategies import SearchStrategy from pydantic import AwareDatetime, BaseModel, ConfigDict, ValidationError @@ -293,3 +295,45 @@ def fixture_id(case_input: FixtureInput, prefix: str) -> str: input_json: Final = canonical_json(case_input.canonical_input()) digest: Final = hashlib.sha256(input_json.encode("utf-8")).hexdigest()[:8] return f"{prefix}-{digest}" + + +def parametrize_recorded_fixtures( + metafunc: pytest.Metafunc, + *, + fixture_name: str, + case_type: type[CaseT], + env_var: str, + default_directory: Path, + regeneration_command: str, + id_builder: Callable[[CaseT], str], +) -> None: + if fixture_name not in metafunc.fixturenames: + return + configured: Final = os.environ.get(env_var) + if configured == "": + raise pytest.UsageError(f"{env_var} is set but empty") + directory: Final = Path(configured).expanduser() if configured is not None else default_directory + try: + fixtures: Final = recorded_fixtures(directory, case_type) + except (ValidationError, ValueError) as error: + raise pytest.UsageError( + f"Invalid parity fixture bundle at {directory}. " + "Each fixture must use the current versioned envelope. " + f"Record fresh fixtures in an empty directory with: `{regeneration_command}`. " + f"Validation details: {error}" + ) from error + if fixtures: + metafunc.parametrize(fixture_name, fixtures, ids=tuple(id_builder(fixture) for fixture in fixtures)) + return + if configured is not None: + raise pytest.UsageError(f"no recorded fixtures in {directory}") + metafunc.parametrize( + fixture_name, + ( + pytest.param( + None, + marks=pytest.mark.skip(reason=f"no recorded fixtures in {directory}"), + id="no-recorded-fixtures", + ), + ), + ) diff --git a/tests/test_litellm/ocr/conftest.py b/tests/test_litellm/ocr/conftest.py index ba2bd84a9d8..367e2f02405 100644 --- a/tests/test_litellm/ocr/conftest.py +++ b/tests/test_litellm/ocr/conftest.py @@ -1,27 +1,16 @@ from __future__ import annotations -import os from pathlib import Path from typing import Final import pytest -from pydantic import ValidationError -from tests.test_litellm._fixture_recorder import fixture_id, recorded_fixtures +from tests.test_litellm._fixture_recorder import fixture_id, parametrize_recorded_fixtures from tests.test_litellm.ocr.fixture_models import OcrParityCase FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" -def _fixture_directory() -> Path: - if FIXTURE_DIR_ENV not in os.environ: - return Path(__file__).with_name(".fixtures") - configured: Final = os.environ[FIXTURE_DIR_ENV] - if not configured: - raise pytest.UsageError(f"{FIXTURE_DIR_ENV} is set but empty") - return Path(configured).expanduser() - - def _fixture_id(fixture: OcrParityCase) -> str: case_input: Final = fixture.litellm_input model_id: Final = case_input.model.replace("/", "-") @@ -31,31 +20,16 @@ def _fixture_id(fixture: OcrParityCase) -> str: def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: - if "ocr_fixture" not in metafunc.fixturenames: - return - directory: Final = _fixture_directory() - try: - fixtures: Final = recorded_fixtures(directory, OcrParityCase) - except (ValidationError, ValueError) as error: - raise pytest.UsageError( - f"Invalid OCR parity fixture bundle at {directory}. " - "Each fixture must use the current versioned envelope. " - "Record fresh fixtures in an empty directory with: " - f"`uv run python tests/test_litellm/ocr/generate_fixtures.py --fixture-dir {directory}`. " - f"Validation details: {error}" - ) from error - if not fixtures: - if FIXTURE_DIR_ENV in os.environ: - raise pytest.UsageError(f"no recorded OCR fixtures in {directory}") - metafunc.parametrize( - "ocr_fixture", - ( - pytest.param( - None, - marks=pytest.mark.skip(reason=f"no recorded OCR fixtures in {directory}"), - id="no-recorded-fixtures", - ), - ), - ) - return - metafunc.parametrize("ocr_fixture", fixtures, ids=tuple(_fixture_id(fixture) for fixture in fixtures)) + default_directory: Final = Path(__file__).with_name(".fixtures") + parametrize_recorded_fixtures( + metafunc, + fixture_name="ocr_fixture", + case_type=OcrParityCase, + env_var=FIXTURE_DIR_ENV, + default_directory=default_directory, + regeneration_command=( + "uv run python tests/test_litellm/ocr/generate_fixtures.py " + f"--fixture-dir {default_directory}" + ), + id_builder=_fixture_id, + ) diff --git a/tests/test_litellm/ocr/fixture_models.py b/tests/test_litellm/ocr/fixture_models.py index 67070c33fde..9494c3690b1 100644 --- a/tests/test_litellm/ocr/fixture_models.py +++ b/tests/test_litellm/ocr/fixture_models.py @@ -2,39 +2,36 @@ from __future__ import annotations import base64 import binascii -from typing import Annotated, Literal, cast +from typing import Annotated, Literal -from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator +from pydantic import Field, field_validator, model_validator from typing_extensions import Self -from tests.test_litellm._recorded_http import RecordedResponse +from tests.test_litellm._fixture_models import ( + FixtureModel, + JsonObject, + JsonSchemaDefinition, + JsonSchemaResponseFormat, + ParityCase, + SdkInputBase, +) -JsonObject = dict[str, JsonValue] +__all__ = ("JsonSchemaDefinition", "JsonSchemaResponseFormat", "OcrParityCase", "OcrSdkInputBase") + +OcrSdkInputBase = SdkInputBase -class _FixtureModel(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True, serialize_by_alias=True) - - -class OcrSdkInputBase(_FixtureModel): - def as_sdk_kwargs(self) -> dict[str, object]: - return cast(dict[str, object], self.model_dump(mode="python", exclude_unset=True)) - - def canonical_input(self) -> dict[str, object]: - return cast(dict[str, object], self.model_dump(mode="json", exclude_unset=True)) - - -class MistralImageUrlValue(_FixtureModel): +class MistralImageUrlValue(FixtureModel): url: str detail: Literal["low", "auto", "high"] | None = None -class MistralImageUrlDocument(_FixtureModel): +class MistralImageUrlDocument(FixtureModel): type: Literal["image_url"] image_url: str | MistralImageUrlValue -class MistralDocumentUrlDocument(_FixtureModel): +class MistralDocumentUrlDocument(FixtureModel): type: Literal["document_url"] document_url: str document_name: str | None = None @@ -46,18 +43,6 @@ MistralDocument = Annotated[ ] -class JsonSchemaDefinition(_FixtureModel): - name: str - description: str | None = None - schema_definition: JsonObject = Field(alias="schema") - strict: bool = False - - -class JsonSchemaResponseFormat(_FixtureModel): - type: Literal["json_schema"] - json_schema: JsonSchemaDefinition - - MistralModel = Literal[ "mistral/mistral-ocr-2512", "mistral/mistral-ocr-4-0", @@ -117,7 +102,7 @@ def _validate_reducto_source(source: str) -> str: return source -class ReductoImageUrlDocument(_FixtureModel): +class ReductoImageUrlDocument(FixtureModel): type: Literal["image_url"] image_url: str @@ -127,7 +112,7 @@ class ReductoImageUrlDocument(_FixtureModel): return _validate_reducto_source(value) -class ReductoDocumentUrlDocument(_FixtureModel): +class ReductoDocumentUrlDocument(FixtureModel): type: Literal["document_url"] document_url: str @@ -167,7 +152,7 @@ ReductoBlockType = Literal[ ] -class ReductoFormatting(_FixtureModel): +class ReductoFormatting(FixtureModel): add_page_markers: bool = False table_output_format: ReductoTableOutputFormat = "dynamic" merge_tables: bool = False @@ -181,7 +166,7 @@ class ReductoFormatting(_FixtureModel): return value -class ReductoChunking(_FixtureModel): +class ReductoChunking(FixtureModel): chunk_mode: Literal["variable", "section", "page", "disabled", "block", "page_sections"] = "disabled" chunk_size: int | None = None chunk_overlap: int = Field(default=0, ge=0) @@ -195,7 +180,7 @@ class ReductoChunking(_FixtureModel): return self -class ReductoRetrieval(_FixtureModel): +class ReductoRetrieval(FixtureModel): chunking: ReductoChunking = Field(default_factory=ReductoChunking) filter_blocks: list[ReductoBlockType] = Field(default_factory=list) embedding_optimized: bool = False @@ -208,7 +193,7 @@ class ReductoRetrieval(_FixtureModel): return value -class ReductoPageRange(_FixtureModel): +class ReductoPageRange(FixtureModel): start: int | None = Field(default=None, ge=1) end: int | None = Field(default=None, ge=1) @@ -219,19 +204,19 @@ class ReductoPageRange(_FixtureModel): return self -class ReductoTenantThrottling(_FixtureModel): +class ReductoTenantThrottling(FixtureModel): tenant_id: str = Field(min_length=1, max_length=256) max_share: float = Field(default=0.5, gt=0, le=1) -class ReductoHybridVpcSettings(_FixtureModel): +class ReductoHybridVpcSettings(FixtureModel): environment: str | None = None ReductoPageSelection = ReductoPageRange | list[ReductoPageRange] | list[int] | list[str] -class ReductoSettings(_FixtureModel): +class ReductoSettings(FixtureModel): ocr_system: Literal["standard", "legacy"] = "standard" extraction_mode: Literal["ocr", "hybrid"] = "hybrid" force_url_result: bool = False @@ -285,6 +270,5 @@ class ReductoParseLegacySdkInput(OcrSdkInputBase): return self -class OcrParityCase(_FixtureModel): - litellm_input: MistralOcrSdkInput - provider_response: RecordedResponse +class OcrParityCase(ParityCase[MistralOcrSdkInput]): + pass diff --git a/tests/test_litellm/ocr/generate_fixtures.py b/tests/test_litellm/ocr/generate_fixtures.py index 9a63b533b5d..531d44feb00 100644 --- a/tests/test_litellm/ocr/generate_fixtures.py +++ b/tests/test_litellm/ocr/generate_fixtures.py @@ -1,10 +1,8 @@ from __future__ import annotations -import argparse import logging import os -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass +from collections.abc import Callable, Mapping from pathlib import Path from typing import Final, cast from urllib.parse import quote @@ -15,11 +13,17 @@ from hypothesis.strategies import DrawFn, SearchStrategy import litellm from litellm.rust_bridge.ocr import use_litellm_rust +from tests.test_litellm._fixture_generator import ( + FixtureTarget, + generate_target_fixtures, + parse_generator_args, +) +from tests.test_litellm._fixture_generator import ( + require_targets as require_fixture_targets, +) from tests.test_litellm._fixture_recorder import ( ProviderSpec, fixture_directory, - generate_case_inputs, - record_cases, ) from tests.test_litellm.ocr.fixture_models import ( JsonSchemaDefinition, @@ -38,26 +42,13 @@ from tests.test_litellm.ocr.fixture_models import ( ) FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" -LOGGER: Final = logging.getLogger(__name__) _TEXT: Final = st.just("invoice 123") _VALUE_TEXT: Final = st.just("case-1") _FONT_SIZE: Final = st.just(24) _MISTRAL_MODEL: Final = "mistral/mistral-ocr-latest" -@dataclass(frozen=True, slots=True) -class GeneratorArgs: - concurrency: int - examples: int - fixture_dir: Path | None - - -@dataclass(frozen=True, slots=True) -class OcrFixtureTarget: - name: str - provider_spec: ProviderSpec - strategy: SearchStrategy[OcrSdkInputBase] - invoke: Callable[[str, OcrSdkInputBase], object] +OcrFixtureTarget = FixtureTarget[OcrSdkInputBase] def _image_document(text: str, font_size: int) -> MistralImageUrlDocument: @@ -249,22 +240,7 @@ def _generate_examples( examples: int, concurrency: int, ) -> None: - case_inputs: Final = generate_case_inputs(target.strategy, examples) - results: Final = record_cases( - target.provider_spec, - root, - case_inputs, - target.invoke, - OcrParityCase, - concurrency, - ) - for result in results: - LOGGER.info( - "%s %s %s", - "cached" if result.cache_hit else "recorded", - target.name, - result.case.litellm_input.model, - ) + generate_target_fixtures(target, root, examples, concurrency, OcrParityCase) def _mistral_upstream_base(environ: Mapping[str, str]) -> str: @@ -300,21 +276,9 @@ def discover_targets( def require_targets(targets: tuple[OcrFixtureTarget, ...]) -> tuple[OcrFixtureTarget, ...]: - if targets: - return targets - raise SystemExit("No OCR fixture providers are configured. Set MISTRAL_API_KEY") - - -def parse_generator_args(argv: Sequence[str] | None = None) -> GeneratorArgs: - parser: Final = argparse.ArgumentParser() - parser.add_argument("--concurrency", type=int, default=4) - parser.add_argument("--examples", type=int, default=4) - parser.add_argument("--fixture-dir", type=Path) - namespace: Final = parser.parse_args(argv) - return GeneratorArgs( - concurrency=cast(int, namespace.concurrency), - examples=cast(int, namespace.examples), - fixture_dir=cast(Path | None, namespace.fixture_dir), + return require_fixture_targets( + targets, + "No OCR fixture providers are configured. Set MISTRAL_API_KEY", ) diff --git a/tests/test_litellm/ocr/test_sdk_parity.py b/tests/test_litellm/ocr/test_sdk_parity.py index 01ce7af10b1..7cc8c05f533 100644 --- a/tests/test_litellm/ocr/test_sdk_parity.py +++ b/tests/test_litellm/ocr/test_sdk_parity.py @@ -15,10 +15,10 @@ from tests.test_litellm.ocr.fixture_models import MistralOcrSdkInput, OcrParityC from tests.test_litellm.parity.compare import assert_parity from tests.test_litellm.parity.models import SDKCommand, SDKReport, WorkerFailure, WorkerResult, WorkerSuccess from tests.test_litellm.parity.runner import ( - WORKER_RESULT_PREFIX, PythonScriptRunner, PythonScriptWorker, - execution_worker, + execution_worker_pair, + parity_worker_main, run_execution, ) @@ -36,7 +36,7 @@ def _call_kwargs(sdk_input: MistralOcrSdkInput, mock_url: str, route: SDKRoute) **sdk_input.as_sdk_kwargs(), "api_base": mock_url, "api_key": API_KEY, - "extra_headers": {"x-ocr-parity-route": route.value}, + "extra_headers": {"x-litellm-parity-route": route.value}, } @@ -52,10 +52,10 @@ def _execute_sdk_case( if route is SDKRoute.OCR: sync_route: Final = cast(Callable[..., OCRResponse], litellm.ocr) response: Final = sync_route(**call_kwargs) - return SDKReport(response=response) + return SDKReport(response=response.model_dump(mode="json")) async_route: Final = cast(Callable[..., Coroutine[object, object, OCRResponse]], litellm.aocr) async_response: Final = event_loop.run_until_complete(async_route(**call_kwargs)) - return SDKReport(response=async_response) + return SDKReport(response=async_response.model_dump(mode="json")) @pytest.fixture(scope="module") @@ -64,10 +64,10 @@ def sdk_workers() -> Generator[tuple[PythonScriptWorker, PythonScriptWorker]]: entrypoint=Path(__file__), rust_env_var="LITELLM_USE_RUST_OCR", python_user_agent=PYTHON_HTTP_SENTINEL, + route_label="OCR", ) - with execution_worker(runner, rust_enabled=False) as python_worker: - with execution_worker(runner, rust_enabled=True) as rust_worker: - yield python_worker, rust_worker + with execution_worker_pair(runner) as workers: + yield workers @pytest.mark.parametrize("route", tuple(SDKRoute), ids=tuple(route.value for route in SDKRoute)) @@ -112,19 +112,7 @@ def _execute_worker_command( return WorkerFailure(error=traceback.format_exc()) -def _worker_main(mock_url: str) -> None: - event_loop: Final = asyncio.new_event_loop() - try: - for line in sys.stdin: - sys.stdout.write( - f"{WORKER_RESULT_PREFIX}{_execute_worker_command(line, mock_url, event_loop).model_dump_json()}\n" - ) - sys.stdout.flush() - finally: - event_loop.close() - - if __name__ == "__main__": if len(sys.argv) != 3 or sys.argv[1] != "--parity-worker": raise SystemExit("usage: test_sdk_parity.py --parity-worker MOCK_URL") - _worker_main(sys.argv[2]) + parity_worker_main(_execute_worker_command, sys.argv[2]) diff --git a/tests/test_litellm/parity/compare.py b/tests/test_litellm/parity/compare.py index 89e07ba717c..7386684c331 100644 --- a/tests/test_litellm/parity/compare.py +++ b/tests/test_litellm/parity/compare.py @@ -5,23 +5,23 @@ from typing import Final from tests.test_litellm.parity.models import CapturedRequest, Execution -def validate_harness(python: Execution, rust: Execution, python_user_agent: str) -> None: +def validate_harness(python: Execution, accelerated: Execution, python_user_agent: str) -> None: if python.request.user_agent != python_user_agent: raise AssertionError( f"Python provider request did not carry fallback sentinel user-agent {python_user_agent!r}: " f"{python.request.user_agent!r}" ) - if rust.request.user_agent == python_user_agent: - raise AssertionError("Rust OCR fell back to the Python HTTP implementation") + if accelerated.request.user_agent == python_user_agent: + raise AssertionError("accelerated route fell back to the Python HTTP implementation") def _request_after_transformation(request: CapturedRequest) -> CapturedRequest: return request.model_copy(update={"user_agent": None}) -def assert_parity(python: Execution, rust: Execution, python_user_agent: str) -> None: - validate_harness(python, rust, python_user_agent) +def assert_parity(python: Execution, accelerated: Execution, python_user_agent: str) -> None: + validate_harness(python, accelerated, python_user_agent) python_request: Final = _request_after_transformation(python.request) - rust_request: Final = _request_after_transformation(rust.request) - assert python_request == rust_request - assert python.report.response == rust.report.response + accelerated_request: Final = _request_after_transformation(accelerated.request) + assert python_request == accelerated_request + assert python.report.response == accelerated.report.response diff --git a/tests/test_litellm/parity/models.py b/tests/test_litellm/parity/models.py index 5747110b7e0..a776936832e 100644 --- a/tests/test_litellm/parity/models.py +++ b/tests/test_litellm/parity/models.py @@ -4,9 +4,6 @@ from typing import Annotated, Literal from pydantic import BaseModel, ConfigDict, Field, JsonValue -from litellm.llms.base_llm.ocr.transformation import OCRResponse - - class CapturedRequest(BaseModel): model_config = ConfigDict(frozen=True) @@ -20,7 +17,7 @@ class CapturedRequest(BaseModel): class SDKReport(BaseModel): model_config = ConfigDict(frozen=True) - response: OCRResponse + response: JsonValue class Execution(BaseModel): diff --git a/tests/test_litellm/parity/replay.py b/tests/test_litellm/parity/replay.py index c2228c70be6..e9cf70e009a 100644 --- a/tests/test_litellm/parity/replay.py +++ b/tests/test_litellm/parity/replay.py @@ -20,7 +20,7 @@ EXCLUDED_REQUEST_HEADERS: Final = frozenset( "connection", "accept-encoding", "user-agent", - "x-ocr-parity-route", + "x-litellm-parity-route", } ) EXCLUDED_RESPONSE_HEADERS: Final = frozenset({"content-length", "transfer-encoding", "connection"}) diff --git a/tests/test_litellm/parity/runner.py b/tests/test_litellm/parity/runner.py index eaec9bcb05d..d37c738627e 100644 --- a/tests/test_litellm/parity/runner.py +++ b/tests/test_litellm/parity/runner.py @@ -1,10 +1,11 @@ from __future__ import annotations +import asyncio import os import subprocess import sys from collections import deque -from collections.abc import Generator +from collections.abc import Callable, Generator from concurrent.futures import ThreadPoolExecutor, TimeoutError from contextlib import contextmanager from dataclasses import dataclass @@ -32,6 +33,7 @@ class PythonScriptRunner: entrypoint: Path rust_env_var: str python_user_agent: str + route_label: str def command(self, provider_url: str) -> tuple[str, ...]: return ( @@ -53,6 +55,7 @@ class PythonScriptWorker: "PYTHONPATH": os.pathsep.join(path for path in (project_root, existing_pythonpath) if path), } self.mode: Final = "Rust" if rust_enabled else "Python" + self.route_label: Final = runner.route_label self.provider: Final = provider self.process: Final = subprocess.Popen( runner.command(provider.url), @@ -74,7 +77,7 @@ class PythonScriptWorker: ) -> Execution: stdin: Final = self.process.stdin if stdin is None or self.process.poll() is not None: - raise AssertionError(f"{self.mode} OCR worker exited before processing {case_file}") + raise AssertionError(f"{self.mode} {self.route_label} worker exited before processing {case_file}") self.provider.enqueue_response(response) command: Final = SDKCommand(case_file=str(case_file), route=route) try: @@ -84,7 +87,9 @@ class PythonScriptWorker: except TimeoutError as error: self.provider.reset() self.close() - raise AssertionError(f"{self.mode} OCR worker timed out after 60s while processing {case_file}") from error + raise AssertionError( + f"{self.mode} {self.route_label} worker timed out after 60s while processing {case_file}" + ) from error except AssertionError: self.provider.reset() raise @@ -93,7 +98,9 @@ class PythonScriptWorker: raise AssertionError(self._failure_message(f"worker pipe failed while processing {case_file}")) from error if isinstance(result, WorkerFailure): self.provider.reset() - raise AssertionError(f"{self.mode} OCR worker failed while processing {case_file}:\n{result.error}") + raise AssertionError( + f"{self.mode} {self.route_label} worker failed while processing {case_file}:\n{result.error}" + ) assert isinstance(result, WorkerSuccess) try: return Execution(request=self.provider.take_request(), report=result.report) @@ -121,7 +128,8 @@ class PythonScriptWorker: def _failure_message(self, message: str) -> str: output: Final = "\n".join(self.recent_output) - return f"{self.mode} OCR {message}" if not output else f"{self.mode} OCR {message}\noutput:\n{output}" + prefix: Final = f"{self.mode} {self.route_label}" + return f"{prefix} {message}" if not output else f"{prefix} {message}\noutput:\n{output}" def close(self) -> None: stdin: Final = self.process.stdin @@ -155,3 +163,27 @@ def run_execution( response: RecordedResponse, ) -> Execution: return worker.execute(case_file, route, response) + + +@contextmanager +def execution_worker_pair( + runner: PythonScriptRunner, +) -> Generator[tuple[PythonScriptWorker, PythonScriptWorker]]: + with execution_worker(runner, rust_enabled=False) as python_worker: + with execution_worker(runner, rust_enabled=True) as accelerated_worker: + yield python_worker, accelerated_worker + + +def parity_worker_main( + execute_command: Callable[[str, str, asyncio.AbstractEventLoop], WorkerResult], + mock_url: str, +) -> None: + event_loop: Final = asyncio.new_event_loop() + try: + for line in sys.stdin: + sys.stdout.write( + f"{WORKER_RESULT_PREFIX}{execute_command(line, mock_url, event_loop).model_dump_json()}\n" + ) + sys.stdout.flush() + finally: + event_loop.close() diff --git a/tests/test_litellm/parity/test_parity.py b/tests/test_litellm/parity/test_parity.py index c67938c366f..03de1d76f43 100644 --- a/tests/test_litellm/parity/test_parity.py +++ b/tests/test_litellm/parity/test_parity.py @@ -5,28 +5,23 @@ from typing import Final import pytest from pydantic import JsonValue -from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse from tests.test_litellm.parity.compare import assert_parity from tests.test_litellm.parity.models import CapturedRequest, Execution, SDKReport -SENTINEL: Final = "python-ocr-parity-fallback" +SENTINEL: Final = "python-parity-fallback" def _execution(*, body: JsonValue = None, markdown: str = "same", user_agent: str | None = None) -> Execution: return Execution( request=CapturedRequest( method="POST", - path="/v1/ocr?mode=test", + path="/v1/test-route?mode=test", headers=(("authorization", "Bearer test-key"), ("content-type", "application/json")), - body={"model": "mistral-ocr-latest"} if body is None else body, + body={"model": "test-model"} if body is None else body, user_agent=user_agent, ), report=SDKReport( - response=OCRResponse( - pages=[OCRPage(index=0, markdown=markdown)], - model="mistral-ocr-latest", - object="ocr", - ) + response={"items": [{"text": markdown}], "model": "test-model"} ), )