diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 3b5980023c7..292bde44202 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -134,26 +134,28 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrailed_texts = guardrailed_inputs.get("texts", []) guardrailed_tool_calls = guardrailed_inputs.get("tool_calls", []) guardrailed_tools = guardrailed_inputs.get("tools") + guardrailed_structured_messages = guardrailed_inputs.get("structured_messages") if guardrailed_tools is not None: data["tools"] = guardrailed_tools - # Step 3: Map guardrail responses back to original message structure - if guardrailed_texts and texts_to_check: - await self._apply_guardrail_responses_to_input_texts( - messages=messages, - responses=guardrailed_texts, - task_mappings=text_task_mappings, - ) + if guardrailed_structured_messages is not None: + data["messages"] = guardrailed_structured_messages + else: + # Step 3: Map guardrail responses back to original message structure + if guardrailed_texts and texts_to_check: + await self._apply_guardrail_responses_to_input_texts( + messages=messages, + responses=guardrailed_texts, + task_mappings=text_task_mappings, + ) - # Step 4: Apply guardrailed tool calls back to messages - if guardrailed_tool_calls: - # Note: The guardrail may modify tool_calls_to_check in place - # or we may need to handle returned tool calls differently - await self._apply_guardrail_responses_to_input_tool_calls( - messages=messages, - tool_calls=guardrailed_tool_calls, # type: ignore - task_mappings=tool_call_task_mappings, - ) + # Step 4: Apply guardrailed tool calls back to messages + if guardrailed_tool_calls: + await self._apply_guardrail_responses_to_input_tool_calls( + messages=messages, + tool_calls=guardrailed_tool_calls, # type: ignore + task_mappings=tool_call_task_mappings, + ) verbose_proxy_logger.debug( "OpenAI Chat Completions: Processed input messages: %s", messages diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 4c268d9dfc9..6f1868ccc1c 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1137,3 +1137,78 @@ class TestGetStructuredMessages: if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +class MockStructuredMessagesGuardrail(CustomGuardrail): + """Mock guardrail that returns compressed structured_messages.""" + + def __init__(self, compressed_messages: list): + super().__init__(guardrail_name="test-compression") + self.compressed_messages = compressed_messages + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + result = GenericGuardrailAPIInputs(texts=inputs.get("texts", [])) + result["structured_messages"] = self.compressed_messages # type: ignore + return result + + +class TestStructuredMessagesResponse: + """Test that a guardrail returning structured_messages replaces data['messages'].""" + + @pytest.mark.asyncio + async def test_structured_messages_replaces_data_messages(self): + compressed = [ + {"role": "user", "content": "Analyze this."}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t1", + "content": '{"rows":"[40]{date,revenue}\\n2024-01,120000"}', + } + ], + }, + ] + guardrail = MockStructuredMessagesGuardrail(compressed_messages=compressed) + handler = OpenAIChatCompletionsHandler() + + data = { + "messages": [ + {"role": "user", "content": "Analyze this."}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t1", + "content": '{"rows": [{"date": "2024-01", "revenue": 120000}' + ', {"date": "2024-01", "revenue": 120000}' * 39 + ']}', + } + ], + }, + ] + } + + result = await handler.process_input_messages(data, guardrail) + assert result["messages"] == compressed + + @pytest.mark.asyncio + async def test_texts_not_applied_when_structured_messages_returned(self): + """When structured_messages is returned, texts are NOT positionally applied.""" + original_text = "original user message" + compressed = [{"role": "user", "content": "compressed replacement"}] + + guardrail = MockStructuredMessagesGuardrail(compressed_messages=compressed) + handler = OpenAIChatCompletionsHandler() + + data = {"messages": [{"role": "user", "content": original_text}]} + + result = await handler.process_input_messages(data, guardrail) + assert result["messages"] == compressed + assert result["messages"][0]["content"] == "compressed replacement"