fix(guardrails): hand the input scan context to the logging_only response scan

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/integrations/custom_guardrail.py

@@ -16,6 +16,7 @@ from litellm.litellm_core_utils.core_helpers import (
     get_or_create_metadata_bucket,
     redact_nested_match_and_regex_keys,
 )
+from litellm.llms.base_llm.guardrail_translation.base_translation import REQUEST_SCAN_CONTEXT_KEY
 from litellm.secret_managers.main import str_to_bool
 from litellm.types.guardrails import (
     DynamicGuardrailParams,
@@ -906,11 +907,10 @@ class CustomGuardrail(CustomLogger):
         response: Final = (
             kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result
         )
-        from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler
         from litellm.types.utils import ModelResponse
 
         output_translation: Final = (
-            OpenAIChatCompletionsHandler(request_scoping=translation)
+            get_guardrail_translation_mapping(CallTypes.acompletion)()
             if isinstance(response, ModelResponse)
             else translation
         )
@@ -950,9 +950,31 @@ class CustomGuardrail(CustomLogger):
         await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
         if response is None:
             return
-        await output_translation.process_output_response(
-            response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request
+        output_request: Final = (
+            scratch_request
+            if type(output_translation) is type(translation)
+            else self._chat_shaped_request(scratch_request, translation)
         )
+        await output_translation.process_output_response(
+            response=copy.deepcopy(response), guardrail_to_apply=self, request_data=output_request
+        )
+
+    def _chat_shaped_request(
+        self,
+        scratch_request: Mapping[str, object],
+        translation: "BaseTranslation",
+    ) -> dict[str, object]:  # mutable-ok: BaseTranslation.process_output_response contract
+        """The logged request in OpenAI chat shape, for an output scan whose translation differs from the input's."""
+        context: Final = translation.request_scan_context(
+            dict(scratch_request),  # mutable-ok: BaseTranslation.request_scan_context requires a dict
+            self,
+        )
+        return {
+            **scratch_request,
+            "messages": list(context.structured_messages),
+            "tools": list(context.tools),
+            REQUEST_SCAN_CONTEXT_KEY: context,
+        }
 
     def supports_scan_only_tool_results(self) -> bool:
         """Whether this guardrail can scan tool-result content.

litellm/llms/base_llm/guardrail_translation/base_translation.py

@@ -56,6 +56,9 @@ class RequestScanContext:
         )
 
 
+REQUEST_SCAN_CONTEXT_KEY: Final = "litellm_request_scan_context"
+
+
 @dataclass(slots=True)
 class StreamTransformSink:
     """Out-parameter used by ``process_output_streaming_response`` to hand the
@@ -313,7 +316,12 @@ class BaseTranslation(ABC):
         """``inputs`` plus the scoped request conversation, closed by the scanned reply, and the request tools."""
         if request_data is None:
             return inputs
-        context: Final = self.request_scan_context(request_data, guardrail_to_apply)
+        precomputed: Final = request_data.get(REQUEST_SCAN_CONTEXT_KEY)
+        context: Final = (
+            precomputed
+            if isinstance(precomputed, RequestScanContext)
+            else self.request_scan_context(request_data, guardrail_to_apply)
+        )
         if not context.conversation_supplied:
             return inputs
         assistant_turn: Final = response_assistant_turn(inputs.get("texts") or (), inputs.get("tool_calls") or ())

litellm/llms/openai/chat/guardrail_translation/handler.py

@@ -26,7 +26,6 @@ import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.llms.base_llm.guardrail_translation.base_translation import (
     BaseTranslation,
-    RequestScanContext,
     StreamingScanKey,
     StreamTransformSink,
 )
@@ -85,9 +84,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
     delivers_ended_stream_rewrites = True
     assembles_streamed_response = True
 
-    def __init__(self, request_scoping: BaseTranslation | None = None) -> None:
-        self._request_scoping: Final = request_scoping
-
     def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
         """
         Convert chat completions request data to OpenAI-spec structured messages.
@@ -99,12 +95,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             return None
         return cast(list[AllMessageValues], messages)
 
-    def request_scan_context(self, data: dict, guardrail_to_apply: "CustomGuardrail") -> RequestScanContext:
-        """Scoped by the translation the request arrived in, so a chat-shaped reply scan sees the request's own scope."""
-        if self._request_scoping is None:
-            return super().request_scan_context(data, guardrail_to_apply)
-        return self._request_scoping.request_scan_context(data, guardrail_to_apply)
-
     async def process_input_messages(
         self,
         data: dict,

tests/test_litellm/integrations/test_custom_guardrail.py

@@ -2716,7 +2716,7 @@ class TestLoggingOnlyApplyGuardrail:
         assert guardrail.calls == [("response", [{"role": "assistant", "content": "general kenobi"}], None)]
 
     @pytest.mark.asyncio
-    async def test_anthropic_messages_response_scan_keeps_midturn_system_turns_under_skip_system(self):
+    async def test_anthropic_messages_response_scan_keeps_midturn_system_when_skip_system(self):
         class _ContextObserver(_ApplyOnlyObserver):
             @log_guardrail_information
             async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
@@ -2727,7 +2727,8 @@ class TestLoggingOnlyApplyGuardrail:
         guardrail.skip_system_message_in_guardrail = True
         kwargs, response = _logged_call(
             [
-                {"role": "system", "content": "Mid-turn operator note"},
+                {"role": "user", "content": "hi"},
+                {"role": "system", "content": "mid-turn note"},
                 {"role": "user", "content": "What is the capital of France?"},
             ]
         )
@@ -2735,8 +2736,8 @@ class TestLoggingOnlyApplyGuardrail:
         await guardrail.async_logging_hook(kwargs, response, CallTypes.anthropic_messages.value)
 
         assert guardrail.calls == [
-            ("request", ["system", "user"]),
-            ("response", ["system", "user", "assistant"]),
+            ("request", ["user", "system", "user"]),
+            ("response", ["user", "system", "user", "assistant"]),
         ]
 
     @pytest.mark.asyncio