diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index fb886e42780..bac95918599 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -4,11 +4,13 @@ import json from collections.abc import Callable, Iterator, Mapping, Sequence from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles +from openai.types.responses import ResponseFunctionWebSearch from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage +from litellm.types.responses.main import with_hosted_web_tool_arguments _DECODED_JSON: Final = TypeAdapter(object) @@ -135,6 +137,30 @@ def stream_item_field(item: object, field: str) -> object | None: return getattr(item, field, None) +def web_search_call(item: object) -> ResponseFunctionWebSearch | None: + if isinstance(item, ResponseFunctionWebSearch): + return item + if stream_item_field(item, "type") != "web_search_call": + return None + try: + return ResponseFunctionWebSearch.model_validate(item.model_dump() if isinstance(item, BaseModel) else item) + except ValidationError: + return None + + +def with_web_search_call_arguments(item: object, arguments: str) -> object: + search: Final = web_search_call(item) + if search is None: + return item + action: Final = with_hosted_web_tool_arguments(search.action, arguments) + if isinstance(item, ResponseFunctionWebSearch): + return item.model_copy(update={"action": action}) + if isinstance(item, BaseModel): + return item.model_copy(update={"action": action.model_dump()}) + fields: Final = cast("Mapping[str, object]", item) # cast-ok: a validated web_search_call dict is str-keyed JSON + return {**fields, "action": action.model_dump()} + + def stream_item_fingerprint(item: object) -> str: plain: Final = item.model_dump() if isinstance(item, BaseModel) else item return json.dumps(plain, sort_keys=True, default=str) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 9ec8923b4cc..3d30b909da8 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -38,7 +38,6 @@ from itertools import accumulate, chain, groupby, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeVar, Union, cast -from openai.types.responses import ResponseFunctionWebSearch from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict @@ -67,6 +66,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_fingerprint, stream_item_items, unappliable_request_rewrite, + web_search_call, + with_web_search_call_arguments, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import ( append_returned_tools, @@ -111,7 +112,6 @@ from litellm.types.llms.openai import ( from litellm.types.responses.main import ( GenericResponseOutputItem, OutputText, - with_hosted_web_tool_arguments, ) from litellm.types.utils import GenericGuardrailAPIInputs @@ -334,19 +334,8 @@ def _is_tool_call_item(item: object) -> bool: return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES -def _web_search_call(item: object) -> ResponseFunctionWebSearch | None: - if isinstance(item, ResponseFunctionWebSearch): - return item - if stream_item_field(item, "type") != _WEB_SEARCH_CALL_TYPE: - return None - try: - return ResponseFunctionWebSearch.model_validate(item.model_dump() if isinstance(item, BaseModel) else item) - except ValidationError: - return None - - def _claude_web_search_tool_call(item: object) -> ChatCompletionAssistantToolCall | None: - search: Final = _web_search_call(item) + search: Final = web_search_call(item) return None if search is None else claude_web_search_call_tool_call(search) @@ -405,16 +394,7 @@ def _scanned_tool_call(item: Mapping[str, object], index: int) -> ChatCompletion def _rewritten_web_search_call(item: object, payload: str | None) -> object: - search: Final = _web_search_call(item) - if payload is None or search is None: - return item - action: Final = with_hosted_web_tool_arguments(search.action, payload) - if isinstance(item, ResponseFunctionWebSearch): - return item.model_copy(update={"action": action}) - if isinstance(item, BaseModel): - return item.model_copy(update={"action": action.model_dump()}) - fields: Final = _tool_call_output_item_mapping(item) - return item if fields is None else {**fields, "action": action.model_dump()} + return item if payload is None else with_web_search_call_arguments(item, payload) def _tool_search_rewrite(payload: str) -> Mapping[str, object]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index 972160bf3ae..8724093a7a3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -43,6 +43,10 @@ from litellm.llms.anthropic.common_utils import ( with_server_tool_call_arguments, write_anthropic_server_tool_calls, ) +from litellm.llms.base_llm.guardrail_translation.utils import ( + web_search_call, + with_web_search_call_arguments, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -52,6 +56,7 @@ from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.anthropic import claude_hosted_web_tool_call from litellm.types.utils import ( Choices, Delta, @@ -129,6 +134,12 @@ class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" +def _claude_hosted_web_tool_arguments(item: object) -> str | None: + search: Final = web_search_call(item) + hosted_call: Final = None if search is None else claude_hosted_web_tool_call(search) + return None if hosted_call is None else hosted_call[2] + + def _streamed_choice_index(choice: object) -> int: index: Final[object] = getattr(choice, "index", None) return index if isinstance(index, int) else 0 @@ -1691,7 +1702,7 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_responses_api_output( self, - output_items: list, + output_items: list[object], sanitized_text: str | None, sanitized_messages: list[dict[str, object]] | None, ) -> bool: @@ -1703,7 +1714,10 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): if not replacement_text: return False applied = False - for item in output_items: + for index, item in enumerate(output_items): + if _claude_hosted_web_tool_arguments(item): + output_items[index] = with_web_search_call_arguments(item, "{}") # rebind-ok: in-place redaction + applied = True content = getattr(item, "content", None) or (item.get("content") if isinstance(item, dict) else None) if isinstance(content, list): for part in content: @@ -2064,8 +2078,9 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): output_items: Final = getattr(response, "output", None) if not isinstance(output_items, list): return [] + typed_output_items: Final[Sequence[object]] = output_items output_parts: Final[list[str]] = [] - for item in output_items: + for item in typed_output_items: get = item.get if isinstance(item, dict) else (lambda k: getattr(item, k, None)) for part in get("content") or []: pget = part.get if isinstance(part, dict) else (lambda k: getattr(part, k, None)) @@ -2079,6 +2094,8 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): direct = get("text") if isinstance(direct, str) and direct: output_parts.append(direct) + if hosted_arguments := _claude_hosted_web_tool_arguments(item): + output_parts.append(hosted_arguments) joined = " ".join(output_parts) return [{"role": "assistant", "content": joined}] if joined else [] diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index bb205c5e251..365bb0ebfef 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -7,7 +7,7 @@ with guardrail transformations, including tool calls. import json from collections.abc import Mapping -from typing import Any, Literal, Optional +from typing import Any, Final, Literal, Optional import pytest @@ -163,8 +163,8 @@ def _hosted_web_call_inputs(web_search_calls) -> list[tuple[str, str, Optional[s ] -_SEARCH_INPUT = {"query": "jane@example.com"} -_FETCH_INPUT = {"url": "https://example.com/?email=jane@example.com"} +_SEARCH_INPUT: Final = {"query": "jane@example.com"} +_FETCH_INPUT: Final = {"url": "https://example.com/?email=jane@example.com"} class QueryRedactingGuardrail(CustomGuardrail): diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py index 3506afbaa5c..740c3c90758 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py @@ -404,6 +404,33 @@ class TestCiscoAIDefenseChatMode: ) +_CLAUDE_HOSTED_WEB_ACTIONS = ( + {"type": "search", "query": "ssn 123-45-6789"}, + {"type": "open_page", "url": "https://example.com/?ssn=123-45-6789"}, +) + + +def _claude_hosted_web_call_response(action: dict): + from openai.types.responses import ResponseFunctionWebSearch + + from litellm.types.llms.openai import ResponsesAPIResponse + + return ResponsesAPIResponse( + id="resp_1", + created_at=0, + output=[ + ResponseFunctionWebSearch( + id="ws_srvtoolu_01", type="web_search_call", status="completed", action=action + ) + ], + parallel_tool_calls=False, + tool_choice=None, + tools=None, + top_p=None, + usage=None, + ) + + class TestCiscoAIDefenseResponsesAPIOutput: @staticmethod @@ -505,6 +532,23 @@ class TestCiscoAIDefenseResponsesAPIOutput: joined = " ".join(m.get("content", "") for m in (sent.get("messages") or [])) assert "4111-1111-1111-1111" in joined + @pytest.mark.parametrize("action", _CLAUDE_HOSTED_WEB_ACTIONS) + @pytest.mark.asyncio + async def test_post_call_scans_claude_hosted_web_call(self, action): + g = _make_guardrail(event_hook="post_call") + post_mock = AsyncMock(return_value=_safe_response()) + + with _patch_inspection_post(g, post_mock): + await g.async_post_call_success_hook( + data={"input": [{"role": "user", "content": "look it up"}]}, + user_api_key_dict=UserAPIKeyAuth(), + response=_claude_hosted_web_call_response(action), + ) + + sent = post_mock.call_args.kwargs["json"] + joined = " ".join(m.get("content", "") for m in (sent.get("messages") or [])) + assert "123-45-6789" in joined + @pytest.mark.asyncio async def test_post_call_responses_api_violation_is_blocked(self): g = _make_guardrail(event_hook="post_call") @@ -573,6 +617,24 @@ class TestCiscoAIDefenseResponsesAPIOutputRedaction: ) + @pytest.mark.parametrize("action", _CLAUDE_HOSTED_WEB_ACTIONS) + @pytest.mark.asyncio + async def test_redact_clears_claude_hosted_web_call(self, action): + g = _make_guardrail(event_hook="post_call", on_flagged_action="monitor") + cisco_resp = _redact_response(sanitized_text="[REDACTED]") + + with _patch_inspection_post(g, AsyncMock(return_value=cisco_resp)): + result = await g.async_post_call_success_hook( + data={"input": [{"role": "user", "content": "look it up"}]}, + user_api_key_dict=UserAPIKeyAuth(), + response=_claude_hosted_web_call_response(action), + ) + + call = result.output[0] + assert (call.id, call.action.type) == ("ws_srvtoolu_01", action["type"]) + assert "123-45-6789" not in call.action.model_dump_json() + + class TestCiscoAIDefenseResponsesAPIInputRedaction: @pytest.mark.parametrize(