fix(guardrails): walk custom_tool_call_output items in _content_utils (#32969)

* fix(guardrails): walk custom_tool_call_output items in _content_utils

* Change _OUTPUT_ITEM_TYPES to Frozenset type

* fix(guardrails): use builtin frozenset generic for _OUTPUT_ITEM_TYPES annotation

Frozenset is not a defined name (typing exports FrozenSet, the builtin is
frozenset), so module import raised NameError and broke every proxy test
suite. The builtin generic is valid on the supported python floor (3.10)
and keeps the UP006 ruff-strict budget at its ceiling, which the typing
alias would exceed
This commit is contained in:
yucheng-berri 2026-07-13 09:27:57 -07:00 • committed by GitHub
parent c2141b1113
commit ff2b690dd4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 47 additions and 2 deletions

View file

@ -34,6 +34,12 @@ def is_text_content_call_type(call_type: str) -> bool:
TEXT_PART_TYPES: FrozenSet[str] = frozenset({"text", "input_text", "output_text"})
# Responses-API item types whose ``output`` field carries user/tool text
# that guardrails should inspect. ``function_call_output`` is the
# built-in shape; ``custom_tool_call_output`` is the custom-tool
# counterpart (see ``ChatCompletionCustomToolCallOutput``).
_OUTPUT_ITEM_TYPES: frozenset[str] = frozenset({"function_call_output", "custom_tool_call_output"})
def _iter_text_parts_in_content(content: Any) -> Iterator[str]:
"""Yield text fragments from a ``message.content`` value (string or
@ -72,7 +78,7 @@ def _coerce_input_to_messages(input_value: Any) -> List[Dict[str, Any]]:
messages.append({"role": item.get("role") or "user", "content": [item]})
elif "content" in item:
messages.append({"role": item.get("role") or "user", "content": item["content"]})
elif item.get("type") == "function_call_output" and "output" in item:
elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item:
messages.append({"role": item.get("role") or "tool", "content": item["output"]})
return messages
@ -157,7 +163,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int:
input_value[idx] = {**item, "text": visit(item["text"])}
elif "content" in item:
item["content"] = _rewrite_content(item["content"])
elif item.get("type") == "function_call_output" and "output" in item:
elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item:
item["output"] = _rewrite_content(item["output"])
return visited

View file

@ -524,3 +524,42 @@ def test_apply_redacted_messages_back_skips_input_when_not_string():
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)