From 45258ecb8574c4b1c97724e7883f69e49ea6dbad Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 16:51:53 -0700 Subject: [PATCH] test(ocr): add shared Mistral transformation contracts --- .../ocr/test_mistral_ocr_transformation.py | 236 -------- tests/transform_contracts/__init__.py | 1 + .../cases/mistral/ocr.json | 530 ++++++++++++++++++ tests/transform_contracts/conftest.py | 19 + tests/transform_contracts/loader.py | 38 ++ tests/transform_contracts/mistral_ocr.py | 101 ++++ tests/transform_contracts/registry.py | 31 + tests/transform_contracts/schema.py | 124 ++++ tests/transform_contracts/test_contracts.py | 10 + tests/transform_contracts/test_foundation.py | 79 +++ 10 files changed, 933 insertions(+), 236 deletions(-) delete mode 100644 tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py create mode 100644 tests/transform_contracts/__init__.py create mode 100644 tests/transform_contracts/cases/mistral/ocr.json create mode 100644 tests/transform_contracts/conftest.py create mode 100644 tests/transform_contracts/loader.py create mode 100644 tests/transform_contracts/mistral_ocr.py create mode 100644 tests/transform_contracts/registry.py create mode 100644 tests/transform_contracts/schema.py create mode 100644 tests/transform_contracts/test_contracts.py create mode 100644 tests/transform_contracts/test_foundation.py diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py deleted file mode 100644 index 8ce8b777f71..00000000000 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py +++ /dev/null @@ -1,236 +0,0 @@ -""" -Unit tests for MistralOCRConfig transformation. - -Tests the supported OCR parameters and their mapping behaviour. -No real API calls are made — all tests are fully mocked/local. -""" - -import httpx -import pytest - -from litellm.llms.mistral.ocr.transformation import MistralOCRConfig - - -@pytest.fixture -def config() -> MistralOCRConfig: - return MistralOCRConfig() - - -MODEL = "mistral-ocr-latest" - - -class TestGetSupportedOcrParams: - def test_extract_header_in_supported_params(self, config: MistralOCRConfig) -> None: - """extract_header must be in the Mistral OCR supported params list.""" - supported = config.get_supported_ocr_params(model=MODEL) - assert "extract_header" in supported - - def test_extract_footer_in_supported_params(self, config: MistralOCRConfig) -> None: - """extract_footer must be in the Mistral OCR supported params list.""" - supported = config.get_supported_ocr_params(model=MODEL) - assert "extract_footer" in supported - - def test_existing_params_still_present(self, config: MistralOCRConfig) -> None: - """Ensure the previously supported params were not accidentally removed.""" - supported = config.get_supported_ocr_params(model=MODEL) - for param in [ - "pages", - "include_image_base64", - "image_limit", - "image_min_size", - "bbox_annotation_format", - "document_annotation_format", - ]: - assert param in supported, f"Previously supported param '{param}' is missing" - - -class TestMapOcrParams: - def test_extract_header_passed_through(self, config: MistralOCRConfig) -> None: - """extract_header=True must survive the map_ocr_params filter.""" - result = config.map_ocr_params( - non_default_params={"extract_header": True}, - optional_params={}, - model=MODEL, - ) - assert result == {"extract_header": True} - - def test_extract_footer_passed_through(self, config: MistralOCRConfig) -> None: - """extract_footer=True must survive the map_ocr_params filter.""" - result = config.map_ocr_params( - non_default_params={"extract_footer": True}, - optional_params={}, - model=MODEL, - ) - assert result == {"extract_footer": True} - - def test_extract_header_and_footer_together(self, config: MistralOCRConfig) -> None: - """Both params can be passed together and are both forwarded.""" - result = config.map_ocr_params( - non_default_params={"extract_header": True, "extract_footer": False}, - optional_params={}, - model=MODEL, - ) - assert result == {"extract_header": True, "extract_footer": False} - - def test_unknown_param_is_dropped(self, config: MistralOCRConfig) -> None: - """Parameters not in the supported list must be silently dropped.""" - result = config.map_ocr_params( - non_default_params={"extract_header": True, "unsupported_param": "value"}, - optional_params={}, - model=MODEL, - ) - assert "extract_header" in result - assert "unsupported_param" not in result - - -class TestNewSupportedParams: - """Verify the newly added params are in the supported list.""" - - @pytest.mark.parametrize( - "param_name", - [ - "table_format", - "confidence_scores_granularity", - "document_annotation_prompt", - "include_blocks", - "id", - ], - ) - def test_new_param_in_supported_list(self, config: MistralOCRConfig, param_name: str) -> None: - supported = config.get_supported_ocr_params(model=MODEL) - assert param_name in supported - - -class TestNewParamsMapOcr: - """Verify the newly added params survive map_ocr_params.""" - - @pytest.mark.parametrize( - "param_name,param_value", - [ - ("table_format", "html"), - ("table_format", "markdown"), - ("confidence_scores_granularity", "word"), - ("confidence_scores_granularity", "page"), - ("document_annotation_prompt", "Extract all invoice line items"), - ("include_blocks", True), - ("id", "req-123"), - ], - ) - def test_new_param_passed_through(self, config: MistralOCRConfig, param_name: str, param_value: str) -> None: - result = config.map_ocr_params( - non_default_params={param_name: param_value}, - optional_params={}, - model=MODEL, - ) - assert result == {param_name: param_value} - - -class TestTransformOcrRequest: - """Verify params end up in the final request body via transform_ocr_request.""" - - SAMPLE_DOCUMENT = { - "type": "document_url", - "document_url": "https://example.com/doc.pdf", - } - - @pytest.mark.parametrize( - "param_name,param_value", - [ - ("table_format", "html"), - ("confidence_scores_granularity", "word"), - ("document_annotation_prompt", "Extract all invoice line items"), - ("id", "req-123"), - ("extract_header", True), - ("include_blocks", True), - ("pages", [0, 1]), - ], - ) - def test_param_included_in_request_body(self, config: MistralOCRConfig, param_name: str, param_value) -> None: - result = config.transform_ocr_request( - model=MODEL, - document=self.SAMPLE_DOCUMENT, - optional_params={param_name: param_value}, - headers={}, - ) - assert result.data[param_name] == param_value - assert result.data["model"] == MODEL - assert result.data["document"] == self.SAMPLE_DOCUMENT - assert result.files is None - - def test_multiple_new_params_together(self, config: MistralOCRConfig) -> None: - """Multiple new params can be passed together in a single request.""" - optional_params = { - "table_format": "html", - "confidence_scores_granularity": "page", - "extract_header": True, - } - result = config.transform_ocr_request( - model=MODEL, - document=self.SAMPLE_DOCUMENT, - optional_params=optional_params, - headers={}, - ) - for key, value in optional_params.items(): - assert result.data[key] == value - - -class TestTransformOcrResponseOcr4Fields: - """OCR 4 adds blocks, confidence_scores, tables, hyperlinks, header and footer - to each page. These must survive transform_ocr_response so callers actually - receive the new structured output rather than having it silently dropped.""" - - def _response(self, page: dict) -> httpx.Response: - return httpx.Response( - 200, - json={ - "pages": [page], - "model": "mistral-ocr-4-0", - "usage_info": {"pages_processed": 1}, - }, - ) - - def test_blocks_and_confidence_scores_preserved(self, config: MistralOCRConfig) -> None: - page = { - "index": 0, - "markdown": "# Invoice", - "blocks": [ - { - "type": "title", - "top_left_x": 10, - "top_left_y": 20, - "bottom_right_x": 300, - "bottom_right_y": 60, - "content": "Invoice", - } - ], - "confidence_scores": {"page": 0.98}, - } - - result = config.transform_ocr_response( - model="mistral-ocr-4-0", - raw_response=self._response(page), - logging_obj=None, - ) - - assert result.pages[0].blocks == page["blocks"] - assert result.pages[0].confidence_scores == page["confidence_scores"] - - def test_ocr4_fields_survive_model_dump(self, config: MistralOCRConfig) -> None: - page = { - "index": 0, - "markdown": "table page", - "tables": [{"rows": 2, "cols": 3}], - "hyperlinks": ["https://example.com"], - "header": "Acme Corp", - "footer": "Page 1", - } - - result = config.transform_ocr_response( - model="mistral-ocr-4-0", - raw_response=self._response(page), - logging_obj=None, - ) - - dumped_page = result.model_dump()["pages"][0] - for field in ("tables", "hyperlinks", "header", "footer"): - assert dumped_page[field] == page[field] diff --git a/tests/transform_contracts/__init__.py b/tests/transform_contracts/__init__.py new file mode 100644 index 00000000000..a7cba34bf52 --- /dev/null +++ b/tests/transform_contracts/__init__.py @@ -0,0 +1 @@ +"""Shared provider transformation contracts.""" diff --git a/tests/transform_contracts/cases/mistral/ocr.json b/tests/transform_contracts/cases/mistral/ocr.json new file mode 100644 index 00000000000..38ae9b1cc0f --- /dev/null +++ b/tests/transform_contracts/cases/mistral/ocr.json @@ -0,0 +1,530 @@ +{ + "schema_version": 1, + "cases": [ + { + "id": "mistral.ocr.get_supported_ocr_params.latest", + "operation": "mistral.ocr.get_supported_ocr_params", + "input": { + "model": "mistral-ocr-latest" + }, + "expected": [ + "pages", + "include_image_base64", + "image_limit", + "image_min_size", + "bbox_annotation_format", + "document_annotation_format", + "document_annotation_prompt", + "extract_header", + "extract_footer", + "table_format", + "confidence_scores_granularity", + "include_blocks", + "id" + ] + }, + { + "id": "mistral.ocr.map_ocr_params.extract_header_true", + "operation": "mistral.ocr.map_ocr_params", + "input": { + "model": "mistral-ocr-latest", + "non_default_params": { + "extract_header": true + }, + "optional_params": {} + }, + "expected": { + "extract_header": true + } + }, + { + "id": "mistral.ocr.map_ocr_params.extract_footer_true", + "operation": "mistral.ocr.map_ocr_params", + "input": { + "model": "mistral-ocr-latest", + "non_default_params": { + "extract_footer": true + }, + "optional_params": {} + }, + "expected": { + "extract_footer": true + } + }, + { + "id": "mistral.ocr.map_ocr_params.extract_header_and_footer", + "operation": "mistral.ocr.map_ocr_params", + "input": { + "model": "mistral-ocr-latest", + "non_default_params": { + "extract_header": true, + "extract_footer": false + }, + "optional_params": {} + }, + "expected": { + "extract_header": true, + "extract_footer": false + } + }, + { + "id": "mistral.ocr.map_ocr_params.unsupported_param_dropped", + "operation": "mistral.ocr.map_ocr_params", + "input": { + "model": "mistral-ocr-latest", + "non_default_params": { + "extract_header": true, + "unsupported_param": "value" + }, + "optional_params": {} + }, + "expected": { + "extract_header": true + } + }, + { + "id": "mistral.ocr.map_ocr_params.table_format_html", + "operation": "mistral.ocr.map_ocr_params", + "input": { + "model": "mistral-ocr-latest", + "non_default_params": { + "table_format": "html" + }, + "optional_params": {} + }, + "expected": { + "table_format": "html" + } + }, + { + "id": "mistral.ocr.map_ocr_params.table_format_markdown", + "operation": "mistral.ocr.map_ocr_params", + "input": { + "model": "mistral-ocr-latest", + "non_default_params": { + "table_format": "markdown" + }, + "optional_params": {} + }, + "expected": { + "table_format": "markdown" + } + }, + { + "id": "mistral.ocr.map_ocr_params.confidence_scores_word", + "operation": "mistral.ocr.map_ocr_params", + "input": { + "model": "mistral-ocr-latest", + "non_default_params": { + "confidence_scores_granularity": "word" + }, + "optional_params": {} + }, + "expected": { + "confidence_scores_granularity": "word" + } + }, + { + "id": "mistral.ocr.map_ocr_params.confidence_scores_page", + "operation": "mistral.ocr.map_ocr_params", + "input": { + "model": "mistral-ocr-latest", + "non_default_params": { + "confidence_scores_granularity": "page" + }, + "optional_params": {} + }, + "expected": { + "confidence_scores_granularity": "page" + } + }, + { + "id": "mistral.ocr.map_ocr_params.document_annotation_prompt", + "operation": "mistral.ocr.map_ocr_params", + "input": { + "model": "mistral-ocr-latest", + "non_default_params": { + "document_annotation_prompt": "Extract all invoice line items" + }, + "optional_params": {} + }, + "expected": { + "document_annotation_prompt": "Extract all invoice line items" + } + }, + { + "id": "mistral.ocr.map_ocr_params.include_blocks_true", + "operation": "mistral.ocr.map_ocr_params", + "input": { + "model": "mistral-ocr-latest", + "non_default_params": { + "include_blocks": true + }, + "optional_params": {} + }, + "expected": { + "include_blocks": true + } + }, + { + "id": "mistral.ocr.map_ocr_params.request_id", + "operation": "mistral.ocr.map_ocr_params", + "input": { + "model": "mistral-ocr-latest", + "non_default_params": { + "id": "req-123" + }, + "optional_params": {} + }, + "expected": { + "id": "req-123" + } + }, + { + "id": "mistral.ocr.transform_ocr_request.table_format_html", + "operation": "mistral.ocr.transform_ocr_request", + "input": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "optional_params": { + "table_format": "html" + }, + "headers": {} + }, + "expected": { + "data": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "table_format": "html" + }, + "files": null + } + }, + { + "id": "mistral.ocr.transform_ocr_request.confidence_scores_word", + "operation": "mistral.ocr.transform_ocr_request", + "input": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "optional_params": { + "confidence_scores_granularity": "word" + }, + "headers": {} + }, + "expected": { + "data": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "confidence_scores_granularity": "word" + }, + "files": null + } + }, + { + "id": "mistral.ocr.transform_ocr_request.document_annotation_prompt", + "operation": "mistral.ocr.transform_ocr_request", + "input": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "optional_params": { + "document_annotation_prompt": "Extract all invoice line items" + }, + "headers": {} + }, + "expected": { + "data": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "document_annotation_prompt": "Extract all invoice line items" + }, + "files": null + } + }, + { + "id": "mistral.ocr.transform_ocr_request.request_id", + "operation": "mistral.ocr.transform_ocr_request", + "input": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "optional_params": { + "id": "req-123" + }, + "headers": {} + }, + "expected": { + "data": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "id": "req-123" + }, + "files": null + } + }, + { + "id": "mistral.ocr.transform_ocr_request.extract_header_true", + "operation": "mistral.ocr.transform_ocr_request", + "input": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "optional_params": { + "extract_header": true + }, + "headers": {} + }, + "expected": { + "data": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "extract_header": true + }, + "files": null + } + }, + { + "id": "mistral.ocr.transform_ocr_request.include_blocks_true", + "operation": "mistral.ocr.transform_ocr_request", + "input": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "optional_params": { + "include_blocks": true + }, + "headers": {} + }, + "expected": { + "data": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "include_blocks": true + }, + "files": null + } + }, + { + "id": "mistral.ocr.transform_ocr_request.pages", + "operation": "mistral.ocr.transform_ocr_request", + "input": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "optional_params": { + "pages": [0, 1] + }, + "headers": {} + }, + "expected": { + "data": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "pages": [0, 1] + }, + "files": null + } + }, + { + "id": "mistral.ocr.transform_ocr_request.multiple_optional_params", + "operation": "mistral.ocr.transform_ocr_request", + "input": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "optional_params": { + "table_format": "html", + "confidence_scores_granularity": "page", + "extract_header": true + }, + "headers": {} + }, + "expected": { + "data": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "table_format": "html", + "confidence_scores_granularity": "page", + "extract_header": true + }, + "files": null + } + }, + { + "id": "mistral.ocr.transform_ocr_response.blocks_and_confidence_scores", + "operation": "mistral.ocr.transform_ocr_response", + "input": { + "model": "mistral-ocr-4-0", + "response": { + "status_code": 200, + "body": { + "pages": [ + { + "index": 0, + "markdown": "# Invoice", + "blocks": [ + { + "type": "title", + "top_left_x": 10, + "top_left_y": 20, + "bottom_right_x": 300, + "bottom_right_y": 60, + "content": "Invoice" + } + ], + "confidence_scores": { + "page": 0.98 + } + } + ], + "model": "mistral-ocr-4-0", + "usage_info": { + "pages_processed": 1 + } + } + } + }, + "expected": { + "pages": [ + { + "index": 0, + "markdown": "# Invoice", + "images": null, + "dimensions": null, + "blocks": [ + { + "type": "title", + "top_left_x": 10, + "top_left_y": 20, + "bottom_right_x": 300, + "bottom_right_y": 60, + "content": "Invoice" + } + ], + "confidence_scores": { + "page": 0.98 + } + } + ], + "model": "mistral-ocr-4-0", + "document_annotation": null, + "usage_info": { + "pages_processed": 1, + "credits": null, + "doc_size_bytes": null + }, + "content": null, + "tables": null, + "keyValuePairs": null, + "object": "ocr" + } + }, + { + "id": "mistral.ocr.transform_ocr_response.ocr4_page_fields", + "operation": "mistral.ocr.transform_ocr_response", + "input": { + "model": "mistral-ocr-4-0", + "response": { + "status_code": 200, + "body": { + "pages": [ + { + "index": 0, + "markdown": "table page", + "tables": [ + { + "rows": 2, + "cols": 3 + } + ], + "hyperlinks": [ + "https://example.com" + ], + "header": "Invoice header", + "footer": "Page 1" + } + ], + "model": "mistral-ocr-4-0", + "usage_info": { + "pages_processed": 1 + } + } + } + }, + "expected": { + "pages": [ + { + "index": 0, + "markdown": "table page", + "images": null, + "dimensions": null, + "tables": [ + { + "rows": 2, + "cols": 3 + } + ], + "hyperlinks": [ + "https://example.com" + ], + "header": "Invoice header", + "footer": "Page 1" + } + ], + "model": "mistral-ocr-4-0", + "document_annotation": null, + "usage_info": { + "pages_processed": 1, + "credits": null, + "doc_size_bytes": null + }, + "content": null, + "tables": null, + "keyValuePairs": null, + "object": "ocr" + } + } + ] +} diff --git a/tests/transform_contracts/conftest.py b/tests/transform_contracts/conftest.py new file mode 100644 index 00000000000..623e94f8509 --- /dev/null +++ b/tests/transform_contracts/conftest.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from typing import Final + +import pytest + +from tests.transform_contracts.loader import load_contract_cases + + +def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: + if "contract_case" not in metafunc.fixturenames: + return + cases: Final = load_contract_cases() + metafunc.parametrize("contract_case", cases, ids=tuple(case.id for case in cases)) + + +@pytest.fixture(autouse=True) +def use_local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") diff --git a/tests/transform_contracts/loader.py b/tests/transform_contracts/loader.py new file mode 100644 index 00000000000..3d83115345c --- /dev/null +++ b/tests/transform_contracts/loader.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from collections import Counter +from pathlib import Path +from typing import Final + +from pydantic import ValidationError + +from tests.transform_contracts.schema import CONTRACT_SUITE_ADAPTER, ContractSuiteV1, TransformationCase + +CONTRACTS_ROOT: Final = Path(__file__).resolve().parent +CASES_ROOT: Final = CONTRACTS_ROOT / "cases" + + +def discover_contract_paths(root: Path = CASES_ROOT) -> tuple[Path, ...]: + if not root.is_dir(): + raise FileNotFoundError(f"transformation contract directory does not exist: {root}") + paths: Final = tuple(sorted(root.rglob("*.json"))) + if paths: + return paths + raise FileNotFoundError(f"no transformation contract files found under: {root}") + + +def load_contract_file(path: Path) -> ContractSuiteV1: + try: + return CONTRACT_SUITE_ADAPTER.validate_json(path.read_text(encoding="utf-8")) + except (OSError, ValidationError) as exc: + raise ValueError(f"invalid transformation contract file: {path}") from exc + + +def load_contract_cases(root: Path = CASES_ROOT) -> tuple[TransformationCase, ...]: + cases: Final = tuple(case for path in discover_contract_paths(root) for case in load_contract_file(path).cases) + duplicates: Final = tuple( + sorted(case_id for case_id, count in Counter(case.id for case in cases).items() if count > 1) + ) + if duplicates: + raise ValueError(f"duplicate transformation contract case ids: {duplicates}") + return tuple(sorted(cases, key=lambda case: case.id)) diff --git a/tests/transform_contracts/mistral_ocr.py b/tests/transform_contracts/mistral_ocr.py new file mode 100644 index 00000000000..262c00c7bac --- /dev/null +++ b/tests/transform_contracts/mistral_ocr.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from typing import Final, Protocol, cast + +import httpx +from pydantic import BaseModel, TypeAdapter + +from tests.transform_contracts.schema import ( + GetSupportedOCRParamsCase, + JsonObject, + JsonValue, + MapOCRParamsCase, + TransformOCRRequestCase, + TransformOCRResponseCase, + TransformationCase, +) + +_JSON_OBJECT_ADAPTER: Final[TypeAdapter[JsonObject]] = TypeAdapter(JsonObject) +_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_DOCUMENT_ADAPTER: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) + + +class _MistralOCRConfig(Protocol): + def get_supported_ocr_params(self, model: str) -> object: ... + + def map_ocr_params( + self, + non_default_params: JsonObject, + optional_params: JsonObject, + model: str, + ) -> object: ... + + def transform_ocr_request( + self, + model: str, + document: dict[str, str], + optional_params: JsonObject, + headers: dict[str, str], + ) -> BaseModel: ... + + def transform_ocr_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: object | None, + ) -> BaseModel: ... + + +def _config() -> _MistralOCRConfig: + from litellm.llms.mistral.ocr.transformation import MistralOCRConfig + + return cast(_MistralOCRConfig, MistralOCRConfig()) + + +def _model_dump(model: BaseModel) -> JsonObject: + return _JSON_OBJECT_ADAPTER.validate_python(cast(object, model.model_dump(mode="json"))) + + +def run_get_supported_ocr_params(case: TransformationCase) -> JsonValue: + if not isinstance(case, GetSupportedOCRParamsCase): + raise TypeError(f"invalid case type for {case.operation}: {type(case).__name__}") + return _JSON_ADAPTER.validate_python(_config().get_supported_ocr_params(model=case.input.model)) + + +def run_map_ocr_params(case: TransformationCase) -> JsonValue: + if not isinstance(case, MapOCRParamsCase): + raise TypeError(f"invalid case type for {case.operation}: {type(case).__name__}") + result: Final = _config().map_ocr_params( + non_default_params=case.input.non_default_params, + optional_params=case.input.optional_params, + model=case.input.model, + ) + return _JSON_OBJECT_ADAPTER.validate_python(result) + + +def run_transform_ocr_request(case: TransformationCase) -> JsonValue: + if not isinstance(case, TransformOCRRequestCase): + raise TypeError(f"invalid case type for {case.operation}: {type(case).__name__}") + result: Final = _config().transform_ocr_request( + model=case.input.model, + document=_DOCUMENT_ADAPTER.validate_python(cast(object, case.input.document.model_dump(mode="json"))), + optional_params=case.input.optional_params, + headers=case.input.headers, + ) + return _model_dump(result) + + +def run_transform_ocr_response(case: TransformationCase) -> JsonValue: + if not isinstance(case, TransformOCRResponseCase): + raise TypeError(f"invalid case type for {case.operation}: {type(case).__name__}") + response: Final = httpx.Response( + status_code=case.input.response.status_code, + json=case.input.response.body, + headers=case.input.response.headers, + ) + result: Final = _config().transform_ocr_response( + model=case.input.model, + raw_response=response, + logging_obj=None, + ) + return _model_dump(result) diff --git a/tests/transform_contracts/registry.py b/tests/transform_contracts/registry.py new file mode 100644 index 00000000000..105df774825 --- /dev/null +++ b/tests/transform_contracts/registry.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final + +from tests.transform_contracts.mistral_ocr import ( + run_get_supported_ocr_params, + run_map_ocr_params, + run_transform_ocr_request, + run_transform_ocr_response, +) +from tests.transform_contracts.schema import JsonValue, TransformationCase + +ContractOperation = Callable[[TransformationCase], JsonValue] + +OPERATIONS: Final[Mapping[str, ContractOperation]] = MappingProxyType( + { + "mistral.ocr.get_supported_ocr_params": run_get_supported_ocr_params, + "mistral.ocr.map_ocr_params": run_map_ocr_params, + "mistral.ocr.transform_ocr_request": run_transform_ocr_request, + "mistral.ocr.transform_ocr_response": run_transform_ocr_response, + } +) + + +def run_contract_case(case: TransformationCase) -> JsonValue: + operation: Final = OPERATIONS.get(case.operation) + if operation is None: + raise ValueError(f"unsupported transformation contract operation: {case.operation}") + return operation(case) diff --git a/tests/transform_contracts/schema.py b/tests/transform_contracts/schema.py new file mode 100644 index 00000000000..12e1a16bd42 --- /dev/null +++ b/tests/transform_contracts/schema.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias, cast + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, model_validator + +if TYPE_CHECKING: + JsonValue: TypeAlias = list["JsonValue"] | dict[str, "JsonValue"] | str | bool | int | float | None +else: + from pydantic import JsonValue + +JsonObject = dict[str, JsonValue] + + +class _ContractModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + +class DocumentUrl(_ContractModel): + type: Literal["document_url"] + document_url: str + + +class ImageUrl(_ContractModel): + type: Literal["image_url"] + image_url: str + + +OCRDocument = Annotated[DocumentUrl | ImageUrl, Field(discriminator="type")] + + +class GetSupportedOCRParamsInput(_ContractModel): + model: str + + +class MapOCRParamsInput(_ContractModel): + model: str + non_default_params: JsonObject + optional_params: JsonObject + + +class TransformOCRRequestInput(_ContractModel): + model: str + document: OCRDocument + optional_params: JsonObject + headers: dict[str, str] + + +class HTTPResponseInput(_ContractModel): + status_code: int = Field(ge=100, le=599) + body: JsonObject + headers: dict[str, str] = Field(default_factory=dict) + + +class TransformOCRResponseInput(_ContractModel): + model: str + response: HTTPResponseInput + + +class OCRRequestOutput(_ContractModel): + data: JsonObject + files: JsonObject | None + + +class _Case(_ContractModel): + id: str = Field(pattern=r"^[a-z0-9][a-z0-9._-]*$") + + +class GetSupportedOCRParamsCase(_Case): + operation: Literal["mistral.ocr.get_supported_ocr_params"] + input: GetSupportedOCRParamsInput + expected: tuple[str, ...] + + +class MapOCRParamsCase(_Case): + operation: Literal["mistral.ocr.map_ocr_params"] + input: MapOCRParamsInput + expected: JsonObject + + +class TransformOCRRequestCase(_Case): + operation: Literal["mistral.ocr.transform_ocr_request"] + input: TransformOCRRequestInput + expected: OCRRequestOutput + + +class TransformOCRResponseCase(_Case): + operation: Literal["mistral.ocr.transform_ocr_response"] + input: TransformOCRResponseInput + expected: JsonObject + + +TransformationCase = Annotated[ + GetSupportedOCRParamsCase | MapOCRParamsCase | TransformOCRRequestCase | TransformOCRResponseCase, + Field(discriminator="operation"), +] + +_JSON_OBJECT_ADAPTER: TypeAdapter[JsonObject] = TypeAdapter(JsonObject) + + +def expected_output(case: TransformationCase) -> JsonValue: + if isinstance(case, GetSupportedOCRParamsCase): + return list(case.expected) + if isinstance(case, TransformOCRRequestCase): + return _JSON_OBJECT_ADAPTER.validate_python(cast(object, case.expected.model_dump(mode="json"))) + return case.expected + + +class ContractSuiteV1(_ContractModel): + schema_version: Literal[1] + cases: tuple[TransformationCase, ...] = Field(min_length=1) + + @model_validator(mode="after") + def validate_id_namespaces(self) -> ContractSuiteV1: + invalid: tuple[TransformationCase, ...] = tuple( + case for case in self.cases if not case.id.startswith(f"{case.operation}.") + ) + if not invalid: + return self + case = invalid[0] + raise ValueError(f"case id must start with '{case.operation}.'") + + +CONTRACT_SUITE_ADAPTER: TypeAdapter[ContractSuiteV1] = TypeAdapter(ContractSuiteV1) diff --git a/tests/transform_contracts/test_contracts.py b/tests/transform_contracts/test_contracts.py new file mode 100644 index 00000000000..0c7c722e6e0 --- /dev/null +++ b/tests/transform_contracts/test_contracts.py @@ -0,0 +1,10 @@ +from typing import Final + +from tests.transform_contracts.registry import run_contract_case +from tests.transform_contracts.schema import JsonValue, TransformationCase, expected_output + + +def test_transformation_contract(contract_case: TransformationCase) -> None: + actual: Final[JsonValue] = run_contract_case(contract_case) + expected: Final[JsonValue] = expected_output(contract_case) + assert actual == expected diff --git a/tests/transform_contracts/test_foundation.py b/tests/transform_contracts/test_foundation.py new file mode 100644 index 00000000000..732f2caa991 --- /dev/null +++ b/tests/transform_contracts/test_foundation.py @@ -0,0 +1,79 @@ +from pathlib import Path +from typing import Final + +import pytest + +from tests.transform_contracts.loader import discover_contract_paths, load_contract_cases + +_VALID_CASE: Final = """ +{ + "schema_version": 1, + "cases": [ + { + "id": "mistral.ocr.get_supported_ocr_params.latest", + "operation": "mistral.ocr.get_supported_ocr_params", + "input": {"model": "mistral-ocr-latest"}, + "expected": ["pages"] + } + ] +} +""" + + +def _write(path: Path, contents: str) -> None: + path.write_text(contents, encoding="utf-8") + + +def test_contract_discovery_is_sorted(tmp_path: Path) -> None: + first: Final = tmp_path / "a.json" + second: Final = tmp_path / "nested" / "b.json" + second.parent.mkdir() + _write(second, _VALID_CASE) + _write(first, _VALID_CASE) + assert discover_contract_paths(tmp_path) == (first, second) + + +def test_invalid_json_fails_loudly(tmp_path: Path) -> None: + _write(tmp_path / "invalid.json", "{") + with pytest.raises(ValueError, match="invalid transformation contract file"): + load_contract_cases(tmp_path) + + +def test_unsupported_schema_version_fails_loudly(tmp_path: Path) -> None: + _write(tmp_path / "future.json", _VALID_CASE.replace('"schema_version": 1', '"schema_version": 2')) + with pytest.raises(ValueError, match="invalid transformation contract file"): + load_contract_cases(tmp_path) + + +def test_duplicate_case_ids_fail_loudly(tmp_path: Path) -> None: + _write(tmp_path / "first.json", _VALID_CASE) + _write(tmp_path / "second.json", _VALID_CASE) + with pytest.raises(ValueError, match="duplicate transformation contract case ids"): + load_contract_cases(tmp_path) + + +def test_missing_required_field_fails_loudly(tmp_path: Path) -> None: + _write( + tmp_path / "missing.json", + _VALID_CASE.replace('"input": {"model": "mistral-ocr-latest"},', '"input": {},'), + ) + with pytest.raises(ValueError, match="invalid transformation contract file"): + load_contract_cases(tmp_path) + + +def test_unsupported_operation_fails_loudly(tmp_path: Path) -> None: + _write( + tmp_path / "unsupported.json", + _VALID_CASE.replace("mistral.ocr.get_supported_ocr_params", "mistral.ocr.unsupported"), + ) + with pytest.raises(ValueError, match="invalid transformation contract file"): + load_contract_cases(tmp_path) + + +def test_case_id_must_match_operation_namespace(tmp_path: Path) -> None: + _write( + tmp_path / "unstable-id.json", + _VALID_CASE.replace("mistral.ocr.get_supported_ocr_params.latest", "mistral.ocr.other.latest"), + ) + with pytest.raises(ValueError, match="invalid transformation contract file"): + load_contract_cases(tmp_path)