add support for pages param
This commit is contained in:
shivam
parent
bdb4f396bb
commit
b6d5728134
@ -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)
|
||||
|
||||
|
||||
@ -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"}
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user