From b6d5728134c2d7f9410139628bae7708cfe62f82 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 16 Apr 2026 19:07:02 -0700 Subject: [PATCH] add support for pages param --- .../document_intelligence/transformation.py | 103 +++++++++++++- .../test_ocr_azure_document_intelligence.py | 126 ++++++++++++++++++ 2 files changed, 223 insertions(+), 6 deletions(-) diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index 6ef309ca67..f2888b1cd2 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -54,10 +54,87 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): """ Get supported OCR parameters for Azure Document Intelligence. - Azure DI has minimal optional parameters compared to Mistral OCR. - Most Mistral-specific params are ignored during transformation. + Azure DI exposes a `pages` query parameter on the analyze endpoint + (1-based, e.g. "1-3,5,7-9"). To keep the public request shape + aligned with Mistral OCR, callers pass `pages` using Mistral + semantics — a list of 0-based integers — or a pre-formatted + Azure-style string. Other Mistral-specific params (e.g. + `include_image_base64`) are not supported by Azure DI and are + ignored during transformation. """ - return [] + return ["pages"] + + def map_ocr_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + ) -> dict: + """ + Map OCR params to Azure DI format. + + Translates Mistral-style `pages` (list[int], 0-based) into Azure's + `pages` query string (1-based, e.g. "1,2,3" or "1-3,5"). A raw + string that already matches Azure's format is passed through + unchanged. + """ + pages = non_default_params.get("pages") + if pages is None: + return optional_params + + normalized = self._normalize_pages_param(pages) + if normalized: + optional_params["pages"] = normalized + return optional_params + + @staticmethod + def _normalize_pages_param(pages: Any) -> str: + """ + Convert a caller-provided `pages` value to Azure DI's query-string + form. Azure expects 1-based page numbers, grammar: `^(\\d+(-\\d+)?)(,\\s*(\\d+(-\\d+)?))*$`. + + Accepted inputs: + - list[int]: Mistral-style 0-based indices. Converted to 1-based + and joined (e.g. [0,1,2] -> "1,2,3"). + - list[str]: tokens like "1" or "3-5". Validated, joined as-is + (treated as Azure-native, i.e. 1-based). + - str: already in Azure format. Validated and whitespace-stripped. + """ + pages_pattern = re.compile(r"^\s*\d+(-\d+)?(\s*,\s*\d+(-\d+)?)*\s*$") + + if isinstance(pages, str): + if not pages_pattern.match(pages): + raise ValueError( + f"Invalid `pages` string for Azure Document Intelligence: " + f"{pages!r}. Expected format like '1-3,5,7-9'." + ) + return pages.replace(" ", "") + + if isinstance(pages, list): + if len(pages) == 0: + return "" + if all(isinstance(p, bool) for p in pages): + raise ValueError("`pages` must be integers, not booleans") + if all(isinstance(p, int) for p in pages): + if any(p < 0 for p in pages): + raise ValueError( + "`pages` integers must be >= 0 (Mistral 0-based indices)" + ) + # Mistral 0-based -> Azure 1-based. + return ",".join(str(p + 1) for p in sorted(set(pages))) + if all(isinstance(p, str) for p in pages): + joined = ",".join(p.strip() for p in pages) + if not pages_pattern.match(joined): + raise ValueError( + f"Invalid `pages` list for Azure Document Intelligence: " + f"{pages!r}. Expected tokens like '1' or '3-5'." + ) + return joined + + raise ValueError( + "`pages` must be a list[int] (0-based, Mistral-style) or a " + "string like '1-3,5,7-9'." + ) def validate_environment( self, @@ -141,7 +218,20 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): # Azure Document Intelligence analyze endpoint # Note: API version 2024-11-30+ uses /documentintelligence/ (not /formrecognizer/) - return f"{api_base}/documentintelligence/documentModels/{model_id}:analyze?api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" + url = ( + f"{api_base}/documentintelligence/documentModels/{model_id}:analyze" + f"?api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" + ) + + # Azure DI accepts `pages` as a query param (1-based, e.g. "1-3,5"). + # `optional_params` has already been normalized in `map_ocr_params`. + pages = optional_params.get("pages") if optional_params else None + if pages: + from urllib.parse import quote + + url += f"&pages={quote(str(pages), safe=',-')}" + + return url def _extract_base64_from_data_uri(self, data_uri: str) -> str: """ @@ -233,8 +323,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): data["urlSource"] = document_url verbose_logger.debug("Using urlSource for Azure Document Intelligence") - # Azure DI doesn't support most Mistral-specific params - # Ignore pages, include_image_base64, etc. + # Azure DI: `pages` is a query param (wired in get_complete_url), + # not a body field. Other Mistral-specific params (e.g. + # include_image_base64, image_limit) are unsupported and ignored. return OCRRequestData(data=data, files=None) diff --git a/tests/ocr_tests/test_ocr_azure_document_intelligence.py b/tests/ocr_tests/test_ocr_azure_document_intelligence.py index 9c1c9e134d..c67cfcbc98 100644 --- a/tests/ocr_tests/test_ocr_azure_document_intelligence.py +++ b/tests/ocr_tests/test_ocr_azure_document_intelligence.py @@ -9,6 +9,10 @@ import os import pytest from base_ocr_unit_tests import BaseOCRTest +from litellm.constants import AZURE_DOCUMENT_INTELLIGENCE_API_VERSION +from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( + AzureDocumentIntelligenceOCRConfig, +) class TestAzureDocumentIntelligenceOCR(BaseOCRTest): @@ -42,3 +46,125 @@ class TestAzureDocumentIntelligenceOCR(BaseOCRTest): "api_base": endpoint, } + +class TestAzureDocumentIntelligencePagesParam: + """ + Unit tests for the Mistral-compatible `pages` parameter translation to + Azure Document Intelligence's `pages` query string. + + These tests exercise the transformation layer directly and do not + require Azure credentials or a network call. + """ + + @pytest.fixture + def cfg(self) -> AzureDocumentIntelligenceOCRConfig: + return AzureDocumentIntelligenceOCRConfig() + + def test_get_supported_ocr_params_includes_pages(self, cfg): + assert cfg.get_supported_ocr_params("prebuilt-layout") == ["pages"] + + def test_map_ocr_params_mistral_zero_based_int_list(self, cfg): + mapped = cfg.map_ocr_params({"pages": [0, 1, 2]}, {}, "prebuilt-layout") + assert mapped == {"pages": "1,2,3"} + + def test_map_ocr_params_dedupes_and_sorts(self, cfg): + mapped = cfg.map_ocr_params({"pages": [2, 0, 0, 1]}, {}, "prebuilt-layout") + assert mapped == {"pages": "1,2,3"} + + def test_map_ocr_params_empty_list_omits_pages(self, cfg): + mapped = cfg.map_ocr_params({"pages": []}, {}, "prebuilt-layout") + assert mapped == {} + + def test_map_ocr_params_azure_native_string_range(self, cfg): + mapped = cfg.map_ocr_params({"pages": "3-9"}, {}, "prebuilt-layout") + assert mapped == {"pages": "3-9"} + + def test_map_ocr_params_azure_native_string_with_spaces_stripped(self, cfg): + mapped = cfg.map_ocr_params({"pages": "1-3, 5"}, {}, "prebuilt-layout") + assert mapped == {"pages": "1-3,5"} + + def test_map_ocr_params_list_of_string_tokens(self, cfg): + mapped = cfg.map_ocr_params({"pages": ["1", "3-5"]}, {}, "prebuilt-layout") + assert mapped == {"pages": "1,3-5"} + + def test_map_ocr_params_invalid_string_raises(self, cfg): + with pytest.raises(ValueError, match="Invalid `pages` string"): + cfg.map_ocr_params({"pages": "a,b"}, {}, "prebuilt-layout") + + def test_map_ocr_params_negative_index_raises(self, cfg): + with pytest.raises(ValueError, match="must be >= 0"): + cfg.map_ocr_params({"pages": [-1]}, {}, "prebuilt-layout") + + def test_map_ocr_params_bool_list_raises(self, cfg): + with pytest.raises(ValueError, match="must be integers, not booleans"): + cfg.map_ocr_params({"pages": [True, False]}, {}, "prebuilt-layout") + + def test_map_ocr_params_unsupported_type_raises(self, cfg): + with pytest.raises(ValueError): + cfg.map_ocr_params({"pages": 5}, {}, "prebuilt-layout") + + def test_get_complete_url_appends_pages_query(self, cfg): + url = cfg.get_complete_url( + api_base="https://example.cognitiveservices.azure.com/", + model="azure_ai/doc-intelligence/prebuilt-layout", + optional_params={"pages": "1-3,5"}, + ) + assert ( + f"api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" in url + ), url + assert "pages=1-3,5" in url, url + assert "/documentintelligence/documentModels/prebuilt-layout:analyze" in url + + def test_get_complete_url_no_pages_when_optional_params_empty(self, cfg): + url = cfg.get_complete_url( + api_base="https://example.cognitiveservices.azure.com", + model="prebuilt-layout", + optional_params={}, + ) + assert "pages=" not in url + + def test_transform_ocr_request_does_not_put_pages_in_body(self, cfg): + req = cfg.transform_ocr_request( + model="prebuilt-layout", + document={ + "type": "document_url", + "document_url": "https://example.com/x.pdf", + }, + optional_params={"pages": "1,2,3"}, + headers={}, + ) + assert req.data is not None + assert "pages" not in req.data + assert req.data.get("urlSource") == "https://example.com/x.pdf" + + def test_end_to_end_mistral_shape_to_azure_query(self, cfg): + """ + Caller sends Mistral-style `pages: [2,3,4,5,6,7,8]` (0-based, + meaning human pages 3-9). LiteLLM should turn that into Azure's + `&pages=3,4,5,6,7,8,9` on the analyze URL, and the body should + still only contain urlSource. + """ + non_default_params = {"pages": [2, 3, 4, 5, 6, 7, 8]} + optional_params = cfg.map_ocr_params( + non_default_params=non_default_params, + optional_params={}, + model="prebuilt-layout", + ) + url = cfg.get_complete_url( + api_base="https://example.cognitiveservices.azure.com", + model="prebuilt-layout", + optional_params=optional_params, + ) + req = cfg.transform_ocr_request( + model="prebuilt-layout", + document={ + "type": "document_url", + "document_url": "https://example.com/x.pdf", + }, + optional_params=optional_params, + headers={}, + ) + + assert "pages=3,4,5,6,7,8,9" in url + assert req.data == {"urlSource": "https://example.com/x.pdf"} +