scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 66

refactor(tests): generalize route parity harness

Browse source
This commit is contained in:
Yujong Lee 2026-08-29 16:49:49 -07:00 • committed by GitHub
parent ce0dda405c
commit 4246fafcb9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 276 additions and 181 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

76
tests/test_litellm/_fixture_generator.py Normal file
View file

@ -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),
)

41
tests/test_litellm/_fixture_models.py Normal file
View file

@ -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

44
tests/test_litellm/_fixture_recorder.py
View file

@ -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",
),
),
)

54
tests/test_litellm/ocr/conftest.py
View file

@ -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,
)

70
tests/test_litellm/ocr/fixture_models.py
View file

@ -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

64
tests/test_litellm/ocr/generate_fixtures.py
View file

@ -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",
)

30
tests/test_litellm/ocr/test_sdk_parity.py
View file

@ -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])

16
tests/test_litellm/parity/compare.py
View file

@ -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

5
tests/test_litellm/parity/models.py
View file

@ -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):

2
tests/test_litellm/parity/replay.py
View file

@ -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"})

42
tests/test_litellm/parity/runner.py
View file

@ -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()

13
tests/test_litellm/parity/test_parity.py
View file

@ -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"}
),
)