Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions litellm/llms/mistral/ocr/transformation.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""
Mistral OCR transformation implementation.
"""

from typing import Any, Dict, Optional

import httpx
Expand Down Expand Up @@ -36,8 +37,12 @@ def get_supported_ocr_params(self, model: str) -> list:
- 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",
Expand All @@ -46,8 +51,12 @@ def get_supported_ocr_params(self, model: str) -> list:
"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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
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
Expand Down Expand Up @@ -39,7 +40,9 @@ def test_existing_params_still_present(self, config: MistralOCRConfig) -> None:
"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:
Expand Down Expand Up @@ -79,3 +82,97 @@ def test_unknown_param_is_dropped(self, config: MistralOCRConfig) -> None:
)
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
Loading