fix(guardrails): apply structured_messages from guardrail response to data[messages]

process_input_messages was ignoring structured_messages in guardrailed_inputs;
it now replaces data['messages'] directly when the guardrail returns them

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

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"