diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index 64b92f7d847..3d1d0fcd6c3 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -43,6 +43,7 @@ jobs: tests/test_litellm/proxy/video_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/image_endpoints + tests/test_litellm/proxy/ocr_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/a2a diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index b95fa20c41e..e7b94b3812b 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -11,6 +11,7 @@ The operation location must be polled until the analysis completes. import asyncio import re import time +from collections.abc import Mapping from typing import Any, Final from urllib.parse import quote @@ -23,15 +24,19 @@ from litellm.constants import ( AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI, AZURE_OPERATION_POLLING_TIMEOUT, ) +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin, encode_url_path_segment from litellm.llms.base_llm.ocr.transformation import ( + OCR_REQUEST_FORMAT_PARAM, BaseOCRConfig, DocumentType, OCRPage, OCRPageDimensions, OCRRequestData, + OCRRequestFormat, OCRResponse, OCRUsageInfo, + parse_ocr_request_format, ) from litellm.secret_managers.main import get_secret_str @@ -97,8 +102,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): comma-separated string. Other Mistral-specific params (e.g. `include_image_base64`) are not supported by Azure DI and are ignored during transformation. + + `req_format` selects the response shape: "litellm" (default) returns + the normalized OCR schema, "native" returns Azure DI's own analyze + operation payload as-is. """ - return ["pages", "features"] + return ["pages", "features", OCR_REQUEST_FORMAT_PARAM] def map_ocr_params( self, @@ -117,14 +126,27 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): """ pages: Final = non_default_params.get("pages") features: Final = non_default_params.get("features") + request_format: Final = non_default_params.get(OCR_REQUEST_FORMAT_PARAM) normalized_pages: Final = self._normalize_pages_param(pages) if pages is not None else "" normalized_features: Final = 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 {}), + **( + {OCR_REQUEST_FORMAT_PARAM: self._parse_request_format(request_format, model)} + if request_format is not None + else {} + ), } + @staticmethod + def _parse_request_format(request_format: object, model: str) -> OCRRequestFormat: + try: + return parse_ocr_request_format(request_format) + except ValueError as e: + raise UnsupportedParamsError(message=f"{e}", model=model, llm_provider="azure_ai") from e + @staticmethod def _normalize_pages_param(pages: Any) -> str: """ @@ -594,14 +616,33 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): 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: + @staticmethod + def _get_request_format(optional_params: object) -> OCRRequestFormat: + if not isinstance(optional_params, dict): + return "litellm" + request_format: Final = optional_params.get(OCR_REQUEST_FORMAT_PARAM) + if request_format is None: + return "litellm" + return parse_ocr_request_format(request_format) + + def _transform_completed_response( + self, + model: str, + raw_response: httpx.Response, + request_format: OCRRequestFormat, + ) -> 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. + + When `request_format` is "native", the untouched Azure operation + payload is attached to the response's hidden params so the proxy can + return it verbatim while cost tracking still reads `usage_info`. """ - operation: Final = AzureDocumentIntelligenceOperation.model_validate(raw_response.json()) + raw_operation: Final[Mapping[str, object]] = raw_response.json() + operation: Final = AzureDocumentIntelligenceOperation.model_validate(raw_operation) verbose_logger.debug("Azure Document Intelligence response status: %s", operation.status) @@ -614,7 +655,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): mistral_pages: Final = [self._transform_azure_page(azure_page) for azure_page in analyze_result.pages] usage_info: Final = OCRUsageInfo(pages_processed=len(mistral_pages), doc_size_bytes=None) - return OCRResponse( + response: Final = OCRResponse( pages=mistral_pages, model=model, usage_info=usage_info, @@ -624,6 +665,11 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): keyValuePairs=analyze_result.keyValuePairs, ) + if request_format == "native": + response.set_provider_native_response(raw_operation) + + return response + def transform_ocr_response( self, model: str, @@ -681,8 +727,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): Returns: OCRResponse in Mistral format """ + request_format: Final = self._get_request_format(kwargs.get("optional_params")) + if raw_response.status_code != 202: - return self._transform_completed_response(model=model, raw_response=raw_response) + return self._transform_completed_response( + model=model, raw_response=raw_response, request_format=request_format + ) verbose_logger.debug("Azure DI returned 202 Accepted, polling operation...") operation_url, poll_headers = self._get_polling_target(raw_response) @@ -691,7 +741,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): headers=poll_headers, timeout_secs=AZURE_OPERATION_POLLING_TIMEOUT, ) - return self._transform_completed_response(model=model, raw_response=completed_response) + return self._transform_completed_response( + model=model, raw_response=completed_response, request_format=request_format + ) async def async_transform_ocr_response( self, @@ -714,8 +766,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): Returns: OCRResponse in Mistral format """ + request_format: Final = self._get_request_format(kwargs.get("optional_params")) + if raw_response.status_code != 202: - return self._transform_completed_response(model=model, raw_response=raw_response) + return self._transform_completed_response( + model=model, raw_response=raw_response, request_format=request_format + ) verbose_logger.debug("Azure DI returned 202 Accepted, polling operation (async)...") operation_url, poll_headers = self._get_polling_target(raw_response) @@ -724,4 +780,6 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): headers=poll_headers, timeout_secs=AZURE_OPERATION_POLLING_TIMEOUT, ) - return self._transform_completed_response(model=model, raw_response=completed_response) + return self._transform_completed_response( + model=model, raw_response=completed_response, request_format=request_format + ) diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 96f86bc8dc0..d1c77186ea8 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -2,7 +2,8 @@ Base OCR transformation configuration. """ -from typing import TYPE_CHECKING, Any +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Literal import httpx from pydantic import PrivateAttr @@ -21,6 +22,26 @@ else: # File-type inputs are preprocessed to this format in litellm/ocr/main.py. DocumentType = dict[str, str] +OCRRequestFormat = Literal["litellm", "native"] + +OCR_REQUEST_FORMATS: Final[tuple[OCRRequestFormat, ...]] = ("litellm", "native") + +OCR_REQUEST_FORMAT_PARAM: Final = "req_format" + +OCR_REQUEST_FORMAT_HEADER: Final = "x-req-format" + +PROVIDER_NATIVE_RESPONSE_KEY: Final = "provider_native_response" + + +def parse_ocr_request_format(value: object) -> OCRRequestFormat: + if value == "litellm": + return "litellm" + if value == "native": + return "native" + raise ValueError( + f"Invalid `{OCR_REQUEST_FORMAT_PARAM}`: {value!r}. Expected one of {', '.join(OCR_REQUEST_FORMATS)}." + ) + class OCRPageDimensions(LiteLLMPydanticObjectBase): """Page dimensions from OCR response.""" @@ -80,6 +101,15 @@ class OCRResponse(LiteLLMPydanticObjectBase): # Define private attributes using PrivateAttr _hidden_params: dict = PrivateAttr(default_factory=dict) + def set_provider_native_response(self, native_response: Mapping[str, object]) -> None: + """Keep the provider's own response payload alongside the normalized one.""" + self._hidden_params[PROVIDER_NATIVE_RESPONSE_KEY] = native_response + + def get_provider_native_response(self) -> Mapping[str, object] | None: + """The provider's own response payload, when `req_format=native` was requested.""" + native_response: Final = self._hidden_params.get(PROVIDER_NATIVE_RESPONSE_KEY) + return native_response if isinstance(native_response, dict) else None + class OCRRequestData(LiteLLMPydanticObjectBase): """OCR request data structure.""" diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index e22ec89847e..d67497dd4da 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1556,12 +1556,14 @@ class BaseLLMHTTPHandler: model: str, response: httpx.Response, logging_obj: LiteLLMLoggingObj, + optional_params: Mapping[str, object], ) -> OCRResponse: """Shared logic for transforming OCR responses.""" return provider_config.transform_ocr_response( model=model, raw_response=response, logging_obj=logging_obj, + optional_params=optional_params, ) def ocr( @@ -1637,6 +1639,7 @@ class BaseLLMHTTPHandler: model=model, response=response, logging_obj=logging_obj, + optional_params=optional_params, ) async def async_ocr( @@ -1699,6 +1702,7 @@ class BaseLLMHTTPHandler: model=model, raw_response=response, logging_obj=logging_obj, + optional_params=optional_params, ) def search( diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index d02adca8a6d..b918f013700 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -21,7 +21,12 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.azure_ai.ocr.common_utils import ( is_azure_document_intelligence_model, ) -from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse +from litellm.llms.base_llm.ocr.transformation import ( + OCR_REQUEST_FORMAT_PARAM, + BaseOCRConfig, + OCRResponse, + parse_ocr_request_format, +) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import ocr as rust_ocr_bridge from litellm.types.router import GenericLiteLLMParams @@ -124,6 +129,24 @@ def _prepare_ocr_request( litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs) supported_params: Final = ocr_provider_config.get_supported_ocr_params(model=model) + requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM) + if requested_format is not None: + try: + parsed_format: Final = parse_ocr_request_format(requested_format) + except ValueError as e: + raise litellm.exceptions.UnsupportedParamsError( + message=f"{e}", model=model, llm_provider=custom_llm_provider + ) from e + if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native": + raise litellm.exceptions.UnsupportedParamsError( + message=( + f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, " + f"model: {model}" + ), + model=model, + llm_provider=custom_llm_provider, + ) + non_default_params: Final = {} for param in supported_params: if param in kwargs: @@ -166,6 +189,8 @@ def _prepare_ocr_request( def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: + if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native": + return False return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index eb2f132456f..ebf4d988fdd 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -1,13 +1,20 @@ #### OCR Endpoints ##### import json +from collections.abc import Mapping from typing import Any, Final, cast import orjson -from fastapi import APIRouter, Depends, Request, Response, UploadFile +from fastapi import APIRouter, Depends, HTTPException, Request, Response, UploadFile from fastapi.responses import ORJSONResponse from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.ocr.transformation import ( + OCR_REQUEST_FORMAT_HEADER, + OCR_REQUEST_FORMAT_PARAM, + OCRResponse, + parse_ocr_request_format, +) from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth @@ -41,6 +48,48 @@ def _build_document_from_upload( ) +def _with_request_format(data: Mapping[str, Any], request: Request) -> Mapping[str, Any]: + """ + Resolve the requested response format from the body or the `x-req-format` header. + + An explicit `req_format` in the body wins over the header. + """ + body_value: Final = data.get(OCR_REQUEST_FORMAT_PARAM) + header_value: Final = request.headers.get(OCR_REQUEST_FORMAT_HEADER) + raw_value: Final = body_value if body_value is not None else header_value + if raw_value is None: + return data + try: + request_format: Final = parse_ocr_request_format( + raw_value.strip().lower() if isinstance(raw_value, str) else raw_value + ) + except ValueError as e: + raise HTTPException(status_code=400, detail={"error": f"{e}"}) + return {**data, OCR_REQUEST_FORMAT_PARAM: request_format} + + +def _native_response(response: object, fastapi_response: Response) -> Response | None: + """ + Return the provider's native payload when the caller asked for + `req_format=native` and the provider config captured it, carrying over the + LiteLLM response headers (cost, call id, etc.) built for the normalized response. + """ + if not isinstance(response, OCRResponse): + return None + native_payload: Final = response.get_provider_native_response() + if native_payload is None: + return None + return Response( + content=orjson.dumps(native_payload), + media_type="application/json", + headers={ + key: value + for key, value in fastapi_response.headers.items() + if key.lower() not in ("content-length", "content-type") + }, + ) + + async def _parse_multipart_form(request: Request) -> dict[str, Any]: """ Extract OCR data from a multipart form request. @@ -105,7 +154,12 @@ async def _parse_multipart_form(request: Request) -> dict[str, Any]: return data -async def _parse_ocr_request(request: Request) -> dict[str, Any]: +async def _parse_ocr_request(request: Request) -> Mapping[str, Any]: + """Parse an OCR request and apply the `x-req-format` header, if any.""" + return _with_request_format(await _parse_ocr_request_body(request), request) + + +async def _parse_ocr_request_body(request: Request) -> dict[str, Any]: """ Parse an OCR request, supporting both JSON and multipart form data. @@ -238,6 +292,11 @@ async def ocr( -F "model=mistral-ocr" \ -F "file=@document.pdf" ``` + + Response format is normalized to the LiteLLM OCR schema by default. Providers + that support it (Azure Document Intelligence) can return their own payload + instead, with cost tracking unchanged, via `x-req-format: native` (or + `"req_format": "native"` in the body). """ from litellm.proxy.proxy_server import ( general_settings, @@ -256,12 +315,12 @@ async def ocr( data: dict = {} try: # Parse request body (JSON or multipart form) - data = await _parse_ocr_request(request) + data = dict(await _parse_ocr_request(request)) # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) - return await processor.base_process_llm_request( + response: Final = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -279,6 +338,8 @@ async def ocr( user_api_base=user_api_base, version=version, ) + + return _native_response(response, fastapi_response) or response except Exception as e: processor = ProxyBaseLLMRequestProcessing(data=data) raise await processor._handle_llm_api_exception( 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 39d6f1dc355..66f4f432eb8 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 @@ -3,6 +3,8 @@ from unittest.mock import MagicMock import httpx import pytest +from litellm.exceptions import UnsupportedParamsError + from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( AzureDocumentIntelligenceOCRConfig, ) @@ -174,7 +176,101 @@ def test_transform_ocr_response_non_succeeded_status_raises(): def test_get_supported_ocr_params_includes_features(): config = AzureDocumentIntelligenceOCRConfig() - assert config.get_supported_ocr_params("prebuilt-layout") == ["pages", "features"] + assert config.get_supported_ocr_params("prebuilt-layout") == ["pages", "features", "req_format"] + + +AZURE_ANALYZE_WITH_NATIVE_ONLY_FIELDS = { + **AZURE_ANALYZE_SUCCEEDED, + "analyzeResult": { + **AZURE_ANALYZE_SUCCEEDED["analyzeResult"], + "paragraphs": [{"content": "Invoice", "spans": [{"offset": 0, "length": 7}]}], + "pages": [ + { + **AZURE_ANALYZE_SUCCEEDED["analyzeResult"]["pages"][0], + "angle": 0.13, + "spans": [{"offset": 0, "length": 44}], + "words": [{"content": "Invoice", "confidence": 0.994, "polygon": [1, 2, 3, 4]}], + } + ], + }, +} + + +def test_transform_ocr_response_native_format_carries_raw_operation(): + config = AzureDocumentIntelligenceOCRConfig() + + result = config.transform_ocr_response( + model="azure_ai/doc-intelligence/prebuilt-layout", + raw_response=_completed_response(AZURE_ANALYZE_WITH_NATIVE_ONLY_FIELDS), + logging_obj=MagicMock(), + optional_params={"req_format": "native"}, + ) + + assert result.get_provider_native_response() == AZURE_ANALYZE_WITH_NATIVE_ONLY_FIELDS + # cost tracking reads usage_info off the normalized response, so it must survive native mode + assert result.usage_info is not None + assert result.usage_info.pages_processed == 1 + _assert_native_fields_preserved(result.model_dump()) + + +@pytest.mark.asyncio +async def test_async_transform_ocr_response_native_format_carries_raw_operation(): + config = AzureDocumentIntelligenceOCRConfig() + + result = await config.async_transform_ocr_response( + model="azure_ai/doc-intelligence/prebuilt-layout", + raw_response=_completed_response(AZURE_ANALYZE_WITH_NATIVE_ONLY_FIELDS), + logging_obj=MagicMock(), + optional_params={"req_format": "native"}, + ) + + assert result.get_provider_native_response() == AZURE_ANALYZE_WITH_NATIVE_ONLY_FIELDS + assert result.usage_info is not None + assert result.usage_info.pages_processed == 1 + + +@pytest.mark.parametrize("optional_params", [{}, {"req_format": "litellm"}]) +def test_transform_ocr_response_default_format_omits_raw_operation(optional_params): + config = AzureDocumentIntelligenceOCRConfig() + + result = config.transform_ocr_response( + model="azure_ai/doc-intelligence/prebuilt-layout", + raw_response=_completed_response(AZURE_ANALYZE_WITH_NATIVE_ONLY_FIELDS), + logging_obj=MagicMock(), + optional_params=optional_params, + ) + + assert result.get_provider_native_response() is None + _assert_native_fields_preserved(result.model_dump()) + + +@pytest.mark.parametrize("req_format", ["native", "litellm"]) +def test_map_ocr_params_passes_through_req_format(req_format): + config = AzureDocumentIntelligenceOCRConfig() + + assert config.map_ocr_params({"req_format": req_format}, {}, "prebuilt-layout") == {"req_format": req_format} + + +def test_map_ocr_params_rejects_unknown_req_format_as_bad_request(): + config = AzureDocumentIntelligenceOCRConfig() + + with pytest.raises(UnsupportedParamsError, match="Invalid `req_format`") as exc_info: + config.map_ocr_params({"req_format": "azure"}, {}, "prebuilt-layout") + + assert exc_info.value.status_code == 400 + + +def test_get_complete_url_omits_req_format_query_param(): + config = AzureDocumentIntelligenceOCRConfig() + + url = config.get_complete_url( + api_base="https://example.cognitiveservices.azure.com", + model="prebuilt-layout", + optional_params={"req_format": "native"}, + litellm_params={}, + ) + + assert "req_format" not in url @pytest.mark.parametrize( diff --git a/tests/test_litellm/ocr/test_ocr_native_format.py b/tests/test_litellm/ocr/test_ocr_native_format.py new file mode 100644 index 00000000000..463213a2071 --- /dev/null +++ b/tests/test_litellm/ocr/test_ocr_native_format.py @@ -0,0 +1,65 @@ +""" +Tests for the OCR `req_format` option in the SDK request path: +providers that don't support a native response must reject it, and the Rust +bridge (which only returns the normalized shape) must not serve native requests. +""" + +from unittest.mock import MagicMock + +import pytest + +import litellm +from litellm.ocr.main import _PreparedOCRRequest, _rust_ocr_supported + +DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} + + +def _prepared(optional_params: dict[str, object]) -> _PreparedOCRRequest: + return _PreparedOCRRequest( + model="doc-intelligence/prebuilt-layout", + document=dict(DOCUMENT), + api_key="fake-key", + api_base="https://example.cognitiveservices.azure.com", + custom_llm_provider="azure_ai", + extra_headers=None, + provider_config=MagicMock(), + optional_params=optional_params, + litellm_params={}, + effective_timeout=60.0, + litellm_logging_obj=MagicMock(), + ) + + +@pytest.mark.parametrize("optional_params", [{}, {"req_format": "litellm"}]) +def test_rust_ocr_serves_default_format(optional_params): + assert _rust_ocr_supported(_prepared(optional_params)) is True + + +def test_rust_ocr_skipped_for_native_format(): + assert _rust_ocr_supported(_prepared({"req_format": "native"})) is False + + +@pytest.mark.asyncio +async def test_native_format_rejected_for_provider_without_support_as_bad_request(): + with pytest.raises(litellm.BadRequestError, match="not supported for provider") as exc_info: + await litellm.aocr( + model="mistral/mistral-ocr-latest", + document=DOCUMENT, + api_key="fake-key", + req_format="native", + ) + + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_unknown_format_rejected_for_provider_without_support_as_bad_request(): + with pytest.raises(litellm.BadRequestError, match="Invalid `req_format`") as exc_info: + await litellm.aocr( + model="mistral/mistral-ocr-latest", + document=DOCUMENT, + api_key="fake-key", + req_format="raw", + ) + + assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/proxy/ocr_endpoints/__init__.py b/tests/test_litellm/proxy/ocr_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py b/tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py new file mode 100644 index 00000000000..153e8c36eda --- /dev/null +++ b/tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py @@ -0,0 +1,111 @@ +""" +Tests for the proxy OCR endpoint helpers that select the response format +(`x-req-format: native | litellm`) and return the provider's native payload. +""" + +from unittest.mock import AsyncMock, MagicMock + +import orjson +import pytest +from fastapi import HTTPException + +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse +from litellm.proxy.ocr_endpoints.endpoints import _native_response, _parse_ocr_request + +AZURE_NATIVE_OPERATION = { + "status": "succeeded", + "createdDateTime": "2026-07-02T00:00:00Z", + "analyzeResult": { + "content": "Invoice", + "pages": [{"pageNumber": 1, "words": [{"content": "Invoice", "confidence": 0.99}]}], + "paragraphs": [{"content": "Invoice"}], + }, +} + + +def _json_request(body: dict, headers: dict[str, str]) -> MagicMock: + request = MagicMock() + request.headers = {"content-type": "application/json", **headers} + request.body = AsyncMock(return_value=orjson.dumps(body)) + request._form = None + return request + + +def _ocr_response(native_payload: dict[str, object] | None) -> OCRResponse: + response = OCRResponse(pages=[OCRPage(index=0, markdown="Invoice")], model="azure-prebuilt-layout") + if native_payload is not None: + response.set_provider_native_response(native_payload) + return response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("header_value", ["native", "NATIVE", " native "]) +async def test_should_read_req_format_from_header(header_value): + request = _json_request( + {"model": "azure-prebuilt-layout", "document": {"type": "document_url", "document_url": "https://x/y.pdf"}}, + {"x-req-format": header_value}, + ) + + assert (await _parse_ocr_request(request))["req_format"] == "native" + + +@pytest.mark.asyncio +async def test_should_prefer_body_req_format_over_header(): + request = _json_request( + { + "model": "azure-prebuilt-layout", + "document": {"type": "document_url", "document_url": "https://x/y.pdf"}, + "req_format": "litellm", + }, + {"x-req-format": "native"}, + ) + + assert (await _parse_ocr_request(request))["req_format"] == "litellm" + + +@pytest.mark.asyncio +async def test_should_omit_req_format_when_header_absent(): + request = _json_request( + {"model": "azure-prebuilt-layout", "document": {"type": "document_url", "document_url": "https://x/y.pdf"}}, + {}, + ) + + assert "req_format" not in await _parse_ocr_request(request) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body_format, headers", + [ + (None, {"x-req-format": "azure"}), + ("azure", {}), + ("azure", {"x-req-format": "native"}), + ], +) +async def test_should_reject_unknown_req_format(body_format, headers): + body = {"model": "azure-prebuilt-layout", "document": {"type": "document_url", "document_url": "https://x/y.pdf"}} + request = _json_request( + body if body_format is None else {**body, "req_format": body_format}, + headers, + ) + + with pytest.raises(HTTPException) as exc_info: + await _parse_ocr_request(request) + + assert exc_info.value.status_code == 400 + assert "Invalid `req_format`" in f"{exc_info.value.detail}" + + +def test_should_return_native_payload_with_litellm_response_headers(): + fastapi_response = MagicMock() + fastapi_response.headers = {"x-litellm-response-cost": "0.0015"} + + native = _native_response(_ocr_response(AZURE_NATIVE_OPERATION), fastapi_response) + + assert native is not None + assert orjson.loads(native.body) == AZURE_NATIVE_OPERATION + assert native.headers["x-litellm-response-cost"] == "0.0015" + + +def test_should_return_normalized_response_when_no_native_payload(): + assert _native_response(_ocr_response(None), MagicMock()) is None diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 42ac146ec02..4da0cada6bb 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8625,6 +8625,11 @@ export interface paths { * ```bash * curl -X POST "http://localhost:4000/v1/ocr" -H "Authorization: Bearer sk-1234" -F "model=mistral-ocr" -F "file=@document.pdf" * ``` + * + * Response format is normalized to the LiteLLM OCR schema by default. Providers + * that support it (Azure Document Intelligence) can return their own payload + * instead, with cost tracking unchanged, via `x-req-format: native` (or + * `"req_format": "native"` in the body). */ post: operations["ocr_ocr_post"]; delete?: never; @@ -17777,6 +17782,11 @@ export interface paths { * ```bash * curl -X POST "http://localhost:4000/v1/ocr" -H "Authorization: Bearer sk-1234" -F "model=mistral-ocr" -F "file=@document.pdf" * ``` + * + * Response format is normalized to the LiteLLM OCR schema by default. Providers + * that support it (Azure Document Intelligence) can return their own payload + * instead, with cost tracking unchanged, via `x-req-format: native` (or + * `"req_format": "native"` in the body). */ post: operations["ocr_v1_ocr_post"]; delete?: never;