fix(azure_ai): cost OCR relays per page so non-chat relays debit budgets

litellm/llms/azure_ai/passthrough/transformation.py

@@ -1,25 +1,28 @@
 from __future__ import annotations
 
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final
 
 import httpx
 from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 
+from litellm._logging import verbose_logger
 from litellm.llms.azure_ai.common_utils import (
     AzureFoundryModelInfo,
     api_key_header_for_base,
     get_azure_ai_auth_headers,
 )
+from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config
 from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, strip_leading_model_segment
 from litellm.types.llms.openai import AllMessageValues
-from litellm.types.utils import StandardPassThroughResponseObject
+from litellm.types.utils import CallTypes, StandardPassThroughResponseObject
 
 if TYPE_CHECKING:
     from httpx import URL, Response
 
     from litellm.litellm_core_utils.litellm_logging import Logging
+    from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse
     from litellm.types.utils import CostResponseTypes
 
 
@@ -75,6 +78,10 @@ def relayed_body(httpx_response: Response) -> str | dict:
 
 
 class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig):
+    def __init__(self, ocr_config_for: Callable[[str], BaseOCRConfig | None] = get_azure_ai_ocr_config) -> None:
+        super().__init__()
+        self.ocr_config_for: Final = ocr_config_for
+
     def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool:
         return bool(request_data.get("stream"))
 
@@ -123,7 +130,7 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig):
         request_data: Mapping[str, object],
         logging_obj: Logging,
         endpoint: str,
-    ) -> CostResponseTypes | StandardPassThroughResponseObject | None:
+    ) -> CostResponseTypes | OCRResponse | StandardPassThroughResponseObject | None:
         from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig
 
         chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response(  # pyright: ignore[reportUnknownMemberType]  # the Azure config still types request_data as a bare dict
@@ -136,8 +143,40 @@ class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig):
         )
         if chat_result is not None:
             return chat_result
+        ocr_result: Final = self.logged_ocr_response(model, httpx_response, logging_obj, endpoint)
+        if ocr_result is not None:
+            return ocr_result
         return StandardPassThroughResponseObject(response=relayed_body(httpx_response))
 
+    def logged_ocr_response(
+        self, model: str, httpx_response: Response, logging_obj: Logging, endpoint: str
+    ) -> OCRResponse | None:
+        ocr_config: Final = self.ocr_config_for(model)
+        if ocr_config is None or httpx_response.status_code != 200:
+            return None
+        relayed_url: Final = httpx_response.request.url
+        relayed_origin: Final = str(relayed_url.copy_with(path="/", query=None, fragment=None)).rstrip("/")
+        ocr_url: Final = httpx.URL(
+            ocr_config.get_complete_url(
+                api_base=relayed_origin,
+                model=model,
+                optional_params={},  # mutable-ok: BaseOCRConfig wants a dict
+            )
+        )
+        known_prefixes: Final = (model, model_group_from(logging_obj.litellm_params))
+        native_endpoint: Final = strip_leading_model_segment(endpoint, known_prefixes)
+        if f"/{native_endpoint.strip('/')}" != ocr_url.path:
+            return None
+        try:
+            ocr_response: Final = ocr_config.transform_ocr_response(
+                model=model, raw_response=httpx_response, logging_obj=logging_obj
+            )
+        except (ValueError, AttributeError) as error:
+            verbose_logger.warning("azure_ai passthrough: OCR body from %s is not costable: %s", ocr_url, error)
+            return None
+        logging_obj.call_type = CallTypes.aocr.value  # rebind-ok: routes cost calculation to the per-page OCR path
+        return ocr_response
+
     def handle_logging_collected_chunks(
         self,
         all_chunks: Sequence[str],

litellm/llms/base_llm/passthrough/transformation.py

@@ -11,6 +11,7 @@ if TYPE_CHECKING:
     from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject
 
     from ..chat.transformation import BaseLLMException
+    from ..ocr.transformation import OCRResponse
 
 
 def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str:
@@ -120,7 +121,7 @@ class BasePassthroughConfig(BaseLLMModelInfo):
         request_data: dict,
         logging_obj: "LiteLLMLoggingObj",
         endpoint: str,
-    ) -> Optional["CostResponseTypes | StandardPassThroughResponseObject"]:
+    ) -> Optional["CostResponseTypes | OCRResponse | StandardPassThroughResponseObject"]:
         pass
 
     def handle_logging_collected_chunks(

tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py

@@ -1,11 +1,14 @@
 import json
+from datetime import datetime
 from unittest.mock import MagicMock
 
 import httpx
 import pytest
 
 import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig
+from litellm.llms.base_llm.ocr.transformation import OCRResponse
 from litellm.types.utils import LlmProviders, ModelResponse
 from litellm.utils import ProviderConfigManager
 
@@ -238,16 +241,135 @@ def _non_chat_logging_result(content: bytes, content_type: str):
     )
 
 
-def test_non_chat_relay_logs_the_parsed_body_so_spend_tracking_sees_the_call():
-    result = _non_chat_logging_result(b'{"id":"parse-1","pages":[],"meta":{"billed_units":{"pages":1}}}', "application/json")
-
-    assert result == {"response": {"id": "parse-1", "pages": [], "meta": {"billed_units": {"pages": 1}}}}
-
-
 def test_non_chat_relay_with_a_non_json_body_logs_the_raw_text():
     assert _non_chat_logging_result(b"page one", "text/plain") == {"response": "page one"}
 
 
+def _relay_logging_obj(model: str, api_base: str) -> Logging:
+    logging_obj = Logging(
+        model=model,
+        messages=[],
+        stream=False,
+        call_type="allm_passthrough_route",
+        start_time=datetime.now(),
+        litellm_call_id="call-1",
+        function_id="fn-1",
+    )
+    logging_obj.update_environment_variables(
+        model=model,
+        litellm_params={"api_base": api_base, "custom_llm_provider": "azure_ai"},
+        optional_params={},
+        custom_llm_provider="azure_ai",
+    )
+    return logging_obj
+
+
+def _relay_logging_result(
+    config: AzureAIPassthroughConfig, model: str, native_path: str, body, api_base: str = FOUNDRY_BASE, status_code: int = 200
+):
+    relayed_url = f"{FOUNDRY_BASE}/{native_path}?api-version=2024-05-01-preview"
+    logging_obj = _relay_logging_obj(model, api_base)
+    response = httpx.Response(
+        status_code=status_code,
+        headers={"content-type": "application/json"},
+        content=json.dumps(body).encode("utf-8"),
+        request=httpx.Request("POST", relayed_url),
+    )
+    result = config.logging_non_streaming_response(
+        model=model,
+        custom_llm_provider="azure_ai",
+        httpx_response=response,
+        request_data={"model": model},
+        logging_obj=logging_obj,
+        endpoint=f"{model}/{native_path}",
+    )
+    return result, logging_obj
+
+
+MISTRAL_OCR_BODY = {
+    "pages": [{"index": 0, "markdown": "page one"}, {"index": 1, "markdown": "page two"}],
+    "model": "mistral-document-ai-2512",
+    "usage_info": {"pages_processed": 2, "doc_size_bytes": 4321},
+}
+
+
+def test_mistral_document_ai_relay_is_costed_per_page():
+    result, logging_obj = _relay_logging_result(
+        AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", MISTRAL_OCR_BODY
+    )
+    per_page = litellm.get_model_info("azure_ai/mistral-document-ai-2512")["ocr_cost_per_page"]
+
+    assert isinstance(result, OCRResponse)
+    assert result.usage_info.pages_processed == 2
+    assert per_page > 0
+    assert logging_obj._response_cost_calculator(result=result) == pytest.approx(2 * per_page)
+
+
+def test_ocr_route_under_a_models_api_base_is_still_recognised():
+    result, _ = _relay_logging_result(
+        AzureAIPassthroughConfig(),
+        "mistral-document-ai-2512",
+        "providers/mistral/azure/ocr",
+        MISTRAL_OCR_BODY,
+        api_base=f"{FOUNDRY_BASE}/models",
+    )
+
+    assert isinstance(result, OCRResponse)
+
+
+def test_relay_to_a_non_ocr_route_keeps_the_passthrough_object_and_call_type():
+    result, logging_obj = _relay_logging_result(
+        AzureAIPassthroughConfig(), "mistral-document-ai-2512", "models/info", {"name": "mistral-document-ai-2512"}
+    )
+
+    assert result == {"response": {"name": "mistral-document-ai-2512"}}
+    assert logging_obj.call_type == "allm_passthrough_route"
+
+
+COHERE_PARSE_BODY = {"id": "parse-1", "pages": [], "meta": {"billed_units": {"pages": 3}}}
+
+
+def test_cohere_parse_relay_is_costed_per_billed_page():
+    result, logging_obj = _relay_logging_result(
+        AzureAIPassthroughConfig(), "Cohere-parse-v5", "providers/cohere/v2/parse", COHERE_PARSE_BODY
+    )
+    per_page = litellm.get_model_info("azure_ai/Cohere-parse-v5")["ocr_cost_per_page"]
+
+    assert isinstance(result, OCRResponse)
+    assert result.usage_info.pages_processed == 3
+    assert logging_obj.call_type == "aocr"
+    assert per_page > 0
+    assert logging_obj._response_cost_calculator(result=result) == pytest.approx(3 * per_page)
+
+
+def test_deployment_without_an_ocr_config_is_never_costed_as_ocr():
+    config = AzureAIPassthroughConfig(ocr_config_for=lambda model: None)
+    result, logging_obj = _relay_logging_result(
+        config, "mistral-document-ai-2512", "providers/mistral/azure/ocr", MISTRAL_OCR_BODY
+    )
+
+    assert result == {"response": MISTRAL_OCR_BODY}
+    assert logging_obj.call_type == "allm_passthrough_route"
+
+
+def test_accepted_ocr_job_without_a_result_body_is_not_costed():
+    result, logging_obj = _relay_logging_result(
+        AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", {"status": "running"}, status_code=202
+    )
+
+    assert result == {"response": {"status": "running"}}
+    assert logging_obj.call_type == "allm_passthrough_route"
+
+
+def test_unparseable_ocr_body_falls_back_to_the_passthrough_object():
+    result, logging_obj = _relay_logging_result(
+        AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", ["not", "an", "ocr", "body"]
+    )
+
+    assert result == {"response": '["not", "an", "ocr", "body"]'}
+    assert logging_obj.call_type == "allm_passthrough_route"
+
+
 def test_streaming_chat_completion_chunks_are_costed_like_azure():
     head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"}
     chunks = [