diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index c67703f64d7..7d915892a28 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -15,6 +15,7 @@ from typing import Any, Dict from urllib.parse import quote import httpx +from pydantic import BaseModel from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin @@ -38,6 +39,30 @@ from litellm.secret_managers.main import get_secret_str AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" +class AzureDocumentIntelligenceLine(BaseModel): + content: str | None = None + + +class AzureDocumentIntelligencePage(BaseModel): + pageNumber: int | None = None + width: float | None = None + height: float | None = None + unit: str | None = None + lines: tuple[AzureDocumentIntelligenceLine, ...] = () + + +class AzureDocumentIntelligenceAnalyzeResult(BaseModel): + content: str | None = None + pages: tuple[AzureDocumentIntelligencePage, ...] = () + tables: list[dict[str, object]] | None = None + keyValuePairs: list[dict[str, object]] | None = None + + +class AzureDocumentIntelligenceOperation(BaseModel): + status: str | None = None + analyzeResult: AzureDocumentIntelligenceAnalyzeResult | None = None + + class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): """ Azure Document Intelligence OCR transformation configuration. @@ -67,11 +92,14 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): (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. + Azure-style string. Azure DI also exposes a `features` query + parameter enabling add-on capabilities (e.g. "keyValuePairs", + "languages"), passed as a list of feature names or a + comma-separated string. Other Mistral-specific params (e.g. `include_image_base64`) are not supported by Azure DI and are ignored during transformation. """ - return ["pages"] + return ["pages", "features"] def map_ocr_params( self, @@ -85,16 +113,18 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): 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. + unchanged. `features` (list[str] or comma-separated string) is + normalized into Azure's comma-joined `features` query string. """ 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 + features = non_default_params.get("features") + normalized_pages = self._normalize_pages_param(pages) if pages is not None else "" + normalized_features = self._normalize_features_param(features) if features is not None else "" + return { + **optional_params, + **({"pages": normalized_pages} if normalized_pages else {}), + **({"features": normalized_features} if normalized_features else {}), + } @staticmethod def _normalize_pages_param(pages: Any) -> str: @@ -140,6 +170,39 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): raise ValueError("`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'.") + @staticmethod + def _normalize_features_param(features: object) -> str: + """ + Convert a caller-provided `features` value to Azure DI's query-string + form (comma-joined feature names, e.g. "keyValuePairs,languages"). + + Accepted inputs: + - list[str]: feature names like ["keyValuePairs", "languages"]. + - str: a single feature name or comma-separated names. + """ + invalid_features_error = ValueError( + f"Invalid `features` for Azure Document Intelligence: {features!r}. " + f"Expected a list of feature names or a comma-separated string like " + f"'keyValuePairs' or 'keyValuePairs,languages'." + ) + + if isinstance(features, str): + raw_tokens = features.split(",") + elif isinstance(features, list): + if len(features) == 0: + return "" + raw_tokens = [feature for feature in features if isinstance(feature, str)] + if len(raw_tokens) != len(features): + raise invalid_features_error + else: + raise invalid_features_error + + tokens = tuple(token.strip() for token in raw_tokens) + feature_pattern = re.compile(r"^[A-Za-z][A-Za-z0-9]*$") + if not all(feature_pattern.match(token) for token in tokens): + raise invalid_features_error + return ",".join(tokens) + def validate_environment( self, headers: Dict, @@ -228,13 +291,15 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): f"?api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" ) - # Azure DI accepts `pages` as a query param (1-based, e.g. "1-3,5"). + # Azure DI accepts `pages` (1-based, e.g. "1-3,5") and `features` + # (comma-joined names, e.g. "keyValuePairs") as query params. # `optional_params` has already been normalized in `map_ocr_params`. pages = optional_params.get("pages") if optional_params else None - if pages: - url += f"&pages={quote(str(pages), safe=',-')}" + features = optional_params.get("features") if optional_params else None + pages_query = f"&pages={quote(str(pages), safe=',-')}" if pages else "" + features_query = f"&features={quote(str(features), safe=',')}" if features else "" - return url + return f"{url}{pages_query}{features_query}" def _extract_base64_from_data_uri(self, data_uri: str) -> str: """ @@ -328,27 +393,15 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): return OCRRequestData(data=data, files=None) - def _extract_page_markdown(self, page_data: Dict[str, Any]) -> str: - """ - Extract text from Azure DI page and format as markdown. - - Azure DI provides text in 'lines' array. We concatenate them with newlines. - - Args: - page_data: Azure DI page object - - Returns: - Markdown-formatted text - """ - lines = page_data.get("lines", []) - if not lines: - return "" - - # Extract text content from each line - text_lines = [line.get("content", "") for line in lines] - - # Join with newlines to preserve structure - return "\n".join(text_lines) + def _transform_azure_page(self, azure_page: AzureDocumentIntelligencePage) -> OCRPage: + page_number = azure_page.pageNumber if azure_page.pageNumber is not None else 1 + markdown = "\n".join(line.content or "" for line in azure_page.lines) + dimensions = self._convert_dimensions( + width=azure_page.width if azure_page.width is not None else 8.5, + height=azure_page.height if azure_page.height is not None else 11, + unit=azure_page.unit if azure_page.unit is not None else "inch", + ) + return OCRPage(index=page_number - 1, markdown=markdown, dimensions=dimensions) def _convert_dimensions(self, width: float, height: float, unit: str) -> OCRPageDimensions: """ @@ -526,6 +579,52 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): retry_after = self._get_retry_after(response=response) await asyncio.sleep(retry_after) + def _get_polling_target(self, raw_response: httpx.Response) -> tuple[str, Dict[str, str]]: + operation_url = raw_response.headers.get("Operation-Location") + if not operation_url: + raise ValueError("Azure Document Intelligence returned 202 but no Operation-Location header found") + + # Reject cross-origin polling URLs — the auth headers + # below would otherwise leak to whatever URL the upstream + # (or an attacker-controlled upstream) returns. VERIA-51. + try: + assert_same_origin(operation_url, str(raw_response.request.url)) + except SSRFError as ssrf_err: + raise ValueError(f"Azure Document Intelligence: rejected polling URL ({ssrf_err})") + + poll_headers = {"Ocp-Apim-Subscription-Key": raw_response.request.headers.get("Ocp-Apim-Subscription-Key", "")} + return operation_url, poll_headers + + def _transform_completed_response(self, model: str, raw_response: httpx.Response) -> OCRResponse: + """ + Transform a completed Azure Document Intelligence analyze operation + into the Mistral OCR response shape, preserving Azure-native + `analyzeResult` fields (`content`, `tables`, `keyValuePairs`) as + top-level response fields. + """ + operation = AzureDocumentIntelligenceOperation.model_validate(raw_response.json()) + + verbose_logger.debug(f"Azure Document Intelligence response status: {operation.status}") + + if operation.status != "succeeded": + raise ValueError(f"Azure Document Intelligence analysis failed with status: {operation.status}") + + analyze_result = ( + operation.analyzeResult if operation.analyzeResult is not None else AzureDocumentIntelligenceAnalyzeResult() + ) + mistral_pages = [self._transform_azure_page(azure_page) for azure_page in analyze_result.pages] + usage_info = OCRUsageInfo(pages_processed=len(mistral_pages), doc_size_bytes=None) + + return OCRResponse( + pages=mistral_pages, + model=model, + usage_info=usage_info, + object="ocr", + content=analyze_result.content, + tables=analyze_result.tables, + keyValuePairs=analyze_result.keyValuePairs, + ) + def transform_ocr_response( self, model: str, @@ -552,11 +651,13 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): "unit": "inch", "lines": [{"content": "text", "boundingBox": [...]}] } - ] + ], + "tables": [...], + "keyValuePairs": [...] } } - Mistral OCR format: + Mistral OCR format (with Azure-native fields preserved): { "pages": [ { @@ -567,7 +668,10 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ], "model": "azure_ai/doc-intelligence/prebuilt-layout", "usage_info": {"pages_processed": 1}, - "object": "ocr" + "object": "ocr", + "content": "Full document text...", + "tables": [...], + "keyValuePairs": [...] } Args: @@ -578,86 +682,17 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): Returns: OCRResponse in Mistral format """ - try: - # Check if we got 202 Accepted (async operation started) - if raw_response.status_code == 202: - verbose_logger.debug("Azure DI returned 202 Accepted, polling operation...") + if raw_response.status_code != 202: + return self._transform_completed_response(model=model, raw_response=raw_response) - # Get Operation-Location header - operation_url = raw_response.headers.get("Operation-Location") - if not operation_url: - raise ValueError("Azure Document Intelligence returned 202 but no Operation-Location header found") - - # Reject cross-origin polling URLs — the auth headers - # below would otherwise leak to whatever URL the upstream - # (or an attacker-controlled upstream) returns. VERIA-51. - try: - assert_same_origin(operation_url, str(raw_response.request.url)) - except SSRFError as ssrf_err: - raise ValueError(f"Azure Document Intelligence: rejected polling URL ({ssrf_err})") - - # Get headers for polling (need auth) - poll_headers = { - "Ocp-Apim-Subscription-Key": raw_response.request.headers.get("Ocp-Apim-Subscription-Key", "") - } - - # Get timeout from kwargs or use default - timeout_secs = AZURE_OPERATION_POLLING_TIMEOUT - - # Poll until operation completes - raw_response = self._poll_operation_sync( - operation_url=operation_url, - headers=poll_headers, - timeout_secs=timeout_secs, - ) - - # Now parse the completed response - response_json = raw_response.json() - - verbose_logger.debug(f"Azure Document Intelligence response status: {response_json.get('status')}") - - # Check if request succeeded - status = response_json.get("status") - if status != "succeeded": - raise ValueError(f"Azure Document Intelligence analysis failed with status: {status}") - - # Extract analyze result - analyze_result = response_json.get("analyzeResult", {}) - azure_pages = analyze_result.get("pages", []) - - # Transform pages to Mistral format - mistral_pages = [] - for azure_page in azure_pages: - page_number = azure_page.get("pageNumber", 1) - index = page_number - 1 # Convert to 0-based index - - # Extract markdown text - markdown = self._extract_page_markdown(azure_page) - - # Convert dimensions - width = azure_page.get("width", 8.5) - height = azure_page.get("height", 11) - unit = azure_page.get("unit", "inch") - dimensions = self._convert_dimensions(width=width, height=height, unit=unit) - - # Build OCR page - ocr_page = OCRPage(index=index, markdown=markdown, dimensions=dimensions) - mistral_pages.append(ocr_page) - - # Build usage info - usage_info = OCRUsageInfo(pages_processed=len(mistral_pages), doc_size_bytes=None) - - # Return Mistral OCR response - return OCRResponse( - pages=mistral_pages, - model=model, - usage_info=usage_info, - object="ocr", - ) - - except Exception as e: - verbose_logger.error(f"Error parsing Azure Document Intelligence response: {e}") - raise e + verbose_logger.debug("Azure DI returned 202 Accepted, polling operation...") + operation_url, poll_headers = self._get_polling_target(raw_response) + completed_response = self._poll_operation_sync( + operation_url=operation_url, + headers=poll_headers, + timeout_secs=AZURE_OPERATION_POLLING_TIMEOUT, + ) + return self._transform_completed_response(model=model, raw_response=completed_response) async def async_transform_ocr_response( self, @@ -680,81 +715,14 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): Returns: OCRResponse in Mistral format """ - try: - # Check if we got 202 Accepted (async operation started) - if raw_response.status_code == 202: - verbose_logger.debug("Azure DI returned 202 Accepted, polling operation (async)...") + if raw_response.status_code != 202: + return self._transform_completed_response(model=model, raw_response=raw_response) - # Get Operation-Location header - operation_url = raw_response.headers.get("Operation-Location") - if not operation_url: - raise ValueError("Azure Document Intelligence returned 202 but no Operation-Location header found") - - # Reject cross-origin polling URLs (see sync path). VERIA-51. - try: - assert_same_origin(operation_url, str(raw_response.request.url)) - except SSRFError as ssrf_err: - raise ValueError(f"Azure Document Intelligence: rejected polling URL ({ssrf_err})") - - # Get headers for polling (need auth) - poll_headers = { - "Ocp-Apim-Subscription-Key": raw_response.request.headers.get("Ocp-Apim-Subscription-Key", "") - } - - # Get timeout from kwargs or use default - timeout_secs = AZURE_OPERATION_POLLING_TIMEOUT - - # Poll until operation completes (async) - raw_response = await self._poll_operation_async( - operation_url=operation_url, - headers=poll_headers, - timeout_secs=timeout_secs, - ) - - # Now parse the completed response - response_json = raw_response.json() - - verbose_logger.debug(f"Azure Document Intelligence response status: {response_json.get('status')}") - - # Check if request succeeded - status = response_json.get("status") - if status != "succeeded": - raise ValueError(f"Azure Document Intelligence analysis failed with status: {status}") - - # Extract analyze result - analyze_result = response_json.get("analyzeResult", {}) - azure_pages = analyze_result.get("pages", []) - - # Transform pages to Mistral format - mistral_pages = [] - for azure_page in azure_pages: - page_number = azure_page.get("pageNumber", 1) - index = page_number - 1 # Convert to 0-based index - - # Extract markdown text - markdown = self._extract_page_markdown(azure_page) - - # Convert dimensions - width = azure_page.get("width", 8.5) - height = azure_page.get("height", 11) - unit = azure_page.get("unit", "inch") - dimensions = self._convert_dimensions(width=width, height=height, unit=unit) - - # Build OCR page - ocr_page = OCRPage(index=index, markdown=markdown, dimensions=dimensions) - mistral_pages.append(ocr_page) - - # Build usage info - usage_info = OCRUsageInfo(pages_processed=len(mistral_pages), doc_size_bytes=None) - - # Return Mistral OCR response - return OCRResponse( - pages=mistral_pages, - model=model, - usage_info=usage_info, - object="ocr", - ) - - except Exception as e: - verbose_logger.error(f"Error parsing Azure Document Intelligence response (async): {e}") - raise e + verbose_logger.debug("Azure DI returned 202 Accepted, polling operation (async)...") + operation_url, poll_headers = self._get_polling_target(raw_response) + completed_response = await self._poll_operation_async( + operation_url=operation_url, + headers=poll_headers, + timeout_secs=AZURE_OPERATION_POLLING_TIMEOUT, + ) + return self._transform_completed_response(model=model, raw_response=completed_response) diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index a38e5bfdcd6..0d878bd308c 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -70,6 +70,9 @@ class OCRResponse(LiteLLMPydanticObjectBase): model: str document_annotation: Any | None = None usage_info: OCRUsageInfo | None = None + content: str | None = None + tables: list[dict[str, object]] | None = None + keyValuePairs: list[dict[str, object]] | None = None object: str = "ocr" model_config = {"extra": "allow"} diff --git a/tests/ocr_tests/test_ocr_azure_document_intelligence.py b/tests/ocr_tests/test_ocr_azure_document_intelligence.py index 7269890b7b6..09c21842ad7 100644 --- a/tests/ocr_tests/test_ocr_azure_document_intelligence.py +++ b/tests/ocr_tests/test_ocr_azure_document_intelligence.py @@ -61,8 +61,8 @@ class TestAzureDocumentIntelligencePagesParam: 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_get_supported_ocr_params_includes_pages_and_features(self, cfg): + assert cfg.get_supported_ocr_params("prebuilt-layout") == ["pages", "features"] def test_map_ocr_params_mistral_zero_based_int_list(self, cfg): mapped = cfg.map_ocr_params({"pages": [0, 1, 2]}, {}, "prebuilt-layout") diff --git a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py index e638be68ec0..39d6f1dc355 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py @@ -1,3 +1,6 @@ +from unittest.mock import MagicMock + +import httpx import pytest from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( @@ -31,3 +34,217 @@ def test_should_reject_dot_segment_azure_document_intelligence_model_id(): optional_params={}, litellm_params={}, ) + + +AZURE_TABLES = [ + { + "rowCount": 2, + "columnCount": 2, + "cells": [ + {"kind": "columnHeader", "rowIndex": 0, "columnIndex": 0, "content": "Item"}, + {"kind": "columnHeader", "rowIndex": 0, "columnIndex": 1, "content": "Price"}, + {"rowIndex": 1, "columnIndex": 0, "content": "Widget"}, + {"rowIndex": 1, "columnIndex": 1, "content": "$100.00"}, + ], + }, + { + "rowCount": 1, + "columnCount": 1, + "cells": [{"rowIndex": 0, "columnIndex": 0, "content": "Totals"}], + }, +] + +AZURE_KEY_VALUE_PAIRS = [ + {"key": {"content": "Invoice No"}, "value": {"content": "INV-12345"}, "confidence": 0.98}, + {"key": {"content": "Total"}, "value": {"content": "$100.00"}, "confidence": 0.95}, +] + +AZURE_ANALYZE_SUCCEEDED = { + "status": "succeeded", + "createdDateTime": "2026-07-02T00:00:00Z", + "lastUpdatedDateTime": "2026-07-02T00:00:05Z", + "analyzeResult": { + "apiVersion": "2024-11-30", + "modelId": "prebuilt-layout", + "content": "Invoice\nInvoice No: INV-12345\nTotal: $100.00", + "pages": [ + { + "pageNumber": 1, + "width": 8.5, + "height": 11, + "unit": "inch", + "lines": [ + {"content": "Invoice"}, + {"content": "Invoice No: INV-12345"}, + {"content": "Total: $100.00"}, + ], + } + ], + "tables": AZURE_TABLES, + "keyValuePairs": AZURE_KEY_VALUE_PAIRS, + }, +} + + +def _completed_response(payload: dict) -> httpx.Response: + return httpx.Response( + status_code=200, + json=payload, + request=httpx.Request("GET", "https://example.cognitiveservices.azure.com/analyzeResults/xyz"), + ) + + +def _assert_native_fields_preserved(serialized: dict) -> None: + assert serialized["content"] == "Invoice\nInvoice No: INV-12345\nTotal: $100.00" + assert serialized["tables"] == AZURE_TABLES + assert serialized["keyValuePairs"] == AZURE_KEY_VALUE_PAIRS + assert serialized["object"] == "ocr" + assert serialized["usage_info"]["pages_processed"] == 1 + assert serialized["pages"][0]["index"] == 0 + assert serialized["pages"][0]["markdown"] == "Invoice\nInvoice No: INV-12345\nTotal: $100.00" + assert serialized["pages"][0]["dimensions"] == {"width": 816, "height": 1056, "dpi": 96} + + +def test_transform_ocr_response_preserves_azure_native_fields(): + config = AzureDocumentIntelligenceOCRConfig() + + result = config.transform_ocr_response( + model="azure_ai/doc-intelligence/prebuilt-layout", + raw_response=_completed_response(AZURE_ANALYZE_SUCCEEDED), + logging_obj=MagicMock(), + ) + + _assert_native_fields_preserved(result.model_dump()) + + +@pytest.mark.asyncio +async def test_async_transform_ocr_response_preserves_azure_native_fields(): + config = AzureDocumentIntelligenceOCRConfig() + + result = await config.async_transform_ocr_response( + model="azure_ai/doc-intelligence/prebuilt-layout", + raw_response=_completed_response(AZURE_ANALYZE_SUCCEEDED), + logging_obj=MagicMock(), + ) + + _assert_native_fields_preserved(result.model_dump()) + + +def test_transform_ocr_response_tolerates_missing_native_fields(): + config = AzureDocumentIntelligenceOCRConfig() + payload = { + "status": "succeeded", + "analyzeResult": { + "pages": [ + { + "pageNumber": 1, + "width": 8.5, + "height": 11, + "unit": "inch", + "lines": [{"content": "hello"}], + } + ], + }, + } + + result = config.transform_ocr_response( + model="azure_ai/doc-intelligence/prebuilt-read", + raw_response=_completed_response(payload), + logging_obj=MagicMock(), + ) + + serialized = result.model_dump() + assert serialized["pages"][0]["markdown"] == "hello" + assert serialized["content"] is None + assert serialized["tables"] is None + assert serialized["keyValuePairs"] is None + + +def test_transform_ocr_response_non_succeeded_status_raises(): + config = AzureDocumentIntelligenceOCRConfig() + + with pytest.raises(ValueError, match="failed with status: failed"): + config.transform_ocr_response( + model="azure_ai/doc-intelligence/prebuilt-layout", + raw_response=_completed_response({"status": "failed"}), + logging_obj=MagicMock(), + ) + + +def test_get_supported_ocr_params_includes_features(): + config = AzureDocumentIntelligenceOCRConfig() + + assert config.get_supported_ocr_params("prebuilt-layout") == ["pages", "features"] + + +@pytest.mark.parametrize( + "features,expected", + [ + (["keyValuePairs"], "keyValuePairs"), + (["keyValuePairs", "languages"], "keyValuePairs,languages"), + ("keyValuePairs", "keyValuePairs"), + ("keyValuePairs,languages", "keyValuePairs,languages"), + ("keyValuePairs, languages", "keyValuePairs,languages"), + ], +) +def test_map_ocr_params_features(features, expected): + config = AzureDocumentIntelligenceOCRConfig() + + mapped = config.map_ocr_params({"features": features}, {}, "prebuilt-layout") + + assert mapped == {"features": expected} + + +def test_map_ocr_params_empty_features_list_omitted(): + config = AzureDocumentIntelligenceOCRConfig() + + assert config.map_ocr_params({"features": []}, {}, "prebuilt-layout") == {} + + +@pytest.mark.parametrize( + "features", + [ + "keyValuePairs&pages=9", + "key value pairs", + "", + [1, 2], + [["keyValuePairs"]], + {"feature": "keyValuePairs"}, + 5, + ], +) +def test_map_ocr_params_invalid_features_raises(features): + config = AzureDocumentIntelligenceOCRConfig() + + with pytest.raises(ValueError, match="Invalid `features`"): + config.map_ocr_params({"features": features}, {}, "prebuilt-layout") + + +def test_get_complete_url_appends_features_query(): + config = AzureDocumentIntelligenceOCRConfig() + + url = config.get_complete_url( + api_base="https://example.cognitiveservices.azure.com", + model="azure_ai/doc-intelligence/prebuilt-layout", + optional_params={"features": "keyValuePairs"}, + ) + + assert "&features=keyValuePairs" in url + + +def test_get_complete_url_combines_pages_and_features(): + config = AzureDocumentIntelligenceOCRConfig() + + optional_params = config.map_ocr_params( + {"pages": [0, 1, 2], "features": ["keyValuePairs", "languages"]}, + {}, + "prebuilt-layout", + ) + url = config.get_complete_url( + api_base="https://example.cognitiveservices.azure.com", + model="prebuilt-layout", + optional_params=optional_params, + ) + + assert "&pages=1,2,3" in url + assert "&features=keyValuePairs,languages" in url