diff --git a/tests/test_litellm/_fixture_recorder.py b/tests/test_litellm/_fixture_recorder.py index 94fecac85a1..85d9d9e9c96 100644 --- a/tests/test_litellm/_fixture_recorder.py +++ b/tests/test_litellm/_fixture_recorder.py @@ -1,6 +1,5 @@ from __future__ import annotations -import argparse import hashlib import queue import threading @@ -18,7 +17,7 @@ from pydantic import ValidationError from tests.test_litellm._json_fs_cache import JsonFileCache, canonical_json from tests.test_litellm.ocr.fixture_models import ( HttpHeader, - MistralOcrParityInput, + LiteLLMOcrInput, OcrParityCase, RecordedHttpResponse, ) @@ -50,14 +49,6 @@ class RecorderResult: cache_hit: bool -@dataclass(frozen=True, slots=True) -class GeneratorArgs: - concurrency: int - examples: int - fixture_dir: Path | None - model: str - - def _excluded_headers(headers: tuple[tuple[str, str], ...]) -> frozenset[str]: connection_values: Final = tuple(value for name, value in headers if name.lower() == "connection") connection_headers: Final = frozenset( @@ -154,14 +145,14 @@ def _recording_provider(spec: ProviderSpec) -> Generator[_RecordingProvider]: thread.join(timeout=5) -def fixture_cache_key(case_input: MistralOcrParityInput) -> dict[str, object]: +def fixture_cache_key(case_input: LiteLLMOcrInput) -> dict[str, object]: return case_input.canonical_input() def record_case( spec: ProviderSpec, root: Path, - case_input: MistralOcrParityInput, + case_input: LiteLLMOcrInput, sdk_call: Callable[..., object], ) -> RecorderResult: cache: Final = JsonFileCache(root) @@ -182,7 +173,7 @@ def record_case( def record_cases( spec: ProviderSpec, root: Path, - case_inputs: tuple[MistralOcrParityInput, ...], + case_inputs: tuple[LiteLLMOcrInput, ...], sdk_call: Callable[..., object], max_concurrency: int, ) -> tuple[RecorderResult, ...]: @@ -198,21 +189,6 @@ def record_cases( return tuple(future.result() for future in futures) -def parse_generator_args() -> 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) - parser.add_argument("--model", default="mistral/mistral-ocr-latest") - namespace: Final = parser.parse_args() - return GeneratorArgs( - concurrency=cast(int, namespace.concurrency), - examples=cast(int, namespace.examples), - fixture_dir=cast(Path | None, namespace.fixture_dir), - model=cast(str, namespace.model), - ) - - def fixture_directory(configured: Path | None, env_value: str | None, default: Path) -> Path: return (configured or Path(env_value or default)).expanduser() @@ -234,4 +210,7 @@ def recorded_fixtures(directory: Path) -> tuple[OcrParityCase, ...]: def fixture_id(fixture: OcrParityCase) -> str: input_json: Final = canonical_json(fixture.litellm_input.canonical_input()) digest: Final = hashlib.sha256(input_json.encode("utf-8")).hexdigest()[:8] - return f"mistral-{fixture.litellm_input.model.rsplit('/', 1)[-1]}-{digest}" + model_id: Final = fixture.litellm_input.model.replace("/", "-") + provider: Final = fixture.litellm_input.custom_llm_provider + prefix: Final = f"{provider}-{model_id}" if provider and not model_id.startswith(f"{provider}-") else model_id + return f"{prefix}-{digest}" diff --git a/tests/test_litellm/ocr/fixture_models.py b/tests/test_litellm/ocr/fixture_models.py index ab4432cc349..bd07de2ce73 100644 --- a/tests/test_litellm/ocr/fixture_models.py +++ b/tests/test_litellm/ocr/fixture_models.py @@ -3,7 +3,9 @@ from __future__ import annotations import base64 from typing import Annotated, Literal, cast -from pydantic import BaseModel, ConfigDict, Field, JsonValue, NonNegativeInt, PositiveInt +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +JsonObject = dict[str, JsonValue] class _FixtureModel(BaseModel): @@ -20,36 +22,18 @@ class DocumentUrlDocument(_FixtureModel): document_url: str -MistralOcrDocument = Annotated[ImageUrlDocument | DocumentUrlDocument, Field(discriminator="type")] +OcrDocument = Annotated[ImageUrlDocument | DocumentUrlDocument, Field(discriminator="type")] -class JsonSchemaDefinition(_FixtureModel): - name: str - schema_value: JsonValue = Field(alias="schema") - strict: bool | None = None +class LiteLLMOcrInput(_FixtureModel): + model_config = ConfigDict(frozen=True, extra="allow", populate_by_name=True, serialize_by_alias=True) - -class AnnotationFormat(_FixtureModel): - type: Literal["json_schema"] - json_schema: JsonSchemaDefinition - - -class MistralOcrParityInput(_FixtureModel): + __pydantic_extra__: JsonObject = Field( # pyright: ignore[reportIncompatibleVariableOverride] # Pydantic typed extras + init=False + ) model: str - document: MistralOcrDocument - pages: list[NonNegativeInt] | None = None - include_image_base64: bool | None = None - image_limit: PositiveInt | None = None - image_min_size: NonNegativeInt | None = None - bbox_annotation_format: AnnotationFormat | None = None - document_annotation_format: AnnotationFormat | None = None - document_annotation_prompt: str | None = None - extract_header: bool | None = None - extract_footer: bool | None = None - table_format: Literal["markdown", "html"] | None = None - confidence_scores_granularity: Literal["word", "page", "block"] | None = None - include_blocks: bool | None = None - id: str | None = None + document: OcrDocument + custom_llm_provider: str | None = None def as_sdk_kwargs(self) -> dict[str, object]: return cast(dict[str, object], self.model_dump(mode="python", exclude_unset=True)) @@ -87,5 +71,5 @@ class RecordedHttpResponse(_FixtureModel): class OcrParityCase(_FixtureModel): - litellm_input: MistralOcrParityInput + litellm_input: LiteLLMOcrInput provider_response: RecordedHttpResponse diff --git a/tests/test_litellm/ocr/generate_fixtures.py b/tests/test_litellm/ocr/generate_fixtures.py index 0cdd9237680..a15b80c48ca 100644 --- a/tests/test_litellm/ocr/generate_fixtures.py +++ b/tests/test_litellm/ocr/generate_fixtures.py @@ -1,9 +1,11 @@ from __future__ import annotations +import argparse import logging import os import queue from collections.abc import Callable +from dataclasses import dataclass from pathlib import Path from typing import Final, cast from urllib.parse import quote @@ -18,12 +20,11 @@ from litellm.rust_bridge.ocr import use_litellm_rust from tests.test_litellm._fixture_recorder import ( ProviderSpec, fixture_directory, - parse_generator_args, record_cases, ) from tests.test_litellm.ocr.fixture_models import ( ImageUrlDocument, - MistralOcrParityInput, + LiteLLMOcrInput, ) FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" @@ -32,12 +33,20 @@ _TEXT: Final = st.from_regex(r"[A-Za-z0-9 ]{1,24}", fullmatch=True) _VALUE_TEXT: Final = st.text(alphabet="abcdefghijklmnopqrstuvwxyz0123456789 -_", min_size=1, max_size=32) +@dataclass(frozen=True, slots=True) +class GeneratorArgs: + concurrency: int + examples: int + fixture_dir: Path | None + model: str + + def _image_document(text: str, font_size: int) -> ImageUrlDocument: url: Final = f"https://dummyjson.com/image/800x300/ffffff/000000?text={quote(text)}&fontSize={font_size}" return ImageUrlDocument(type="image_url", image_url=url) -def _input_strategy(model: str) -> SearchStrategy[MistralOcrParityInput]: +def _mistral_input_strategy(model: str) -> SearchStrategy[LiteLLMOcrInput]: document_strategy: Final = st.builds(_image_document, _TEXT, st.integers(min_value=12, max_value=36)) input_values: Final = st.fixed_dictionaries( {"model": st.just(model), "document": document_strategy}, @@ -52,7 +61,7 @@ def _input_strategy(model: str) -> SearchStrategy[MistralOcrParityInput]: "id": _VALUE_TEXT, }, ) - return input_values.map(MistralOcrParityInput.model_validate) + return input_values.map(LiteLLMOcrInput.model_validate) def _generate_examples( @@ -62,11 +71,11 @@ def _generate_examples( concurrency: int, sdk_call: Callable[..., object], ) -> None: - generated: Final[queue.SimpleQueue[MistralOcrParityInput | None]] = queue.SimpleQueue() + generated: Final[queue.SimpleQueue[LiteLLMOcrInput | None]] = queue.SimpleQueue() @settings(max_examples=examples, deadline=None, derandomize=True) - @given(case_input=_input_strategy(spec.model)) - def generate_case(case_input: MistralOcrParityInput) -> None: + @given(case_input=_mistral_input_strategy(spec.model)) + def generate_case(case_input: LiteLLMOcrInput) -> None: generated.put(case_input) generate_case() @@ -82,10 +91,25 @@ def _mistral_upstream_base() -> str: return configured.removesuffix("/v1") +def _parse_args() -> 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) + parser.add_argument("--model", default="mistral/mistral-ocr-latest") + namespace: Final = parser.parse_args() + return GeneratorArgs( + concurrency=cast(int, namespace.concurrency), + examples=cast(int, namespace.examples), + fixture_dir=cast(Path | None, namespace.fixture_dir), + model=cast(str, namespace.model), + ) + + def main() -> None: logging.basicConfig(level=logging.INFO, format="%(message)s") load_dotenv() - args: Final = parse_generator_args() + args: Final = _parse_args() api_key: Final = os.environ.get("MISTRAL_API_KEY") or os.environ.get("LITELLM_API_KEY") if api_key is None: raise SystemExit("MISTRAL_API_KEY is required") diff --git a/tests/test_litellm/ocr/test_fixture_models.py b/tests/test_litellm/ocr/test_fixture_models.py new file mode 100644 index 00000000000..fe452833c64 --- /dev/null +++ b/tests/test_litellm/ocr/test_fixture_models.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +from typing import Final + +import pytest +from pydantic import ValidationError + +from tests.test_litellm.ocr.fixture_models import LiteLLMOcrInput + + +@pytest.mark.parametrize( + ("raw_input", "expected_params"), + ( + ( + { + "model": "azure_ai/doc-intelligence/prebuilt-layout", + "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + "pages": "1-3,5", + "features": ["keyValuePairs", "languages"], + }, + {"pages": "1-3,5", "features": ["keyValuePairs", "languages"]}, + ), + ( + { + "model": "reducto/parse-v3", + "document": {"type": "document_url", "document_url": "reducto://fixture-file"}, + "formatting": {"table_output_format": "html"}, + "retrieval": {"chunking": {"chunk_mode": "variable"}}, + }, + { + "formatting": {"table_output_format": "html"}, + "retrieval": {"chunking": {"chunk_mode": "variable"}}, + }, + ), + ), +) +def test_litellm_ocr_input_preserves_provider_params( + raw_input: dict[str, object], expected_params: dict[str, object] +) -> None: + fixture_input: Final = LiteLLMOcrInput.model_validate(raw_input) + sdk_kwargs: Final = fixture_input.as_sdk_kwargs() + canonical_input: Final = fixture_input.canonical_input() + + assert {name: sdk_kwargs[name] for name in expected_params} == expected_params + assert {name: canonical_input[name] for name in expected_params} == expected_params + + +def test_litellm_ocr_input_rejects_non_json_provider_params() -> None: + with pytest.raises(ValidationError): + LiteLLMOcrInput.model_validate( + { + "model": "reducto/parse-v3", + "document": {"type": "document_url", "document_url": "reducto://fixture-file"}, + "settings": object(), + } + ) diff --git a/tests/test_litellm/ocr/test_sdk_parity.py b/tests/test_litellm/ocr/test_sdk_parity.py index 9b855e08bef..4b3ce140007 100644 --- a/tests/test_litellm/ocr/test_sdk_parity.py +++ b/tests/test_litellm/ocr/test_sdk_parity.py @@ -11,7 +11,7 @@ from typing import Final, cast import pytest from litellm.llms.base_llm.ocr.transformation import OCRResponse -from tests.test_litellm.ocr.fixture_models import MistralOcrParityInput, OcrParityCase +from tests.test_litellm.ocr.fixture_models import LiteLLMOcrInput, OcrParityCase 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 ( @@ -31,7 +31,7 @@ class SDKRoute(str, Enum): AOCR = "aocr" -def _call_kwargs(sdk_input: MistralOcrParityInput, mock_url: str, route: SDKRoute) -> dict[str, object]: +def _call_kwargs(sdk_input: LiteLLMOcrInput, mock_url: str, route: SDKRoute) -> dict[str, object]: return { **sdk_input.as_sdk_kwargs(), "api_base": mock_url, @@ -41,7 +41,7 @@ def _call_kwargs(sdk_input: MistralOcrParityInput, mock_url: str, route: SDKRout def _execute_sdk_case( - sdk_input: MistralOcrParityInput, + sdk_input: LiteLLMOcrInput, route: SDKRoute, mock_url: str, event_loop: asyncio.AbstractEventLoop, diff --git a/tests/test_litellm/test__fixture_recorder.py b/tests/test_litellm/test__fixture_recorder.py index f457f97ee31..1f058bb2a2b 100644 --- a/tests/test_litellm/test__fixture_recorder.py +++ b/tests/test_litellm/test__fixture_recorder.py @@ -10,7 +10,7 @@ from typing import Final, cast import httpx from tests.test_litellm._fixture_recorder import ProviderSpec, record_cases -from tests.test_litellm.ocr.fixture_models import ImageUrlDocument, MistralOcrParityInput +from tests.test_litellm.ocr.fixture_models import ImageUrlDocument, LiteLLMOcrInput class _ControlledUpstream(ThreadingHTTPServer): @@ -78,11 +78,13 @@ def _controlled_upstream() -> Generator[_ControlledUpstream]: thread.join(timeout=5) -def _case(identifier: str) -> MistralOcrParityInput: - return MistralOcrParityInput( - model="mistral/mistral-ocr-latest", - document=ImageUrlDocument(type="image_url", image_url="https://example.com/image.png"), - id=identifier, +def _case(identifier: str) -> LiteLLMOcrInput: + return LiteLLMOcrInput.model_validate( + { + "model": "mistral/mistral-ocr-latest", + "document": ImageUrlDocument(type="image_url", image_url="https://example.com/image.png"), + "id": identifier, + } )