mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(guardrails): scan reasoning_content in chat completions history
The chat completions guardrail handler only handed message content and tool calls to guardrails, so an assistant message's reasoning_content reached the provider exactly as the client sent it. Clients that replay reasoning, which providers with thinking plus tool calls require, sent any PII or blocked text in it straight past masking and moderation guardrails reasoning_content is now extracted as one more text of its message and the guardrail's rewrite is written back to that field, including on assistant turns whose content is null
This commit is contained in:
parent
f2118fc767
commit
49b5a7d005
2 changed files with 55 additions and 4 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")],
|
||||
},
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue