Merge 8cc213247e into a9b181489b

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

@@ -19,7 +19,7 @@ import time
 import uuid
 from collections.abc import Iterator, Mapping, Sequence
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Union, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast
 
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 
@@ -121,7 +121,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         texts_to_check: Final[list[str]] = []
         images_to_check: Final[list[str]] = []
         tool_calls_to_check: Final[list[ChatCompletionToolParam]] = []
-        text_task_mappings: Final[list[tuple[int, int | None]]] = []
+        text_task_mappings: Final[list[tuple[int, int | Literal["reasoning_content"] | None]]] = []
         tool_call_task_mappings: Final[list[tuple[int, int]]] = []
 
         # Step 1: Extract all text content, images, and tool calls
@@ -280,7 +280,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         texts_to_check: list[str],
         images_to_check: list[str],
         tool_calls_to_check: list[ChatCompletionToolParam],
-        text_task_mappings: list[tuple[int, int | None]],
+        text_task_mappings: list[tuple[int, int | Literal["reasoning_content"] | None]],
         tool_call_task_mappings: list[tuple[int, int]],
         skip_system_message: bool = False,
         skip_tool_message: bool = False,
@@ -325,6 +325,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                         elif isinstance(image_url, str):
                             images_to_check.append(image_url)
 
+        reasoning_content: Final = message.get("reasoning_content", None)
+        if isinstance(reasoning_content, str):
+            texts_to_check.append(reasoning_content)
+            text_task_mappings.append((msg_idx, "reasoning_content"))
+
         # Extract tool calls (typically in assistant messages)
         tool_calls: Final = message.get("tool_calls", None)
         if tool_calls is not None and isinstance(tool_calls, list):
@@ -338,7 +343,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         self,
         messages: list[dict[str, object]],
         responses: list[str],
-        task_mappings: list[tuple[int, int | None]],
+        task_mappings: list[tuple[int, int | Literal["reasoning_content"] | None]],
     ) -> None:
         """
         Apply guardrail responses back to input message text content.
@@ -348,6 +353,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         for task_idx, guardrail_response in enumerate(responses):
             mapping = task_mappings[task_idx]
             msg_idx = cast(int, mapping[0])
+            if mapping[1] == "reasoning_content":
+                message = messages[msg_idx]
+                message["reasoning_content"] = guardrail_response
+                continue
             content_idx_optional = cast(int | None, mapping[1])
 
             # Handle content

tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py

@@ -2365,3 +2365,45 @@ class TestStreamingScanKey:
         for released in (text_only, finished_tool_call):
             ended = handler.released_stream_as_ended(released)
             assert len(ended) == len(released) and all(a is b for a, b in zip(ended, released, strict=True))
+
+
+def _weather_tool_call(location: str) -> dict:
+    return {
+        "id": "call_1",
+        "type": "function",
+        "function": {"name": "get_weather", "arguments": json.dumps({"location": location})},
+    }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("content", [None, "Let me check."])
+async def test_history_reasoning_content_is_guardrailed_and_written_back(content: Optional[str]):
+    guardrail = MockGuardrail()
+    data = {
+        "messages": [
+            {"role": "user", "content": "weather in boston?"},
+            {
+                "role": "assistant",
+                "content": content,
+                "reasoning_content": "the user wants boston weather",
+                "tool_calls": [_weather_tool_call("boston")],
+            },
+        ]
+    }
+
+    await OpenAIChatCompletionsHandler().process_input_messages(data, guardrail)
+
+    assert guardrail.last_inputs["texts"] == [
+        "weather in boston?",
+        *([] if content is None else [content]),
+        "the user wants boston weather",
+    ]
+    assert data["messages"] == [
+        {"role": "user", "content": "WEATHER IN BOSTON?"},
+        {
+            "role": "assistant",
+            "content": None if content is None else content.upper(),
+            "reasoning_content": "THE USER WANTS BOSTON WEATHER",
+            "tool_calls": [_weather_tool_call("BOSTON")],
+        },
+    ]