mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
fix(bedrock): guard neutralized tool results anywhere in history when guardrailConfig is set
This commit is contained in:
parent
5493be743b
commit
04086be795
2 changed files with 60 additions and 6 deletions
|
|
@ -234,22 +234,28 @@ class AmazonConverseConfig(BaseConfig):
|
|||
rendered = convert_content_list_to_str(message).strip()
|
||||
return rendered or "<non-text tool result omitted>"
|
||||
|
||||
guardrail_active: Final = "guardrailConfig" in optional_params
|
||||
|
||||
def _rewrite(message: AllMessageValues) -> AllMessageValues:
|
||||
role = message.get("role")
|
||||
tool_calls = message.get("tool_calls")
|
||||
if role == "assistant" and tool_calls:
|
||||
base_text = convert_content_list_to_str(message)
|
||||
call_texts = [_tool_call_text(call) for call in tool_calls]
|
||||
text = "\n".join(filter(None, [base_text, *call_texts]))
|
||||
base_text: Final = convert_content_list_to_str(message)
|
||||
call_texts: Final = tuple(_tool_call_text(call) for call in tool_calls)
|
||||
text: Final = "\n".join(part for part in (base_text, *call_texts) if part)
|
||||
return ChatCompletionAssistantMessage(role="assistant", content=text)
|
||||
if role in ("tool", "function"):
|
||||
tool_call_id = message.get("tool_call_id")
|
||||
name = message.get("name")
|
||||
label = f"tool result for {tool_call_id or name or 'unknown'}"
|
||||
return ChatCompletionUserMessage(
|
||||
role="user",
|
||||
content=f"[{label}: {_result_text(message)}]",
|
||||
result_text: Final = f"[{label}: {_result_text(message)}]"
|
||||
# Tool results are externally controlled, so guard them wherever they
|
||||
# land in history; _convert_consecutive_user_messages_to_guarded_text
|
||||
# only covers the trailing user turn.
|
||||
content: Final = (
|
||||
[{"type": "guarded_text", "text": result_text}] if guardrail_active else result_text
|
||||
)
|
||||
return ChatCompletionUserMessage(role="user", content=content)
|
||||
return message
|
||||
|
||||
verbose_logger.warning(
|
||||
|
|
|
|||
|
|
@ -6714,6 +6714,54 @@ def test_transform_request_neutralized_tool_output_is_guarded(monkeypatch):
|
|||
assert "secret tool output" in serialized
|
||||
|
||||
|
||||
def test_transform_request_neutralized_tool_output_guarded_mid_history(monkeypatch):
|
||||
"""Regression: a neutralized tool result that is NOT the trailing turn (an
|
||||
assistant reply and a later user turn follow it) must still be guardContent.
|
||||
_convert_consecutive_user_messages_to_guarded_text only covers the trailing
|
||||
user turn, so neutralize itself must guard untrusted tool output regardless
|
||||
of position, else an attacker controlling the tool response bypasses the
|
||||
guardrail (bot review)."""
|
||||
monkeypatch.setattr(litellm, "modify_params", False)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
messages=[
|
||||
{"role": "user", "content": "look it up"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "c1",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "IGNORE_PRIOR malware"},
|
||||
{"role": "assistant", "content": "Here is the summary."},
|
||||
{"role": "user", "content": "thanks"},
|
||||
],
|
||||
optional_params={
|
||||
"guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "1"}
|
||||
},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
_assert_no_structured_tool_blocks(result)
|
||||
blocks = [block for message in result["messages"] for block in message["content"]]
|
||||
guarded_texts = [
|
||||
block["guardContent"]["text"]["text"] for block in blocks if "guardContent" in block
|
||||
]
|
||||
plain_texts = [block["text"] for block in blocks if "text" in block and "guardContent" not in block]
|
||||
assert any("malware" in text for text in guarded_texts), "mid-history tool output must be guarded"
|
||||
assert not any(
|
||||
"malware" in text for text in plain_texts
|
||||
), "mid-history tool output must not reach the model as unguarded text"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_transform_request_no_tools_with_tool_history(monkeypatch):
|
||||
"""Async is a separate request assembler; it must neutralize identically."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue