litellm/tests/test_litellm/proxy/guardrails/test_content_utils.py

619 lines
22 KiB
Python

"""Tests for the shared guardrail content extraction helpers."""
from litellm.proxy.guardrails._content_utils import (
apply_redacted_messages_back,
build_inspection_messages,
has_non_string_content,
iter_message_text,
walk_user_text,
)
# ── iter_message_text ────────────────────────────────────────────────────────────
def test_iter_message_text_string_messages():
data = {
"messages": [
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "hi"},
]
}
assert list(iter_message_text(data)) == ["hello", "hi"]
def test_iter_message_text_multimodal_list_content():
"""VERIA-11: list-format content must be inspected, not silently skipped."""
data = {
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "AWS_KEY=AKIA..."},
{"type": "image_url", "image_url": {"url": "..."}},
{"type": "text", "text": "more secrets"},
],
}
]
}
assert list(iter_message_text(data)) == ["AWS_KEY=AKIA...", "more secrets"]
def test_iter_message_text_responses_api_string_input():
"""fniVO9-F: Responses-API ``input`` must be inspectable when ``messages`` absent."""
data = {"input": "tell me a secret"}
assert list(iter_message_text(data)) == ["tell me a secret"]
def test_iter_message_text_responses_api_list_input_messages():
data = {
"input": [
{"role": "user", "content": "first"},
{"role": "user", "content": "second"},
]
}
assert list(iter_message_text(data)) == ["first", "second"]
def test_iter_message_text_responses_api_list_input_content_parts():
data = {
"input": [
{"type": "text", "text": "alpha"},
{"type": "image_url", "image_url": {"url": "..."}},
{"type": "text", "text": "beta"},
]
}
assert list(iter_message_text(data)) == ["alpha", "beta"]
def test_iter_message_text_responses_api_list_input_mixed_dicts_and_strings():
"""Greptile P2: mixed-list ``input`` with content-part dicts AND bare
strings must yield every text fragment — read helpers used to truncate
the bare strings."""
data = {
"input": [
{"type": "text", "text": "from-dict"},
"from-bare-string",
{"type": "image_url", "image_url": {"url": "..."}},
"another-bare-string",
]
}
assert list(iter_message_text(data)) == [
"from-dict",
"from-bare-string",
"another-bare-string",
]
def test_iter_message_text_walks_messages_and_input_independently():
"""When both are present (rare), every fragment from either field is
inspected — a stricter guarantee than "first one wins"."""
data = {
"messages": [{"role": "user", "content": "msg-content"}],
"input": "input-content",
}
assert list(iter_message_text(data)) == ["msg-content", "input-content"]
def test_iter_message_text_empty_data():
assert list(iter_message_text({})) == []
assert list(iter_message_text({"messages": []})) == []
assert list(iter_message_text({"input": ""})) == []
def test_iter_message_text_responses_api_input_text_and_output_text_parts():
"""LIT-4294: Responses-API content parts use ``input_text`` (request) and
``output_text`` (assistant); reading only ``type == "text"`` skipped every
``/v1/responses`` body and every text guardrail was a no-op on that path."""
data = {
"input": [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "user text"}],
},
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": "assistant text"}],
},
]
}
assert list(iter_message_text(data)) == ["user text", "assistant text"]
def test_iter_message_text_responses_api_tool_call_taxonomy():
"""LIT-4294: a Responses ``input`` list freely mixes message items,
``function_call`` (no ``role``), and ``function_call_output`` items. The
old ``all(item has 'role')`` gate wrapped the whole list as one blob and
yielded nothing; every text fragment must be visited independently."""
data = {
"input": [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "hello"}],
},
{
"type": "function_call",
"call_id": "c1",
"name": "get_weather",
"arguments": "{}",
},
{
"type": "function_call_output",
"call_id": "c1",
"output": [{"type": "input_text", "text": "sunny"}],
},
]
}
assert list(iter_message_text(data)) == ["hello", "sunny"]
def test_iter_message_text_inspects_reasoning_content_and_summary():
"""VERIA: reasoning items forwarded as ``reasoning_content`` must be
inspected, including ``summary`` blocks the bridge reads as a fallback."""
data = {
"input": [
{
"type": "reasoning",
"id": "rs_1",
"content": [{"type": "summary_text", "text": "content secret"}],
"summary": [{"type": "summary_text", "text": "summary secret"}],
}
]
}
assert list(iter_message_text(data)) == ["content secret", "summary secret"]
# ── walk_user_text ────────────────────────────────────────────────────────────
def test_walk_user_text_redacts_string_messages_in_place():
data = {
"messages": [
{"role": "user", "content": "leak: AKIAEXAMPLE"},
{"role": "assistant", "content": "ok"},
]
}
visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
assert visited == 2
assert data["messages"][0]["content"] == "leak: [REDACTED]"
assert data["messages"][1]["content"] == "ok"
def test_walk_user_text_redacts_multimodal_text_parts():
"""VERIA-11: list-content text parts must be mutable for in-place redaction."""
data = {
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "AKIAEXAMPLE here"},
{"type": "image_url", "image_url": {"url": "..."}},
{"type": "text", "text": "no secret"},
],
}
]
}
visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
assert visited == 2
parts = data["messages"][0]["content"]
assert parts[0] == {"type": "text", "text": "[REDACTED] here"}
# Non-text part must be left untouched.
assert parts[1] == {"type": "image_url", "image_url": {"url": "..."}}
assert parts[2] == {"type": "text", "text": "no secret"}
def test_walk_user_text_redacts_responses_api_string_input():
data = {"input": "leak AKIAEXAMPLE"}
visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
assert visited == 1
assert data["input"] == "leak [REDACTED]"
def test_walk_user_text_redacts_responses_api_list_input():
data = {
"input": [
{"type": "text", "text": "AKIAEXAMPLE"},
{"type": "image_url", "image_url": {"url": "..."}},
]
}
visited = walk_user_text(data, lambda s: f"[redacted]{s}[/]")
assert visited == 1
assert data["input"][0] == {"type": "text", "text": "[redacted]AKIAEXAMPLE[/]"}
assert data["input"][1] == {"type": "image_url", "image_url": {"url": "..."}}
def test_walk_user_text_redacts_responses_input_text_and_output_text_parts():
"""LIT-4294: ``walk_user_text`` must recognise the Responses text-part
variants so masking guardrails (secret detection, PII) actually redact
``/v1/responses`` bodies instead of no-op'ing on them."""
data = {
"input": [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "AKIAEXAMPLE"}],
},
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": "AKIAEXAMPLE too"}],
},
]
}
visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
assert visited == 2
assert data["input"][0]["content"][0] == {
"type": "input_text",
"text": "[REDACTED]",
}
assert data["input"][1]["content"][0] == {
"type": "output_text",
"text": "[REDACTED] too",
}
def test_walk_user_text_redacts_function_call_output_text():
"""LIT-4294: tool-call round-trips carry secrets in
``function_call_output.output``; the redact walker must descend into it
while leaving ``function_call`` items (call_id, arguments) untouched."""
data = {
"input": [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "AKIAEXAMPLE user"}],
},
{
"type": "function_call",
"call_id": "c1",
"name": "get_weather",
"arguments": '{"AKIAEXAMPLE": 1}',
},
{
"type": "function_call_output",
"call_id": "c1",
"output": [{"type": "input_text", "text": "AKIAEXAMPLE tool"}],
},
]
}
visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
assert visited == 2
assert data["input"][0]["content"][0]["text"] == "[REDACTED] user"
assert data["input"][1] == {
"type": "function_call",
"call_id": "c1",
"name": "get_weather",
"arguments": '{"AKIAEXAMPLE": 1}',
}
assert data["input"][2]["output"][0]["text"] == "[REDACTED] tool"
def test_walk_user_text_redacts_function_call_output_string_output():
"""LIT-4294: ``function_call_output.output`` is also a plain string in
OpenAI's Responses spec; the redact walker must handle both forms."""
data = {
"input": [
{
"type": "function_call_output",
"call_id": "c1",
"output": "AKIAEXAMPLE tool",
},
]
}
visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
assert visited == 1
assert data["input"][0]["output"] == "[REDACTED] tool"
def test_walk_user_text_redacts_mixed_list_input():
"""Read and write helpers must agree on coverage — bare strings inside
a mixed ``input`` list are inspected by both."""
data = {
"input": [
{"type": "text", "text": "secret-one"},
"secret-two",
{"type": "image_url", "image_url": {"url": "..."}},
]
}
visited = walk_user_text(data, lambda s: f"<{s}>")
assert visited == 2
assert data["input"][0] == {"type": "text", "text": "<secret-one>"}
assert data["input"][1] == "<secret-two>"
assert data["input"][2] == {"type": "image_url", "image_url": {"url": "..."}}
def test_walk_user_text_redacts_reasoning_content_and_summary():
"""VERIA: in-place redaction must cover both plaintext shapes the bridge
forwards from a reasoning item."""
data = {
"input": [
{
"type": "reasoning",
"id": "rs_1",
"content": [{"type": "summary_text", "text": "AKIAEXAMPLE content"}],
"summary": [{"type": "summary_text", "text": "AKIAEXAMPLE summary"}],
}
]
}
visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]"))
assert visited == 2
item = data["input"][0]
assert item["content"][0]["text"] == "[REDACTED] content"
assert item["summary"][0]["text"] == "[REDACTED] summary"
assert item["id"] == "rs_1"
# ── build_inspection_messages ─────────────────────────────────────────────────
def test_build_inspection_messages_chat_completion_passthrough():
data = {
"messages": [
{"role": "system", "content": "be helpful"},
{"role": "user", "content": "hi"},
]
}
assert build_inspection_messages(data) == [
{"role": "system", "content": "be helpful"},
{"role": "user", "content": "hi"},
]
def test_build_inspection_messages_joins_multimodal_text_parts():
data = {
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "first part"},
{"type": "image_url", "image_url": {"url": "..."}},
{"type": "text", "text": "second part"},
],
}
]
}
assert build_inspection_messages(data) == [{"role": "user", "content": "first part\nsecond part"}]
def test_build_inspection_messages_lifts_responses_api_input():
"""fniVO9-F: ``input`` must be visible to hooks that POST messages to a remote API."""
data = {"input": "responses-api content"}
assert build_inspection_messages(data) == [{"role": "user", "content": "responses-api content"}]
def test_build_inspection_messages_drops_messages_with_no_text():
data = {
"messages": [
{"role": "user", "content": ""},
{
"role": "user",
"content": [{"type": "image_url", "image_url": {"url": "..."}}],
},
{"role": "user", "content": "kept"},
]
}
assert build_inspection_messages(data) == [{"role": "user", "content": "kept"}]
def test_build_inspection_messages_responses_api_tool_call_taxonomy():
"""LIT-4294: mixed Responses ``input`` (message + function_call +
function_call_output) must produce a non-empty inspection list. The
customer's writeup reproduced a 422 from AIM's ``/fw/v1/analyze``
(``No messages in the request``) when this synthesised list came back
empty; every other guardrail silently scanned nothing on the same
input."""
data = {
"input": [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "hello"}],
},
{
"type": "function_call",
"call_id": "c1",
"name": "get_weather",
"arguments": "{}",
},
{
"type": "function_call_output",
"call_id": "c1",
"output": [{"type": "input_text", "text": "sunny"}],
},
]
}
assert build_inspection_messages(data) == [
{"role": "user", "content": "hello"},
{"role": "tool", "content": "sunny"},
]
def test_build_inspection_messages_function_call_output_defaults_to_tool():
"""LIT-4294: a Responses ``function_call_output`` item is the semantic
equivalent of a chat-completions ``role: "tool"`` message, so the shared
helper synthesises ``role: "tool"`` when the item has no explicit role.
AIM's schema-safe coercion happens at the AIM call site, not here."""
data = {
"input": [
{
"type": "function_call_output",
"call_id": "c1",
"output": [{"type": "input_text", "text": "tool text"}],
},
]
}
assert build_inspection_messages(data) == [{"role": "tool", "content": "tool text"}]
def test_build_inspection_messages_function_call_output_preserves_explicit_role():
"""When ``function_call_output`` carries a caller-supplied ``role`` the
shared helper preserves it rather than synthesising ``tool``."""
data = {
"input": [
{
"type": "function_call_output",
"role": "assistant",
"call_id": "c1",
"output": [{"type": "input_text", "text": "tool text"}],
},
]
}
assert build_inspection_messages(data) == [{"role": "assistant", "content": "tool text"}]
def test_build_inspection_messages_bare_content_part_preserves_explicit_role():
"""A bare content-part dict with an explicit ``role`` keeps it. Only
absent roles get defaulted to ``user``."""
data = {
"input": [
{"type": "input_text", "text": "no role"},
{"type": "output_text", "role": "assistant", "text": "with role"},
]
}
assert build_inspection_messages(data) == [
{"role": "user", "content": "no role"},
{"role": "assistant", "content": "with role"},
]
def test_build_inspection_messages_message_item_preserves_role():
"""Responses message items carry a role explicitly; the shared helper
passes it through untouched."""
data = {
"input": [
{"type": "message", "role": "system", "content": [{"type": "input_text", "text": "sys"}]},
{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "asst"}]},
]
}
assert build_inspection_messages(data) == [
{"role": "system", "content": "sys"},
{"role": "assistant", "content": "asst"},
]
def test_build_inspection_messages_empty_data():
assert build_inspection_messages({}) == []
assert build_inspection_messages({"messages": []}) == []
assert build_inspection_messages({"input": ""}) == []
def test_build_inspection_messages_includes_reasoning_summary():
"""VERIA: remote guardrail APIs must see reasoning summaries even when
the reasoning item has no ``content`` field."""
data = {
"input": [
{
"type": "reasoning",
"id": "rs_1",
"summary": [{"type": "summary_text", "text": "secret summary"}],
}
]
}
assert build_inspection_messages(data) == [
{"role": "assistant", "content": "secret summary"}
]
# ── has_non_string_content ────────────────────────────────────────────────────
def test_has_non_string_content_string_messages():
data = {"messages": [{"role": "user", "content": "hello"}]}
assert has_non_string_content(data) is False
def test_has_non_string_content_multimodal_messages():
data = {"messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]}
assert has_non_string_content(data) is True
def test_has_non_string_content_responses_api_string_input():
assert has_non_string_content({"input": "plain string"}) is False
def test_has_non_string_content_responses_api_list_input():
assert has_non_string_content({"input": ["a", "b"]}) is True
def test_has_non_string_content_empty_data():
assert has_non_string_content({}) is False
assert has_non_string_content({"messages": []}) is False
assert has_non_string_content({"input": ""}) is False
# ── apply_redacted_messages_back ──────────────────────────────────────────────
def test_apply_redacted_messages_back_chat_completion():
data = {"messages": [{"role": "user", "content": "secret"}]}
apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}])
assert data["messages"] == [{"role": "user", "content": "[REDACTED]"}]
assert "input" not in data
def test_apply_redacted_messages_back_responses_api_string_input():
"""A Responses-API request reads ``data["input"]``; writing only to
``messages`` would let unredacted text reach the LLM."""
data = {"input": "secret payload"}
apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}])
assert data["input"] == "[REDACTED]"
def test_apply_redacted_messages_back_both_fields():
"""Defensive: when both fields are present, both are updated."""
data = {
"messages": [{"role": "user", "content": "old"}],
"input": "old",
}
apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}])
assert data["messages"] == [{"role": "user", "content": "[REDACTED]"}]
assert data["input"] == "[REDACTED]"
def test_apply_redacted_messages_back_skips_input_when_not_string():
"""List ``input`` (multimodal Responses-API) is left alone — the
multimodal-degrades-to-block guard runs upstream."""
data = {"input": [{"type": "text", "text": "leak"}]}
apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}])
assert data["input"] == [{"type": "text", "text": "leak"}]
# -------------------------------------------------------------------
# LIT-4302: custom_tool_call_output walking
# -------------------------------------------------------------------
def test_iter_message_text_walks_custom_tool_call_output():
"""custom_tool_call_output items should yield their output text."""
data = {
"input": [
{"type": "custom_tool_call_output", "output": "tool-secret"},
]
}
from litellm.proxy.guardrails._content_utils import iter_message_text
texts = list(iter_message_text(data))
assert "tool-secret" in texts
def test_walk_user_text_redacts_custom_tool_call_output():
"""walk_user_text should rewrite text inside custom_tool_call_output."""
data = {
"input": [
{"type": "custom_tool_call_output", "output": "PII-data"},
]
}
count = walk_user_text(data, lambda t: t.replace("PII-data", "[MASKED]"))
assert count >= 1
assert data["input"][0]["output"] == "[MASKED]"
def test_build_inspection_messages_custom_tool_call_output():
"""build_inspection_messages should include custom_tool_call_output text."""
data = {
"input": [
{"type": "custom_tool_call_output", "output": "custom-tool-leak"},
]
}
msgs = build_inspection_messages(data)
assert any("custom-tool-leak" in m["content"] for m in msgs)