diff --git a/tests/test_litellm/_fixture_recorder.py b/tests/test_litellm/_fixture_recorder.py index 85d9d9e9c96..e40a915e6f3 100644 --- a/tests/test_litellm/_fixture_recorder.py +++ b/tests/test_litellm/_fixture_recorder.py @@ -17,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, - LiteLLMOcrInput, + MistralOcrSdkInput, OcrParityCase, RecordedHttpResponse, ) @@ -145,14 +145,14 @@ def _recording_provider(spec: ProviderSpec) -> Generator[_RecordingProvider]: thread.join(timeout=5) -def fixture_cache_key(case_input: LiteLLMOcrInput) -> dict[str, object]: +def fixture_cache_key(case_input: MistralOcrSdkInput) -> dict[str, object]: return case_input.canonical_input() def record_case( spec: ProviderSpec, root: Path, - case_input: LiteLLMOcrInput, + case_input: MistralOcrSdkInput, sdk_call: Callable[..., object], ) -> RecorderResult: cache: Final = JsonFileCache(root) @@ -173,7 +173,7 @@ def record_case( def record_cases( spec: ProviderSpec, root: Path, - case_inputs: tuple[LiteLLMOcrInput, ...], + case_inputs: tuple[MistralOcrSdkInput, ...], sdk_call: Callable[..., object], max_concurrency: int, ) -> tuple[RecorderResult, ...]: diff --git a/tests/test_litellm/ocr/fixture_models.py b/tests/test_litellm/ocr/fixture_models.py index bd07de2ce73..dc247e6a145 100644 --- a/tests/test_litellm/ocr/fixture_models.py +++ b/tests/test_litellm/ocr/fixture_models.py @@ -1,9 +1,11 @@ from __future__ import annotations import base64 +import binascii from typing import Annotated, Literal, cast -from pydantic import BaseModel, ConfigDict, Field, JsonValue +from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator +from typing_extensions import Self JsonObject = dict[str, JsonValue] @@ -12,29 +14,7 @@ class _FixtureModel(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True, serialize_by_alias=True) -class ImageUrlDocument(_FixtureModel): - type: Literal["image_url"] - image_url: str - - -class DocumentUrlDocument(_FixtureModel): - type: Literal["document_url"] - document_url: str - - -OcrDocument = Annotated[ImageUrlDocument | DocumentUrlDocument, Field(discriminator="type")] - - -class LiteLLMOcrInput(_FixtureModel): - model_config = ConfigDict(frozen=True, extra="allow", populate_by_name=True, serialize_by_alias=True) - - __pydantic_extra__: JsonObject = Field( # pyright: ignore[reportIncompatibleVariableOverride] # Pydantic typed extras - init=False - ) - model: str - document: OcrDocument - custom_llm_provider: str | None = None - +class OcrSdkInputBase(_FixtureModel): def as_sdk_kwargs(self) -> dict[str, object]: return cast(dict[str, object], self.model_dump(mode="python", exclude_unset=True)) @@ -42,6 +22,267 @@ class LiteLLMOcrInput(_FixtureModel): return cast(dict[str, object], self.model_dump(mode="json", exclude_unset=True)) +class MistralImageUrlValue(_FixtureModel): + url: str + detail: Literal["low", "auto", "high"] | None = None + + +class MistralImageUrlDocument(_FixtureModel): + type: Literal["image_url"] + image_url: str | MistralImageUrlValue + + +class MistralDocumentUrlDocument(_FixtureModel): + type: Literal["document_url"] + document_url: str + document_name: str | None = None + + +MistralDocument = Annotated[ + MistralImageUrlDocument | MistralDocumentUrlDocument, + Field(discriminator="type"), +] + + +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", + "mistral/mistral-ocr-4-1", + "mistral/mistral-ocr-4", + "mistral/mistral-ocr-latest", + "mistral-ocr-2512", + "mistral-ocr-4-0", + "mistral-ocr-4-1", + "mistral-ocr-4", + "mistral-ocr-latest", +] + + +class MistralOcrSdkInput(OcrSdkInputBase): + model: MistralModel + document: MistralDocument + custom_llm_provider: Literal["mistral"] | None = None + pages: str | list[int] | None = None + include_image_base64: bool | None = None + image_limit: int | None = None + image_min_size: int | None = None + bbox_annotation_format: JsonSchemaResponseFormat | None = None + document_annotation_format: JsonSchemaResponseFormat | None = None + document_annotation_prompt: str | None = None + extract_header: bool = False + extract_footer: bool = False + table_format: Literal["markdown", "html"] | None = None + confidence_scores_granularity: Literal["page", "word", "block"] | None = None + include_blocks: bool = True + id: str | None = None + + @model_validator(mode="after") + def validate_provider_routing(self) -> Self: + if not self.model.startswith("mistral/") and self.custom_llm_provider != "mistral": + raise ValueError("unqualified Mistral models require custom_llm_provider='mistral'") + if self.document_annotation_prompt is not None and self.document_annotation_format is None: + raise ValueError("document_annotation_prompt requires document_annotation_format") + return self + + +def _validate_reducto_source(source: str) -> str: + if source.startswith("reducto://"): + return source + if not source.startswith("data:"): + raise ValueError("Reducto documents require a reducto:// id or base64 data URI") + try: + header, encoded = source.split(",", 1) + except ValueError as error: + raise ValueError("invalid Reducto data URI") from error + if ";base64" not in header: + raise ValueError("Reducto data URIs must be base64 encoded") + try: + base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError) as error: + raise ValueError("invalid Reducto base64 payload") from error + return source + + +class ReductoImageUrlDocument(_FixtureModel): + type: Literal["image_url"] + image_url: str + + @field_validator("image_url") + @classmethod + def validate_image_url(cls, value: str) -> str: + return _validate_reducto_source(value) + + +class ReductoDocumentUrlDocument(_FixtureModel): + type: Literal["document_url"] + document_url: str + + @field_validator("document_url") + @classmethod + def validate_document_url(cls, value: str) -> str: + return _validate_reducto_source(value) + + +ReductoDocument = Annotated[ + ReductoImageUrlDocument | ReductoDocumentUrlDocument, + Field(discriminator="type"), +] + +ReductoTableOutputFormat = Literal["html", "json", "md", "jsonbbox", "dynamic", "csv"] +ReductoFormattingInclude = Literal[ + "change_tracking", + "highlight", + "comments", + "hyperlinks", + "signatures", + "ignore_watermarks", +] +ReductoBlockType = Literal[ + "Header", + "Footer", + "Title", + "Section Header", + "Page Number", + "List Item", + "Figure", + "Table", + "Key Value", + "Text", + "Comment", + "Signature", +] + + +class ReductoFormatting(_FixtureModel): + add_page_markers: bool = False + table_output_format: ReductoTableOutputFormat = "dynamic" + merge_tables: bool = False + include: list[ReductoFormattingInclude] = Field(default_factory=list) + + @field_validator("include") + @classmethod + def validate_unique_include(cls, value: list[ReductoFormattingInclude]) -> list[ReductoFormattingInclude]: + if len(value) != len(set(value)): + raise ValueError("formatting.include entries must be unique") + return value + + +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) + + @model_validator(mode="after") + def validate_chunking(self) -> Self: + if self.chunk_size is not None and self.chunk_size <= 0: + raise ValueError("chunk_size must be positive") + if self.chunk_size is not None and self.chunk_overlap >= self.chunk_size: + raise ValueError("chunk_overlap must be less than chunk_size") + return self + + +class ReductoRetrieval(_FixtureModel): + chunking: ReductoChunking = Field(default_factory=ReductoChunking) + filter_blocks: list[ReductoBlockType] = Field(default_factory=list) + embedding_optimized: bool = False + + @field_validator("filter_blocks") + @classmethod + def validate_unique_blocks(cls, value: list[ReductoBlockType]) -> list[ReductoBlockType]: + if len(value) != len(set(value)): + raise ValueError("retrieval.filter_blocks entries must be unique") + return value + + +class ReductoPageRange(_FixtureModel): + start: int | None = Field(default=None, ge=1) + end: int | None = Field(default=None, ge=1) + + @model_validator(mode="after") + def validate_range(self) -> Self: + if self.start is not None and self.end is not None and self.end < self.start: + raise ValueError("page range end must be greater than or equal to start") + return self + + +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): + environment: str | None = None + + +ReductoPageSelection = ReductoPageRange | list[ReductoPageRange] | list[int] | list[str] + + +class ReductoSettings(_FixtureModel): + ocr_system: Literal["standard", "legacy"] = "standard" + extraction_mode: Literal["ocr", "hybrid"] = "hybrid" + force_url_result: bool = False + force_file_extension: str | None = None + return_ocr_data: bool = False + return_images: list[Literal["figure", "table", "page"]] = Field(default_factory=list) + embed_pdf_metadata: bool = False + embed_pdf_metadata_dpi: int = Field(default=100, ge=50, le=250) + persist_results: bool = False + tenant_throttling: ReductoTenantThrottling | None = None + timeout: float | None = Field(default=None, gt=0) + page_range: ReductoPageSelection | None = None + document_password: str | None = None + hybrid_vpc: ReductoHybridVpcSettings = Field(default_factory=ReductoHybridVpcSettings) + + @field_validator("return_images") + @classmethod + def validate_unique_images( + cls, value: list[Literal["figure", "table", "page"]] + ) -> list[Literal["figure", "table", "page"]]: + if len(value) != len(set(value)): + raise ValueError("settings.return_images entries must be unique") + return value + + +class ReductoParseV3SdkInput(OcrSdkInputBase): + model: Literal["reducto/parse-v3", "parse-v3"] + document: ReductoDocument + custom_llm_provider: Literal["reducto"] | None = None + formatting: ReductoFormatting = Field(default_factory=ReductoFormatting) + retrieval: ReductoRetrieval = Field(default_factory=ReductoRetrieval) + settings: ReductoSettings = Field(default_factory=ReductoSettings) + + @model_validator(mode="after") + def validate_provider_routing(self) -> Self: + if self.model == "parse-v3" and self.custom_llm_provider != "reducto": + raise ValueError("unqualified Reducto models require custom_llm_provider='reducto'") + return self + + +class ReductoParseLegacySdkInput(OcrSdkInputBase): + model: Literal["reducto/parse-legacy", "parse-legacy"] + document: ReductoDocument + custom_llm_provider: Literal["reducto"] | None = None + enhance: JsonObject | None = None + + @model_validator(mode="after") + def validate_provider_routing(self) -> Self: + if self.model == "parse-legacy" and self.custom_llm_provider != "reducto": + raise ValueError("unqualified Reducto models require custom_llm_provider='reducto'") + return self + + class HttpHeader(_FixtureModel): name: str value: str @@ -71,5 +312,5 @@ class RecordedHttpResponse(_FixtureModel): class OcrParityCase(_FixtureModel): - litellm_input: LiteLLMOcrInput + litellm_input: MistralOcrSdkInput provider_response: RecordedHttpResponse diff --git a/tests/test_litellm/ocr/generate_fixtures.py b/tests/test_litellm/ocr/generate_fixtures.py index a15b80c48ca..173f4ef3966 100644 --- a/tests/test_litellm/ocr/generate_fixtures.py +++ b/tests/test_litellm/ocr/generate_fixtures.py @@ -13,7 +13,7 @@ from urllib.parse import quote from dotenv import load_dotenv from hypothesis import given, settings from hypothesis import strategies as st -from hypothesis.strategies import SearchStrategy +from hypothesis.strategies import DrawFn, SearchStrategy import litellm from litellm.rust_bridge.ocr import use_litellm_rust @@ -23,14 +23,32 @@ from tests.test_litellm._fixture_recorder import ( record_cases, ) from tests.test_litellm.ocr.fixture_models import ( - ImageUrlDocument, - LiteLLMOcrInput, + JsonSchemaDefinition, + JsonSchemaResponseFormat, + MistralImageUrlDocument, + MistralModel, + MistralOcrSdkInput, + ReductoChunking, + ReductoDocumentUrlDocument, + ReductoFormatting, + ReductoParseLegacySdkInput, + ReductoParseV3SdkInput, + ReductoRetrieval, + ReductoSettings, ) FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" LOGGER: Final = logging.getLogger(__name__) -_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) +_TEXT: Final = st.just("invoice 123") +_VALUE_TEXT: Final = st.just("case-1") +_FONT_SIZE: Final = st.just(24) +_MISTRAL_MODELS: Final = ( + "mistral/mistral-ocr-2512", + "mistral/mistral-ocr-4-0", + "mistral/mistral-ocr-4-1", + "mistral/mistral-ocr-4", + "mistral/mistral-ocr-latest", +) @dataclass(frozen=True, slots=True) @@ -41,27 +59,187 @@ class GeneratorArgs: model: str -def _image_document(text: str, font_size: int) -> ImageUrlDocument: +def _image_document(text: str, font_size: int) -> MistralImageUrlDocument: url: Final = f"https://dummyjson.com/image/800x300/ffffff/000000?text={quote(text)}&fontSize={font_size}" - return ImageUrlDocument(type="image_url", image_url=url) + return MistralImageUrlDocument(type="image_url", image_url=url) -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}, - optional={ - "include_image_base64": st.booleans(), - "image_limit": st.integers(min_value=1, max_value=100), - "image_min_size": st.integers(min_value=0, max_value=10_000), - "extract_header": st.booleans(), - "extract_footer": st.booleans(), - "table_format": st.sampled_from(("markdown", "html")), - "include_blocks": st.booleans(), - "id": _VALUE_TEXT, - }, +def _annotation_format(name: str) -> JsonSchemaResponseFormat: + return JsonSchemaResponseFormat( + type="json_schema", + json_schema=JsonSchemaDefinition( + name=name, + description="Extract the visible document fields", + schema={ + "type": "object", + "properties": {"title": {"type": "string"}}, + "required": ["title"], + "additionalProperties": False, + }, + strict=True, + ), + ) + + +def _mistral_confidence_strategy() -> SearchStrategy[str]: + return st.just("page") + + +@st.composite +def mistral_input_strategy(draw: DrawFn, model: str) -> MistralOcrSdkInput: + document: Final = draw(st.builds(_image_document, _TEXT, _FONT_SIZE)) + annotation: Final = draw( + st.sampled_from( + ( + {}, + {"document_annotation_format": _annotation_format("document_title")}, + { + "document_annotation_format": _annotation_format("prompted_document_title"), + "document_annotation_prompt": "Extract the visible title", + }, + ) + ) + ) + optional_params: Final = draw( + st.fixed_dictionaries( + {}, + optional={ + "pages": st.just([0]), + "include_image_base64": st.booleans(), + "image_limit": st.just(1), + "image_min_size": st.just(300), + "bbox_annotation_format": st.just(_annotation_format("bounding_boxes")), + "extract_header": st.booleans(), + "extract_footer": st.booleans(), + "table_format": st.just("markdown"), + "confidence_scores_granularity": _mistral_confidence_strategy(), + "include_blocks": st.booleans(), + "id": _VALUE_TEXT, + }, + ) + ) + return MistralOcrSdkInput.model_validate( + { + "model": model, + "document": document, + **annotation, + **optional_params, + } + ) + + +def _reducto_formatting_strategy() -> SearchStrategy[ReductoFormatting]: + return st.builds( + ReductoFormatting, + add_page_markers=st.booleans(), + table_output_format=st.sampled_from(("html", "json", "md", "jsonbbox", "dynamic", "csv")), + merge_tables=st.booleans(), + include=st.sampled_from( + ( + [], + ["hyperlinks"], + ["change_tracking", "highlight", "comments"], + ["signatures", "ignore_watermarks"], + ) + ), + ) + + +def _reducto_chunking_strategy() -> SearchStrategy[ReductoChunking]: + return st.one_of( + st.builds( + ReductoChunking, + chunk_mode=st.sampled_from(("section", "page", "disabled", "block", "page_sections")), + chunk_size=st.just(None), + chunk_overlap=st.just(0), + ), + st.builds( + ReductoChunking, + chunk_mode=st.just("variable"), + chunk_size=st.sampled_from((250, 1000, 1500)), + chunk_overlap=st.sampled_from((0, 32, 128)), + ), + ) + + +def _reducto_retrieval_strategy() -> SearchStrategy[ReductoRetrieval]: + return st.builds( + ReductoRetrieval, + chunking=_reducto_chunking_strategy(), + filter_blocks=st.sampled_from( + ( + [], + ["Header"], + ["Header", "Footer", "Page Number"], + ["Figure", "Table", "Key Value"], + ) + ), + embedding_optimized=st.booleans(), + ) + + +def _reducto_settings_strategy() -> SearchStrategy[ReductoSettings]: + return st.builds( + ReductoSettings, + ocr_system=st.sampled_from(("standard", "legacy")), + extraction_mode=st.sampled_from(("ocr", "hybrid")), + force_url_result=st.booleans(), + return_ocr_data=st.booleans(), + return_images=st.sampled_from(([], ["figure"], ["table"], ["page"], ["figure", "table", "page"])), + embed_pdf_metadata=st.booleans(), + embed_pdf_metadata_dpi=st.sampled_from((50, 100, 250)), + persist_results=st.just(False), + timeout=st.sampled_from((None, 300.0, 900.0)), + page_range=st.sampled_from((None, [1], [1, 2], ["Sheet1"])), + ) + + +@st.composite +def reducto_v3_input_strategy(draw: DrawFn) -> ReductoParseV3SdkInput: + model, custom_llm_provider = draw(st.sampled_from((("reducto/parse-v3", None), ("parse-v3", "reducto")))) + options: Final = draw( + st.fixed_dictionaries( + {}, + optional={ + "formatting": _reducto_formatting_strategy(), + "retrieval": _reducto_retrieval_strategy(), + "settings": _reducto_settings_strategy(), + }, + ) + ) + return ReductoParseV3SdkInput.model_validate( + { + "model": model, + "custom_llm_provider": custom_llm_provider, + "document": ReductoDocumentUrlDocument( + type="document_url", + document_url="reducto://fixture-document.pdf", + ), + **options, + } + ) + + +def reducto_legacy_input_strategy() -> SearchStrategy[ReductoParseLegacySdkInput]: + return st.sampled_from( + ( + ReductoParseLegacySdkInput( + model="reducto/parse-legacy", + document=ReductoDocumentUrlDocument( + type="document_url", + document_url="reducto://fixture-document.pdf", + ), + ), + ReductoParseLegacySdkInput( + model="parse-legacy", + custom_llm_provider="reducto", + document=ReductoDocumentUrlDocument( + type="document_url", + document_url="reducto://fixture-document.pdf", + ), + ), + ) ) - return input_values.map(LiteLLMOcrInput.model_validate) def _generate_examples( @@ -71,11 +249,11 @@ def _generate_examples( concurrency: int, sdk_call: Callable[..., object], ) -> None: - generated: Final[queue.SimpleQueue[LiteLLMOcrInput | None]] = queue.SimpleQueue() + generated: Final[queue.SimpleQueue[MistralOcrSdkInput | None]] = queue.SimpleQueue() @settings(max_examples=examples, deadline=None, derandomize=True) - @given(case_input=_mistral_input_strategy(spec.model)) - def generate_case(case_input: LiteLLMOcrInput) -> None: + @given(case_input=mistral_input_strategy(spec.model)) + def generate_case(case_input: MistralOcrSdkInput) -> None: generated.put(case_input) generate_case() @@ -96,13 +274,13 @@ def _parse_args() -> GeneratorArgs: 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") + parser.add_argument("--model", choices=_MISTRAL_MODELS, 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), + model=cast(MistralModel, namespace.model), ) diff --git a/tests/test_litellm/ocr/test_fixture_models.py b/tests/test_litellm/ocr/test_fixture_models.py index fe452833c64..70966e47bda 100644 --- a/tests/test_litellm/ocr/test_fixture_models.py +++ b/tests/test_litellm/ocr/test_fixture_models.py @@ -1,56 +1,231 @@ from __future__ import annotations -from typing import Final +from collections.abc import Callable +from typing import Final, cast import pytest +from hypothesis import given, settings from pydantic import ValidationError -from tests.test_litellm.ocr.fixture_models import LiteLLMOcrInput +from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig +from litellm.llms.mistral.ocr.transformation import MistralOCRConfig +from litellm.llms.reducto.ocr.transformation import ReductoParseLegacyConfig, ReductoParseV3Config +from tests.test_litellm.ocr.fixture_models import ( + JsonSchemaDefinition, + JsonSchemaResponseFormat, + MistralDocumentUrlDocument, + MistralImageUrlDocument, + MistralImageUrlValue, + MistralOcrSdkInput, + OcrSdkInputBase, + ReductoChunking, + ReductoDocumentUrlDocument, + ReductoFormatting, + ReductoPageRange, + ReductoParseLegacySdkInput, + ReductoParseV3SdkInput, + ReductoRetrieval, + ReductoSettings, +) +from tests.test_litellm.ocr.generate_fixtures import ( + mistral_input_strategy, + reducto_legacy_input_strategy, + reducto_v3_input_strategy, +) + +COMMON_FIELDS: Final = frozenset({"model", "document", "custom_llm_provider"}) + + +def _provider_fields(model: type[OcrSdkInputBase]) -> set[str]: + return set(model.model_fields) - COMMON_FIELDS + + +def _supported_params(config: BaseOCRConfig, model: str) -> set[str]: + get_supported_params: Final = cast(Callable[[str], list[str]], config.get_supported_ocr_params) + return set(get_supported_params(model)) + + +def _mistral_input(**params: object) -> MistralOcrSdkInput: + return MistralOcrSdkInput.model_validate( + { + "model": "mistral/mistral-ocr-latest", + "document": {"type": "image_url", "image_url": "https://example.com/image.png"}, + **params, + } + ) + + +def _reducto_document() -> ReductoDocumentUrlDocument: + return ReductoDocumentUrlDocument( + type="document_url", + document_url="reducto://fixture-document.pdf", + ) + + +def test_mistral_fixture_fields_match_provider_config() -> None: + assert _provider_fields(MistralOcrSdkInput) == _supported_params(MistralOCRConfig(), "mistral-ocr-latest") + + +def test_reducto_fixture_fields_match_provider_configs() -> None: + assert _provider_fields(ReductoParseV3SdkInput) == _supported_params(ReductoParseV3Config(), "parse-v3") + assert _provider_fields(ReductoParseLegacySdkInput) == _supported_params(ReductoParseLegacyConfig(), "parse-legacy") + + +def test_mistral_input_preserves_omission_and_explicit_boolean_values() -> None: + omitted: Final = _mistral_input().as_sdk_kwargs() + explicit: Final = _mistral_input(extract_header=False, include_blocks=True).as_sdk_kwargs() + + assert "extract_header" not in omitted + assert "include_blocks" not in omitted + assert explicit["extract_header"] is False + assert explicit["include_blocks"] is True + + +def test_mistral_input_supports_document_and_page_variants() -> None: + nested_image: Final = MistralOcrSdkInput( + model="mistral/mistral-ocr-4-1", + document=MistralImageUrlDocument( + type="image_url", + image_url=MistralImageUrlValue(url="https://example.com/image.png", detail="high"), + ), + pages="0,2-4", + ) + named_document: Final = MistralOcrSdkInput( + model="mistral/mistral-ocr-2512", + document=MistralDocumentUrlDocument( + type="document_url", + document_url="https://example.com/document.pdf", + document_name="invoice.pdf", + ), + ) + + assert nested_image.canonical_input()["document"] == { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png", "detail": "high"}, + } + assert nested_image.as_sdk_kwargs()["pages"] == "0,2-4" + assert named_document.canonical_input()["document"] == { + "type": "document_url", + "document_url": "https://example.com/document.pdf", + "document_name": "invoice.pdf", + } + + +def test_mistral_annotation_schema_serializes_provider_alias() -> None: + annotation: Final = JsonSchemaResponseFormat( + type="json_schema", + json_schema=JsonSchemaDefinition( + name="invoice", + schema={"type": "object"}, + ), + ) + sdk_input: Final = _mistral_input( + document_annotation_format=annotation, + document_annotation_prompt="Extract invoice fields", + ) + + assert sdk_input.canonical_input()["document_annotation_format"] == { + "type": "json_schema", + "json_schema": { + "name": "invoice", + "schema": {"type": "object"}, + }, + } + + +def test_mistral_annotation_prompt_requires_format() -> None: + with pytest.raises(ValidationError, match="requires document_annotation_format"): + _mistral_input(document_annotation_prompt="Extract invoice fields") + + +@pytest.mark.parametrize("field", ("extract_header", "extract_footer", "include_blocks")) +def test_mistral_nonnullable_booleans_reject_null(field: str) -> None: + with pytest.raises(ValidationError): + _mistral_input(**{field: None}) + + +def test_unqualified_models_require_explicit_provider() -> None: + with pytest.raises(ValidationError, match="custom_llm_provider='mistral'"): + MistralOcrSdkInput( + model="mistral-ocr-latest", + document=MistralImageUrlDocument(type="image_url", image_url="https://example.com/image.png"), + ) + with pytest.raises(ValidationError, match="custom_llm_provider='reducto'"): + ReductoParseV3SdkInput(model="parse-v3", document=_reducto_document()) + + +def test_reducto_v3_preserves_nested_provider_params() -> None: + sdk_input: Final = ReductoParseV3SdkInput( + model="reducto/parse-v3", + document=_reducto_document(), + formatting=ReductoFormatting(table_output_format="html", include=["hyperlinks"]), + retrieval=ReductoRetrieval(chunking=ReductoChunking(chunk_mode="variable", chunk_size=250, chunk_overlap=32)), + settings=ReductoSettings(embed_pdf_metadata=True, embed_pdf_metadata_dpi=250, page_range=[1, 3]), + ) + + assert sdk_input.as_sdk_kwargs()["formatting"] == { + "table_output_format": "html", + "include": ["hyperlinks"], + } + assert sdk_input.as_sdk_kwargs()["retrieval"] == { + "chunking": {"chunk_mode": "variable", "chunk_size": 250, "chunk_overlap": 32} + } + assert sdk_input.as_sdk_kwargs()["settings"] == { + "embed_pdf_metadata": True, + "embed_pdf_metadata_dpi": 250, + "page_range": [1, 3], + } + + +def test_reducto_optional_objects_reject_explicit_null() -> None: + with pytest.raises(ValidationError): + ReductoParseV3SdkInput.model_validate( + { + "model": "reducto/parse-v3", + "document": _reducto_document(), + "formatting": None, + } + ) @pytest.mark.parametrize( - ("raw_input", "expected_params"), + "source", ( - ( - { - "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"}}, - }, - ), + "https://example.com/document.pdf", + "not-a-document", + "data:application/pdf,not-base64", + "data:application/pdf;base64,not!base64", ), ) -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: +def test_reducto_document_rejects_unsupported_sources(source: str) -> None: with pytest.raises(ValidationError): - LiteLLMOcrInput.model_validate( - { - "model": "reducto/parse-v3", - "document": {"type": "document_url", "document_url": "reducto://fixture-file"}, - "settings": object(), - } - ) + ReductoDocumentUrlDocument(type="document_url", document_url=source) + + +def test_reducto_nested_constraints() -> None: + with pytest.raises(ValidationError, match="less than chunk_size"): + ReductoChunking(chunk_mode="variable", chunk_size=100, chunk_overlap=100) + with pytest.raises(ValidationError, match="greater than or equal to start"): + ReductoPageRange(start=3, end=2) + with pytest.raises(ValidationError): + ReductoSettings(embed_pdf_metadata_dpi=49) + with pytest.raises(ValidationError, match="must be unique"): + ReductoFormatting(include=["hyperlinks", "hyperlinks"]) + + +@settings(max_examples=50, deadline=None) +@given(sdk_input=mistral_input_strategy("mistral/mistral-ocr-4-1")) +def test_mistral_strategy_only_generates_valid_sdk_inputs(sdk_input: MistralOcrSdkInput) -> None: + assert MistralOcrSdkInput.model_validate(sdk_input.canonical_input()) == sdk_input + + +@settings(max_examples=50, deadline=None) +@given(sdk_input=reducto_v3_input_strategy()) +def test_reducto_v3_strategy_only_generates_valid_sdk_inputs(sdk_input: ReductoParseV3SdkInput) -> None: + assert ReductoParseV3SdkInput.model_validate(sdk_input.canonical_input()) == sdk_input + + +@settings(max_examples=10, deadline=None) +@given(sdk_input=reducto_legacy_input_strategy()) +def test_reducto_legacy_strategy_omits_undocumented_enhance(sdk_input: ReductoParseLegacySdkInput) -> None: + assert "enhance" not in sdk_input.as_sdk_kwargs() diff --git a/tests/test_litellm/ocr/test_sdk_parity.py b/tests/test_litellm/ocr/test_sdk_parity.py index 4b3ce140007..324fc90db9d 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 LiteLLMOcrInput, OcrParityCase +from tests.test_litellm.ocr.fixture_models import MistralOcrSdkInput, 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: LiteLLMOcrInput, mock_url: str, route: SDKRoute) -> dict[str, object]: +def _call_kwargs(sdk_input: MistralOcrSdkInput, 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: LiteLLMOcrInput, mock_url: str, route: SDKRoute) -> def _execute_sdk_case( - sdk_input: LiteLLMOcrInput, + sdk_input: MistralOcrSdkInput, 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 1f058bb2a2b..44a0cf88cf9 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, LiteLLMOcrInput +from tests.test_litellm.ocr.fixture_models import MistralImageUrlDocument, MistralOcrSdkInput class _ControlledUpstream(ThreadingHTTPServer): @@ -78,11 +78,11 @@ def _controlled_upstream() -> Generator[_ControlledUpstream]: thread.join(timeout=5) -def _case(identifier: str) -> LiteLLMOcrInput: - return LiteLLMOcrInput.model_validate( +def _case(identifier: str) -> MistralOcrSdkInput: + return MistralOcrSdkInput.model_validate( { "model": "mistral/mistral-ocr-latest", - "document": ImageUrlDocument(type="image_url", image_url="https://example.com/image.png"), + "document": MistralImageUrlDocument(type="image_url", image_url="https://example.com/image.png"), "id": identifier, } )