From 2d925e5dde1aa4186d1fbf690f97bd4c4c3ca4dd Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 17 Sep 2026 05:40:44 +0000 Subject: [PATCH] fix(guardrails): scope the logging_only reply scan with the request's own translation The chat-shaped output handler now takes the input translation as its request scoping, so the logged request is scoped exactly once and with the pre-call semantics of the surface it arrived on. This drops the unscoped chat_shaped_request_conversation detour from af312dc8, which made the Anthropic response scan remove in-sequence system turns under skip_system while the request scan kept them Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 21 ++----------- .../chat/guardrail_translation/handler.py | 31 +++++++++---------- .../guardrail_translation/base_translation.py | 11 ++----- .../chat/guardrail_translation/handler.py | 10 ++++++ .../integrations/test_custom_guardrail.py | 24 ++++++++++++++ .../test_anthropic_guardrail_handler.py | 17 ++++++++++ 6 files changed, 71 insertions(+), 43 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 1ddfee5fd6d..d9e39cb7fc4 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -906,10 +906,11 @@ 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 = ( - get_guardrail_translation_mapping(CallTypes.acompletion)() + OpenAIChatCompletionsHandler(request_scoping=translation) if isinstance(response, ModelResponse) else translation ) @@ -949,26 +950,10 @@ class CustomGuardrail(CustomLogger): await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self) if response is None: return - 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 + response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_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.""" - messages, tools = translation.chat_shaped_request_conversation( - dict(scratch_request) # mutable-ok: BaseTranslation.chat_shaped_request_conversation requires a dict - ) - return {**scratch_request, "messages": list(messages), "tools": list(tools)} - def supports_scan_only_tool_results(self) -> bool: """Whether this guardrail can scan tool-result content. diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index d0288b1b853..eaa522cdb35 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -528,26 +528,23 @@ class AnthropicMessagesHandler(BaseTranslation): ) return result if result else None - def chat_shaped_request_conversation( - self, data: dict - ) -> tuple[tuple[AllMessageValues, ...], tuple[ChatCompletionToolParam, ...]]: - if data.get("messages") is None: - return (), () - translated: Final = self._translate_to_openai( - {key: value for key, value in data.items() if key != "system"} # mutable-ok: API message payload - ) - hoisted_system_message: Final = self._hoisted_top_level_system_message(data) - messages: Final = ( - *(() if hoisted_system_message is None else (hoisted_system_message,)), - *translated["messages"], - ) - tools: Final = tuple(tool for tool in translated.get("tools") or () if not is_provider_native_tool_dict(tool)) - return messages, tools - def request_scan_context(self, data: dict, guardrail_to_apply: "CustomGuardrail") -> RequestScanContext: if data.get("messages") is None: return RequestScanContext() - return RequestScanContext.scoped(*self.chat_shaped_request_conversation(data), guardrail_to_apply) + translated: Final = self._translate_to_openai( + {key: value for key, value in data.items() if key != "system"} # mutable-ok: API message payload + ) + hoisted_system_message: Final = ( + None + if effective_skip_system_message_for_guardrail(guardrail_to_apply) + else self._hoisted_top_level_system_message(data) + ) + return RequestScanContext.scoped( + (*(() if hoisted_system_message is None else (hoisted_system_message,)), *translated["messages"]), + tuple(tool for tool in translated.get("tools") or () if not is_provider_native_tool_dict(tool)), + guardrail_to_apply, + skip_system=False, + ) async def process_input_messages( self, diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index bcffc4777d9..535ab15721a 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -298,16 +298,11 @@ class BaseTranslation(ABC): """ return None - def chat_shaped_request_conversation( - self, data: dict - ) -> tuple[tuple["AllMessageValues", ...], tuple["ChatCompletionToolParam", ...]]: - """The full, unscoped request turns and tool definitions in OpenAI chat shape.""" - return tuple(self.get_structured_messages(data) or ()), tuple(data.get("tools") or ()) - def request_scan_context(self, data: dict, guardrail_to_apply: "CustomGuardrail") -> RequestScanContext: """Override wherever ``process_input_messages`` scopes or translates the request differently.""" - messages, tools = self.chat_shaped_request_conversation(data) - return RequestScanContext.scoped(messages, tools, guardrail_to_apply) + return RequestScanContext.scoped( + self.get_structured_messages(data) or (), data.get("tools") or (), guardrail_to_apply + ) def with_response_context( self, diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index f85d238484e..1961146a88b 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -26,6 +26,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, + RequestScanContext, StreamingScanKey, StreamTransformSink, ) @@ -84,6 +85,9 @@ 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. @@ -95,6 +99,12 @@ 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, diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 66fbc017875..fe7bc8efbad 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2715,6 +2715,30 @@ 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): + class _ContextObserver(_ApplyOnlyObserver): + @log_guardrail_information + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.calls.append((input_type, [m["role"] for m in inputs.get("structured_messages") or []])) + return inputs + + guardrail = _ContextObserver() + guardrail.skip_system_message_in_guardrail = True + kwargs, response = _logged_call( + [ + {"role": "system", "content": "Mid-turn operator note"}, + {"role": "user", "content": "What is the capital of France?"}, + ] + ) + + await guardrail.async_logging_hook(kwargs, response, CallTypes.anthropic_messages.value) + + assert guardrail.calls == [ + ("request", ["system", "user"]), + ("response", ["system", "user", "assistant"]), + ] + @pytest.mark.asyncio async def test_async_success_handler_records_verdict_in_standard_logging_object(self): import datetime as dt diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 5a5e3b22b4f..2f56838cbb2 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2724,6 +2724,23 @@ class TestAnthropicResponseScanCarriesRequestConversation: [(_, inputs)] = guardrail.seen assert [m["role"] for m in inputs["structured_messages"]] == ["user", "assistant", "tool", "assistant"] + @pytest.mark.asyncio + async def test_skip_system_keeps_in_sequence_system_turns_in_the_response_scan(self): + handler = AnthropicMessagesHandler() + guardrail = TypedInputsRecordingGuardrail() + guardrail.skip_system_message_in_guardrail = True + request = { + **self._request(), + "messages": [{"role": "system", "content": "Mid-turn operator note"}, *self._request()["messages"]], + } + + await handler.process_input_messages(data=request, guardrail_to_apply=guardrail) + await handler.process_output_response(self._tool_use_response(), guardrail, request_data=request) + + (_, request_inputs), (_, response_inputs) = guardrail.seen + assert [m["role"] for m in request_inputs["structured_messages"]] == ["system", "user", "assistant", "tool"] + assert response_inputs["structured_messages"][:-1] == request_inputs["structured_messages"] + @staticmethod def _sse_chunks(ended: bool) -> list: events = [