diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index d9e39cb7fc4..47e2564dc0e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/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. diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 535ab15721a..a61ff7b9785 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/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 ()) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 1961146a88b..f85d238484e 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/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, diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index fe7bc8efbad..24696c94cc3 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/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