fix(ocr): run post-call logging hooks (#40154)

litellm/integrations/custom_guardrail.py

@@ -773,7 +773,7 @@ class CustomGuardrail(CustomLogger):
     def uses_apply_guardrail_interface(self) -> bool:
         return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
 
-    def _deployment_pre_call_target(self) -> "CustomLogger":
+    def _deployment_hook_target(self) -> "CustomLogger":
         if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
             return self
         try:
@@ -802,7 +802,7 @@ class CustomGuardrail(CustomLogger):
 
         # CHECK IF GUARDRAIL REJECTS THE REQUEST
         if call_type == CallTypes.completion or call_type == CallTypes.acompletion:
-            target: Final = self._deployment_pre_call_target()
+            target: Final = self._deployment_hook_target()
             if target is not self:
                 kwargs["guardrail_to_apply"] = self
             result: Final = await target.async_pre_call_hook(
@@ -845,7 +845,9 @@ class CustomGuardrail(CustomLogger):
             return None
 
         # CHECK IF GUARDRAIL REJECTS THE REQUEST
-        result: Final = await self.async_post_call_success_hook(
+        target: Final = self._deployment_hook_target()
+        hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data
+        result: Final = await target.async_post_call_success_hook(
             user_api_key_dict=UserAPIKeyAuth(
                 user_id=request_data.get("user_api_key_user_id"),
                 team_id=request_data.get("user_api_key_team_id"),
@@ -853,7 +855,7 @@ class CustomGuardrail(CustomLogger):
                 api_key=request_data.get("user_api_key_hash"),
                 request_route=request_data.get("user_api_key_request_route"),
             ),
-            data=request_data,
+            data=hook_request_data,
             response=response,
         )
 

litellm/llms/custom_httpx/llm_http_handler.py

@@ -1741,6 +1741,12 @@ class BaseLLMHTTPHandler:
         except Exception as e:
             raise self._handle_error(e=e, provider_config=provider_config)
 
+        logging_obj.post_call(
+            api_key=api_key,
+            original_response=response.text,
+            additional_args={"complete_input_dict": data},
+        )
+
         return self._transform_ocr_response(
             provider_config=provider_config,
             model=model,
@@ -1804,6 +1810,12 @@ class BaseLLMHTTPHandler:
         except Exception as e:
             raise self._handle_error(e=e, provider_config=provider_config)
 
+        logging_obj.post_call(
+            api_key=api_key,
+            original_response=response.text,
+            additional_args={"complete_input_dict": data},
+        )
+
         # Use async response transform for async operations
         return await provider_config.async_transform_ocr_response(
             model=model,

tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py

@@ -262,7 +262,7 @@ def test_proxied_traffic_stays_on_native_hooks():
     never sees ``data["prompt"]``."""
     guardrail = _guardrail()
     assert guardrail.uses_apply_guardrail_interface() is True
-    assert guardrail._deployment_pre_call_target() is guardrail
+    assert guardrail._deployment_hook_target() is guardrail
 
 
 @pytest.mark.asyncio

tests/test_litellm/integrations/test_custom_guardrail.py

@@ -2610,3 +2610,36 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook:
         )
 
         assert result is replacement
+
+    @pytest.mark.asyncio
+    async def test_apply_guardrail_interface_modifies_deployment_response(self):
+        from litellm.types.guardrails import GuardrailEventHooks
+        from litellm.types.utils import ModelResponse
+
+        class ReplacingGuardrail(CustomGuardrail):
+            async def apply_guardrail(
+                self,
+                inputs: GenericGuardrailAPIInputs,
+                request_data: dict[str, object],
+                input_type: Literal["request", "response"],
+                logging_obj: Optional["LiteLLMLoggingObj"] = None,
+            ) -> GenericGuardrailAPIInputs:
+                assert input_type == "response"
+                return {**inputs, "texts": ["filtered response"]}
+
+        guardrail = ReplacingGuardrail(
+            guardrail_name="test-guardrail",
+            event_hook=GuardrailEventHooks.post_call,
+        )
+        response = ModelResponse(choices=[{"message": {"role": "assistant", "content": "original response"}}])
+        request_data = {"guardrails": ["test-guardrail"]}
+
+        result = await guardrail.async_post_call_success_deployment_hook(
+            request_data=request_data,
+            response=response,
+            call_type=CallTypes.acompletion,
+        )
+
+        assert result is response
+        assert response.choices[0].message.content == "filtered response"
+        assert request_data == {"guardrails": ["test-guardrail"]}

tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py

@@ -29,6 +29,7 @@ from litellm.llms.custom_httpx.llm_http_handler import (
     _rust_responses_websocket_enabled,
 )
 from litellm.llms.azure.videos.transformation import AzureVideoConfig
+from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
 from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
 from litellm.types.llms.openai import ResponsesAPIResponse
 from litellm.types.router import GenericLiteLLMParams
@@ -37,6 +38,69 @@ from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, Trans
 _ACTIVE_KEY = "_code_interpreter_interception_active"
 _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
 
+OCR_RESPONSE = {
+    "pages": [{"index": 0, "markdown": "OCR output", "images": []}],
+    "model": "mistral-ocr-latest",
+    "usage_info": {"pages_processed": 1},
+}
+
+
+def _ocr_sync_client() -> HTTPHandler:
+    client = HTTPHandler()
+    client.client = httpx.Client(transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE)))
+    return client
+
+
+def _ocr_async_client() -> AsyncHTTPHandler:
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(
+        transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE))
+    )
+    return client
+
+
+def test_ocr_calls_post_call_with_raw_provider_response():
+    logging_obj = Mock()
+
+    response = BaseLLMHTTPHandler().ocr(
+        model="mistral-ocr-latest",
+        document={"type": "document_url", "document_url": "https://example.com/document.pdf"},
+        optional_params={},
+        timeout=5,
+        logging_obj=logging_obj,
+        api_key="test-key",
+        api_base="https://api.mistral.ai/v1/ocr",
+        custom_llm_provider="mistral",
+        client=_ocr_sync_client(),
+        provider_config=MistralOCRConfig(),
+    )
+
+    assert response.pages[0].markdown == "OCR output"
+    logging_obj.post_call.assert_called_once()
+    assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE
+
+
+@pytest.mark.asyncio
+async def test_async_ocr_calls_post_call_with_raw_provider_response():
+    logging_obj = Mock()
+
+    response = await BaseLLMHTTPHandler().async_ocr(
+        model="mistral-ocr-latest",
+        document={"type": "document_url", "document_url": "https://example.com/document.pdf"},
+        optional_params={},
+        timeout=5,
+        logging_obj=logging_obj,
+        api_key="test-key",
+        api_base="https://api.mistral.ai/v1/ocr",
+        custom_llm_provider="mistral",
+        client=_ocr_async_client(),
+        provider_config=MistralOCRConfig(),
+    )
+
+    assert response.pages[0].markdown == "OCR output"
+    logging_obj.post_call.assert_called_once()
+    assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE
+
 
 def test_prepare_fake_stream_request():
     # Initialize the BaseLLMHTTPHandler

tests/test_litellm/proxy/test_proxy_logging_hook_detection.py

@@ -637,14 +637,14 @@ def test_callback_capabilities_excludes_opted_out_guardrail_from_iterator_overri
     assert [cb for cb, _ in caps.iterator_overrides if cb is opted_out] == []
 
 
-def test_deployment_pre_call_target_stays_native_when_opted_out():
+def test_deployment_hook_target_stays_native_when_opted_out():
     """Model-level guardrails resolve their target here rather than through ProxyLogging."""
-    assert _KeepsNativeHooks()._deployment_pre_call_target() is not None
+    assert _KeepsNativeHooks()._deployment_hook_target() is not None
     opted_out = _KeepsNativeHooks()
-    assert opted_out._deployment_pre_call_target() is opted_out
-    assert _AppliesGuardrail()._deployment_pre_call_target() is not None
+    assert opted_out._deployment_hook_target() is opted_out
+    assert _AppliesGuardrail()._deployment_hook_target() is not None
     routed = _AppliesGuardrail()
-    assert routed._deployment_pre_call_target() is not routed
+    assert routed._deployment_hook_target() is not routed
 
 
 @pytest.mark.asyncio