mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(azure_ai): preserve content, tables, and keyValuePairs in doc-intelligence /v1/ocr (#32018)
* fix(azure_ai): preserve content, tables, and keyValuePairs in doc-intelligence /v1/ocr Azure Document Intelligence analyzeResult.content, .tables, and .keyValuePairs were dropped when normalizing to the Mistral OCR schema. They are now passed through verbatim as top-level response fields, and the duplicated sync/async response parsing is consolidated into one pydantic-validated helper. Also adds the Azure DI features query param (list[str] or comma-separated string, e.g. features=keyValuePairs) which Azure requires for keyValuePairs extraction. * test(azure_ai): replace fastapi jsonable_encoder with model_dump in ocr unit tests
This commit is contained in:
Mateo Wang
GitHub
parent
b61fae397b
commit
ae356cf1fa
4 changed files with 384 additions and 196 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue