test(ocr): generalize parity fixture inputs

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

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

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

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(),
+            }
+        )

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,

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,
+        }
     )