diff --git a/tests/route_parity/README.md b/tests/route_parity/README.md index 3fc8e67b972..c063491bd78 100644 --- a/tests/route_parity/README.md +++ b/tests/route_parity/README.md @@ -8,7 +8,16 @@ - The same LiteLLM input is transformed by isolated Python and Rust workers - The resulting provider requests must match in method, path, headers, and body, excluding runtime-specific HTTP metadata - The recorded provider response is then replayed unchanged to both workers -- Each worker serializes its normalized LiteLLM SDK response to JSON, and the results must match +- The harness compares the values returned through the Python SDK interface +- Non-streaming responses are compared directly, including their concrete return type and public model fields +- Streaming responses are consumed and compared chunk by chunk, including wrapper type, chunk type and order, termination, and public exception behavior +- Route-specific comparators and chunk normalizers handle differences in each public SDK contract + +## Process isolation + +- SDK object and stream parity runs Python and Rust sequentially in the same process so tests can retain the returned objects +- Every test saves and restores the original bridge state +- A small subprocess smoke test verifies environment-based startup configuration and detects fallback to the Python HTTP implementation ## Hypothesis and property-based testing diff --git a/tests/route_parity/compare.py b/tests/route_parity/compare.py index e76079b6cb9..4fe9fdea9b4 100644 --- a/tests/route_parity/compare.py +++ b/tests/route_parity/compare.py @@ -1,9 +1,13 @@ from __future__ import annotations -from typing import Final +from typing import Final, TypeVar + +from pydantic import BaseModel from tests.route_parity.models import CapturedRequest, Execution +ModelT = TypeVar("ModelT", bound=BaseModel) + def validate_harness(python: Execution, accelerated: Execution, python_user_agent: str) -> None: if python.request.user_agent != python_user_agent: @@ -19,9 +23,24 @@ def _request_after_transformation(request: CapturedRequest) -> CapturedRequest: return request.model_copy(update={"user_agent": None}) +def assert_request_parity(python: CapturedRequest, accelerated: CapturedRequest) -> None: + python_request: Final = _request_after_transformation(python) + accelerated_request: Final = _request_after_transformation(accelerated) + assert python_request == accelerated_request + + +def public_model_copy(model: ModelT) -> ModelT: + copied: Final = model.model_copy(deep=True) + object.__setattr__(copied, "__pydantic_private__", None) + return copied + + +def assert_model_parity(python: BaseModel, accelerated: BaseModel) -> None: + assert type(python) is type(accelerated) + assert public_model_copy(python) == public_model_copy(accelerated) + + 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) - accelerated_request: Final = _request_after_transformation(accelerated.request) - assert python_request == accelerated_request + assert_request_parity(python.request, accelerated.request) assert python.report.response == accelerated.report.response diff --git a/tests/route_parity/inprocess.py b/tests/route_parity/inprocess.py new file mode 100644 index 00000000000..74abbae9b35 --- /dev/null +++ b/tests/route_parity/inprocess.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Generic, TypeVar + +from tests.route_parity.models import CapturedRequest +from tests.route_parity.recorded_http import RecordedResponse +from tests.route_parity.replay import ReplayServer + +ResponseT = TypeVar("ResponseT") + + +@dataclass(frozen=True, slots=True) +class InProcessExecution(Generic[ResponseT]): + request: CapturedRequest + response: ResponseT + + +def run_in_process( + provider: ReplayServer, + recorded_response: RecordedResponse, + call: Callable[[str], ResponseT], +) -> InProcessExecution[ResponseT]: + provider.enqueue_response(recorded_response) + try: + response = call(provider.url) + return InProcessExecution(request=provider.take_request(), response=response) + except Exception: + provider.reset() + raise diff --git a/tests/route_parity/stream.py b/tests/route_parity/stream.py new file mode 100644 index 00000000000..21997d76003 --- /dev/null +++ b/tests/route_parity/stream.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +from collections.abc import AsyncIterable, Awaitable, Callable, Iterable +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import BaseModel + +from tests.route_parity.compare import public_model_copy + + +@dataclass(frozen=True, slots=True) +class StreamCompleted: + kind: Literal["completed"] = "completed" + + +@dataclass(frozen=True, slots=True) +class StreamFailed: + phase: Literal["creation", "iteration"] + exception_type: type[BaseException] + status_code: int | None + llm_provider: str | None + model: str | None + kind: Literal["failed"] = "failed" + + +StreamTerminal: TypeAlias = StreamCompleted | StreamFailed + + +@dataclass(frozen=True, slots=True) +class StreamOutcome: + wrapper_type: type[object] | None + supports_sync_iteration: bool | None + supports_async_iteration: bool | None + chunks: tuple[object, ...] + chunk_types: tuple[type[object], ...] + terminal: StreamTerminal + + +ChunkNormalizer: TypeAlias = Callable[[object], object] + + +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) + model: Final = getattr(error, "model", None) + return StreamFailed( + phase=phase, + exception_type=type(error), + status_code=status_code if isinstance(status_code, int) else None, + llm_provider=llm_provider if isinstance(llm_provider, str) else None, + model=model if isinstance(model, str) else None, + ) + + +def consume_sync_stream(create: Callable[[], Iterable[object]]) -> StreamOutcome: + try: + stream = create() + except Exception as error: + return StreamOutcome( + wrapper_type=None, + supports_sync_iteration=None, + supports_async_iteration=None, + chunks=(), + chunk_types=(), + terminal=_failed("creation", error), + ) + + chunks: list[object] = [] # mutable-ok: iterator consumption builds an ordered trace + try: + for chunk in stream: + chunks.append(chunk) # noqa: PERF402 # partial trace is required if iteration raises + except Exception as error: + recorded: Final = tuple(chunks) + return StreamOutcome( + wrapper_type=type(stream), + supports_sync_iteration=hasattr(stream, "__iter__"), + supports_async_iteration=hasattr(stream, "__aiter__"), + chunks=recorded, + chunk_types=tuple(type(chunk) for chunk in recorded), + terminal=_failed("iteration", error), + ) + completed_chunks: Final = tuple(chunks) + return StreamOutcome( + wrapper_type=type(stream), + supports_sync_iteration=hasattr(stream, "__iter__"), + supports_async_iteration=hasattr(stream, "__aiter__"), + chunks=completed_chunks, + chunk_types=tuple(type(chunk) for chunk in completed_chunks), + terminal=StreamCompleted(), + ) + + +async def consume_async_stream(create: Callable[[], Awaitable[AsyncIterable[object]]]) -> StreamOutcome: + try: + stream = await create() + except Exception as error: + return StreamOutcome( + wrapper_type=None, + supports_sync_iteration=None, + supports_async_iteration=None, + chunks=(), + chunk_types=(), + terminal=_failed("creation", error), + ) + + chunks: list[object] = [] # mutable-ok: iterator consumption builds an ordered trace + try: + async for chunk in stream: + chunks.append(chunk) + except Exception as error: + recorded: Final = tuple(chunks) + return StreamOutcome( + wrapper_type=type(stream), + supports_sync_iteration=hasattr(stream, "__iter__"), + supports_async_iteration=hasattr(stream, "__aiter__"), + chunks=recorded, + chunk_types=tuple(type(chunk) for chunk in recorded), + terminal=_failed("iteration", error), + ) + completed_chunks: Final = tuple(chunks) + return StreamOutcome( + wrapper_type=type(stream), + supports_sync_iteration=hasattr(stream, "__iter__"), + supports_async_iteration=hasattr(stream, "__aiter__"), + chunks=completed_chunks, + chunk_types=tuple(type(chunk) for chunk in completed_chunks), + terminal=StreamCompleted(), + ) + + +def normalize_chunk(chunk: object) -> object: + if isinstance(chunk, BaseModel): + return public_model_copy(chunk) + return chunk + + +def assert_stream_parity( + python: StreamOutcome, + accelerated: StreamOutcome, + *, + normalize: ChunkNormalizer = normalize_chunk, +) -> None: + assert python.wrapper_type is accelerated.wrapper_type + assert python.supports_sync_iteration is accelerated.supports_sync_iteration + assert python.supports_async_iteration is accelerated.supports_async_iteration + assert python.chunk_types == accelerated.chunk_types + assert len(python.chunks) == len(accelerated.chunks) + for python_chunk, accelerated_chunk in zip(python.chunks, accelerated.chunks, strict=True): + assert normalize(python_chunk) == normalize(accelerated_chunk) + assert python.terminal == accelerated.terminal diff --git a/tests/route_parity/test_parity.py b/tests/route_parity/test_parity.py index 3deeec47032..c92b7bc0298 100644 --- a/tests/route_parity/test_parity.py +++ b/tests/route_parity/test_parity.py @@ -3,14 +3,26 @@ from __future__ import annotations from typing import Final import pytest -from pydantic import JsonValue +from pydantic import BaseModel, JsonValue, PrivateAttr -from tests.route_parity.compare import assert_parity +from tests.route_parity.compare import assert_model_parity, assert_parity from tests.route_parity.models import CapturedRequest, Execution, SDKReport SENTINEL: Final = "python-parity-fallback" +class _ComparableResponse(BaseModel): + value: str + _hidden_params: dict[str, object] = PrivateAttr(default_factory=dict) + + def set_hidden_param(self, key: str, value: object) -> None: + self._hidden_params[key] = value + + +class _DifferentResponse(BaseModel): + value: str + + def _execution(*, body: JsonValue = None, markdown: str = "same", user_agent: str | None = None) -> Execution: return Execution( request=CapturedRequest( @@ -46,3 +58,25 @@ def test_parity_rejects_rust_fallback() -> None: with pytest.raises(AssertionError, match="fell back"): assert_parity(python, rust, SENTINEL) + + +def test_model_parity_compares_public_values_and_ignores_private_attrs() -> None: + python: Final = _ComparableResponse(value="same") + rust: Final = _ComparableResponse(value="same") + python.set_hidden_param("litellm_call_id", "python-id") + rust.set_hidden_param("litellm_call_id", "rust-id") + + assert_model_parity(python, rust) + + +def test_model_parity_rejects_public_value_difference() -> None: + python: Final = _ComparableResponse(value="python") + rust: Final = _ComparableResponse(value="rust") + + with pytest.raises(AssertionError): + assert_model_parity(python, rust) + + +def test_model_parity_rejects_type_difference() -> None: + with pytest.raises(AssertionError): + assert_model_parity(_ComparableResponse(value="same"), _DifferentResponse(value="same")) diff --git a/tests/route_parity/test_stream.py b/tests/route_parity/test_stream.py new file mode 100644 index 00000000000..7cb6796c074 --- /dev/null +++ b/tests/route_parity/test_stream.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator, Iterator +from typing import Final + +import pytest +from pydantic import BaseModel, PrivateAttr + +from tests.route_parity.stream import ( + StreamFailed, + assert_stream_parity, + consume_async_stream, + consume_sync_stream, +) + + +class _Chunk(BaseModel): + value: str + _hidden_params: dict[str, object] = PrivateAttr(default_factory=dict) + + def set_hidden_param(self, key: str, value: object) -> None: + self._hidden_params[key] = value + + +class _SyncStream: + def __init__(self, chunks: tuple[_Chunk, ...], error: BaseException | None = None) -> None: + self.chunks: Final = chunks + self.error: Final = error + + def __iter__(self) -> Iterator[object]: + yield from self.chunks + if self.error is not None: + raise self.error + + +class _AsyncStream: + def __init__(self, chunks: tuple[_Chunk, ...], error: BaseException | None = None) -> None: + self.chunks: Final = chunks + self.error: Final = error + + async def __aiter__(self) -> AsyncIterator[object]: + for chunk in self.chunks: + yield chunk + if self.error is not None: + raise self.error + + +class _PublicStreamError(Exception): + def __init__( + self, message: str = "runtime-specific message", *, status_code: int, llm_provider: str, model: str + ) -> None: + super().__init__(message) + self.status_code: Final = status_code + self.llm_provider: Final = llm_provider + self.model: Final = model + + +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: + return _AsyncStream(chunks, error) + + +def test_sync_stream_parity_compares_chunks_and_ignores_private_attrs() -> None: + python_chunk: Final = _Chunk(value="same") + accelerated_chunk: Final = _Chunk(value="same") + python_chunk.set_hidden_param("request_id", "python") + accelerated_chunk.set_hidden_param("request_id", "accelerated") + python: Final = consume_sync_stream(lambda: _SyncStream((python_chunk,))) + accelerated: Final = consume_sync_stream(lambda: _SyncStream((accelerated_chunk,))) + + assert python.supports_sync_iteration is True + assert python.supports_async_iteration is False + assert_stream_parity(python, accelerated) + + +def test_stream_parity_rejects_extra_chunk() -> None: + python: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="one"),))) + accelerated: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="one"), _Chunk(value="two")))) + + with pytest.raises(AssertionError): + assert_stream_parity(python, accelerated) + + +def test_stream_parity_rejects_chunk_value_difference() -> None: + python: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="python"),))) + accelerated: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="accelerated"),))) + + with pytest.raises(AssertionError): + assert_stream_parity(python, accelerated) + + +def test_stream_outcome_distinguishes_creation_and_iteration_errors() -> None: + creation: Final = consume_sync_stream(_creation_error) + iteration: Final = consume_sync_stream( + lambda: _SyncStream( + (_Chunk(value="before-error"),), + _PublicStreamError(status_code=429, llm_provider="test", model="test-model"), + ) + ) + + assert isinstance(creation.terminal, StreamFailed) + assert creation.terminal.phase == "creation" + assert creation.chunks == () + assert isinstance(iteration.terminal, StreamFailed) + assert iteration.terminal.phase == "iteration" + assert len(iteration.chunks) == 1 + with pytest.raises(AssertionError): + assert_stream_parity(creation, iteration) + + +@pytest.mark.asyncio +async def test_async_stream_uses_same_trace_contract() -> None: + python_error: Final = _PublicStreamError( + "python runtime detail", status_code=500, llm_provider="test", model="test-model" + ) + accelerated_error: Final = _PublicStreamError( + "rust runtime detail", status_code=500, llm_provider="test", model="test-model" + ) + python: Final = await consume_async_stream(lambda: _async_stream((_Chunk(value="same"),), python_error)) + accelerated: Final = await consume_async_stream(lambda: _async_stream((_Chunk(value="same"),), accelerated_error)) + + assert python.supports_sync_iteration is False + assert python.supports_async_iteration is True + assert_stream_parity(python, accelerated) + + +def test_stream_parity_accepts_route_specific_chunk_normalizer() -> None: + python: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="python-generated-id"),))) + accelerated: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="rust-generated-id"),))) + + assert_stream_parity(python, accelerated, normalize=lambda chunk: type(chunk)) diff --git a/tests/test_litellm/ocr/test_sdk_parity.py b/tests/test_litellm/ocr/test_sdk_parity.py index 577d8b48522..97d836e849c 100644 --- a/tests/test_litellm/ocr/test_sdk_parity.py +++ b/tests/test_litellm/ocr/test_sdk_parity.py @@ -1,9 +1,11 @@ from __future__ import annotations import asyncio +import os import sys import traceback -from collections.abc import Callable, Coroutine, Generator +from collections.abc import Awaitable, Callable, Coroutine, Generator +from contextlib import contextmanager from enum import Enum from pathlib import Path from typing import Final, cast @@ -11,8 +13,14 @@ from typing import Final, cast import pytest from litellm.llms.base_llm.ocr.transformation import OCRResponse -from tests.route_parity.compare import assert_parity +from litellm.rust_bridge import get_native_bridge +from litellm.rust_bridge import ocr as rust_ocr_bridge +from litellm.rust_bridge.ocr import RustAocr, RustOcr +from tests.route_parity.compare import assert_model_parity, assert_parity, assert_request_parity +from tests.route_parity.fixture_recorder import recorded_fixtures +from tests.route_parity.inprocess import run_in_process from tests.route_parity.models import SDKCommand, SDKReport, WorkerFailure, WorkerResult, WorkerSuccess +from tests.route_parity.replay import replay_server from tests.route_parity.runner import ( PythonScriptRunner, PythonScriptWorker, @@ -24,6 +32,7 @@ from tests.test_litellm.ocr.fixture_models import MistralOcrSdkInput, OcrParityC API_KEY: Final = "test-key" PYTHON_HTTP_SENTINEL: Final = "python-ocr-parity-fallback" +FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" class SDKRoute(str, Enum): @@ -58,6 +67,96 @@ def _execute_sdk_case( return SDKReport(response=async_response.model_dump(mode="json")) +def _call_sdk_case(sdk_input: MistralOcrSdkInput, route: SDKRoute, mock_url: str) -> OCRResponse: + import litellm + + call_kwargs: Final = _call_kwargs(sdk_input, mock_url, route) + if route is SDKRoute.OCR: + sync_route: Final = cast(Callable[..., OCRResponse], litellm.ocr) + return sync_route(**call_kwargs) + async_route: Final = cast(Callable[..., Coroutine[object, object, OCRResponse]], litellm.aocr) + return asyncio.run(async_route(**call_kwargs)) + + +class _RustOcrSpy: + def __init__(self, delegate: RustOcr) -> None: + self.delegate: Final = delegate + self.calls = 0 + + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + self.calls += 1 + return self.delegate( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_seconds, + ) + + +class _RustAocrSpy: + def __init__(self, delegate: RustAocr) -> None: + self.delegate: Final = delegate + self.calls = 0 + + async def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + self.calls += 1 + result: Final[Awaitable[dict[str, object]]] = self.delegate( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_seconds, + ) + return await result + + +@contextmanager +def _restore_rust_ocr_state() -> Generator[None]: + enabled: Final = rust_ocr_bridge._rust_ocr_enabled # pyright: ignore[reportPrivateUsage] # restore test state + ocr_impl: Final = rust_ocr_bridge._rust_ocr_impl # pyright: ignore[reportPrivateUsage] # restore test state + aocr_impl: Final = rust_ocr_bridge._rust_aocr_impl # pyright: ignore[reportPrivateUsage] # restore test state + try: + yield + finally: + rust_ocr_bridge.use_litellm_rust(enabled, ocr=ocr_impl, aocr=aocr_impl) + + +def _native_spies() -> tuple[_RustOcrSpy, _RustAocrSpy]: + native_bridge: Final = get_native_bridge() + if native_bridge is None: + pytest.fail("native Rust bridge is required for OCR parity testing") + sync_spy: Final = _RustOcrSpy(cast(RustOcr, getattr(native_bridge, "ocr"))) + async_spy: Final = _RustAocrSpy(cast(RustAocr, getattr(native_bridge, "aocr"))) + return sync_spy, async_spy + + @pytest.fixture(scope="module") def sdk_workers() -> Generator[tuple[PythonScriptWorker, PythonScriptWorker]]: runner: Final = PythonScriptRunner( @@ -70,28 +169,65 @@ def sdk_workers() -> Generator[tuple[PythonScriptWorker, PythonScriptWorker]]: yield workers +@pytest.fixture(scope="module") +def startup_ocr_fixture() -> OcrParityCase: + default_directory: Final = Path(__file__).with_name("fixtures") + configured: Final = os.environ.get(FIXTURE_DIR_ENV) + directory: Final = Path(configured).expanduser() if configured is not None else default_directory + fixtures: Final = recorded_fixtures(directory, OcrParityCase) + if not fixtures: + pytest.skip(f"no recorded fixtures in {directory}") + return fixtures[0] + + @pytest.mark.parametrize("route", tuple(SDKRoute), ids=tuple(route.value for route in SDKRoute)) def test_recorded_ocr_sdk_parity( ocr_fixture: OcrParityCase, route: SDKRoute, +) -> None: + sync_spy, async_spy = _native_spies() + with _restore_rust_ocr_state(), replay_server() as provider: + rust_ocr_bridge.use_litellm_rust(False, ocr=sync_spy, aocr=async_spy) + python: Final = run_in_process( + provider, + ocr_fixture.provider_response, + lambda mock_url: _call_sdk_case(ocr_fixture.litellm_input, route, mock_url), + ) + assert sync_spy.calls == 0 + assert async_spy.calls == 0 + + rust_ocr_bridge.use_litellm_rust(True, ocr=sync_spy, aocr=async_spy) + rust: Final = run_in_process( + provider, + ocr_fixture.provider_response, + lambda mock_url: _call_sdk_case(ocr_fixture.litellm_input, route, mock_url), + ) + + assert sync_spy.calls == (1 if route is SDKRoute.OCR else 0) + assert async_spy.calls == (1 if route is SDKRoute.AOCR else 0) + assert_request_parity(python.request, rust.request) + assert_model_parity(python.response, rust.response) + + +def test_ocr_subprocess_startup_smoke( + startup_ocr_fixture: OcrParityCase, tmp_path: Path, sdk_workers: tuple[PythonScriptWorker, PythonScriptWorker], ) -> None: - case_file: Final = tmp_path / f"{route.value}-ocr-parity-case.json" - case_file.write_text(ocr_fixture.model_dump_json(indent=2, exclude_unset=True), encoding="utf-8") - response: Final = ocr_fixture.provider_response + case_file: Final = tmp_path / "ocr-startup-smoke.json" + case_file.write_text(startup_ocr_fixture.model_dump_json(indent=2, exclude_unset=True), encoding="utf-8") python_worker, rust_worker = sdk_workers python: Final = run_execution( python_worker, case_file, - route.value, - response, + SDKRoute.OCR.value, + startup_ocr_fixture.provider_response, ) rust: Final = run_execution( rust_worker, case_file, - route.value, - response, + SDKRoute.OCR.value, + startup_ocr_fixture.provider_response, ) assert_parity(python, rust, PYTHON_HTTP_SENTINEL)