diff --git a/litellm/llms/mistral/ocr/transformation.py b/litellm/llms/mistral/ocr/transformation.py index 3d5e876302..21e0e27a31 100644 --- a/litellm/llms/mistral/ocr/transformation.py +++ b/litellm/llms/mistral/ocr/transformation.py @@ -1,6 +1,7 @@ """ Mistral OCR transformation implementation. """ + from typing import Any, Dict, Optional import httpx @@ -36,8 +37,12 @@ class MistralOCRConfig(BaseOCRConfig): - image_min_size: Minimum size of images to include - bbox_annotation_format: Format for bounding box annotations - document_annotation_format: Format for document annotations + - document_annotation_prompt: Prompt for document annotation extraction - extract_header: Whether to extract document header - extract_footer: Whether to extract document footer + - table_format: Table output format ("markdown" or "html") + - confidence_scores_granularity: Confidence score level ("word" or "page") + - id: Request identifier """ return [ "pages", @@ -46,8 +51,12 @@ class MistralOCRConfig(BaseOCRConfig): "image_min_size", "bbox_annotation_format", "document_annotation_format", + "document_annotation_prompt", "extract_header", "extract_footer", + "table_format", + "confidence_scores_granularity", + "id", ] def map_ocr_params( 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 index ca823d6fb5..97461561a0 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py @@ -4,6 +4,7 @@ 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 pytest from litellm.llms.mistral.ocr.transformation import MistralOCRConfig @@ -39,7 +40,9 @@ class TestGetSupportedOcrParams: "bbox_annotation_format", "document_annotation_format", ]: - assert param in supported, f"Previously supported param '{param}' is missing" + assert ( + param in supported + ), f"Previously supported param '{param}' is missing" class TestMapOcrParams: @@ -79,3 +82,97 @@ class TestMapOcrParams: ) 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", + "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"), + ("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), + ("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