mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
test(ocr): add shared Mistral transformation contracts
This commit is contained in:
Yujong Lee
parent
3829418878
commit
45258ecb85
10 changed files with 933 additions and 236 deletions
|
|
@ -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]
|
||||
1
tests/transform_contracts/__init__.py
Normal file
1
tests/transform_contracts/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
"""Shared provider transformation contracts."""
|
||||
530
tests/transform_contracts/cases/mistral/ocr.json
Normal file
530
tests/transform_contracts/cases/mistral/ocr.json
Normal file
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
19
tests/transform_contracts/conftest.py
Normal file
19
tests/transform_contracts/conftest.py
Normal file
|
|
@ -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")
|
||||
38
tests/transform_contracts/loader.py
Normal file
38
tests/transform_contracts/loader.py
Normal file
|
|
@ -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))
|
||||
101
tests/transform_contracts/mistral_ocr.py
Normal file
101
tests/transform_contracts/mistral_ocr.py
Normal file
|
|
@ -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)
|
||||
31
tests/transform_contracts/registry.py
Normal file
31
tests/transform_contracts/registry.py
Normal file
|
|
@ -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)
|
||||
124
tests/transform_contracts/schema.py
Normal file
124
tests/transform_contracts/schema.py
Normal file
|
|
@ -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)
|
||||
10
tests/transform_contracts/test_contracts.py
Normal file
10
tests/transform_contracts/test_contracts.py
Normal file
|
|
@ -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
|
||||
79
tests/transform_contracts/test_foundation.py
Normal file
79
tests/transform_contracts/test_foundation.py
Normal file
|
|
@ -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)
|
||||
Loading…
Add table
Reference in a new issue