mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
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:
parent
c2141b1113
commit
ff2b690dd4
2 changed files with 47 additions and 2 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue