diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 1d31e33c77c..6cdf49818e6 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -32,6 +32,9 @@ def is_text_content_call_type(call_type: str) -> bool: return call_type in TEXT_CONTENT_CALL_TYPES +TEXT_PART_TYPES: FrozenSet[str] = frozenset({"text", "input_text", "output_text"}) + + def _iter_text_parts_in_content(content: Any) -> Iterator[str]: """Yield text fragments from a ``message.content`` value (string or multimodal list). Non-text parts (images, audio, …) are skipped.""" @@ -48,7 +51,7 @@ def _iter_text_parts_in_content(content: Any) -> Iterator[str]: continue if not isinstance(part, dict): continue - if part.get("type") == "text": + if part.get("type") in TEXT_PART_TYPES: text = part.get("text") if isinstance(text, str) and text: yield text @@ -58,14 +61,20 @@ def _coerce_input_to_messages(input_value: Any) -> List[Dict[str, Any]]: """Coerce a Responses-API ``data["input"]`` value into chat-style messages.""" if isinstance(input_value, str): return [{"role": "user", "content": input_value}] - if isinstance(input_value, list): - if input_value and all(isinstance(item, dict) and "role" in item for item in input_value): - return list(input_value) - # Mixed lists (content-part dicts + bare strings) and pure - # string/dict lists all become a single user message; the content - # iterator below handles each element type uniformly. - return [{"role": "user", "content": input_value}] - return [] + if not isinstance(input_value, list): + return [] + messages: List[Dict[str, Any]] = [] + for item in input_value: + if isinstance(item, str): + messages.append({"role": "user", "content": item}) + elif isinstance(item, dict): + if item.get("type") in TEXT_PART_TYPES: + 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: + messages.append({"role": item.get("role") or "tool", "content": item["output"]}) + return messages def _iter_inspection_messages(data: Dict[str, Any]) -> Iterator[Dict[str, Any]]: @@ -112,7 +121,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: new_parts.append(visit(part)) elif ( isinstance(part, dict) - and part.get("type") == "text" + and part.get("type") in TEXT_PART_TYPES and isinstance(part.get("text"), str) and part["text"] ): @@ -136,25 +145,20 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: data["input"] = visit(input_value) return visited if isinstance(input_value, list): - # List of full messages: rewrite each message's content. - if input_value and all(isinstance(item, dict) and "role" in item for item in input_value): - for item in input_value: - if "content" in item: - item["content"] = _rewrite_content(item["content"]) - return visited - # List of content parts and/or bare strings: rewrite in place. for idx, item in enumerate(input_value): - if isinstance(item, str) and item: - visited += 1 - input_value[idx] = visit(item) - elif ( - isinstance(item, dict) - and item.get("type") == "text" - and isinstance(item.get("text"), str) - and item["text"] - ): - visited += 1 - input_value[idx] = {**item, "text": visit(item["text"])} + if isinstance(item, str): + if item: + visited += 1 + input_value[idx] = visit(item) + elif isinstance(item, dict): + if item.get("type") in TEXT_PART_TYPES: + if isinstance(item.get("text"), str) and item["text"]: + visited += 1 + 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: + item["output"] = _rewrite_content(item["output"]) return visited return visited diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index e7d9406ae3b..d22243cbe88 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -93,11 +93,10 @@ class AimGuardrail(CustomGuardrail): user_email=user_email, litellm_call_id=call_id, ) - # Covers multimodal list content + Responses-API input. response = await self.async_handler.post( f"{self.api_base}/fw/v1/analyze", headers=headers, - json={"messages": build_inspection_messages(data)}, + json={"messages": self._build_aim_inspection_messages(data)}, ) response.raise_for_status() res = response.json() @@ -116,6 +115,15 @@ class AimGuardrail(CustomGuardrail): verbose_proxy_logger.error(f"Aim: {action_type} action") return data + @staticmethod + def _build_aim_inspection_messages(data: dict) -> list[dict[str, str]]: + """AIM validates against the OpenAI chat schema. Bare ``role: "tool"`` + without ``tool_call_id`` and bare ``role: "function"`` without ``name`` + are rejected; the flatten drops those fields, so any role outside + ``{system, user, assistant}`` collapses to ``user`` for the AIM POST.""" + safe_roles = {"system", "user", "assistant"} + return [{**m, "role": "user"} if m["role"] not in safe_roles else m for m in build_inspection_messages(data)] + @staticmethod def _rejection(message: str, *, openai_code: str | None = None) -> ProxyException: return ProxyException( @@ -177,7 +185,10 @@ class AimGuardrail(CustomGuardrail): user_email=user_email, litellm_call_id=call_id, ), - json={"messages": build_inspection_messages(request_data) + [{"role": "assistant", "content": output}]}, + json={ + "messages": self._build_aim_inspection_messages(request_data) + + [{"role": "assistant", "content": output}] + }, ) response.raise_for_status() res = response.json() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_aim.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_aim.py new file mode 100644 index 00000000000..2e83422074e --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_aim.py @@ -0,0 +1,88 @@ +"""Tests for the AIM guardrail's inspection-payload construction.""" + +from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail + + +def test_aim_inspection_messages_coerces_chat_completions_tool_role_to_user(): + """LIT-4294: A valid chat-completions ``role: "tool"`` message carries a + ``tool_call_id``, but the inspection flatten drops every field except + ``role`` and ``content``. A bare ``tool`` message without ``tool_call_id`` + is schema-invalid per the OpenAI chat schema, and the customer's writeup + reproduced AIM's ``/fw/v1/analyze`` returning 422 on exactly that shape. + The AIM POST collapses the role to ``user``; the outbound request to the + LLM is untouched.""" + data = { + "messages": [ + {"role": "user", "content": "weather in SF"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "sunny"}, + ] + } + assert AimGuardrail._build_aim_inspection_messages(data) == [ + {"role": "user", "content": "weather in SF"}, + {"role": "user", "content": "sunny"}, + ] + + +def test_aim_inspection_messages_coerces_non_standard_caller_role_to_user(): + """LIT-4294: A caller-supplied role outside {system, user, assistant} + (e.g. ``developer``, ``function``) is coerced to ``user`` for the AIM + POST, since AIM validates the payload against the OpenAI chat schema + and rejects unknown roles the same way it rejects bare ``tool``.""" + data = { + "messages": [ + {"role": "developer", "content": "system-ish instruction"}, + {"role": "user", "content": "normal user text"}, + ] + } + assert AimGuardrail._build_aim_inspection_messages(data) == [ + {"role": "user", "content": "system-ish instruction"}, + {"role": "user", "content": "normal user text"}, + ] + + +def test_aim_inspection_messages_coerces_responses_function_call_output_role(): + """LIT-4294: the shared helper synthesises ``role: "tool"`` for a + Responses ``function_call_output`` item (semantic equivalent of + chat-completions tool messages). AIM's schema-validating POST cannot + carry ``tool_call_id`` in the flat inspection payload, so AIM collapses + that ``tool`` role to ``user`` locally before POSTing.""" + data = { + "input": [ + { + "type": "function_call_output", + "call_id": "c1", + "output": [{"type": "input_text", "text": "sunny"}], + }, + ] + } + assert AimGuardrail._build_aim_inspection_messages(data) == [ + {"role": "user", "content": "sunny"}, + ] + + +def test_aim_inspection_messages_preserves_safe_roles(): + """Safe roles pass through untouched — the coercion only fires for + roles the OpenAI chat schema flatten cannot represent standalone.""" + data = { + "messages": [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + } + assert AimGuardrail._build_aim_inspection_messages(data) == [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index 099fca78a62..34d92505359 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -8,7 +8,6 @@ from litellm.proxy.guardrails._content_utils import ( walk_user_text, ) - # ── iter_message_text ──────────────────────────────────────────────────────────── @@ -101,6 +100,55 @@ def test_iter_message_text_empty_data(): 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"] + + # ── walk_user_text ──────────────────────────────────────────────────────────── @@ -160,6 +208,89 @@ def test_walk_user_text_redacts_responses_api_list_input(): 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.""" @@ -206,17 +337,13 @@ def test_build_inspection_messages_joins_multimodal_text_parts(): } ] } - assert build_inspection_messages(data) == [ - {"role": "user", "content": "first part\nsecond 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"} - ] + assert build_inspection_messages(data) == [{"role": "user", "content": "responses-api content"}] def test_build_inspection_messages_drops_messages_with_no_text(): @@ -233,6 +360,102 @@ def test_build_inspection_messages_drops_messages_with_no_text(): 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": []}) == []