diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index f719dacd148..fd912839c95 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.redact_messages import ( should_redact_message_logging, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall, StandardLoggingPayload if TYPE_CHECKING: @@ -663,12 +664,11 @@ def _to_plain_dict(value): def _get_tool_calls(message) -> list | None: - """Return ``message.tool_calls`` only when it's a non-empty list. - - Works for dicts and Pydantic message objects via ``_safe_get``. + """Return the message's tool calls, Claude's server tool calls first, or + ``None`` when there are none. Works for dicts and Pydantic message objects. """ - tool_calls: Final = _safe_get(message, "tool_calls") - return tool_calls if isinstance(tool_calls, list) and tool_calls else None + tool_calls: Final = tool_calls_with_server_tool_calls(message) + return list(tool_calls) if tool_calls else None class _NormalizedToolCallFunction(TypedDict): diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index a72023eaedf..79788d35854 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -44,6 +44,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -315,7 +316,10 @@ def _to_dd_tool_calls(message: Mapping[str, object]) -> tuple[ToolCall, ...]: tool_id=_text_field(tool_call, "id"), type=_text_field(tool_call, "type", "function"), ) - for tool_call in (raw_tool_calls if isinstance(raw_tool_calls, list) else ()) + for tool_call in ( + *anthropic_server_tool_call_params(message), + *(raw_tool_calls if isinstance(raw_tool_calls, list) else ()), + ) if isinstance(tool_call, dict) for function in [_mapping_field(tool_call, "function")] ) diff --git a/litellm/integrations/helicone.py b/litellm/integrations/helicone.py index 63acc562d98..a0052069271 100644 --- a/litellm/integrations/helicone.py +++ b/litellm/integrations/helicone.py @@ -10,6 +10,23 @@ from litellm.integrations.helicone_mock_client import ( create_mock_helicone_client, should_use_helicone_mock, ) +from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls +from litellm.types.utils import ChatCompletionMessageCustomToolCall, ChatCompletionMessageToolCall + + +def _tool_use_block( + tool_call: ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall, +) -> dict[str, object]: + match tool_call: + case ChatCompletionMessageCustomToolCall(custom=custom): + return {"type": "tool_use", "id": tool_call.id, "name": custom.name, "input": custom.input} + case _: + return { + "type": "tool_use", + "id": tool_call.id, + "name": tool_call.function.name, + "input": tool_call.function.arguments, + } class HeliconeLogger: @@ -57,23 +74,8 @@ class HeliconeLogger: message = choice["message"] content = [] - if "tool_calls" in message and message["tool_calls"]: - for tool_call in message["tool_calls"]: - function = tool_call.get("function") - custom = tool_call.get("custom") - if not function and not custom: - continue - name, tool_input = ( - (function["name"], function["arguments"]) if function else (custom["name"], custom["input"]) - ) - content.append( - { - "type": "tool_use", - "id": tool_call["id"], - "name": name, - "input": tool_input, - } - ) + if tool_calls := tool_calls_with_server_tool_calls(message): + content = [_tool_use_block(tool_call) for tool_call in tool_calls] elif "content" in message and message["content"]: content = [{"type": "text", "text": message["content"]}] diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index ee75e6501e8..0117348da66 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -26,6 +26,7 @@ from litellm.integrations.opentelemetry import ( from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.langfuse_otel import ( LangfuseSpanAttributes, @@ -605,8 +606,8 @@ def _extract_choices_output(response_obj: _Gettable) -> str | None: message: Final = _as_gettable(first_choice.get("message", {})) if message is None: return None - tool_calls: Final = message.get("tool_calls") - if isinstance(tool_calls, list) and tool_calls: + tool_calls: Final = tool_calls_with_server_tool_calls(message) + if tool_calls: transformed_tool_calls: Final = [ entry for tool_call in tool_calls if (entry := _transformed_tool_call(response_obj, tool_call)) is not None ] diff --git a/litellm/integrations/lunary.py b/litellm/integrations/lunary.py index 6c020c78c32..08acb2c80dc 100644 --- a/litellm/integrations/lunary.py +++ b/litellm/integrations/lunary.py @@ -6,6 +6,7 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, Final import packaging +from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls if TYPE_CHECKING: from openai.types.chat.chat_completion_message_custom_tool_call import Custom @@ -67,8 +68,8 @@ def parse_messages(input): } # Only add tool_calls and function_call to res if they are set - if message.get("tool_calls"): - serialized["tool_calls"] = parse_tool_calls(message.get("tool_calls")) + if tool_calls := tool_calls_with_server_tool_calls(message): + serialized["tool_calls"] = parse_tool_calls(tool_calls) return serialized diff --git a/litellm/integrations/newrelic/newrelic.py b/litellm/integrations/newrelic/newrelic.py index 3623335f599..efc796825e4 100644 --- a/litellm/integrations/newrelic/newrelic.py +++ b/litellm/integrations/newrelic/newrelic.py @@ -56,6 +56,7 @@ from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.litellm_core_utils.redact_messages import should_redact_message_logging +from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.newrelic import NewRelicInitParams from litellm.types.utils import Message, ModelResponse, StandardLoggingPayload @@ -479,11 +480,11 @@ class NewRelicLogger(CustomLogger): content: Final = message.get("content") # Handle tool calls - if message.get("tool_calls"): + if tool_calls := tool_calls_with_server_tool_calls(message): try: - return json.dumps(message["tool_calls"]) + return json.dumps([tool_call.model_dump() for tool_call in tool_calls]) except Exception: - return str(message["tool_calls"]) + return str(tool_calls) # Handle None or missing content if content is None: diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index a9bd93e6420..39c485cb31f 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -42,6 +42,7 @@ from litellm.litellm_core_utils.service_tier_utils import ( get_requested_service_tier, get_served_service_tier, ) +from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params, tool_calls_with_server_tool_calls from litellm.secret_managers.main import get_secret_bool, str_to_bool from litellm.types.integrations.otel_span_attributes import SpanAttributes from litellm.types.services import ServiceLoggerPayload @@ -2736,9 +2737,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): for idx, choice in enumerate(response_obj.get("choices")): if choice.get("finish_reason"): message = choice.get("message") - tool_calls = message.get("tool_calls") - if tool_calls: - kv_pairs = OpenTelemetry._tool_calls_kv_pair(tool_calls) + if tool_calls := tool_calls_with_server_tool_calls(message): + kv_pairs = OpenTelemetry._tool_calls_kv_pair(list(tool_calls)) for key, value in kv_pairs.items(): self.safe_set_attribute( span=span, @@ -2856,7 +2856,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): transformed_msg = {"role": role, "parts": parts} if "id" in msg: transformed_msg["id"] = msg["id"] - if "tool_calls" in msg: + if server_tool_calls := anthropic_server_tool_call_params(cast(object, msg)): + transformed_msg["tool_calls"] = [*server_tool_calls, *(msg.get("tool_calls") or [])] + elif "tool_calls" in msg: transformed_msg["tool_calls"] = msg["tool_calls"] if "tool_call_id" in msg: transformed_msg["tool_call_id"] = msg["tool_call_id"] diff --git a/litellm/integrations/otel/mappers/utils.py b/litellm/integrations/otel/mappers/utils.py index 80ef2dd553a..bfd50fea5d5 100644 --- a/litellm/integrations/otel/mappers/utils.py +++ b/litellm/integrations/otel/mappers/utils.py @@ -12,6 +12,8 @@ from typing import Final from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue from litellm.integrations.otel.model.payloads import LLMCallSpanData, ToolDefinition +from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of +from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params @dataclass(frozen=True, slots=True) @@ -141,12 +143,13 @@ def output_messages(data: LLMCallSpanData) -> list: def message_tool_calls(message: object) -> tuple[MessageToolCall, ...]: + tool_call_values: Final = ( + *anthropic_server_tool_call_params(message), + *parts_of(message_field(message, "tool_calls")), + ) if not isinstance(message, dict): return () - tool_calls: Final = message.get("tool_calls") - if not isinstance(tool_calls, (list, tuple)) or not tool_calls: - return () - return tuple(tool_call for value in tool_calls if (tool_call := _message_tool_call(value)) is not None) + return tuple(tool_call for value in tool_call_values if (tool_call := _message_tool_call(value)) is not None) def _message_tool_call(value: object) -> MessageToolCall | None: diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 7a62ed6aed9..4238d72fc5e 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -10,6 +10,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, Literal, Protocol, cast from urllib.parse import urlsplit +from openai.types.responses.response_function_web_search import ResponseFunctionWebSearch +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm.integrations.otel.model.metadata import ( @@ -36,6 +38,7 @@ from litellm.integrations.otel.model.utils import ( as_str_tuple, ) from litellm.integrations.otel.routing import RoutingAttributeValue, routing_decision_attributes +from litellm.types.llms.anthropic import claude_hosted_web_tool_call # ``RequestIdentity`` and the request-metadata translation now live in # :mod:`metadata`; re-exported here so existing ``model.payloads`` imports keep @@ -809,6 +812,7 @@ class _Choice(TypedDict): _RESPONSES_TOOL_CALL_TYPES: Final = frozenset({"function_call", "custom_tool_call"}) +_WEB_SEARCH_CALL: Final = TypeAdapter(ResponseFunctionWebSearch) def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: @@ -816,9 +820,13 @@ def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: items: Final = _dicts(response.get("output")) messages: Final = tuple(item for item in items if item.get("type") == "message") parts: Final = tuple(part for item in messages for part in _dicts(item.get("content"))) - tool_calls: Final = tuple( + client_tool_calls: Final = tuple( _responses_tool_call(item) for item in items if item.get("type") in _RESPONSES_TOOL_CALL_TYPES ) + tool_calls: Final = ( + *(call for item in items if (call := _claude_web_search_tool_call(item)) is not None), + *client_tool_calls, + ) if not messages and not tool_calls: return () message: Final[_AssistantMessage] = { @@ -827,7 +835,10 @@ def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: "refusal": _responses_parts_text(parts, "refusal", "refusal"), "tool_calls": tool_calls or None, } - choice: Final[_Choice] = {"message": message, "finish_reason": _responses_finish_reason(response, bool(tool_calls))} + choice: Final[_Choice] = { + "message": message, + "finish_reason": _responses_finish_reason(response, bool(client_tool_calls)), + } return (choice,) @@ -981,6 +992,20 @@ def _responses_tool_call(item: Mapping[str, object]) -> _ToolCall: return tool_call +def _claude_web_search_tool_call(item: Mapping[str, object]) -> _ToolCall | None: + if item.get("type") != "web_search_call": + return None + try: + search: Final = _WEB_SEARCH_CALL.validate_python(item) + except ValidationError: + return None + claude_call: Final = claude_hosted_web_tool_call(search) + if claude_call is None: + return None + tool_id, name, arguments = claude_call + return {"id": tool_id, "type": "function", "function": {"name": name, "arguments": arguments}} + + def _responses_finish_reason(response: Mapping[str, object], has_tool_calls: bool) -> str | None: status: Final = as_str(response.get("status")) if status == "completed": diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index f683108a51a..2c3d8a731e1 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -39,6 +39,7 @@ from litellm.litellm_core_utils.llm_judge import ( parse_json_verdict, ) from litellm.litellm_core_utils.redact_messages import should_redact_message_logging +from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.redis.stores.spend_counters import ( shadow_eval_counter, @@ -205,7 +206,7 @@ def _chat_final_text(response_obj: object) -> str: if read is None: return "" prose: Final = extract_text_from_content(read("content")) - if not (read("tool_calls") or read("function_call")): + if not _tool_calls_list(read): return prose serialized: Final = _serialize_tool_calls(read) return f"{prose} {serialized}".strip() if prose else serialized @@ -221,8 +222,9 @@ _RESPONSES_TOOL_CALL_TYPES: Final = frozenset(("function_call", "custom_tool_cal def _tool_calls_list(read: Callable[[str], object]) -> tuple[object, ...]: - calls: Final = read("tool_calls") - listed: Final = tuple(calls) if isinstance(calls, Sequence) and not isinstance(calls, str) else () + listed: Final = tool_calls_with_server_tool_calls( + {"tool_calls": read("tool_calls"), "provider_specific_fields": read("provider_specific_fields")} + ) single: Final = read("function_call") return listed if listed else ((single,) if single is not None else ()) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 8f874edfddf..ec857218c79 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -604,6 +604,8 @@ def get_content_from_model_response(response: ModelResponse | dict) -> str: else: new_response = response + from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls + content = "" for choice in new_response.choices: @@ -611,9 +613,8 @@ def get_content_from_model_response(response: ModelResponse | dict) -> str: content += choice.message.content if choice.message.content else "" if choice.message.function_call: content += choice.message.function_call.model_dump_json() - if choice.message.tool_calls: - for tc in choice.message.tool_calls: - content += tc.model_dump_json() + for tc in tool_calls_with_server_tool_calls(choice.message): + content += tc.model_dump_json() elif isinstance(choice, StreamingChoices): content += getattr(choice, "delta", {}).get("content", "") or "" return content diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 5c18fc15bcb..17a1fa7d81b 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -2012,12 +2012,39 @@ def _tool_call_ids_with_a_kept_server_result(message: object) -> frozenset[str]: return frozenset() return frozenset( tool_call_id - for tool_call in parts_of(message_field(message, "tool_calls")) + for tool_call in ( + *parts_of(message_field(message, "tool_calls")), + *parts_of(message_field(provider_specific_fields, "server_tool_calls")), + ) if isinstance(tool_call_id := message_field(tool_call, "id"), str) and tool_call_is_rebuilt_as_server_tool_use(tool_call_id, provider_specific_fields) ) +def _assistant_tool_calls_with_server_tool_calls( + message: ChatCompletionAssistantMessage, +) -> Sequence[ChatCompletionAssistantToolCall] | None: + from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params + + tool_calls: Final = message.get("tool_calls") + server_tool_calls: Final = anthropic_server_tool_call_params(message) + if not server_tool_calls: + return tool_calls + client_tool_call_ids: Final = frozenset(tool_call["id"] for tool_call in tool_calls or ()) + return ( + *(tool_call for tool_call in server_tool_calls if tool_call["id"] not in client_tool_call_ids), + *(tool_calls or ()), + ) + + +def _server_tool_call_ids(message: object) -> frozenset[str]: + from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params + + return frozenset( + tool_call["id"] for tool_call in anthropic_server_tool_call_params(message) if tool_call["id"] is not None + ) + + def _server_tool_answers_replaced_after(messages: Sequence[object], index: int) -> frozenset[str]: kept_server_result_ids: Final = _tool_call_ids_with_a_kept_server_result(messages[index]) if not kept_server_result_ids: @@ -2102,6 +2129,12 @@ def _unrun_server_tool_use_ids(blocks: Sequence[Mapping[str, object]]) -> frozen ) +def _paused_server_tool_use_ids(blocks: Sequence[Mapping[str, object]]) -> frozenset[object]: + if any(block.get("type") == "tool_use" for block in blocks): + return frozenset() + return _unrun_server_tool_use_ids(blocks) + + def _last_block_id_among(blocks: Sequence[Mapping[str, object]], ids: frozenset[object]) -> object: return next(block_id for block in reversed(blocks) if (block_id := _block_id(block)) in ids) @@ -2115,7 +2148,7 @@ def _paused_turn_replay( tuple[_AnthropicToolInvokeItem, ...], tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...], ]: - unrun_ids: Final = _unrun_server_tool_use_ids(tool_invokes) + unrun_ids: Final = _paused_server_tool_use_ids(tool_invokes) if not unrun_ids: return tuple(tool_invokes), tuple(thinking_blocks) dropped_ids: Final = ( @@ -2125,10 +2158,11 @@ def _paused_turn_replay( def convert_to_anthropic_tool_invoke( - tool_calls: list[ChatCompletionAssistantToolCall], + tool_calls: Sequence[ChatCompletionAssistantToolCall], web_search_results: Sequence[object] | None = None, tool_results: Sequence[object] | None = None, answered_tool_ids: Container[str] = frozenset(), + server_tool_call_ids: Container[str] = frozenset(), ) -> list[_AnthropicToolInvokeItem]: """ OpenAI tool invokes: @@ -2206,14 +2240,11 @@ def convert_to_anthropic_tool_invoke( anthropic_tool_invoke.append(server_tool_result) elif is_anthropic_tool_search_use(tool_id, tool_name) and tool_id not in answered_tool_ids: continue - elif is_anthropic_server_tool_use(tool_id) and tool_id not in answered_tool_ids: + elif ( + tool_id in server_tool_call_ids or is_anthropic_server_tool_use(tool_id) + ) and tool_id not in answered_tool_ids: anthropic_tool_invoke.append( - { - "type": "server_tool_use", - "id": tool_id, - "name": tool_name, - "input": tool_input, - } + {"type": "server_tool_use", "id": tool_id, "name": tool_name, "input": tool_input} ) else: sanitized_tool_id = sanitize_anthropic_tool_use_id(tool_id) @@ -3174,7 +3205,7 @@ def anthropic_messages_pt( # If so, we need to interleave thinking blocks with tool call groups # to preserve the original content block ordering. # Fixes: https://github.com/BerriAI/litellm/issues/23047 - assistant_tool_calls = assistant_content_block.get("tool_calls") + assistant_tool_calls = _assistant_tool_calls_with_server_tool_calls(assistant_content_block) _has_server_tool_calls = False if assistant_tool_calls is not None: for _tc in assistant_tool_calls: @@ -3198,11 +3229,12 @@ def anthropic_messages_pt( web_search_results=_provider_specific_fields.get("web_search_results"), tool_results=_provider_specific_fields.get("tool_results"), answered_tool_ids=client_answered_tool_ids, + server_tool_call_ids=_server_tool_call_ids(assistant_content_block), ) if assistant_tool_calls is not None else () ) - turn_paused = bool(_unrun_server_tool_use_ids(tool_invoke_results)) + turn_paused = bool(_paused_server_tool_use_ids(tool_invoke_results)) replayed_tool_invokes, replayed_thinking_blocks = _paused_turn_replay( tool_invoke_results, thinking_blocks or (), is_final_message=msg_i == len(messages) - 1 ) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index ac5c0720eff..a5e79d7b2b4 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -14,7 +14,7 @@ from collections.abc import Iterable, Mapping, Sequence from dataclasses import replace from typing import TYPE_CHECKING, Any, Final -from pydantic import JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError import litellm from litellm.constants import REDACTED_BY_LITELLM @@ -301,9 +301,63 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ ) if key in fields }, + **( + {"server_tool_calls": _redacted_tool_call_dicts(fields["server_tool_calls"])} + if "server_tool_calls" in fields + else {} + ), + **( + {"web_search_calls": _redacted_web_search_calls(fields["web_search_calls"])} + if "web_search_calls" in fields + else {} + ), } +def _redacted_tool_call_dict(call: object) -> object: + if (mapping := _as_mapping(call)) is None or (function := _as_mapping(mapping.get("function"))) is None: + return call + return {**mapping, "function": {**function, "arguments": REDACTED_BY_LITELLM}} + + +def _redacted_tool_call_dicts(calls: object) -> object: + return [_redacted_tool_call_dict(call) for call in calls] if isinstance(calls, list) else calls + + +_WEB_ACTION_TEXT_FIELDS: Final = ("query", "url", "pattern") +_WEB_ACTION_LIST_FIELDS: Final = ("queries", "sources") + + +def _web_action_redactions(action: Mapping[str, object]) -> Mapping[str, object]: + return { + **{key: REDACTED_BY_LITELLM for key in _WEB_ACTION_TEXT_FIELDS if action.get(key) is not None}, + **{key: None for key in _WEB_ACTION_LIST_FIELDS if action.get(key) is not None}, + } + + +def _redacted_web_search_action(action: object) -> object: + if isinstance(action, BaseModel): + return action.model_copy(update=dict(_web_action_redactions(action.model_dump()))) + mapping: Final = _as_mapping(action) + return action if mapping is None else {**mapping, **_web_action_redactions(mapping)} + + +def _redacted_web_search_call(item: object) -> object: + if isinstance(item, BaseModel): + action: Final = getattr(item, "action", None) + return item if action is None else item.model_copy(update={"action": _redacted_web_search_action(action)}) + mapping: Final = _as_mapping(item) + if mapping is None or mapping.get("action") is None: + return item + return {**mapping, "action": _redacted_web_search_action(mapping["action"])} + + +def _redacted_web_search_calls(calls: object) -> object: + if (mapping := _as_mapping(calls)) is not None: + return {key: _redacted_web_search_call(call) for key, call in mapping.items()} + return [_redacted_web_search_call(call) for call in calls] if isinstance(calls, list) else calls + + def _redact_field_when_present(item: object, field: str) -> None: if getattr(item, field, None) is not None: setattr(item, field, REDACTED_BY_LITELLM) @@ -335,6 +389,8 @@ def _redact_responses_api_output(output_items): output_item.arguments = REDACTED_BY_LITELLM if output_type == "custom_tool_call" and hasattr(output_item, "input"): output_item.input = REDACTED_BY_LITELLM + if output_type == "web_search_call" and getattr(output_item, "action", None) is not None: + output_item.action = _redacted_web_search_action(output_item.action) _RESPONSE_ITEMS: Final = TypeAdapter(tuple[object, ...]) @@ -389,6 +445,11 @@ def _redacted_responses_api_output_item(item: object, redacted_str: str) -> obje if item_mapping.get("type") == "custom_tool_call" and "input" in item_mapping else {} ), + **( + {"action": _redacted_web_search_action(item_mapping["action"])} + if item_mapping.get("type") == "web_search_call" and item_mapping.get("action") is not None + else {} + ), } diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 6436cd832c7..8f1a08d2964 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -16,7 +16,7 @@ import json from collections.abc import Iterable, Mapping, MutableSequence, Sequence from copy import deepcopy from dataclasses import dataclass -from itertools import chain, repeat +from itertools import chain, filterfalse, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast, overload, runtime_checkable @@ -26,6 +26,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.common_utils import anthropic_server_tool_calls from litellm.llms.anthropic.messages.adapters.transformation import ( OPENAI_MAX_TOOL_NAME_LENGTH, LiteLLMAnthropicMessagesAdapter, @@ -206,17 +207,15 @@ class _WebSearchString: value: str -def _streamed_web_search_result_blocks(response: ModelResponse) -> tuple[Mapping[str, object], ...]: - choice: Final = response.choices[0] if response.choices else None - fields: Final = as_str_mapping(choice.message.provider_specific_fields) if isinstance(choice, Choices) else None - blocks: Final = ( - (as_str_mapping(item) for item in iter_items(fields.get("web_search_results"))) - if fields is not None and isinstance(fields.get("web_search_results"), list) - else () - ) - return tuple( - deepcopy(block) for block in blocks if block is not None and block.get("type") == "web_search_tool_result" - ) +def _legacy_hook_block(block: object) -> object: + if not isinstance(block, Mapping): + return block + fields: Final = _as_str_mapping(block) + return {**fields, "type": "tool_use"} if fields.get("type") == "server_tool_use" else fields + + +def _is_web_search_result_block(block: object) -> bool: + return isinstance(block, Mapping) and _as_str_mapping(block).get("type") == "web_search_tool_result" def _web_search_results(block: Mapping[str, object]) -> Iterable[object]: @@ -519,6 +518,13 @@ def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewrit return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}} +def _in_stream_block_order( + tool_calls: Sequence[ChatCompletionMessageToolCall], blocks: Sequence[tuple[int, Mapping[str, object]]] +) -> tuple[ChatCompletionMessageToolCall, ...]: + positions: Final = MappingProxyType({block.get("id"): position for position, (_, block) in enumerate(blocks)}) + return tuple(sorted(tool_calls, key=lambda tool_call: positions.get(tool_call.id, len(positions)))) + + def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]: """The guardrail-visible shape of each tool call, whether the guardrail handed back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts.""" @@ -586,11 +592,14 @@ class AnthropicMessagesHandler(BaseTranslation): if not isinstance(response, ModelResponse): return response translated: Final = self.adapter.translate_openai_response_to_anthropic(response) - result_blocks: Final = _streamed_web_search_result_blocks(response) - if not result_blocks: - return translated - content: Final = translated.get("content") or () - return {**translated, "content": [*result_blocks, *content]} + blocks: Final = tuple(map(_legacy_hook_block, translated.get("content") or ())) + return { + **translated, + "content": [ + *filter(_is_web_search_result_block, blocks), + *filterfalse(_is_web_search_result_block, blocks), + ], + } @staticmethod def _build_streaming_usage_response( @@ -1570,15 +1579,18 @@ class AnthropicMessagesHandler(BaseTranslation): if built_response is not None and hasattr(built_response, "choices") and built_response.choices: model_response: Final = cast(ModelResponse, built_response) first_choice: Final = cast(Choices, model_response.choices[0]) - tool_calls_list: Final = cast( - list[ChatCompletionMessageToolCall] | None, - first_choice.message.tool_calls, + tool_call_blocks: Final = self._streamed_tool_call_blocks(responses_so_far) + tool_calls_list: Final = _in_stream_block_order( + ( + *anthropic_server_tool_calls(first_choice.message), + *cast(list[ChatCompletionMessageToolCall], first_choice.message.tool_calls or []), + ), + tool_call_blocks, ) string_so_far = first_choice.message.content pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ()) guardrail_inputs: Final = GenericGuardrailAPIInputs() leading_texts: Final = (string_so_far,) if isinstance(string_so_far, str) and string_so_far else () - tool_call_blocks: Final = self._streamed_tool_call_blocks(responses_so_far) streamed_query_strings: Final = _streamed_query_strings(tool_call_blocks, pre_guardrail_tool_calls) streamed_result_strings: Final = self._streamed_web_search_result_strings(responses_so_far) streamed_web_search_strings: Final = (*streamed_result_strings, *streamed_query_strings) @@ -1586,7 +1598,7 @@ class AnthropicMessagesHandler(BaseTranslation): if texts_to_check: guardrail_inputs["texts"] = texts_to_check if tool_calls_list: - guardrail_inputs["tool_calls"] = tool_calls_list + guardrail_inputs["tool_calls"] = list(tool_calls_list) try: prepared_request_data = self._prepare_request_data( diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index c273c715ba6..4062a957e00 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -54,6 +54,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallFunctionChunk, ) from litellm.types.responses.main import ( + HOSTED_WEB_TOOL_NAMES, OutputCodeInterpreterCall, build_code_interpreter_log_outputs, build_web_search_call, @@ -761,6 +762,8 @@ class ModelResponseIterator: self._server_tool_inputs: dict[str, object] = {} self.tool_results: list[dict[str, Any]] = [] self._current_server_tool_id: str | None = None + self._current_server_tool_name: str = "" + self._server_tool_calls: tuple[ChatCompletionToolCallChunk, ...] = () self._container_id: str | None = None @property @@ -836,10 +839,9 @@ class ModelResponseIterator: ) self.content_blocks.append(content_block) if "partial_json" in content_block["delta"]: - # Only emit tool calls if we're in a tool_use or server_tool_use block # web_search_tool_result blocks also have input_json_delta but should not be treated as tool calls # See: https://github.com/BerriAI/litellm/issues/17254 - if self.current_content_block_type in ("tool_use", "server_tool_use"): + if self.current_content_block_type == "tool_use": tool_use = cast( ChatCompletionToolCallChunk, { @@ -938,7 +940,7 @@ class ModelResponseIterator: def _web_search_call_snapshot(self) -> dict[str, object]: return dict(self._web_search_calls) - def _refresh_pending_web_search_call(self, tool_use_id: str) -> bool: + def _refresh_pending_web_search_call(self, tool_use_id: str, tool_name: str) -> bool: if tool_use_id not in self._web_search_calls: return False self._web_search_calls[tool_use_id] = build_web_search_call( @@ -946,10 +948,11 @@ class ModelResponseIterator: self._server_tool_inputs.get(tool_use_id, {}), {"content": []}, status="in_progress", + tool_name=tool_name, ) return True - def _complete_web_search_call(self, result: dict[str, object]) -> None: + def _complete_web_search_call(self, result: dict[str, object], tool_name: str) -> None: tool_use_id: Final = result.get("tool_use_id") if not isinstance(tool_use_id, str) or tool_use_id not in self._web_search_calls: return @@ -957,6 +960,7 @@ class ModelResponseIterator: tool_id=tool_use_id, tool_input=self._server_tool_inputs.get(tool_use_id, {}), result=result, + tool_name=tool_name, ) def _build_code_interpreter_results(self) -> list: @@ -1036,7 +1040,6 @@ class ModelResponseIterator: content_block_start["content_block"]["type"] == "tool_use" or content_block_start["content_block"]["type"] == "server_tool_use" ): - self.tool_index += 1 # Reverse-map the (sanitized) tool name back to the # caller's original. No-op when the map is empty. _stream_tool_name = content_block_start["content_block"]["name"] @@ -1045,35 +1048,39 @@ class ModelResponseIterator: # Use empty string for arguments in content_block_start - actual arguments # come in subsequent content_block_delta chunks and get accumulated. # Using str(input) here would prepend '{}' causing invalid JSON accumulation. - tool_use = ChatCompletionToolCallChunk( - id=content_block_start["content_block"]["id"], - type="function", - function=ChatCompletionToolCallFunctionChunk( - name=_stream_tool_name, - arguments="", - ), - index=self.tool_index, - ) + if content_block_start["content_block"]["type"] == "tool_use": + self.tool_index += 1 + tool_use = ChatCompletionToolCallChunk( + id=content_block_start["content_block"]["id"], + type="function", + function=ChatCompletionToolCallFunctionChunk( + name=_stream_tool_name, + arguments="", + ), + index=self.tool_index, + ) # Track server tool use inputs for code_interpreter_results. # The initial input in content_block_start is typically {} # for streaming; the full input arrives via input_json_delta # and is assembled at content_block_stop. if content_block_start["content_block"]["type"] == "server_tool_use": self._current_server_tool_id = content_block_start["content_block"]["id"] + self._current_server_tool_name = _stream_tool_name tool_input: Final = content_block_start["content_block"].get("input", {}) self._server_tool_inputs[self._current_server_tool_id] = tool_input - if _stream_tool_name == "web_search": + if _stream_tool_name in HOSTED_WEB_TOOL_NAMES: self._web_search_calls[self._current_server_tool_id] = build_web_search_call( self._current_server_tool_id, tool_input, {"content": []}, status="in_progress", + tool_name=_stream_tool_name, ) provider_specific_fields["web_search_calls"] = self._web_search_call_snapshot() # Include caller information if present (for programmatic tool calling) if "caller" in content_block_start["content_block"]: caller_data: Final = content_block_start["content_block"]["caller"] - if caller_data: + if caller_data and tool_use is not None: tool_use["caller"] = cast(dict[str, Any], caller_data) elif content_block_start["content_block"]["type"] == "redacted_thinking": ( @@ -1104,7 +1111,7 @@ class ModelResponseIterator: # The full content comes in content_block_start, not in deltas # See: https://github.com/BerriAI/litellm/issues/17737 self.web_search_results.append(content_block_start["content_block"]) - self._complete_web_search_call(content_block_start["content_block"]) + self._complete_web_search_call(content_block_start["content_block"], "web_search") provider_specific_fields["web_search_results"] = self.web_search_results provider_specific_fields["web_search_calls"] = self._web_search_call_snapshot() elif content_type == "web_fetch_tool_result": @@ -1112,7 +1119,9 @@ class ModelResponseIterator: # The full content comes in content_block_start, not in deltas # Fixes: https://github.com/BerriAI/litellm/issues/18137 self.web_search_results.append(content_block_start["content_block"]) + self._complete_web_search_call(content_block_start["content_block"], "web_fetch") provider_specific_fields["web_search_results"] = self.web_search_results + provider_specific_fields["web_search_calls"] = self._web_search_call_snapshot() else: self.tool_results.append(content_block_start["content_block"]) provider_specific_fields["tool_results"] = self.tool_results @@ -1127,7 +1136,7 @@ class ModelResponseIterator: # check if tool call content block - only for tool_use and server_tool_use blocks if self.current_content_block_type in ("tool_use", "server_tool_use"): is_empty: Final = self.check_empty_tool_call_args() - if is_empty: + if is_empty and self.current_content_block_type == "tool_use": tool_use = ChatCompletionToolCallChunk( id=None, type="function", @@ -1151,8 +1160,24 @@ class ModelResponseIterator: self._server_tool_inputs[self._current_server_tool_id] = json.loads(args) except (json.JSONDecodeError, TypeError): pass - if self._refresh_pending_web_search_call(self._current_server_tool_id): + server_tool_input: Final = self._server_tool_inputs.get(self._current_server_tool_id, {}) + if self._refresh_pending_web_search_call( + self._current_server_tool_id, self._current_server_tool_name + ): provider_specific_fields["web_search_calls"] = self._web_search_call_snapshot() + self._server_tool_calls = ( + *self._server_tool_calls, + ChatCompletionToolCallChunk( + id=self._current_server_tool_id, + type="function", + function=ChatCompletionToolCallFunctionChunk( + name=self._current_server_tool_name, + arguments=args or json.dumps(server_tool_input), + ), + index=len(self._server_tool_calls), + ), + ) + provider_specific_fields["server_tool_calls"] = list(self._server_tool_calls) self._current_server_tool_id = None if self.is_response_format_tool: text = self._flush_response_format_buffer() diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index a6f6ec3d598..775eb6eda39 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -6,10 +6,11 @@ from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from functools import partial from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, cast +from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Final, NoReturn, cast import httpx -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from openai.types.responses.response_function_web_search import ResponseFunctionWebSearch +from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -110,11 +111,11 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, ChatCompletionToolParam, - OpenAIChatCompletionFinishReason, OpenAIMcpServerTool, OpenAIWebSearchOptions, ) from litellm.types.responses.main import ( + HOSTED_WEB_TOOL_NAMES, OutputCodeInterpreterCall, build_code_interpreter_log_outputs, build_web_search_call, @@ -150,6 +151,7 @@ from ..common_utils import ( stop_reason_provider_specific_fields, strip_advisor_blocks_from_messages, tools_marked_for_a_tool_config_point, + validated_object_mapping, ) if TYPE_CHECKING: @@ -187,9 +189,9 @@ _ANTHROPIC_TOOL_NAME_MAX_LEN: Final = 128 _DEFER_LOADING_ADAPTER: Final = TypeAdapter(bool) _TOOL_FIELDS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _TOOLS_ADAPTER: Final = TypeAdapter(list[object]) -_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(tuple[object, ...]) _RESPONSE_BLOCKS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) -_PENDING_WEB_SEARCH_RESULT: Final = MappingProxyType({"content": ()}) +_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _REQUEST_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _NATIVE_OUTPUT_STRING_FORMATS: Final = frozenset( ("date-time", "time", "date", "duration", "email", "hostname", "uri", "ipv4", "ipv6", "uuid") @@ -346,6 +348,51 @@ def _validated_defer_loading(value: object) -> bool: raise ValueError("defer_loading must be a boolean") from exc +def _hosted_web_tool_use(block: object) -> tuple[str, str, object] | None: + match block: + case { + "type": "server_tool_use", + "name": str() as name, + "id": str() as call_id, + "input": object() as tool_input, + } if name in HOSTED_WEB_TOOL_NAMES: + return call_id, name, tool_input + case _: + return None + + +def _server_tool_use_id(block: object) -> str | None: + match block: + case {"type": "server_tool_use", "id": str() as call_id}: + return call_id + case _: + return None + + +def _dict_or_empty(value: object) -> object: + return validated_object_mapping(value) or {} + + +_RESPONSE_BODY_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter( + Annotated[Mapping[str, object], BeforeValidator(_dict_or_empty)] +) + + +def _server_tool_use_ids(completion_response: Mapping[str, object]) -> frozenset[str]: + blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(completion_response.get("content") or ()) + return frozenset(filter(None, map(_server_tool_use_id, blocks))) + + +def _hosted_web_call(call_id: str, tool_name: str, tool_input: object, result: object) -> ResponseFunctionWebSearch: + return build_web_search_call( + tool_id=call_id, + tool_input=tool_input, + result=result, + status=None if result is not None else "in_progress", + tool_name=tool_name, + ) + + class _AnthropicUsageIteration(TypedDict, total=False): """One entry of the ``usage.iterations`` array on an Anthropic response.""" @@ -3079,25 +3126,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): self, web_search_results: Sequence[object], completion_response: Mapping[str, object], - ) -> list[object]: - blocks: Final = _RESPONSE_BLOCKS.validate_python(completion_response.get("content") or ()) - results_by_id: Final = { - tool_use_id: result + ) -> list[ResponseFunctionWebSearch]: + blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(completion_response.get("content") or ()) + results_by_call: Final = { + (tool_use_id, result.get("type")): result for result in _RESPONSE_BLOCKS.validate_python(web_search_results) - if result.get("type") == "web_search_tool_result" - and isinstance((tool_use_id := result.get("tool_use_id")), str) + if isinstance((tool_use_id := result.get("tool_use_id")), str) } return [ - build_web_search_call( - tool_id=call_id, - tool_input=block.get("input"), - result=results_by_id.get(call_id, _PENDING_WEB_SEARCH_RESULT), - status=None if call_id in results_by_id else "in_progress", - ) - for block in blocks - if block.get("type") == "server_tool_use" - and block.get("name") == "web_search" - and isinstance((call_id := block.get("id")), str) + _hosted_web_call(call_id, tool_name, tool_input, results_by_call.get((call_id, f"{tool_name}_tool_result"))) + for call_id, tool_name, tool_input in filter(None, map(_hosted_web_tool_use, blocks)) ] def _build_provider_specific_fields( @@ -3109,6 +3147,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): tool_results: Sequence[_AnthropicToolResultBlock] | None, compaction_blocks: Sequence[object] | None, tool_calls: list[ChatCompletionToolCallChunk], + server_tool_use_ids: frozenset[str], ) -> dict[str, object]: provider_specific_fields: Final[dict[str, object]] = { "citations": citations, @@ -3148,6 +3187,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if compaction_blocks is not None: provider_specific_fields["compaction_blocks"] = compaction_blocks + server_tool_calls: Final = [tc for tc in tool_calls if tc["id"] in server_tool_use_ids] + if server_tool_calls: + provider_specific_fields["server_tool_calls"] = server_tool_calls + stop_reason_fields: Final = stop_reason_provider_specific_fields( STOP_REASON_FIELDS.validate_python(completion_response) ) @@ -3207,6 +3250,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ): text_content = prefix_prompt + text_content + server_tool_use_ids: Final = _server_tool_use_ids(_RESPONSE_BODY_ADAPTER.validate_python(completion_response)) + client_tool_calls: Final = [tc for tc in tool_calls if tc["id"] not in server_tool_use_ids] provider_specific_fields: Final = self._build_provider_specific_fields( completion_response, citations, @@ -3215,11 +3260,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): tool_results, compaction_blocks, tool_calls, + server_tool_use_ids, ) json_mode_message, tool_calls_for_message, json_extra_content = self._resolve_json_mode_non_streaming( json_mode=json_mode, - tool_calls=tool_calls, + tool_calls=client_tool_calls, response_format_schema=response_format_schema, ) merged_text = text_content or "" @@ -3258,10 +3304,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): model_response.choices[0].message = _message model_response.hidden_params["original_response"] = completion_response["content"] - model_response.choices[0].finish_reason = cast( - OpenAIChatCompletionFinishReason, - map_finish_reason(completion_response["stop_reason"]), - ) + model_response.choices[0].finish_reason = map_finish_reason(completion_response["stop_reason"]) usage: Final = self.calculate_usage( usage_object=completion_response["usage"], diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 275cc5cbe5c..5d4957fa166 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -6,7 +6,7 @@ import copy import json import os import re -from collections.abc import Container, Mapping, MutableMapping, Sequence +from collections.abc import Callable, Container, Mapping, MutableMapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from itertools import chain @@ -70,10 +70,26 @@ from litellm.types.llms.anthropic import ( AnthropicThinkingParam, ) from litellm.types.llms.base import LiteLLMBaseModel -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionAssistantToolCall, + ChatCompletionToolCallFunctionChunk, +) from litellm.types.proxy.auth.special_headers import SpecialHeaders from litellm.types.proxy.model_listing import ModelInfoResponse -from litellm.types.utils import LlmProviders +from litellm.types.responses.main import ( + ResponseFunctionWebSearch, + web_search_call_tool_id, + with_hosted_web_tool_arguments, +) +from litellm.types.utils import ( + ChatCompletionMessageCustomToolCall, + ChatCompletionMessageToolCall, + Delta, + LlmProviders, + Message, + chat_completion_tool_call_from_dict, +) _MessageT = TypeVar("_MessageT") @@ -126,6 +142,11 @@ _DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)") _CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:" OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_SERVER_TOOL_CALLS_ADAPTER: Final = TypeAdapter(tuple[ChatCompletionAssistantToolCall, ...]) +_SERVER_TOOL_CALL_ADAPTER: Final = TypeAdapter(ChatCompletionAssistantToolCall) +_SERVER_CALL_FIELD_KEYS: Final = ("server_tool_calls", "web_search_calls") +_WEB_SEARCH_CALL_ADAPTER: Final = TypeAdapter(ResponseFunctionWebSearch) +_TOOL_CALL_FIELDS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _DECODED_JSON: Final = TypeAdapter(object) _JSON_VALUE: Final = TypeAdapter(JsonValue) @@ -2290,6 +2311,261 @@ def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None: return None +def anthropic_server_tool_call_params(message: object) -> tuple[ChatCompletionAssistantToolCall, ...]: + fields: Final = validated_object_mapping(message_field(message, "provider_specific_fields")) + try: + return _SERVER_TOOL_CALLS_ADAPTER.validate_python(fields.get("server_tool_calls", ()) if fields else ()) + except ValidationError: + return () + + +def anthropic_server_tool_calls(message: object) -> tuple[ChatCompletionMessageToolCall, ...]: + return tuple(_message_tool_call_from_params(call) for call in anthropic_server_tool_call_params(message)) + + +def tool_calls_with_server_tool_calls( + message: object, +) -> tuple[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall, ...]: + client_tool_calls: Final = tuple( + call + for tool_call in parts_of(message_field(message, "tool_calls")) + if (call := _client_message_tool_call(tool_call)) is not None + ) + return (*anthropic_server_tool_calls(message), *client_tool_calls) + + +def anthropic_server_call_fields(holder: object) -> Mapping[str, object]: + fields: Final = validated_object_mapping(message_field(holder, "provider_specific_fields")) or {} + return MappingProxyType({key: fields[key] for key in _SERVER_CALL_FIELD_KEYS if fields.get(key)}) + + +def message_dict_with_server_tool_calls(message: object) -> object: + mapping: Final = validated_object_mapping(message) + server_tool_calls: Final = anthropic_server_tool_call_params(message) + if mapping is None or not server_tool_calls: + return message + return {**mapping, "tool_calls": [*server_tool_calls, *parts_of(mapping.get("tool_calls"))]} + + +def _server_tool_call_param(tool_call: object) -> ChatCompletionAssistantToolCall | None: + try: + call: Final = _SERVER_TOOL_CALL_ADAPTER.validate_python(tool_call) + except ValidationError: + return None + return call if _server_tool_call_id(call).startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX) else None + + +def message_dict_with_server_tool_calls_split_out(message: object) -> object: + mapping: Final = validated_object_mapping(message) + tool_calls: Final = parts_of(mapping.get("tool_calls")) if mapping is not None else () + server_tool_calls: Final = tuple( + call for tool_call in tool_calls if (call := _server_tool_call_param(tool_call)) is not None + ) + if mapping is None or not server_tool_calls: + return message + fields: Final = validated_object_mapping(mapping.get("provider_specific_fields")) or {} + changed: Final = changed_server_tool_call_arguments(anthropic_server_tool_call_params(message), server_tool_calls) + client_tool_calls: Final = [tool_call for tool_call in tool_calls if _server_tool_call_param(tool_call) is None] + return { + **mapping, + "tool_calls": client_tool_calls or None, + "provider_specific_fields": { + **fields, + "server_tool_calls": list(server_tool_calls), + **( + { + "web_search_calls": _rewritten_web_search_calls( + fields["web_search_calls"], changed, _with_hosted_web_call_arguments + ) + } + if changed and fields.get("web_search_calls") is not None + else {} + ), + }, + } + + +def with_server_tool_call_arguments( + call: ChatCompletionAssistantToolCall, arguments: str +) -> ChatCompletionAssistantToolCall: + return ChatCompletionAssistantToolCall( + id=call["id"], + type=call["type"], + function=ChatCompletionToolCallFunctionChunk(name=call["function"].get("name"), arguments=arguments), + ) + + +def write_anthropic_server_tool_calls( + holder: Message | Delta, rewritten: tuple[ChatCompletionAssistantToolCall, ...] +) -> None: + original: Final = anthropic_server_tool_call_params(holder) + if rewritten == original: + return + holder.provider_specific_fields = { # rebind-ok: a guardrail rewrite must land on the message it scanned + **(holder.provider_specific_fields or {}), + "server_tool_calls": list(rewritten), + } + write_anthropic_hosted_web_call_arguments(holder, changed_server_tool_call_arguments(original, rewritten)) + + +def write_request_server_tool_calls(message: object, rewritten: tuple[ChatCompletionAssistantToolCall, ...]) -> None: + original: Final = anthropic_server_tool_call_params(message) + if rewritten == original: + return + write_request_hosted_web_call_arguments(message, changed_server_tool_call_arguments(original, rewritten)) + fields: Final = validated_object_mapping(message_field(message, "provider_specific_fields")) or {} + if not isinstance(message, dict): + return + message["provider_specific_fields"] = { # rebind-ok: a guardrail rewrite must land on the request it scanned + **fields, + "server_tool_calls": list(rewritten), + } + + +def write_request_hosted_web_call_arguments(message: object, arguments_by_tool_id: Mapping[str, str]) -> None: + fields: Final = validated_object_mapping(message_field(message, "provider_specific_fields")) + if not arguments_by_tool_id or not isinstance(message, dict) or fields is None: + return + if fields.get("web_search_calls") is None: + return + message["provider_specific_fields"] = { # rebind-ok: a rewritten call's hosted copy must not keep its old input + **fields, + "web_search_calls": _rewritten_web_search_calls( + fields["web_search_calls"], arguments_by_tool_id, _jsonable_hosted_web_call + ), + } + + +def _server_tool_call_id(call: ChatCompletionAssistantToolCall) -> str: + return call["id"] or "" + + +def _server_tool_call_arguments(call: ChatCompletionAssistantToolCall) -> str: + return call["function"].get("arguments") or "" + + +def changed_server_tool_call_arguments( + original: Sequence[ChatCompletionAssistantToolCall], rewritten: Sequence[ChatCompletionAssistantToolCall] +) -> Mapping[str, str]: + original_arguments: Final = MappingProxyType( + {_server_tool_call_id(call): _server_tool_call_arguments(call) for call in original} + ) + return MappingProxyType( + { + _server_tool_call_id(call): _server_tool_call_arguments(call) + for call in rewritten + if _server_tool_call_arguments(call) != original_arguments.get(_server_tool_call_id(call)) + } + ) + + +def write_anthropic_hosted_web_call_arguments(holder: Message | Delta, arguments_by_tool_id: Mapping[str, str]) -> None: + fields: Final = holder.provider_specific_fields + if not arguments_by_tool_id or not fields or fields.get("web_search_calls") is None: + return + holder.provider_specific_fields = { # rebind-ok: the hosted copy of a rewritten call must not keep its old input + **fields, + "web_search_calls": _rewritten_web_search_calls( + fields["web_search_calls"], arguments_by_tool_id, _with_hosted_web_call_arguments + ), + } + + +def _rewritten_web_search_calls( + calls: object, + arguments_by_tool_id: Mapping[str, str], + rewrite: Callable[[object, Mapping[str, str]], object], +) -> object: + mapping: Final = validated_object_mapping(calls) + if mapping is not None: + return {key: rewrite(call, arguments_by_tool_id) for key, call in mapping.items()} + if not isinstance(calls, list): + return calls + return [rewrite(call, arguments_by_tool_id) for call in parts_of(calls)] + + +def _jsonable_hosted_web_call(item: object, arguments_by_tool_id: Mapping[str, str]) -> object: + call: Final = _with_hosted_web_call_arguments(item, arguments_by_tool_id) + if not isinstance(call, ResponseFunctionWebSearch): + return call + return _WEB_SEARCH_CALL_ADAPTER.dump_python(call, mode="json", exclude_none=True) + + +def _with_hosted_web_call_arguments(item: object, arguments_by_tool_id: Mapping[str, str]) -> object: + try: + call: Final = _WEB_SEARCH_CALL_ADAPTER.validate_python(item) + except ValidationError: + return item + arguments: Final = arguments_by_tool_id.get(web_search_call_tool_id(call.id) or "") + if arguments is None: + return item + return call.model_copy(update={"action": with_hosted_web_tool_arguments(call.action, arguments)}) + + +def without_anthropic_server_tool_calls(holder: Message, tool_ids: frozenset[str]) -> None: + fields: Final = holder.provider_specific_fields or {} + holder.provider_specific_fields = { # rebind-ok: a denied call must leave every copy of the message it was on + **fields, + "server_tool_calls": [ + call for call in anthropic_server_tool_call_params(holder) if _server_tool_call_id(call) not in tool_ids + ], + **( + {"web_search_calls": _web_search_calls_without(fields["web_search_calls"], tool_ids)} + if fields.get("web_search_calls") is not None + else {} + ), + } + + +def _web_search_calls_without(calls: object, tool_ids: frozenset[str]) -> object: + mapping: Final = validated_object_mapping(calls) + if mapping is not None: + return {key: call for key, call in mapping.items() if _hosted_web_call_tool_id(call) not in tool_ids} + items: Final = validated_object_list(calls) + if items is None: + return calls + return [call for call in items if _hosted_web_call_tool_id(call) not in tool_ids] + + +def _hosted_web_call_tool_id(item: object) -> str | None: + try: + return web_search_call_tool_id(_WEB_SEARCH_CALL_ADAPTER.validate_python(item).id) + except ValidationError: + return None + + +def write_anthropic_server_tool_call_arguments(holder: Message | Delta, index: int, arguments: str) -> None: + write_anthropic_server_tool_calls( + holder, + tuple( + with_server_tool_call_arguments(call, arguments) if call_index == index else call + for call_index, call in enumerate(anthropic_server_tool_call_params(holder)) + ), + ) + + +def _message_tool_call_from_params(call: ChatCompletionAssistantToolCall) -> ChatCompletionMessageToolCall: + return ChatCompletionMessageToolCall(id=call["id"], type=call["type"], function=dict(call["function"])) + + +def _client_message_tool_call( + tool_call: object, +) -> ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall | None: + match tool_call: + case ChatCompletionMessageToolCall() | ChatCompletionMessageCustomToolCall(): + return tool_call + case _: + try: + fields: Final = _TOOL_CALL_FIELDS_ADAPTER.validate_python(tool_call) + return chat_completion_tool_call_from_dict( + MappingProxyType({key: value for key, value in fields.items() if key != "custom"}) + if fields.get("type") == "function" + else fields + ) + except (TypeError, ValidationError): + return None + + def tool_call_is_rebuilt_as_server_tool_use( tool_call_id: object, provider_specific_fields: object, answered_tool_ids: Container[str] = frozenset() ) -> bool: diff --git a/litellm/llms/anthropic/messages/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index 907d4021fd3..43eaa37aba7 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/litellm/llms/anthropic/messages/adapters/transformation.py @@ -46,6 +46,25 @@ def _thought_signature(provider_specific_fields: object) -> str | None: return signature if isinstance(signature, str) else None +def _object_sequence(value: object) -> Sequence[object] | None: + return value if isinstance(value, (list, tuple)) else None + + +def _server_tool_blocks(message: object) -> tuple[Mapping[str, object], ...]: + server_tool_calls: Final = anthropic_server_tool_call_params(message) + if not server_tool_calls: + return () + fields: Final = _as_string_mapping(_optional_attr(message, "provider_specific_fields")) or {} + return tuple( + convert_to_anthropic_tool_invoke( + server_tool_calls, + web_search_results=_object_sequence(fields.get("web_search_results")), + tool_results=_object_sequence(fields.get("tool_results")), + server_tool_call_ids=frozenset(tool_call["id"] for tool_call in server_tool_calls), + ) + ) + + def _compaction_blocks(provider_specific_fields: object) -> tuple[Mapping[str, object], ...]: fields: Final = _as_string_mapping(provider_specific_fields) raw_blocks: Final = fields.get("compaction_blocks") if fields is not None else None @@ -127,6 +146,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.litellm_core_utils.prompt_templates.factory import ( THOUGHT_SIGNATURE_SEPARATOR, + convert_to_anthropic_tool_invoke, ) from litellm.litellm_core_utils.prompt_templates.prompt_cache_breakpoint_utils import ( with_prompt_cache_breakpoint, @@ -145,6 +165,7 @@ from litellm.litellm_core_utils.reasoning_effort_utils import ( reasoning_effort_from_thinking_silently, ) from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, eager_input_streaming_flag, is_empty_unsigned_thinking_block, normalize_anthropic_tool_use_id, @@ -1508,6 +1529,7 @@ class LiteLLMAnthropicMessagesAdapter: ).model_dump() ) + new_content.extend(dict(block) for block in _server_tool_blocks(choice.message)) # Handle text content if choice.message.content is not None and (choice.message.content != "" or not compaction_blocks): new_content.append( diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index fb886e42780..900a9c0712e 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -4,11 +4,14 @@ 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 import claude_hosted_web_tool_call 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 +138,36 @@ 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 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 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/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 5c5204ef2f4..c5488bc8a98 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -21,6 +21,7 @@ from collections.abc import Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast +from pydantic import TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -29,6 +30,13 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i replayed_tool_search_tool_names, ) from litellm.litellm_core_utils.streaming_chunk_builder_utils import stream_chunk_builder +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, + changed_server_tool_call_arguments, + write_anthropic_hosted_web_call_arguments, + write_anthropic_server_tool_calls, + write_request_server_tool_calls, +) from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -50,7 +58,12 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_items, unappliable_request_rewrite, ) -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionAssistantToolCall, + ChatCompletionToolCallFunctionChunk, + ChatCompletionToolParam, +) from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( coerce_stream_holdback_value, ) @@ -139,14 +152,27 @@ class OpenAIChatCompletionsHandler(BaseTranslation): scan_only_tool_results=scan_only_tool_results, ) + structured_messages: Final = self.get_structured_messages(data) + server_tool_call_tasks: Final = tuple( + _input_server_tool_call_tasks( + structured_messages or [], + skip_system_message=skip_system, + skip_tool_message=skip_tool, + scan_only_tool_results=scan_only_tool_results, + ) + ) + tool_calls_with_server_calls: Final = [ + *tool_calls_to_check, + *(_server_tool_call_input(call) for _, call in server_tool_call_tasks), + ] + # Step 2: Apply guardrail to all texts and tool calls in batch - if texts_to_check or tool_calls_to_check: + if texts_to_check or tool_calls_with_server_calls: inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: inputs["images"] = images_to_check - if tool_calls_to_check: - inputs["tool_calls"] = tool_calls_to_check - structured_messages: Final = self.get_structured_messages(data) + if tool_calls_with_server_calls: + inputs["tool_calls"] = tool_calls_with_server_calls scoped_message_indices: Final = scoped_structured_message_indices( structured_messages or [], scan_only_tool_results=scan_only_tool_results, @@ -218,6 +244,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tool_calls=guardrailed_tool_calls, task_mappings=tool_call_task_mappings, ) + if len(guardrailed_tool_calls) == len(tool_calls_with_server_calls): + self._apply_guardrail_responses_to_input_server_tool_calls( + messages=structured_messages or [], + guardrailed_tool_calls=guardrailed_tool_calls[len(tool_calls_to_check) :], + task_mappings=tuple(mapping for mapping, _ in server_tool_call_tasks), + ) elif ( not images_to_check @@ -386,6 +418,22 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Replace the tool call with the guardrailed version message_tool_calls[tool_call_idx] = guardrailed_tool_call + @staticmethod + def _apply_guardrail_responses_to_input_server_tool_calls( + messages: Sequence[object], + guardrailed_tool_calls: Sequence[object], + task_mappings: Sequence[tuple[int, int]], + ) -> None: + guardrailed_by_task: Final = MappingProxyType(dict(zip(task_mappings, guardrailed_tool_calls))) + for msg_idx in sorted(frozenset(msg_idx for msg_idx, _ in task_mappings)): + write_request_server_tool_calls( + messages[msg_idx], + tuple( + _rewritten_server_tool_call(server_tool_call, guardrailed_by_task.get((msg_idx, server_idx))) + for server_idx, server_tool_call in enumerate(anthropic_server_tool_call_params(messages[msg_idx])) + ), + ) + async def process_output_response( self, response: "ModelResponse", @@ -436,9 +484,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): text_task_mappings=text_task_mappings, tool_call_task_mappings=tool_call_task_mappings, ) + server_tool_call_tasks: Final = tuple(_server_tool_call_tasks(response.choices)) + tool_calls_with_server_calls: Final = [ + *tool_calls_to_check, + *(_server_tool_call_input(call) for _, call in server_tool_call_tasks), + ] # Step 2: Apply guardrail to all texts and tool calls in batch - if texts_to_check or tool_calls_to_check: + if texts_to_check or tool_calls_with_server_calls: # Use the real request_data if provided (proxy path), otherwise # create a standalone dict (SDK / direct-call path). if request_data is None: @@ -452,8 +505,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: inputs["images"] = images_to_check - if tool_calls_to_check: - inputs["tool_calls"] = tool_calls_to_check + if tool_calls_with_server_calls: + inputs["tool_calls"] = tool_calls_with_server_calls # Include model information from the response if available if hasattr(response, "model") and response.model: inputs["model"] = response.model @@ -469,8 +522,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") guardrailed_tool_calls: Final[list[dict[str, object]]] = ( cast(list[dict[str, object]], returned_tool_calls) - if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check) - else tool_calls_to_check + if isinstance(returned_tool_calls, list) + and len(returned_tool_calls) == len(tool_calls_with_server_calls) + else tool_calls_with_server_calls ) # Step 3: Map guardrail responses back to original response structure @@ -488,6 +542,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tool_calls=guardrailed_tool_calls, task_mappings=tool_call_task_mappings, ) + self._apply_guardrail_responses_to_output_server_tool_calls( + response=response, + guardrailed_tool_calls=guardrailed_tool_calls[len(tool_calls_to_check) :], + task_mappings=tuple(mapping for mapping, _ in server_tool_call_tasks), + ) verbose_proxy_logger.debug("OpenAI Chat Completions: Processed output response: %s", response) @@ -681,6 +740,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): pre_guardrail_tool_calls=pre_guardrail_tool_calls, guardrail_name=guardrail_to_apply.guardrail_name or "unknown", ) + self._write_ended_stream_server_tool_call_rewrites( + responses_so_far=responses_so_far, + guardrailed_response=model_response, + ) @staticmethod def _rebuild_ended_stream_per_choice( @@ -934,6 +997,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if choice.message.tool_calls and isinstance(choice.message.tool_calls, list): if len(choice.message.tool_calls) > 0: return True + if anthropic_server_tool_call_params(choice.message): + return True elif isinstance(response, ModelResponseStream): for streaming_choice in response.choices: if isinstance(streaming_choice, litellm.StreamingChoices): @@ -944,6 +1009,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if streaming_choice.delta.tool_calls and isinstance(streaming_choice.delta.tool_calls, list): if len(streaming_choice.delta.tool_calls) > 0: return True + if anthropic_server_tool_call_params(streaming_choice.delta): + return True return False def _extract_output_text_images_and_tool_calls( @@ -1094,6 +1161,22 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "name" in func_dict: existing_tool_call.function.name = func_dict["name"] + @staticmethod + def _apply_guardrail_responses_to_output_server_tool_calls( + response: "ModelResponse", + guardrailed_tool_calls: Sequence[object], + task_mappings: Sequence[tuple[int, int]], + ) -> None: + guardrailed_by_task: Final = MappingProxyType(dict(zip(task_mappings, guardrailed_tool_calls))) + for choice_idx, choice in enumerate(response.choices): + write_anthropic_server_tool_calls( + choice.message, + tuple( + _rewritten_server_tool_call(server_tool_call, guardrailed_by_task.get((choice_idx, server_idx))) + for server_idx, server_tool_call in enumerate(anthropic_server_tool_call_params(choice.message)) + ), + ) + @staticmethod def _string_choice_contents(response: "ModelResponse") -> tuple[str | None, ...]: return tuple( @@ -1206,6 +1289,42 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for fragment in tail: fragment.function.arguments = "" + @staticmethod + def _write_ended_stream_server_tool_call_rewrites( + responses_so_far: Sequence["ModelResponseStream"], + guardrailed_response: "ModelResponse", + ) -> None: + guardrailed_by_id: Final = MappingProxyType( + { + (guardrailed_response.choices[choice_position].index, server_tool_call["id"]): server_tool_call + for (choice_position, _), server_tool_call in _server_tool_call_tasks(guardrailed_response.choices) + } + ) + if not guardrailed_by_id: + return + original_by_id: Final = MappingProxyType(dict(_streamed_server_tool_calls(responses_so_far))) + changed_by_choice: Final = MappingProxyType( + { + choice_index: changed_server_tool_call_arguments( + _choice_server_tool_calls(original_by_id, choice_index), + _choice_server_tool_calls(guardrailed_by_id, choice_index), + ) + for choice_index in frozenset(index for index, _ in guardrailed_by_id) + } + ) + for response in responses_so_far: + for choice in response.choices: + write_anthropic_server_tool_calls( + choice.delta, + tuple( + guardrailed_by_id.get((choice.index, server_tool_call["id"]), server_tool_call) + for server_tool_call in anthropic_server_tool_call_params(choice.delta) + ), + ) + write_anthropic_hosted_web_call_arguments( + choice.delta, changed_by_choice.get(choice.index, MappingProxyType({})) + ) + async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], @@ -1395,7 +1514,76 @@ def _stream_chunk_choices(item: object) -> Sequence[object]: def _streamed_delta_tool_calls(delta: object) -> tuple[object, ...]: function_call: Final = stream_item_field(delta, "function_call") legacy: Final = () if function_call is None else (function_call,) - return stream_item_items(delta, "tool_calls") + legacy + return anthropic_server_tool_call_params(delta) + stream_item_items(delta, "tool_calls") + legacy + + +def _input_server_tool_call_tasks( + messages: Sequence[Mapping[str, object]], + *, + skip_system_message: bool, + skip_tool_message: bool, + scan_only_tool_results: bool, +) -> Iterator[tuple[tuple[int, int], ChatCompletionAssistantToolCall]]: + for msg_idx, message in enumerate(messages): + if role_out_of_guardrail_scope( + str(message.get("role") or "").lower(), + skip_system_message=skip_system_message, + skip_tool_message=skip_tool_message, + scan_only_tool_results=scan_only_tool_results, + ): + continue + for server_idx, server_tool_call in enumerate(anthropic_server_tool_call_params(message)): + yield (msg_idx, server_idx), server_tool_call + + +def _server_tool_call_tasks( + choices: Sequence[object], +) -> Iterator[tuple[tuple[int, int], ChatCompletionAssistantToolCall]]: + for choice_idx, choice in enumerate(choices): + if not isinstance(choice, Choices): + continue + for server_idx, server_tool_call in enumerate(anthropic_server_tool_call_params(choice.message)): + yield (choice_idx, server_idx), server_tool_call + + +def _streamed_server_tool_calls( + responses: Sequence["ModelResponseStream"], +) -> Iterator[tuple[tuple[int, str | None], ChatCompletionAssistantToolCall]]: + for response in responses: + for choice in response.choices: + for server_tool_call in anthropic_server_tool_call_params(choice.delta): + yield (choice.index, server_tool_call["id"]), server_tool_call + + +def _choice_server_tool_calls( + calls_by_id: Mapping[tuple[int, str | None], ChatCompletionAssistantToolCall], choice_index: int +) -> tuple[ChatCompletionAssistantToolCall, ...]: + return tuple(call for (index, _), call in calls_by_id.items() if index == choice_index) + + +def _server_tool_call_input(server_tool_call: ChatCompletionAssistantToolCall) -> dict[str, object]: + return { + "id": server_tool_call["id"], + "type": server_tool_call["type"], + "function": dict(server_tool_call["function"]), + } + + +_TOOL_CALL_FUNCTION_ADAPTER: Final = TypeAdapter(ChatCompletionToolCallFunctionChunk) + + +def _rewritten_server_tool_call( + server_tool_call: ChatCompletionAssistantToolCall, guardrailed: object +) -> ChatCompletionAssistantToolCall: + try: + function: Final = _TOOL_CALL_FUNCTION_ADAPTER.validate_python(stream_item_field(guardrailed, "function")) + except ValidationError: + return server_tool_call + return ChatCompletionAssistantToolCall( + id=server_tool_call["id"], + type=server_tool_call["type"], + function={**server_tool_call["function"], **function}, + ) def _released_choices(responses_so_far: Sequence[object]) -> Iterator[object]: diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index a18cbe13925..0ec8497b0c6 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -71,6 +71,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, @@ -86,12 +88,14 @@ from litellm.responses.litellm_completion_transformation.tool_search import is_s from litellm.responses.litellm_completion_transformation.transformation import ( TOOL_CALL_OUTPUT_ITEM_TYPES, LiteLLMCompletionResponsesConfig, + claude_web_search_call_tool_call, tool_output_to_chat_content, ) from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( AllMessageValues, BaseLiteLLMOpenAIResponseObject, + ChatCompletionAssistantToolCall, ChatCompletionToolCallChunk, ChatCompletionToolParam, ContentPartAddedEvent, @@ -266,6 +270,8 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset( _TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call"}) _TOOL_SEARCH_CALL_TYPE: Final = "tool_search_call" +_WEB_SEARCH_CALL_TYPE: Final = "web_search_call" +_ITEM_ID_KEYED_TOOL_CALL_TYPES: Final = frozenset({_TOOL_SEARCH_CALL_TYPE, _WEB_SEARCH_CALL_TYPE}) _OUTPUT_TOOL_CALL_ITEM_TYPES: Final = _TOOL_CALL_ITEM_TYPES | {_TOOL_SEARCH_CALL_TYPE} _SCANNED_TOOL_SEARCH_NAME: Final = "tool_search" _TOOL_SEARCH_ARGUMENTS_ADAPTER: Final = TypeAdapter(dict[str, object]) @@ -335,8 +341,27 @@ def _is_tool_call_item(item: object) -> bool: return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES +def _claude_web_search_tool_call(item: object) -> ChatCompletionAssistantToolCall | None: + search: Final = web_search_call(item) + return None if search is None else claude_web_search_call_tool_call(search) + + +def _claude_web_tool_call(item: object, index: int) -> ChatCompletionToolCallChunk | None: + tool_call: Final = _claude_web_search_tool_call(item) + if tool_call is None: + return None + return ChatCompletionToolCallChunk(id=tool_call["id"], type="function", function=tool_call["function"], index=index) + + +def _is_output_tool_call_item(item: object) -> bool: + return ( + stream_item_field(item, "type") in _OUTPUT_TOOL_CALL_ITEM_TYPES + or _claude_web_search_tool_call(item) is not None + ) + + def _tool_call_output_item_mapping(item: object) -> Mapping[str, object] | None: - if stream_item_field(item, "type") not in _OUTPUT_TOOL_CALL_ITEM_TYPES: + if not _is_output_tool_call_item(item): return None if isinstance(item, Mapping): return cast("Mapping[str, object]", item) # cast-ok: output items are str-keyed JSON objects @@ -348,7 +373,7 @@ def _tool_call_item_key(item: object) -> str | None: if isinstance(call_id, str) and call_id: return call_id item_id: Final = stream_item_field(item, "id") - if stream_item_field(item, "type") == _TOOL_SEARCH_CALL_TYPE and isinstance(item_id, str) and item_id: + if stream_item_field(item, "type") in _ITEM_ID_KEYED_TOOL_CALL_TYPES and isinstance(item_id, str) and item_id: return item_id return None @@ -368,6 +393,17 @@ def _scanned_tool_call_item(item: Mapping[str, object]) -> Mapping[str, object]: return {**item, "name": _SCANNED_TOOL_SEARCH_NAME, "arguments": json.dumps(item.get("arguments") or {})} +def _scanned_tool_call(item: Mapping[str, object], index: int) -> ChatCompletionToolCallChunk: + web_tool_call: Final = _claude_web_tool_call(item, index) + if web_tool_call is not None: + return web_tool_call + return tool_call_dict_from_output_item(_scanned_tool_call_item(item), index) + + +def _rewritten_web_search_call(item: object, payload: str | None) -> object: + return item if payload is None else with_web_search_call_arguments(item, payload) + + def _tool_search_rewrite(payload: str) -> Mapping[str, object]: try: return {"arguments": _TOOL_SEARCH_ARGUMENTS_ADAPTER.validate_json(payload)} @@ -1322,17 +1358,22 @@ class OpenAIResponsesHandler(BaseTranslation): # delivered, so a truncated buffer still reports text undeliverable. # # ------------------------------------------------------------------ # if final_chunk.get("type") == "response.output_item.done": - finished_summary: Final = _output_compaction_summary(stream_item_field(final_chunk, "item")) + finished_item: Final = stream_item_field(final_chunk, "item") + finished_summary: Final = _output_compaction_summary(finished_item) model_response_stream: Final = ( OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(final_chunk) ) - tool_calls: Final = model_response_stream.choices[0].delta.tool_calls + finished_web_tool_call: Final = _claude_web_tool_call(finished_item, 0) + tool_calls: Final = [ + *cast(list[ChatCompletionToolCallChunk], model_response_stream.choices[0].delta.tool_calls or []), + *(() if finished_web_tool_call is None else (finished_web_tool_call,)), + ] if tool_calls or finished_summary is not None: inputs = GenericGuardrailAPIInputs() if finished_summary is not None: inputs["texts"] = [finished_summary] if tool_calls: - inputs["tool_calls"] = cast(list[ChatCompletionToolCallChunk], tool_calls) + inputs["tool_calls"] = tool_calls if hasattr(model_response_stream, "model") and model_response_stream.model: inputs["model"] = model_response_stream.model await guardrail_to_apply.apply_guardrail( @@ -1629,7 +1670,7 @@ class OpenAIResponsesHandler(BaseTranslation): if event_type not in _OUTPUT_ITEM_EVENT_TYPES: return None item: Final = stream_item_field(event, "item") - return _tool_call_item_key(item) if stream_item_field(item, "type") in _OUTPUT_TOOL_CALL_ITEM_TYPES else None + return _tool_call_item_key(item) if _is_output_tool_call_item(item) else None @staticmethod def _write_tool_call_item(item: object, name: str | None, payload: str | None) -> object: @@ -1643,6 +1684,8 @@ class OpenAIResponsesHandler(BaseTranslation): return item if item_type == _TOOL_SEARCH_CALL_TYPE: return _rewritten_tool_search_call(item, payload) + if item_type == _WEB_SEARCH_CALL_TYPE: + return _rewritten_web_search_call(item, payload) if name is not None: OpenAIResponsesHandler._write_event_field(item, "name", name) if payload is not None and isinstance(item_type, str) and item_type in _TOOL_CALL_PAYLOAD_FIELDS: @@ -1698,7 +1741,7 @@ class OpenAIResponsesHandler(BaseTranslation): stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES or ( stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES - and stream_item_field(stream_item_field(event, "item"), "type") in _OUTPUT_TOOL_CALL_ITEM_TYPES + and _is_output_tool_call_item(stream_item_field(event, "item")) ) for event in responses_so_far ) @@ -1706,9 +1749,7 @@ class OpenAIResponsesHandler(BaseTranslation): @staticmethod def _completed_response_scan_key(response: object) -> StreamingScanKey: output_items: Final = stream_item_items(response, "output") - message_items: Final = tuple( - item for item in output_items if stream_item_field(item, "type") not in _OUTPUT_TOOL_CALL_ITEM_TYPES - ) + message_items: Final = tuple(item for item in output_items if not _is_output_tool_call_item(item)) message_parts: Final = chain.from_iterable(stream_item_items(item, "content") for item in message_items) part_texts: Final = tuple( text for part in message_parts if isinstance(text := stream_item_field(part, "text"), str) and text @@ -1716,11 +1757,7 @@ class OpenAIResponsesHandler(BaseTranslation): summaries: Final = tuple(summary for item in output_items if (summary := _output_compaction_summary(item))) return StreamingScanKey( texts=(*part_texts, *summaries), - tool_calls=tuple( - stream_item_fingerprint(item) - for item in output_items - if stream_item_field(item, "type") in _OUTPUT_TOOL_CALL_ITEM_TYPES - ), + tool_calls=tuple(stream_item_fingerprint(item) for item in output_items if _is_output_tool_call_item(item)), stream_ended=True, ) @@ -1833,9 +1870,7 @@ class OpenAIResponsesHandler(BaseTranslation): tool_call_item: Final = _tool_call_output_item_mapping(output_item) if tool_call_item is not None: if tool_calls_to_check is not None: - tool_calls_to_check.append( - tool_call_dict_from_output_item(_scanned_tool_call_item(tool_call_item), output_idx) - ) + tool_calls_to_check.append(_scanned_tool_call(tool_call_item, output_idx)) return summary: Final = _output_compaction_summary(output_item) diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 591ad81a4e5..92181d546eb 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -63,11 +63,26 @@ def response_tool_call_names(completion_response: object) -> tuple[str, ...]: ) return tuple( stripped - for name in chain(called_names, _auto_executed_mcp_tool_names(completion_response)) + for name in chain( + _server_tool_call_names(completion_response), + called_names, + _auto_executed_mcp_tool_names(completion_response), + ) if isinstance(name, str) and (stripped := name.strip()) ) +def _server_tool_call_names(completion_response: object) -> tuple[str | None, ...]: + from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of + from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params + + server_tool_calls: Final = chain.from_iterable( + anthropic_server_tool_call_params(message_field(choice, "message")) + for choice in parts_of(message_field(completion_response, "choices")) + ) + return tuple(server_tool_call["function"].get("name") for server_tool_call in server_tool_calls) + + class _ExecutedTool(BaseModel): model_config = ConfigDict(frozen=True) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 51934c99db6..1a1670715dc 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -18,6 +18,11 @@ from typing import Final from pydantic import JsonValue, TypeAdapter from litellm.litellm_core_utils.core_helpers import iter_items +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, + changed_server_tool_call_arguments, + write_request_hosted_web_call_arguments, +) from litellm.responses.litellm_completion_transformation.compaction_items import ( compaction_summary, input_items, @@ -407,18 +412,27 @@ def walk_user_text( return new_parts return content - messages: Final = data.get("messages") - for message in _object_list(messages) or (): + def _rewrite_call_arguments(calls: object) -> None: + for call in _object_list(calls) or (): + function = call.get("function") if isinstance(call, dict) else None + if isinstance(function, dict) and "arguments" in function: + function["arguments"] = _rewrite_arguments(function["arguments"]) + + messages: Final = _object_list(data.get("messages")) or () + server_tool_calls_before: Final = tuple(anthropic_server_tool_call_params(message) for message in messages) + for message in messages: if (skip_tool_message and _is_tool_message(message)) or not isinstance(message, dict): continue if "content" in message: message["content"] = _rewrite_content(message["content"]) - calls = message.get("tool_calls") - if isinstance(calls, list): - for call in calls: - function = call.get("function") if isinstance(call, dict) else None - if isinstance(function, dict) and "arguments" in function: - function["arguments"] = _rewrite_arguments(function["arguments"]) + _rewrite_call_arguments(message.get("tool_calls")) + provider_fields = message.get("provider_specific_fields") + if isinstance(provider_fields, dict): + _rewrite_call_arguments(provider_fields.get("server_tool_calls")) + for message, original in zip(messages, server_tool_calls_before, strict=True): + write_request_hosted_web_call_arguments( + message, changed_server_tool_call_arguments(original, anthropic_server_tool_call_params(message)) + ) if "system" in data: data["system"] = _rewrite_content(data["system"]) # rebind-ok: preserve the native hook's in-place contract diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index 19b9bccd8aa..c773e88c769 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -9,7 +9,7 @@ import contextlib import json import os import ssl -from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Iterator, Mapping, Sequence from ssl import SSLContext from typing import TYPE_CHECKING, Any, Final @@ -23,7 +23,16 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.litellm_core_utils.core_helpers import iter_items +from litellm.litellm_core_utils.core_helpers import attribute_items, iter_items +from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_calls, + write_anthropic_server_tool_call_arguments, +) +from litellm.llms.base_llm.guardrail_translation.utils import ( + claude_hosted_web_tool_arguments, + with_web_search_call_arguments, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_ssl_configuration, @@ -515,27 +524,37 @@ class CatoNetworksGuardrail(CustomGuardrail): @staticmethod def _output_fragments(message: Message) -> Sequence[tuple[tuple[str, int | None], str]]: """Assistant text the proxy returns to the caller: ``content`` plus every - ``tool_calls[].function.arguments`` string, each tagged with where a - redaction must be written back. ``content`` is only included when present - so a tool-call-only choice keeps its ``None`` content (the text-vs-tool-call - signal downstream consumers rely on) while its arguments are still inspected.""" - fragments: Final[list] = [] - if message.content is not None: - fragments.append((("content", None), message.content)) - for idx, tool_call in enumerate(message.tool_calls or []): - function = getattr(tool_call, "function", None) - arguments = getattr(function, "arguments", None) - if isinstance(arguments, str) and arguments: - fragments.append((("tool_call", idx), arguments)) - return fragments + tool call's ``function.arguments`` string, Claude's server tool calls first, + each tagged with where a redaction must be written back. ``content`` is only + included when present so a tool-call-only choice keeps its ``None`` content + (the text-vs-tool-call signal downstream consumers rely on) while its + arguments are still inspected.""" + content: Final = () if message.content is None else ((("content", None), message.content),) + return ( + *content, + *CatoNetworksGuardrail._tool_call_fragments("server_tool_call", anthropic_server_tool_calls(message)), + *CatoNetworksGuardrail._tool_call_fragments("tool_call", message.tool_calls or ()), + ) + + @staticmethod + def _tool_call_fragments(kind: str, tool_calls: Sequence[object]) -> tuple[tuple[tuple[str, int | None], str], ...]: + return tuple( + ((kind, idx), arguments) + for idx, arguments in enumerate( + getattr(getattr(tool_call, "function", None), "arguments", None) for tool_call in tool_calls + ) + if isinstance(arguments, str) and arguments + ) @staticmethod def _apply_output_fragment(message: Any, target: tuple[str, int | None], redacted: str) -> None: - kind, idx = target - if kind == "content": - message.content = redacted - else: - message.tool_calls[idx].function.arguments = redacted + match target: + case ("content", _): + message.content = redacted + case ("server_tool_call", int() as idx): + write_anthropic_server_tool_call_arguments(message, idx, redacted) + case (_, idx): + message.tool_calls[idx].function.arguments = redacted @staticmethod def _responses_output_field(item: object, key: str) -> str | Sequence[object] | None: @@ -565,12 +584,19 @@ class CatoNetworksGuardrail(CustomGuardrail): return fragments @staticmethod - def _apply_responses_output_fragment(container: object, key: str, redacted: str) -> None: + def _apply_responses_output_fragment(container: object, key: str, redacted: object) -> None: if isinstance(container, dict): container[key] = redacted else: setattr(container, key, redacted) + @staticmethod + def _responses_hosted_web_calls(response: ResponsesAPIResponse) -> Iterator[tuple[object, str]]: + for item in attribute_items(response, "output"): + hosted_arguments = claude_hosted_web_tool_arguments(item) + if hosted_arguments: + yield item, hosted_arguments + async def _inspect_output_text( self, data: dict, @@ -612,6 +638,12 @@ class CatoNetworksGuardrail(CustomGuardrail): redacted_output = await self._inspect_output_text(data, text, user_api_key_dict, user_email) if redacted_output is not None: self._apply_responses_output_fragment(container, key, redacted_output) + for item, hosted_arguments in self._responses_hosted_web_calls(response): + redacted_output = await self._inspect_output_text(data, hosted_arguments, user_api_key_dict, user_email) + if redacted_output is not None: + self._apply_responses_output_fragment( + item, "action", message_field(with_web_search_call_arguments(item, redacted_output), "action") + ) return response async def async_post_call_streaming_iterator_hook( 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 3d555707159..2920c6bec9a 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 @@ -37,6 +37,16 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.core_helpers import iter_items +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, + tool_calls_with_server_tool_calls, + with_server_tool_call_arguments, + write_anthropic_server_tool_calls, +) +from litellm.llms.base_llm.guardrail_translation.utils import ( + claude_hosted_web_tool_arguments, + with_web_search_call_arguments, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -48,7 +58,9 @@ from litellm.proxy.common_utils.callback_utils import ( from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( Choices, + Delta, LLMResponseTypes, + Message, ModelResponse, ModelResponseStream, TextCompletionResponse, @@ -1654,6 +1666,7 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): @classmethod def _clear_tool_call_arguments(cls, message: object) -> None: """Clear tool-call / function-call arguments after Cisco redaction.""" + cls._clear_server_tool_call_arguments(message) tool_calls = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) for tc in tool_calls or []: fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) @@ -1663,9 +1676,26 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): ) cls._clear_arguments_field(function_call) + @staticmethod + def _clear_server_tool_call_arguments(message: object) -> None: + cleared: Final = tuple( + with_server_tool_call_arguments(server_tool_call, "{}") + for server_tool_call in anthropic_server_tool_call_params(message) + ) + match message: + case Message() | Delta(): + write_anthropic_server_tool_calls(message, cleared) + case dict() if cleared: + message["provider_specific_fields"] = { # rebind-ok: a redaction must land on the message it scanned + **(message.get("provider_specific_fields") or {}), + "server_tool_calls": list(cleared), + } + case _: + pass + 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: @@ -1677,7 +1707,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: @@ -2038,8 +2071,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)) @@ -2053,6 +2087,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 [] @@ -2095,8 +2131,7 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): @classmethod def _extract_message_tool_argument_parts(cls, message: object) -> list[str]: parts: Final[list[str]] = [] - tool_calls = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) - for tool_call in tool_calls or []: + for tool_call in tool_calls_with_server_tool_calls(message): args = cls._extract_tool_call_arguments(tool_call) if args: parts.append(args) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index cb1b223fb47..9c8c9cbe96d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -22,6 +22,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -163,11 +164,11 @@ class lakeraAI_Moderation(CustomGuardrail): for message in data["messages"]: role = message.get("role") if role in stringified_roles: - if "tool_calls" in message: - tool_call_messages = [ - *tool_call_messages, - *message["tool_calls"], - ] + tool_call_messages = [ + *tool_call_messages, + *anthropic_server_tool_call_params(message), + *(message.get("tool_calls") or ()), + ] if role == Role.SYSTEM.value: # we need this for later system_message = message continue @@ -185,10 +186,11 @@ class lakeraAI_Moderation(CustomGuardrail): if system_message is not None: if not litellm.add_function_to_prompt: content = system_message.get("content") - function_input: Final = [] - for tool_call in tool_call_messages: - if "function" in tool_call: - function_input.append(tool_call["function"]["arguments"]) + function_input: Final[tuple[str, ...]] = tuple( + tool_call["function"]["arguments"] + for tool_call in tool_call_messages + if "function" in tool_call + ) if len(function_input) > 0: content += " Function Input: " + " ".join(function_input) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 97b35769f63..0530cf22f62 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -40,6 +40,13 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.integrations.custom_guardrail import dc as global_cache from litellm.litellm_core_utils.core_helpers import iter_items +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, + anthropic_server_tool_calls, + with_server_tool_call_arguments, + write_anthropic_server_tool_calls, + write_request_server_tool_calls, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -50,6 +57,7 @@ from litellm.proxy.guardrails._content_utils import ( has_non_string_content, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import ChatCompletionAssistantToolCall class LassoResponse(TypedDict): @@ -259,7 +267,7 @@ class LassoGuardrail(CustomGuardrail): msg = choice.message if msg.content: response_messages.append({"role": "assistant", "content": msg.content}) - for call in getattr(msg, "tool_calls", None) or []: + for call in (*anthropic_server_tool_calls(msg), *(getattr(msg, "tool_calls", None) or ())): call_id, name, input_data = self._extract_tool_call_fields(call) if not call_id or not name: continue @@ -564,6 +572,15 @@ class LassoGuardrail(CustomGuardrail): elif role == "assistant" and not content and orig_msg.get("tool_calls"): msg["tool_calls"] = self._update_tool_calls_from_masked(orig_msg["tool_calls"], masked_tool_use) + if role == "assistant": + write_request_server_tool_calls( + msg, + tuple( + self._masked_server_tool_call(server_tool_call, masked_tool_use) + for server_tool_call in anthropic_server_tool_call_params(orig_msg) + ), + ) + result.append(msg) return result @@ -704,7 +721,7 @@ class LassoGuardrail(CustomGuardrail): expanded.append({"role": role, "content": content}) if role == "assistant": - for call in iter_items(msg.get("tool_calls") or []): + for call in (*anthropic_server_tool_call_params(msg), *iter_items(msg.get("tool_calls") or [])): call_id, name, input_data = self._extract_tool_call_fields(call) if not call_id or not name: verbose_proxy_logger.warning( @@ -962,6 +979,14 @@ class LassoGuardrail(CustomGuardrail): text_cursor += 1 verbose_proxy_logger.debug("Applied masked text content to choice %s", text_cursor) + write_anthropic_server_tool_calls( + msg, + tuple( + self._masked_server_tool_call(server_tool_call, masked_tool_use) + for server_tool_call in anthropic_server_tool_call_params(msg) + ), + ) + for call in getattr(msg, "tool_calls", None) or []: call_id = self._get_field(call, "id") if call_id and call_id in masked_tool_use: @@ -977,6 +1002,16 @@ class LassoGuardrail(CustomGuardrail): func.arguments = json.dumps(masked_input) verbose_proxy_logger.debug("Applied masked tool_call arguments for call_id=%s", call_id) + @staticmethod + def _masked_server_tool_call( + server_tool_call: ChatCompletionAssistantToolCall, + masked_tool_use: Mapping[object, Mapping[str, object]], + ) -> ChatCompletionAssistantToolCall: + masked_input: Final = masked_tool_use.get(server_tool_call["id"], {}).get("input") + if masked_input is None: + return server_tool_call + return with_server_tool_call_arguments(server_tool_call, json.dumps(masked_input)) + @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: from litellm.types.proxy.guardrails.guardrail_hooks.lasso import ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py index d2dafb04fe0..5e7a0240545 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py @@ -24,13 +24,17 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.core_helpers import attribute_items, attribute_or_none, iter_items, require_attributes +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, + write_anthropic_server_tool_call_arguments, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices +from litellm.types.utils import GenericGuardrailAPIInputs, Message, TextChoices if TYPE_CHECKING: from litellm.caching.caching import DualCache @@ -87,6 +91,19 @@ _DEFAULT_TIMEOUT_SECONDS: Final = 10.0 _SHIELD_BODY: Final = TypeAdapter(dict[str, object]) +def _server_tool_call_slots(message: object) -> tuple[Slot, ...]: + if not isinstance(message, Message): + return () + return tuple( + (arguments, functools.partial(write_anthropic_server_tool_call_arguments, message, index)) + for index, arguments in enumerate( + server_tool_call["function"].get("arguments") + for server_tool_call in anthropic_server_tool_call_params(message) + ) + if isinstance(arguments, str) and arguments + ) + + def _string_tool_arguments(calls: Sequence[object]) -> Iterator[tuple[object, str]]: for call in calls: function = read_field(call, "function") @@ -382,6 +399,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): content = getattr(message, "content", None) if isinstance(content, str) and content: pending.append((content, functools.partial(setattr, message, "content"))) + pending.extend(_server_tool_call_slots(message)) for tool_call in attribute_items(message, "tool_calls"): function: object = getattr(tool_call, "function", None) arguments = getattr(function, "arguments", None) if function is not None else None diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py index f83314af548..1956170ba67 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -434,9 +435,7 @@ class PurviewGuardrailBase: args: Final[list[str]] = [] # tool_calls: [{"function": {"arguments": "..."}}] - tool_calls: Final[Sequence[object] | None] = ( - message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) - ) + tool_calls: Final = tool_calls_with_server_tool_calls(message) if tool_calls: for tc in tool_calls: fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py index 72474ba27c3..0b3d56a016b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py @@ -22,6 +22,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.base_llm.guardrail_translation.utils import claude_hosted_web_tool_arguments from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( Choices, @@ -227,6 +228,9 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): if not output: return args for item in output: + hosted_arguments = claude_hosted_web_tool_arguments(item) + if hosted_arguments: + args.append(hosted_arguments) if isinstance(item, dict): item_type = item.get("type") arguments = item.get("arguments") diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py index 2194238cea0..4ec3ed21359 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py @@ -1,5 +1,6 @@ # litellm/proxy/guardrails/guardrail_hooks/pangea.py import os +from collections.abc import Callable from typing import TYPE_CHECKING, Final from fastapi import HTTPException @@ -10,6 +11,11 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.anthropic.common_utils import ( + message_dict_with_server_tool_calls, + message_dict_with_server_tool_calls_split_out, + validated_object_mapping, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -30,6 +36,13 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +def _choice_with_message(choice: object, rewrite: Callable[[object], object]) -> object: + mapping: Final = validated_object_mapping(choice) + if mapping is None or "message" not in mapping: + return choice + return {**mapping, "message": rewrite(mapping["message"])} + + class PangeaGuardrailMissingSecrets(Exception): """Custom exception for missing Pangea secrets.""" @@ -242,9 +255,13 @@ class PangeaHandler(CustomGuardrail): for c in choices: if isinstance(c, Choices): try: - serialized_choices.append(c.model_dump()) + serialized_choices.append( + _choice_with_message(c.model_dump(), message_dict_with_server_tool_calls) + ) except Exception: - serialized_choices.append(c.dict()) + serialized_choices.append( + _choice_with_message(c.dict(), message_dict_with_server_tool_calls) + ) else: serialized_choices.append(c) choices = serialized_choices @@ -269,7 +286,9 @@ class PangeaHandler(CustomGuardrail): return output: Final = ai_guard_response.get("result", {}).get("output", {}) - response.choices = output["choices"] + response.choices = [ # pyright: ignore[reportAttributeAccessIssue] # Pangea returns untyped choice dicts + _choice_with_message(choice, message_dict_with_server_tool_calls_split_out) for choice in output["choices"] + ] return response @log_guardrail_information diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 3cd9dc6a7b4..379c2c4b075 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -27,6 +27,12 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, + tool_calls_with_server_tool_calls, + with_server_tool_call_arguments, + write_anthropic_server_tool_calls, +) from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, @@ -364,10 +370,9 @@ class PanwPrismaAirsHandler(CustomGuardrail): text_parts.append(str(choice.message.content)) # Extract tool call arguments - if hasattr(choice.message, "tool_calls") and choice.message.tool_calls: - for tool_call in choice.message.tool_calls: - if hasattr(tool_call, "function") and hasattr(tool_call.function, "arguments"): - text_parts.append(str(tool_call.function.arguments)) + for tool_call in tool_calls_with_server_tool_calls(choice.message): + if hasattr(tool_call, "function") and hasattr(tool_call.function, "arguments"): + text_parts.append(str(tool_call.function.arguments)) # Extract function call arguments (legacy) if hasattr(choice.message, "function_call") and choice.message.function_call: @@ -756,6 +761,14 @@ class PanwPrismaAirsHandler(CustomGuardrail): elif isinstance(content, list): choice.message.content = self._mask_content_list(content, masked_text) + write_anthropic_server_tool_calls( + choice.message, + tuple( + with_server_tool_call_arguments(server_tool_call, masked_text) + for server_tool_call in anthropic_server_tool_call_params(choice.message) + ), + ) + # Mask tool call arguments if hasattr(choice.message, "tool_calls") and choice.message.tool_calls: for tool_call in choice.message.tool_calls: diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index 3f025cfc8b3..e136b9baaa8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -26,6 +26,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.anthropic.common_utils import message_dict_with_server_tool_calls from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -205,6 +206,7 @@ class PillarGuardrail(CustomGuardrail): include_evidence: bool | None = None, fallback_on_error: str | None = None, timeout: float | None = None, + async_handler: _PillarProtectHTTPClient | None = None, **kwargs, ) -> None: """ @@ -223,8 +225,10 @@ class PillarGuardrail(CustomGuardrail): LiteLLM virtual key context (user_id, team_id, key_alias, etc.) is always automatically passed as X-LiteLLM-* headers to enable application/user tracking. """ - self.async_handler: _PillarProtectHTTPClient = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback + self.async_handler: _PillarProtectHTTPClient = ( + async_handler + if async_handler is not None + else get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) ) self.api_key = api_key or os.environ.get("PILLAR_API_KEY") @@ -430,7 +434,9 @@ class PillarGuardrail(CustomGuardrail): # Extract response messages in the format Pillar expects response_dict = response.model_dump() if hasattr(response, "model_dump") else {} response_messages: Final = [ - choice.get("message") for choice in response_dict.get("choices", []) if choice.get("message") + message_dict_with_server_tool_calls(choice["message"]) + for choice in response_dict.get("choices", []) + if choice.get("message") ] if not response_messages: diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 458528e7e13..1e4621622b3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -12,7 +12,7 @@ import asyncio import json import re import threading -from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Iterator, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime @@ -43,6 +43,11 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, + with_server_tool_call_arguments, + write_anthropic_server_tool_calls, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames from litellm.proxy.guardrails.anthropic_sse import ( @@ -1404,6 +1409,18 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): verbose_proxy_logger.debug("No pii_tokens found in request_data['metadata'] — nothing to unmask") presidio_config: Final = self.get_presidio_settings_from_request_data(request_data or {}) + async def pii_processed_text(text: str) -> str: + match mode: + case "unmask": + return self._unmask_pii_text(text, pii_tokens) + case "mask": + return await self.check_pii( + text=text, + output_parse_pii=False, + presidio_config=presidio_config, + request_data=request_data, + ) + for choice in response.choices: message: Message | None = getattr(choice, "message", None) if message is None: @@ -1412,15 +1429,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # 1. Process content content = getattr(message, "content", None) if isinstance(content, str): - if mode == "unmask": - message.content = self._unmask_pii_text(content, pii_tokens) - elif mode == "mask": - message.content = await self.check_pii( - text=content, - output_parse_pii=False, - presidio_config=presidio_config, - request_data=request_data, - ) + message.content = await pii_processed_text(content) elif isinstance(content, list): for item in content: if not isinstance(item, dict): @@ -1428,54 +1437,43 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): text_value = item.get("text") if text_value is None: continue - if mode == "unmask": - item["text"] = self._unmask_pii_text(text_value, pii_tokens) - elif mode == "mask": - item["text"] = await self.check_pii( - text=text_value, - output_parse_pii=False, - presidio_config=presidio_config, - request_data=request_data, - ) + item["text"] = await pii_processed_text(text_value) + + await self._process_server_tool_calls_for_pii(message, pii_processed_text) # 2. Process tool calls for tool_call in attribute_items(message, "tool_calls"): function = attribute_or_none(tool_call, "function") args = attribute_or_none(function, "arguments") if function else None if isinstance(args, str): - if mode == "unmask": - _write_attribute(function, "arguments", self._unmask_pii_text(args, pii_tokens)) - elif mode == "mask": - _write_attribute( - function, - "arguments", - await self.check_pii( - text=args, - output_parse_pii=False, - presidio_config=presidio_config, - request_data=request_data, - ), - ) + _write_attribute(function, "arguments", await pii_processed_text(args)) # 3. Process legacy function calls function_call = attribute_or_none(message, "function_call") legacy_args = attribute_or_none(function_call, "arguments") if function_call else None if isinstance(legacy_args, str): - if mode == "unmask": - _write_attribute(function_call, "arguments", self._unmask_pii_text(legacy_args, pii_tokens)) - elif mode == "mask": - _write_attribute( - function_call, - "arguments", - await self.check_pii( - text=legacy_args, - output_parse_pii=False, - presidio_config=presidio_config, - request_data=request_data, - ), - ) + _write_attribute(function_call, "arguments", await pii_processed_text(legacy_args)) return response + @staticmethod + async def _process_server_tool_calls_for_pii( + message: Message, + pii_processed_text: Callable[[str], Awaitable[str]], + ) -> None: + server_tool_calls: Final = anthropic_server_tool_call_params(message) + if not server_tool_calls: + return + processed_arguments: Final = tuple( + [ + await pii_processed_text(server_tool_call["function"].get("arguments", "")) + for server_tool_call in server_tool_calls + ] + ) + write_anthropic_server_tool_calls( + message, + tuple(map(with_server_tool_call_arguments, server_tool_calls, processed_arguments)), + ) + async def _mask_output_response( self, response: ModelResponse | EmbeddingResponse | ImageResponse, diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index 6150fccac0c..c76e8e9c084 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -18,6 +18,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -180,8 +181,8 @@ class QualifireGuardrail(CustomGuardrail): api_message["tool_call_id"] = tool_call_id # Handle tool calls if present - tool_calls = msg.get("tool_calls") - if tool_calls and isinstance(tool_calls, list): + tool_calls = [*anthropic_server_tool_call_params(msg), *(msg.get("tool_calls") or [])] + if tool_calls: api_tool_calls = [] for tc in tool_calls: if isinstance(tc, dict): diff --git a/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py b/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py index b1f0f588ade..2e1ac62d850 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py @@ -12,6 +12,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] httpxSpecialProvider, @@ -80,13 +81,11 @@ class RepelloAIGuardrail(CustomGuardrail): def _extract_tool_call_args_from_message(cls, message: object) -> list[str]: args: Final[list[str]] = [] - tool_calls: Final = cls._get_field(message, "tool_calls") - if _is_object_list(tool_calls): - for tool_call in tool_calls: - function = cls._get_field(tool_call, "function") - arguments = cls._get_field(function, "arguments") - if isinstance(arguments, str) and arguments.strip(): - args.append(arguments) + for tool_call in tool_calls_with_server_tool_calls(message): + function = cls._get_field(tool_call, "function") + arguments = cls._get_field(function, "arguments") + if isinstance(arguments, str) and arguments.strip(): + args.append(arguments) function_call: Final = cls._get_field(message, "function_call") arguments = cls._get_field(function_call, "arguments") diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 5f99b1a4611..f2697c0aa11 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -24,6 +24,11 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.tool_fields import declared_function_name, qualified_member_name +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, + anthropic_server_tool_calls, + without_anthropic_server_tool_calls, +) from litellm.llms.base_llm.guardrail_translation.utils import anthropic_tool_names from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import ( @@ -983,7 +988,7 @@ class ToolPermissionGuardrail(CustomGuardrail): for choice_index, choice in enumerate(response.choices): if isinstance(choice, Choices): - for tool in choice.message.tool_calls or []: + for tool in (*(choice.message.tool_calls or ()), *anthropic_server_tool_calls(choice.message)): tool_calls.append(tool) legacy_tool_call = self._legacy_function_call_to_tool_call( getattr(choice.message, "function_call", None), choice_index @@ -1493,6 +1498,23 @@ class ToolPermissionGuardrail(CustomGuardrail): choice.message.tool_calls = filtered_tool_calls if filtered_tool_calls else None + server_error_results = tuple( + error_results.get(self._judged_tool_call(server_tool_call)) + for server_tool_call in anthropic_server_tool_calls(choice.message) + ) + if any(server_error_results): + error_messages.extend(result.content for result in server_error_results if result is not None) + without_anthropic_server_tool_calls( + choice.message, + frozenset( + server_tool_call["id"] or "" + for server_tool_call, result in zip( + anthropic_server_tool_call_params(choice.message), server_error_results + ) + if result is not None + ), + ) + legacy_tool_call = self._legacy_function_call_to_tool_call( getattr(choice.message, "function_call", None), choice_index ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 5ce2dfb3b87..6f6d056423f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -25,6 +25,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings +from litellm.llms.anthropic.common_utils import anthropic_server_call_fields from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( @@ -144,6 +145,15 @@ def _ends_stream(endpoint_translation: object, item: object) -> bool: return has_ended is not None and has_ended((item,)) +def _passthrough_delta(delta: object) -> Delta: + server_call_fields: Final = anthropic_server_call_fields(delta) + role: Final = getattr(delta, "role", None) + tool_calls: Final = getattr(delta, "tool_calls", None) + if not server_call_fields: + return Delta(content=None, role=role, tool_calls=tool_calls) + return Delta(content=None, role=role, tool_calls=tool_calls, provider_specific_fields=dict(server_call_fields)) + + def _held_choices(held_chars_per_choice: Mapping[int, int]) -> frozenset[int]: return frozenset(idx for idx, held in held_chars_per_choice.items() if held > 0) @@ -974,7 +984,7 @@ class UnifiedLLMGuardrails(CustomLogger): def _chunk_has_tool_calls(item: object) -> bool: for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) - if getattr(delta, "tool_calls", None): + if getattr(delta, "tool_calls", None) or anthropic_server_call_fields(delta): return True return False @@ -1042,11 +1052,7 @@ class UnifiedLLMGuardrails(CustomLogger): synthetic_choices.append( StreamingChoices( index=idx, - delta=Delta( - content=None, - role=getattr(delta, "role", None), - tool_calls=getattr(delta, "tool_calls", None), - ), + delta=_passthrough_delta(delta), finish_reason=passthrough_finish, ) ) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 521a698b276..0f8ba6ef98a 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -9,6 +9,7 @@ import litellm from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend +from litellm.llms.anthropic.common_utils import anthropic_server_call_fields, anthropic_server_tool_calls from litellm.llms.anthropic.messages.native.utils import openai_chat_refusal_text from litellm.main import stream_chunk_builder from litellm.responses.litellm_completion_transformation.compaction_items import compaction_output_item @@ -133,6 +134,7 @@ def _delta_carries_output(delta: ChatCompletionDelta) -> bool: getattr(delta, "images", None), getattr(delta, "audio", None), openai_chat_refusal_text(delta), + anthropic_server_call_fields(delta), ) return any(fields) or _delta_has_signed_thinking_block(delta) @@ -228,6 +230,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._queued_web_search_call_ids: set[str] = set() # mutable-ok: emitted call ids self._tool_search = tool_search_routing(request_tools) self._tool_search_executions: Mapping[str, ToolSearchExecution] = MappingProxyType({}) + self._done_tool_search_call_ids: frozenset[str] = frozenset() async def aclose(self) -> None: await self.litellm_custom_stream_wrapper.aclose() @@ -265,6 +268,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _queue_tool_search_call_done( self, call_id: str, output_index: int, final_args: str, execution: ToolSearchExecution ) -> None: + if call_id in self._done_tool_search_call_ids: + return + self._done_tool_search_call_ids = self._done_tool_search_call_ids | {call_id} self._queue_tool_search_call_added(call_id, output_index, "", execution) item: Final = build_tool_search_call( call_id, final_args or self._tool_args_by_call_id.get(call_id, ""), "completed", execution @@ -277,6 +283,18 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) ) + def _queue_streamed_tool_search_calls(self, chunk: ModelResponseStream) -> None: + if not chunk.choices: + return + for call in anthropic_server_tool_calls(chunk.choices[0].delta): + if (execution := self._tool_search_execution(call.id, call.function.name or "")) is None: + continue + if call.id not in self._tool_output_index_by_call_id: + self._queue_held_annotations() + self._queue_tool_search_call_done( + call.id, self._get_or_assign_tool_output_index(call.id), call.function.arguments, execution + ) + def _get_or_assign_tool_output_index(self, call_id: str) -> int: existing: Final = self._tool_output_index_by_call_id.get(call_id) if existing is not None: @@ -509,11 +527,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): try: message: Final = litellm_complete_object.choices[0].message - tool_calls = getattr(message, "tool_calls", None) + server_tool_calls: Final = anthropic_server_tool_calls(message) + tool_calls: Final = (*server_tool_calls, *_tool_call_objects(getattr(message, "tool_calls", None))) except Exception: - tool_calls = None + return - for tc in _tool_call_objects(tool_calls): + if not tool_calls: + return + + server_tool_call_ids: Final = frozenset(tool_call.id for tool_call in server_tool_calls) + + for tc in tool_calls: call_id_raw = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) if not call_id_raw: continue @@ -538,6 +562,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if (search_execution := self._tool_search_execution(call_id, fn_name)) is not None: self._queue_tool_search_call_done(call_id, output_index, fn_args, search_execution) continue + if call_id in server_tool_call_ids: + continue # Track if this is a new tool call that wasn't streamed is_new_tool_call = call_id not in self._tool_item_id_by_call_id @@ -1205,7 +1231,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return # Tool-first - if hasattr(delta, "tool_calls") and delta.tool_calls: + if (hasattr(delta, "tool_calls") and delta.tool_calls) or anthropic_server_call_fields(delta): # Tool calls already handled via _queue_tool_call_delta_events # DO NOT create message item return @@ -1383,6 +1409,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _queue_chunk_events(self, chunk: ModelResponseStream) -> None: self._hold_chunk_annotations(chunk) + self._queue_streamed_tool_search_calls(chunk) response_api_chunk: Final = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) if response_api_chunk: self._pending_response_events.append(response_api_chunk) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 3e605e926f7..5b1964a4b5e 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -73,7 +73,12 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( is_unhashable_value_behind, ) from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of -from litellm.llms.anthropic.common_utils import supports_anthropic_cache_control +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, + anthropic_server_tool_calls, + supports_anthropic_cache_control, + validated_object_mapping, +) from litellm.llms.anthropic.messages.native.utils import anthropic_stop_reason, openai_chat_refusal_text from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import bedrock_declared_cache_point_support @@ -84,11 +89,12 @@ from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler, StoredTurns, ) -from litellm.types.llms.anthropic import AnthropicToolSearchToolRegex +from litellm.types.llms.anthropic import AnthropicToolSearchToolRegex, claude_hosted_web_tool_call from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, + ChatCompletionAssistantToolCall, ChatCompletionAudioObject, ChatCompletionCachedContent, ChatCompletionFileObject, @@ -933,6 +939,22 @@ def _chain_session_metadata(chain_session: str, request_kwargs: Mapping[str, obj return {key: {**_STR_KEY_DICT_ADAPTER.validate_python(metadata), CALLER_SESSION_ID_METADATA_KEY: chain_session}} +def claude_web_search_call_tool_call(search: ResponseFunctionWebSearch) -> ChatCompletionAssistantToolCall | None: + claude_call: Final = claude_hosted_web_tool_call(search) + if claude_call is None: + return None + tool_id, name, arguments = claude_call + return ChatCompletionAssistantToolCall( + id=tool_id, + type="function", + function=ChatCompletionToolCallFunctionChunk(name=name, arguments=arguments), + ) + + +def _pending_claude_server_tool_call(search: ResponseFunctionWebSearch) -> ChatCompletionAssistantToolCall | None: + return claude_web_search_call_tool_call(search) if search.status == "in_progress" else None + + class LiteLLMCompletionResponsesConfig: @staticmethod def get_supported_openai_params(model: str) -> list: @@ -1889,10 +1911,22 @@ class LiteLLMCompletionResponsesConfig: message_cache_control: Final = ( None if isinstance(content, list) else new_cache_control or previous_cache_control ) + new_server_tool_calls: Final = anthropic_server_tool_call_params(new_message) + server_tool_call_fields: Final = ( + { + "provider_specific_fields": { + **(validated_object_mapping(last_message.get("provider_specific_fields")) or {}), + "server_tool_calls": [*anthropic_server_tool_call_params(last_message), *new_server_tool_calls], + } + } + if new_server_tool_calls + else {} + ) merged: Final = { **{key: value for key, value in last_message.items() if key != "cache_control"}, "content": content, **({} if message_cache_control is None else {"cache_control": message_cache_control}), + **server_tool_call_fields, } return cast( # cast-ok: preserves the assistant fields and content blocks ChatCompletionAssistantMessage, merged @@ -2334,10 +2368,15 @@ class LiteLLMCompletionResponsesConfig: content="Hosted web search: " + replay.model_dump_json(exclude_none=True), ) ] + hosted_search_text: Final = "Hosted web search: " + search.model_dump_json(exclude_none=True) + pending_server_tool_call: Final = _pending_claude_server_tool_call(search) + if pending_server_tool_call is None: + return [GenericChatCompletionMessage(role="assistant", content=hosted_search_text)] return [ - GenericChatCompletionMessage( + ChatCompletionResponseMessage( role="assistant", - content="Hosted web search: " + search.model_dump_json(exclude_none=True), + content=hosted_search_text, + provider_specific_fields={"server_tool_calls": [pending_server_tool_call]}, ) ] if LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(input_item): @@ -3085,13 +3124,12 @@ class LiteLLMCompletionResponsesConfig: all_chat_completion_tools: Final[list[ChatCompletionMessageToolCall]] = [] for choice in chat_completion_response.choices: if isinstance(choice, Choices): - if choice.message.tool_calls: - all_chat_completion_tools.extend(choice.message.tool_calls) - for tool_call in choice.message.tool_calls: - TOOL_CALLS_CACHE.set_cache( - key=tool_call.id, - value=tool_call, - ) + for tool_call in (*anthropic_server_tool_calls(choice.message), *(choice.message.tool_calls or ())): + all_chat_completion_tools.append(tool_call) + TOOL_CALLS_CACHE.set_cache( + key=tool_call.id, + value=tool_call, + ) request_tools: Final = responses_api_request.get("tools") if responses_api_request is not None else None custom_tool_names: Final = extract_custom_tool_names(request_tools) @@ -3522,7 +3560,22 @@ class LiteLLMCompletionResponsesConfig: for item in responses_output ] - return responses_output + server_tool_call_ids: Final = LiteLLMCompletionResponsesConfig._server_tool_call_ids(choices) + return [ + item + for item in responses_output + if not ( + getattr(item, "type", None) == "function_call" + and getattr(item, "call_id", None) in server_tool_call_ids + ) + ] + + @staticmethod + def _server_tool_call_ids(choices: Sequence[Choices]) -> frozenset[str]: + server_tool_calls: Final = chain.from_iterable( + anthropic_server_tool_calls(choice.message) for choice in choices + ) + return frozenset(tool_call.id for tool_call in server_tool_calls) @staticmethod def _extract_tool_result_output_items( diff --git a/litellm/tool_loop.py b/litellm/tool_loop.py index d0b2850bb5b..554fdff4fce 100644 --- a/litellm/tool_loop.py +++ b/litellm/tool_loop.py @@ -120,9 +120,12 @@ def _replays_provider_specific_fields( client_tool_calls: tuple[ChatCompletionMessageToolCall, ...], ) -> bool: from litellm.completion_extras.litellm_responses_transformation.transformation import carries_tool_search_items + from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params - return client_tool_calls != tool_calls or carries_tool_search_items( - getattr(message, "provider_specific_fields", None) + return ( + client_tool_calls != tool_calls + or bool(anthropic_server_tool_call_params(message)) + or carries_tool_search_items(getattr(message, "provider_specific_fields", None)) ) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 802f6794839..c322eb25b3e 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -2,10 +2,12 @@ from collections.abc import Iterable, Mapping, Sequence from enum import Enum from typing import Final, Literal, TypeAlias +from openai.types.responses.response_function_web_search import ResponseFunctionWebSearch from pydantic import ConfigDict from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.responses.main import hosted_web_tool_call, web_search_call_tool_id from .openai import ( ChatCompletionCachedContent, @@ -862,6 +864,15 @@ ANTHROPIC_TOOL_SEARCH_TOOL_TYPES: Final = frozenset( ANTHROPIC_TOOL_SEARCH_TOOL_NAMES: Final = frozenset({"tool_search_tool_regex", "tool_search_tool_bm25"}) ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX: Final = "srvtoolu_" + +def claude_hosted_web_tool_call(search: ResponseFunctionWebSearch) -> tuple[str, str, str] | None: + tool_id: Final = web_search_call_tool_id(search.id) + if tool_id is None or not tool_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX): + return None + hosted_call: Final = hosted_web_tool_call(search.action) + return None if hosted_call is None else (tool_id, *hosted_call) + + # Effort beta header constant ANTHROPIC_EFFORT_BETA_HEADER: Final = "effort-2025-11-24" diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py index 86dc857cfd0..7a2f58e9c16 100644 --- a/litellm/types/responses/main.py +++ b/litellm/types/responses/main.py @@ -1,11 +1,18 @@ import builtins +import json from collections.abc import Mapping, Sequence from typing import Final, Literal, Optional, Union from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall -from openai.types.responses.response_function_web_search import ActionSearchSource, ResponseFunctionWebSearch +from openai.types.responses.response_function_web_search import ( + Action, + ActionOpenPage, + ActionSearch, + ActionSearchSource, + ResponseFunctionWebSearch, +) from openai.types.responses.response_output_refusal import ResponseOutputRefusal -from pydantic import ConfigDict, PrivateAttr +from pydantic import ConfigDict, PrivateAttr, TypeAdapter, ValidationError from typing_extensions import TypedDict from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject @@ -43,14 +50,11 @@ class OutputFunctionToolCall(BaseLiteLLMOpenAIResponseObject): phase: Phase = None -def build_web_search_call( - tool_id: str, - tool_input: object, - result: object, - status: Literal["in_progress", "searching", "completed", "failed"] | None = None, -) -> ResponseFunctionWebSearch: +HOSTED_WEB_TOOL_NAMES: Final = frozenset({"web_search", "web_fetch"}) + + +def _web_search_action(tool_input: object, content: object) -> ActionSearch: query: Final = tool_input.get("query", "") if isinstance(tool_input, Mapping) else "" - content: Final = result.get("content") if isinstance(result, Mapping) else None result_items: Final = content if isinstance(content, Sequence) and not isinstance(content, (str, bytes)) else () sources: Final = [ ActionSearchSource(type="url", url=url) @@ -59,20 +63,96 @@ def build_web_search_call( and item.get("type") == "web_search_result" and isinstance((url := item.get("url")), str) ] - failed: Final = isinstance(content, Mapping) and content.get("type") == "web_search_tool_result_error" - return ResponseFunctionWebSearch( - id=f"ws_{tool_id}", - type="web_search_call", - status=status or ("failed" if failed else "completed"), - action={ - "type": "search", - "query": query if isinstance(query, str) else "", - "queries": [query] if isinstance(query, str) and query else [], - "sources": sources, - }, + return ActionSearch( + type="search", + query=query if isinstance(query, str) else "", + queries=[query] if isinstance(query, str) and query else [], + sources=sources, ) +def _web_fetch_action(tool_input: object) -> ActionOpenPage: + match tool_input: + case {"url": str() as url}: + return ActionOpenPage(type="open_page", url=url) + case _: + return ActionOpenPage(type="open_page", url=None) + + +def _result_content(result: object) -> object: + match result: + case {"content": object() as content}: + return content + case _: + return None + + +def _is_tool_result_error(content: object, tool_name: str) -> bool: + match content: + case {"type": str() as content_type}: + return content_type == f"{tool_name}_tool_result_error" + case _: + return False + + +WEB_SEARCH_CALL_ID_PREFIX: Final = "ws_" + + +def build_web_search_call( + tool_id: str, + tool_input: object, + result: object, + status: Literal["in_progress", "searching", "completed", "failed"] | None = None, + tool_name: str = "web_search", +) -> ResponseFunctionWebSearch: + content: Final = _result_content(result) + return ResponseFunctionWebSearch( + id=f"{WEB_SEARCH_CALL_ID_PREFIX}{tool_id}", + type="web_search_call", + status=status or ("failed" if _is_tool_result_error(content, tool_name) else "completed"), + action=_web_fetch_action(tool_input) if tool_name == "web_fetch" else _web_search_action(tool_input, content), + ) + + +def web_search_call_tool_id(call_id: str) -> str | None: + tool_id: Final = call_id.removeprefix(WEB_SEARCH_CALL_ID_PREFIX) + return None if tool_id == call_id else tool_id + + +_HOSTED_WEB_TOOL_ARGUMENTS: Final = TypeAdapter(dict[str, object]) + + +def hosted_web_tool_call(action: object) -> tuple[str, str] | None: + match action: + case ActionSearch(query=str() as query): + return "web_search", json.dumps({"query": query}) + case ActionOpenPage(url=str() as url): + return "web_fetch", json.dumps({"url": url}) + case _: + return None + + +def _decoded_hosted_web_tool_arguments(arguments: str) -> Mapping[str, object]: + try: + return _HOSTED_WEB_TOOL_ARGUMENTS.validate_json(arguments) + except ValidationError: + return {} + + +def with_hosted_web_tool_arguments(action: Action, arguments: str) -> Action: + match action, _decoded_hosted_web_tool_arguments(arguments): + case ActionSearch(), {"query": str() as query}: + return action.model_copy(update={"query": query, "queries": [query] if query else []}) + case ActionSearch(), _: + return action.model_copy(update={"query": "", "queries": []}) + case ActionOpenPage(), {"url": str() as url}: + return action.model_copy(update={"url": url}) + case ActionOpenPage(), _: + return action.model_copy(update={"url": None}) + case _: + return action + + class OutputImageGenerationCall(BaseLiteLLMOpenAIResponseObject): """An image generation call output""" diff --git a/tests/integration/providers/_tool_search_support.py b/tests/integration/providers/_tool_search_support.py index fa6b53040e4..b2a156d238d 100644 --- a/tests/integration/providers/_tool_search_support.py +++ b/tests/integration/providers/_tool_search_support.py @@ -317,6 +317,11 @@ def tool_call_ids(assistant: Event) -> list[JsonValue]: return [JSON_OBJECT.validate_python(call)["id"] for call in calls] +def server_tool_call_ids(assistant: Event) -> list[JsonValue]: + calls: Final = JSON_LIST.validate_python(provider_fields(assistant).get("server_tool_calls") or []) + return [JSON_OBJECT.validate_python(call)["id"] for call in calls] + + def wire_content(message: Event) -> list[JsonValue]: content: Final = message["content"] if isinstance(content, str): diff --git a/tests/integration/providers/test_anthropic_tool_search_result_matrix_wire.py b/tests/integration/providers/test_anthropic_tool_search_result_matrix_wire.py index d4c0e2651d8..adce80d3c75 100644 --- a/tests/integration/providers/test_anthropic_tool_search_result_matrix_wire.py +++ b/tests/integration/providers/test_anthropic_tool_search_result_matrix_wire.py @@ -66,6 +66,7 @@ from integration.providers._tool_search_support import ( response_id, scripted_peer, search_result, + server_tool_call_ids, spend_row, target, target_path, @@ -136,7 +137,8 @@ def test_tool_search_result_round_trips_on_every_anthropic_wire_provider( assert [target_path(request) for request in received] == [target(provider, stream)] * 2, received assert kept_results(assistant) == [SEARCH_RESULT], assistant assert "code_interpreter_results" not in provider_fields(assistant), assistant - assert tool_call_ids(assistant) == [SEARCH_ID, WEATHER_ID], assistant + assert server_tool_call_ids(assistant) == [SEARCH_ID], assistant + assert tool_call_ids(assistant) == [WEATHER_ID], assistant assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body assert answered_ids(received[1]) == [WEATHER_ID], received[1].body for identity in (response_id(first, stream), response_id(second, stream)): @@ -157,7 +159,8 @@ def test_snowflake_claude_keeps_the_tool_search_result_on_the_response(gateway: assistant: Final = assistant_from(first, stream) assert [target_path(request) for request in received] == [target("snowflake", stream)], received assert kept_results(assistant) == [SEARCH_RESULT], assistant - assert tool_call_ids(assistant) == [SEARCH_ID, WEATHER_ID], assistant + assert server_tool_call_ids(assistant) == [SEARCH_ID], assistant + assert tool_call_ids(assistant) == [WEATHER_ID], assistant assert spend_row(response_id(first, stream))["model_group"] == model diff --git a/tests/integration/providers/test_anthropic_tool_search_result_wire.py b/tests/integration/providers/test_anthropic_tool_search_result_wire.py index b46187e3e71..491cd5acbe8 100644 --- a/tests/integration/providers/test_anthropic_tool_search_result_wire.py +++ b/tests/integration/providers/test_anthropic_tool_search_result_wire.py @@ -1,6 +1,7 @@ import json import uuid from collections.abc import Sequence +from functools import reduce from typing import Final from integration._support.anthropic_thinking import sse_chunks @@ -142,13 +143,8 @@ def _streamed_assistant(lines: Sequence[str]) -> dict[str, JsonValue]: entry["id"] = call.get("id") or entry["id"] entry["name"] += call["function"].get("name") or "" entry["arguments"] += call["function"].get("arguments") or "" - fields: Final = next( - ( - delta["provider_specific_fields"] - for delta in reversed(deltas) - if (delta.get("provider_specific_fields") or {}).get("tool_results") - ), - {}, + fields: Final = reduce( + lambda merged, delta: {**merged, **(delta.get("provider_specific_fields") or {})}, deltas, {} ) return { "role": "assistant", diff --git a/tests/integration/providers/test_chat_completions_tool_search_wire.py b/tests/integration/providers/test_chat_completions_tool_search_wire.py index f99db3fe735..51c44996d31 100644 --- a/tests/integration/providers/test_chat_completions_tool_search_wire.py +++ b/tests/integration/providers/test_chat_completions_tool_search_wire.py @@ -352,28 +352,41 @@ def _first_choice(payload: Mapping[str, JsonValue]) -> dict[str, JsonValue]: return rv.ITEMS.validate_python(payload["choices"])[0] -def _streamed_tool_calls(frames: Sequence[Mapping[str, JsonValue]]) -> list[tuple[str, str]]: - deltas: Final = [ - rv.JSON_OBJECT.validate_python(_first_choice(frame)["delta"]) for frame in frames if frame.get("choices") - ] - opened: Final = rv.ITEMS.validate_python( - list(chain.from_iterable(rv.ITEMS.validate_python(delta.get("tool_calls") or []) for delta in deltas)) - ) +def _call_names(calls: Sequence[Mapping[str, JsonValue]]) -> list[tuple[str, str]]: return [ (str(call["id"]), str(rv.JSON_OBJECT.validate_python(call["function"])["name"])) - for call in opened + for call in calls if call.get("id") ] +def _deltas(frames: Sequence[Mapping[str, JsonValue]]) -> list[dict[str, JsonValue]]: + return [rv.JSON_OBJECT.validate_python(_first_choice(frame)["delta"]) for frame in frames if frame.get("choices")] + + +def _server_tool_calls_of(holder: Mapping[str, JsonValue]) -> list[dict[str, JsonValue]]: + fields: Final = rv.JSON_OBJECT.validate_python(holder.get("provider_specific_fields") or {}) + return rv.ITEMS.validate_python(fields.get("server_tool_calls") or []) + + def answered_tool_calls(answer: Answer) -> list[tuple[str, str]]: if answer.frames: - return _streamed_tool_calls(answer.frames) + return _call_names( + list( + chain.from_iterable(rv.ITEMS.validate_python(d.get("tool_calls") or []) for d in _deltas(answer.frames)) + ) + ) message: Final = rv.JSON_OBJECT.validate_python(_first_choice(answer.payload)["message"]) - return [ - (str(call["id"]), str(rv.JSON_OBJECT.validate_python(call["function"])["name"])) - for call in rv.ITEMS.validate_python(message.get("tool_calls") or []) - ] + return _call_names(rv.ITEMS.validate_python(message.get("tool_calls") or [])) + + +def answered_server_tool_calls(answer: Answer) -> list[tuple[str, str]]: + if answer.frames: + streamed: Final = _call_names( + list(chain.from_iterable(_server_tool_calls_of(d) for d in _deltas(answer.frames))) + ) + return list(dict.fromkeys(streamed)) + return _call_names(_server_tool_calls_of(rv.JSON_OBJECT.validate_python(_first_choice(answer.payload)["message"]))) def response_id(answer: Answer) -> str: @@ -382,10 +395,8 @@ def response_id(answer: Answer) -> str: def assert_search_call_answered(answer: Answer, marker: str) -> None: assert answer.status == 200, answer.text - assert answered_tool_calls(answer) == [ - (f"srvtoolu_{marker}", _REGEX_TOOL_NAME), - (f"toolu_{marker}", "get_weather"), - ], answer.text + assert answered_server_tool_calls(answer) == [(f"srvtoolu_{marker}", _REGEX_TOOL_NAME)], answer.text + assert answered_tool_calls(answer) == [(f"toolu_{marker}", "get_weather")], answer.text if answer.frames: return usage: Final = rv.JSON_OBJECT.validate_python(answer.payload["usage"]) diff --git a/tests/integration/providers/test_responses_bridge_tool_search_wire.py b/tests/integration/providers/test_responses_bridge_tool_search_wire.py index 57d9655b7d5..95c05610fba 100644 --- a/tests/integration/providers/test_responses_bridge_tool_search_wire.py +++ b/tests/integration/providers/test_responses_bridge_tool_search_wire.py @@ -998,7 +998,7 @@ def test_a_callers_function_named_like_the_search_tool_stays_a_function_call(gat assert [tool["name"] for tool in _ITEMS.validate_python(body["tools"])] == [_REGEX_TOOL_NAME], body -def test_an_unrequested_server_tool_use_from_the_vendor_stays_a_function_call(gateway: Gateway) -> None: +def test_an_unrequested_server_tool_use_from_the_vendor_is_not_returned_as_a_function_call(gateway: Gateway) -> None: marker: Final = uuid.uuid4().hex with wire_server(_Vendor(searched=True).respond) as wire, gateway.scenario() as scenario: model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) @@ -1007,7 +1007,11 @@ def test_an_unrequested_server_tool_use_from_the_vendor_stays_a_function_call(ga ) assert "tool_search_call" not in _output_types(answer.payload), answer calls: Final = {str(item["call_id"]): str(item["name"]) for item in _tool_items(answer.payload)} - assert calls == {f"srvtoolu_{marker}": _REGEX_TOOL_NAME, f"toolu_{marker}": "get_weather"}, answer + assert calls == {f"toolu_{marker}": "get_weather"}, answer + server_calls: Final = _ITEMS.validate_python( + _JSON_OBJECT.validate_python(answer.payload["provider_specific_fields"])["server_tool_calls"] + ) + assert [call["id"] for call in server_calls] == [f"srvtoolu_{marker}"], answer _, body = _only_request(wire) _assert_dropped_on_wire(body, ("get_weather",)) @@ -1224,15 +1228,17 @@ def test_chat_completions_keep_forwarding_the_regex_tool_and_defer_loading(gatew if stream: assert "data: [DONE]" in response.text and "get_weather" in response.text, response.text else: - calls: Final = _ITEMS.validate_python( - _JSON_OBJECT.validate_python( - _ITEMS.validate_python(_JSON_OBJECT.validate_json(response.content)["choices"])[0]["message"] - )["tool_calls"] + message: Final = _JSON_OBJECT.validate_python( + _ITEMS.validate_python(_JSON_OBJECT.validate_json(response.content)["choices"])[0]["message"] ) - assert [_JSON_OBJECT.validate_python(call["function"])["name"] for call in calls] == [ - _REGEX_TOOL_NAME, - "get_weather", - ], calls + calls: Final = _ITEMS.validate_python(message["tool_calls"]) + server_calls: Final = _ITEMS.validate_python( + _JSON_OBJECT.validate_python(message["provider_specific_fields"])["server_tool_calls"] + ) + assert [_JSON_OBJECT.validate_python(call["function"])["name"] for call in calls] == ["get_weather"], calls + assert [_JSON_OBJECT.validate_python(call["function"])["name"] for call in server_calls] == [ + _REGEX_TOOL_NAME + ], message request, sent = _only_request(wire) _assert_anthropic_wire(request, sent, ("get_weather",)) diff --git a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py index 9816eb91102..3599c938c28 100644 --- a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py @@ -29,6 +29,7 @@ from litellm_enterprise.enterprise_callbacks.secret_detection import ( from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import hash_token +from litellm.types.responses.main import build_web_search_call AWS_KEY = "AKIAIOSFODNN7EXAMPLE" OPENAI_KEY = "sk-test-abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGH" @@ -114,6 +115,51 @@ async def test_deeply_nested_tool_arguments_are_redacted(): assert json.loads(rewritten) == json.loads(arguments.replace(AWS_KEY, "[REDACTED]")) +@pytest.mark.asyncio +async def test_a_replayed_server_tool_call_cannot_forward_unredacted_credentials(): + arguments = json.dumps({"query": f"AWS_ACCESS_KEY_ID={AWS_KEY}"}) + server_call = {"id": "srvtoolu_1", "type": "function", "function": {"name": "web_search", "arguments": arguments}} + data = { + "messages": [ + { + "role": "assistant", + "content": "Searching.", + "provider_specific_fields": {"server_tool_calls": [server_call]}, + } + ] + } + await _guardrail().async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, "acompletion") + rewritten = data["messages"][0]["provider_specific_fields"]["server_tool_calls"][0]["function"]["arguments"] + assert json.loads(rewritten) == {"query": "AWS_ACCESS_KEY_ID=[REDACTED]"} + + +@pytest.mark.asyncio +async def test_a_replayed_server_web_search_keeps_no_credentials_in_its_hosted_copy(): + query = f"AWS_ACCESS_KEY_ID={AWS_KEY}" + server_call = { + "id": "srvtoolu_1", + "type": "function", + "function": {"name": "web_search", "arguments": json.dumps({"query": query})}, + } + data = { + "messages": [ + { + "role": "assistant", + "content": "Searching.", + "provider_specific_fields": { + "server_tool_calls": [server_call], + "web_search_calls": [ + build_web_search_call("srvtoolu_1", {"query": query}, {"content": []}).model_dump() + ], + }, + } + ] + } + await _guardrail().async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, "acompletion") + hosted = json.loads(json.dumps(data["messages"][0]["provider_specific_fields"]["web_search_calls"])) + assert [call["action"]["query"] for call in hosted] == ["AWS_ACCESS_KEY_ID=[REDACTED]"] + + @pytest.mark.parametrize("password", ("hunter2\\brahms9x", "hunter2\tbrahms9x")) def test_tool_payload_redacts_keyword_values_that_need_json_escaping(password): assert _guardrail().redact_tool_payload({"password": password, "path": ".env"}) == { diff --git a/tests/unit/integrations/arize/test_arize_utils.py b/tests/unit/integrations/arize/test_arize_utils.py index 31699359bbe..db7144b125a 100644 --- a/tests/unit/integrations/arize/test_arize_utils.py +++ b/tests/unit/integrations/arize/test_arize_utils.py @@ -725,6 +725,73 @@ def test_arize_output_value_falls_back_to_tool_calls_summary(): assert "SF" in out +def test_arize_server_tool_calls_lead_the_output_summary_and_tool_call_attributes() -> None: + span: Final = TracerProvider().get_tracer(__name__).start_span("litellm_request") + kwargs: Final = { + "model": "claude-opus-5-5", + "messages": [{"role": "user", "content": "persimmon facts?"}], + "standard_logging_object": {"model_parameters": {}, "metadata": {}, "call_type": "completion"}, + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "anthropic"}, + } + response_obj: Final = ModelResponse( + usage={"total_tokens": 10, "completion_tokens": 4, "prompt_tokens": 6}, + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=litellm.Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + } + ], + provider_specific_fields={ + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + } + ] + }, + ), + ) + ], + model="claude-opus-5-5", + id="r-server-tools", + ) + + ArizeLogger.set_arize_attributes(span, kwargs, response_obj) + + tool_calls_prefix: Final = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_TOOL_CALLS}." + tool_call_attributes: Final = { + key: value for key, value in span.attributes.items() if key.startswith(tool_calls_prefix) + } + assert json.loads(span.attributes[SpanAttributes.OUTPUT_VALUE]) == { + "tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + }, + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}}, + ] + } + assert tool_call_attributes == { + f"{tool_calls_prefix}0.{ToolCallAttributes.TOOL_CALL_ID}": "srvtoolu_01", + f"{tool_calls_prefix}0.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}": "web_search", + f"{tool_calls_prefix}0.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}": '{"query": "persimmon"}', + f"{tool_calls_prefix}1.{ToolCallAttributes.TOOL_CALL_ID}": "call_1", + f"{tool_calls_prefix}1.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}": "get_weather", + f"{tool_calls_prefix}1.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}": '{"city": "SF"}', + } + + def test_arize_output_value_unchanged_when_content_present(): """Regression guard: when content is non-empty, OUTPUT_VALUE must be exactly that content (no summary written).""" diff --git a/tests/unit/integrations/datadog/test_datadog_llm_obs.py b/tests/unit/integrations/datadog/test_datadog_llm_obs.py index a7826e11483..3604ed223e1 100644 --- a/tests/unit/integrations/datadog/test_datadog_llm_obs.py +++ b/tests/unit/integrations/datadog/test_datadog_llm_obs.py @@ -110,6 +110,41 @@ def test_output_tool_calls_use_the_datadog_tool_call_schema(logger: DataDogLLMOb assert "function" not in message["tool_calls"][0] + +def test_claude_server_tool_calls_are_reported_before_client_tool_calls(logger: DataDogLLMObsLogger) -> None: + payload: Final = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [ASSISTANT_TOOL_CALL], + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01Search", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "paris weather"}'}, + } + ] + }, + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"] == [ + { + "name": "web_search", + "arguments": {"query": "paris weather"}, + "tool_id": "srvtoolu_01Search", + "type": "function", + }, + { + "name": "get_weather", + "arguments": {"city": "Paris", "unit": "c"}, + "tool_id": "call_abc123", + "type": "function", + }, + ] + def test_tool_call_identifiers_that_are_not_strings_are_blanked_not_stringified(logger: DataDogLLMObsLogger) -> None: payload = build( logger, diff --git a/tests/unit/integrations/newrelic/test_newrelic.py b/tests/unit/integrations/newrelic/test_newrelic.py index d09732c6eb2..312457d07e7 100644 --- a/tests/unit/integrations/newrelic/test_newrelic.py +++ b/tests/unit/integrations/newrelic/test_newrelic.py @@ -1,3 +1,4 @@ +import json import os import sys from datetime import datetime, timezone @@ -45,6 +46,12 @@ NR_ENV = { "NEW_RELIC_APP_NAME": "test-app", } +_SERVER_CALL: Final = { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, +} + def make_logger(**kwargs) -> NewRelicLogger: """Instantiate NewRelicLogger with NR agent calls mocked out.""" @@ -361,6 +368,21 @@ class TestExtractMessageContent: assert "get_weather" in result assert "call_1" in result + @pytest.mark.parametrize( + "message", + [ + {"role": "assistant", "content": None, "provider_specific_fields": {"server_tool_calls": [_SERVER_CALL]}}, + litellm.Message( + role="assistant", content=None, provider_specific_fields={"server_tool_calls": [_SERVER_CALL]} + ), + ], + ids=["dict", "message"], + ) + def test_server_tool_call_only_message_serializes_the_server_call( + self, message: dict[str, object] | litellm.Message + ) -> None: + assert json.loads(self.logger._extract_message_content(message)) == [_SERVER_CALL] + def test_multimodal_list_serialized_as_json(self): msg = { "content": [ diff --git a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py index 2d0b4701c9f..c29da4607eb 100644 --- a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py @@ -48,6 +48,7 @@ from litellm.integrations.otel.model.spans import ( validate_registry, ) from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls +from litellm.types.responses.main import build_web_search_call @pytest.fixture(autouse=True) @@ -798,6 +799,39 @@ def test_responses_tool_calls_fold_into_the_assistant_message_with_tool_calls_fi assert data.finish_reasons == ("tool_calls",) +def test_responses_claude_web_calls_lead_the_tool_calls_without_ending_the_turn_on_tool_calls(): + claude_fetch: Final = ( + build_web_search_call("srvtoolu_01", {"url": "https://example.com/"}, None, tool_name="web_fetch") + ) + openai_search: Final = { + "type": "web_search_call", + "id": "ws_1", + "status": "completed", + "action": {"type": "search", "query": "news"}, + } + client_call: Final = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"} + claude_tool_call: Final = { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_fetch", "arguments": '{"url": "https://example.com/"}'}, + } + + hosted_only: Final = LLMCallSpanData.from_standard_logging_payload( + _responses_payload([_RESPONSES_TEXT_ITEM, claude_fetch.model_dump(), openai_search]), capture_content=True + ) + with_client_call: Final = LLMCallSpanData.from_standard_logging_payload( + _responses_payload([_RESPONSES_TEXT_ITEM, claude_fetch.model_dump(), client_call]), capture_content=True + ) + + assert json.loads(json.dumps(hosted_only.choices_out[0]["message"]["tool_calls"])) == [claude_tool_call] + assert hosted_only.finish_reasons == ("stop",) + assert json.loads(json.dumps(with_client_call.choices_out[0]["message"]["tool_calls"])) == [ + claude_tool_call, + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}}, + ] + assert with_client_call.finish_reasons == ("tool_calls",) + + def test_responses_tool_call_only_output_has_no_content(): data = LLMCallSpanData.from_standard_logging_payload( _responses_payload([{"type": "function_call", "id": "fc_1", "name": "get_weather", "arguments": "{}"}]), diff --git a/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py index 57764bf902b..611e828cf84 100644 --- a/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py +++ b/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py @@ -174,6 +174,45 @@ def test_openinference_output_tool_calls_preserve_calls_in_attributes_and_value( ] +def test_openinference_output_lists_claude_server_tool_calls_before_client_tool_calls(): + data: Final = _llm_call( + choices_out=( + { + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_paris", + "type": "function", + "function": {"name": "lookup_weather", "arguments": '{"city": "Paris"}'}, + } + ], + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01Search", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "paris weather"}'}, + } + ] + }, + }, + }, + ) + ) + attrs: Final = OpenInferenceMapper().map(data) + assert {key: value for key, value in attrs.items() if ".tool_calls." in key} == { + "llm.output_messages.0.message.tool_calls.0.tool_call.id": "srvtoolu_01Search", + "llm.output_messages.0.message.tool_calls.0.tool_call.function.name": "web_search", + "llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments": '{"query": "paris weather"}', + "llm.output_messages.0.message.tool_calls.1.tool_call.id": "call_paris", + "llm.output_messages.0.message.tool_calls.1.tool_call.function.name": "lookup_weather", + "llm.output_messages.0.message.tool_calls.1.tool_call.function.arguments": '{"city": "Paris"}', + } + + def test_openinference_responses_tool_calls_are_emitted_as_output_attributes(): data: Final = LLMCallSpanData.from_standard_logging_payload( _responses_payload( diff --git a/tests/unit/integrations/test_helicone.py b/tests/unit/integrations/test_helicone.py index ce3f61a61eb..a4c23a35b40 100644 --- a/tests/unit/integrations/test_helicone.py +++ b/tests/unit/integrations/test_helicone.py @@ -17,7 +17,8 @@ def _claude_mapping(messages, response_obj): return logger.claude_mapping(model="gpt-5.6", messages=messages, response_obj=response_obj) -def test_claude_mapping_serializes_custom_tool_calls(monkeypatch): +@pytest.fixture() +def anthropic_prompts(monkeypatch: pytest.MonkeyPatch) -> None: """ Stub the anthropic module unconditionally: the SDK may be absent (it lives in the proxy-runtime extra), and the tests/unit/llms/anthropic test package can @@ -27,6 +28,10 @@ def test_claude_mapping_serializes_custom_tool_calls(monkeypatch): stub.HUMAN_PROMPT = "\n\nHuman:" stub.AI_PROMPT = "\n\nAssistant:" monkeypatch.setitem(sys.modules, "anthropic", stub) + + +@pytest.mark.usefixtures("anthropic_prompts") +def test_claude_mapping_serializes_custom_tool_calls(): response_obj = { "id": "chatcmpl-1", "choices": [ @@ -58,6 +63,44 @@ def test_claude_mapping_serializes_custom_tool_calls(monkeypatch): assert {"type": "tool_use", "id": "call_f", "name": "read_file", "input": '{"path": "a.py"}'} in tool_use_blocks +@pytest.mark.usefixtures("anthropic_prompts") +def test_claude_mapping_puts_server_tool_calls_first_as_tool_use_blocks(): + response_obj: Final = { + "id": "chatcmpl-1", + "choices": [ + { + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_f", + "type": "function", + "function": {"name": "read_file", "arguments": '{"path": "a.py"}'}, + }, + ], + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + } + ] + }, + }, + "finish_reason": "tool_calls", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 2}, + } + mapped: Final = _claude_mapping([{"role": "user", "content": "hi"}], response_obj) + assert mapped["content"] == [ + {"type": "tool_use", "id": "srvtoolu_01", "name": "web_search", "input": '{"query": "persimmon"}'}, + {"type": "tool_use", "id": "call_f", "name": "read_file", "input": '{"path": "a.py"}'}, + ] + + @pytest.fixture() def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/integrations/test_langfuse_otel.py b/tests/unit/integrations/test_langfuse_otel.py index a2803719c75..71087b4de51 100644 --- a/tests/unit/integrations/test_langfuse_otel.py +++ b/tests/unit/integrations/test_langfuse_otel.py @@ -1309,6 +1309,55 @@ class TestDerivedTraceFields: assert json.loads(str(attributes["langfuse.observation.output"])) == {"content": "hi"} + def test_server_tool_calls_lead_the_observation_output(self) -> None: + response_obj: Final = ModelResponse( + id="chatcmpl-server", + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=litellm.Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + } + ], + provider_specific_fields={ + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + } + ] + }, + ), + ) + ], + ) + attributes: Final = _emitted({"call_type": "acompletion", "litellm_params": {"metadata": {}}}, response_obj) + + assert json.loads(str(attributes["langfuse.observation.output"])) == [ + { + "id": "chatcmpl-server", + "name": "web_search", + "call_id": "srvtoolu_01", + "type": "function_call", + "arguments": {"query": "persimmon"}, + }, + { + "id": "chatcmpl-server", + "name": "get_weather", + "call_id": "call_1", + "type": "function_call", + "arguments": {"city": "SF"}, + }, + ] + def test_no_messages_leaves_trace_input_unset_but_emits_output(self) -> None: attributes: Final = _emitted( {"call_type": "acompletion", "litellm_params": {"metadata": {}}}, _chat_response("hi") diff --git a/tests/unit/integrations/test_lunary.py b/tests/unit/integrations/test_lunary.py index 6491f5c8b82..51353f938b3 100644 --- a/tests/unit/integrations/test_lunary.py +++ b/tests/unit/integrations/test_lunary.py @@ -1,6 +1,9 @@ -from litellm.integrations.lunary import parse_tool_calls +from typing import Final + +import litellm +from litellm.integrations.lunary import parse_messages, parse_tool_calls from litellm.types.utils import ( ChatCompletionMessageCustomToolCall, ChatCompletionMessageToolCall, @@ -35,3 +38,46 @@ def test_parse_tool_calls_serializes_custom_tool_calls(): def test_parse_tool_calls_none_passthrough(): assert parse_tool_calls(None) is None + + +def test_parse_messages_lists_server_tool_calls_before_client_calls(): + response: Final = litellm.ModelResponse( + choices=[ + litellm.Choices( + index=0, + finish_reason="tool_calls", + message=litellm.Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + } + ], + provider_specific_fields={ + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + } + ] + }, + ), + ) + ], + ) + assert parse_messages(response.choices) == { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "type": "function", + "id": "srvtoolu_01", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + }, + {"type": "function", "id": "call_1", "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}}, + ], + } diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 5f749e951b1..393905da6a2 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -372,6 +372,112 @@ def test_v1_routing_diagnostics_add_scalar_attributes_and_preserve_legacy_metada assert finished.attributes["metadata.routing_decision"] == str(decision) +def test_v1_server_tool_calls_lead_the_completion_function_call_attributes() -> None: + exporter: Final = InMemorySpanExporter() + provider: Final = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + otel: Final = OpenTelemetry(tracer_provider=provider) + response_obj: Final = litellm.ModelResponse( + choices=[ + litellm.Choices( + index=0, + finish_reason="tool_calls", + message=litellm.Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + } + ], + provider_specific_fields={ + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + } + ] + }, + ), + ) + ], + ) + with provider.get_tracer(__name__).start_as_current_span("inference") as span: + otel.set_attributes( + span, + { + "model": "claude-opus-5-5", + "standard_logging_object": {"id": "server-tools-call", "call_type": "acompletion", "metadata": {}}, + }, + response_obj, + ) + (finished,) = exporter.get_finished_spans() + function_call_attributes: Final = { + key: value for key, value in finished.attributes.items() if ".function_call." in key + } + assert function_call_attributes == { + f"{SpanAttributes.LLM_COMPLETIONS.value}.0.function_call.name": "web_search", + f"{SpanAttributes.LLM_COMPLETIONS.value}.0.function_call.arguments": '{"query": "persimmon"}', + f"{SpanAttributes.LLM_COMPLETIONS.value}.1.function_call.name": "get_weather", + f"{SpanAttributes.LLM_COMPLETIONS.value}.1.function_call.arguments": '{"city": "SF"}', + } + + +def test_v1_server_tool_calls_stay_in_the_gen_ai_input_and_output_messages() -> None: + exporter: Final = InMemorySpanExporter() + provider: Final = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + otel: Final = OpenTelemetry(tracer_provider=provider) + weather_call: Final = { + "id": "toolu_01", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + } + search_call: Final = { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + } + replayed_assistant: Final = { + "role": "assistant", + "content": "Searching.", + "tool_calls": [weather_call], + "provider_specific_fields": {"server_tool_calls": [search_call]}, + } + response_obj: Final = litellm.ModelResponse( + choices=[ + litellm.Choices( + index=0, + finish_reason="tool_calls", + message=litellm.Message( + role="assistant", + content=None, + tool_calls=[weather_call], + provider_specific_fields={"server_tool_calls": [search_call]}, + ), + ) + ], + ) + with provider.get_tracer(__name__).start_as_current_span("inference") as span: + otel.set_attributes( + span, + { + "model": "claude-opus-5-5", + "messages": [{"role": "user", "content": "Weather?"}, replayed_assistant], + "standard_logging_object": {"id": "server-tools-call", "call_type": "acompletion", "metadata": {}}, + }, + response_obj, + ) + (finished,) = exporter.get_finished_spans() + input_messages: Final = json.loads(finished.attributes[SpanAttributes.GEN_AI_INPUT_MESSAGES.value]) + output_messages: Final = json.loads(finished.attributes[SpanAttributes.GEN_AI_OUTPUT_MESSAGES.value]) + assert [call["id"] for call in input_messages[1]["tool_calls"]] == ["srvtoolu_01", "toolu_01"] + assert [call["id"] for call in output_messages[0]["tool_calls"]] == ["srvtoolu_01", "toolu_01"] + + def test_v1_pod_routing_none_matches_other_none_metadata_attributes() -> None: exporter: Final = InMemorySpanExporter() provider: Final = TracerProvider() diff --git a/tests/unit/integrations/test_shadow_eval_logger.py b/tests/unit/integrations/test_shadow_eval_logger.py index f802dbfdba9..ab8563c7ced 100644 --- a/tests/unit/integrations/test_shadow_eval_logger.py +++ b/tests/unit/integrations/test_shadow_eval_logger.py @@ -1905,6 +1905,40 @@ class TestShadowPipeline: assert "Let me look that up. [tool call] Read({})" in judge_prompt + async def test_a_server_tool_call_reaches_the_judge_ahead_of_client_tool_calls(self): + router: Final = _shadow_reply_router( + { + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + } + ], + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + } + ] + }, + }, + finish_reason="tool_calls", + ) + await self._judged_shadow_row(router) + + judge_prompt: Final = next( + call.kwargs["messages"][-1]["content"] + for call in router.acompletion.call_args_list + if call.kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) != SHADOW_EVAL_ROUTER_CALL_ORIGIN + ) + + shadow_text: Final = '[tool call] web_search({"query": "persimmon"}), [tool call] get_weather({"city": "SF"})' + assert f":\n{shadow_text}\n\n" in judge_prompt, judge_prompt + async def test_an_empty_shadow_reply_names_the_finish_reason_and_the_routed_model(self): """A reply that really carried no text is diagnosable only if the row says what the arm was doing when it produced none: a truncated turn and a model that answers diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 390da500c32..9d3c19b7460 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -11,6 +11,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_MEDIA_BOUNDARY, TOOL_RESULT_MEDIA_PLACEHOLDER, add_system_prompt_to_messages, + get_content_from_model_response, get_semantic_cache_prompt_from_messages, get_str_from_messages, handle_any_messages_to_chat_completion_str_messages_conversion, @@ -21,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( system_messages_first, ) from litellm.types.llms.openai import ChatCompletionAudioObject, ChatCompletionFileObject, ChatCompletionImageObject +from litellm.types.utils import ChatCompletionMessageToolCall, Choices, Function, Message, ModelResponse _ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' _ARTIFACT_DATA_ID_PATTERN: Final = r"^(?!\.\.?(?:\/|$))[A-Za-z0-9_\-.~:@+]{1,200}$" @@ -983,3 +985,30 @@ def test_content_list_flattening_returns_new_messages_and_leaves_the_callers_unt assert flattened == [{"role": "user", "content": "Hello world"}, {"role": "assistant", "content": "Hi"}] assert messages == messages_before + + +def test_get_content_from_model_response_includes_server_tool_calls_before_client_calls() -> None: + server_call: Final = ChatCompletionMessageToolCall( + id="srvtoolu_01", type="function", function=Function(name="web_search", arguments='{"query": "persimmon"}') + ) + client_call: Final = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="get_weather", arguments='{"city": "SF"}') + ) + response: Final = ModelResponse( + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content="Looking it up.", + tool_calls=[client_call.model_dump()], + provider_specific_fields={"server_tool_calls": [server_call.model_dump()]}, + ), + ) + ], + ) + + assert get_content_from_model_response(response) == ( + "Looking it up." + server_call.model_dump_json() + client_call.model_dump_json() + ) diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 491ea5315ff..8e4bed02d7e 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -11,7 +11,7 @@ import importlib import json import os import threading -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Iterable, Mapping from datetime import datetime from types import MappingProxyType, SimpleNamespace from typing import Final, Optional, Union @@ -20,10 +20,11 @@ from unittest.mock import patch import httpx import pytest import pytest_asyncio -from pydantic import JsonValue +from openai.types.responses.response_function_web_search import ResponseFunctionWebSearch +from pydantic import JsonValue, TypeAdapter import litellm -from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE +from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE, REDACTED_BY_LITELLM from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.redact_messages import ( @@ -37,6 +38,7 @@ from litellm.litellm_core_utils.redact_messages import ( ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.responses.main import mock_responses_api_response +from litellm.types.responses.main import build_web_search_call from litellm.types.router import BaselineRouteStamp from litellm.types.utils import ( ModelResponse, @@ -236,6 +238,116 @@ def test_streaming_redaction_redacts_delta_provider_fields_without_touching_the_ assert fields == _provider_fields() +_SECRET_URL: Final = f"https://example.com/?q={SECRET}" +_WEB_SEARCH_CALL: Final = TypeAdapter(ResponseFunctionWebSearch) + + +def _claude_fetch_call() -> ResponseFunctionWebSearch: + return build_web_search_call( + "srvtoolu_01Fetch", {"url": _SECRET_URL}, {"content": []}, status="in_progress", tool_name="web_fetch" + ) + + +def _web_call_summaries( + calls: Mapping[str, object] | Iterable[object], +) -> list[tuple[str, str | None, dict[str, object]]]: + items: Final = calls.values() if isinstance(calls, dict) else calls + return [ + (call.id, call.status, call.action.model_dump(exclude_none=True)) + for call in (_WEB_SEARCH_CALL.validate_python(item) for item in items) + ] + + +@pytest.mark.parametrize("shape", ["object", "dict"]) +def test_redaction_hides_claude_server_call_inputs_but_keeps_the_calls(shape: str) -> None: + response: Final = ModelResponse( + model="claude-opus-5-5", + choices=[ + { + "message": { + "role": "assistant", + "content": None, + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01Fetch", + "type": "function", + "function": {"name": "web_fetch", "arguments": json.dumps({"url": _SECRET_URL})}, + } + ], + "web_search_calls": [ + _claude_fetch_call(), + build_web_search_call( + "srvtoolu_01Search", {"query": SECRET}, {"content": []}, status="in_progress" + ), + ], + }, + } + } + ], + ) + result: Final = response if shape == "object" else response.model_dump() + + redacted: Final = perform_redaction({}, result) + + message: Final = ( + redacted.choices[0].message if shape == "object" else SimpleNamespace(**redacted["choices"][0]["message"]) + ) + fields: Final = message.provider_specific_fields + assert fields["server_tool_calls"][0]["id"] == "srvtoolu_01Fetch" + assert fields["server_tool_calls"][0]["function"] == {"name": "web_fetch", "arguments": REDACTED_BY_LITELLM} + assert _web_call_summaries(fields["web_search_calls"]) == [ + ("ws_srvtoolu_01Fetch", "in_progress", {"type": "open_page", "url": REDACTED_BY_LITELLM}), + ("ws_srvtoolu_01Search", "in_progress", {"type": "search", "query": REDACTED_BY_LITELLM}), + ] + assert SECRET not in str(redacted), redacted + assert SECRET in str(result) + + +@pytest.mark.parametrize("key", ["complete_streaming_response", "async_complete_streaming_response"]) +def test_streaming_redaction_hides_the_hosted_copy_of_a_claude_web_fetch(key: str) -> None: + response: Final = ModelResponseStream( + choices=[ + litellm.utils.StreamingChoices( + delta=litellm.utils.Delta( + provider_specific_fields={"web_search_calls": {"srvtoolu_01Fetch": _claude_fetch_call()}} + ) + ) + ] + ) + + perform_redaction({"stream": True, key: response}, None) + + assert _web_call_summaries(response.choices[0].delta.provider_specific_fields["web_search_calls"]) == [ + ("ws_srvtoolu_01Fetch", "in_progress", {"type": "open_page", "url": REDACTED_BY_LITELLM}) + ] + + +@pytest.mark.parametrize("shape", ["object", "dict"]) +def test_redaction_hides_web_search_call_actions_in_responses_api_output(shape: str) -> None: + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-opus-5-5", + object="response", + status="completed", + output=[ + build_web_search_call("srvtoolu_01Fetch", {"url": _SECRET_URL}, {"content": []}, tool_name="web_fetch"), + build_web_search_call("srvtoolu_01Search", {"query": SECRET}, {"content": []}), + ], + ) + result: Final = response if shape == "object" else response.model_dump() + + redacted: Final = perform_redaction({}, result) + + output: Final = redacted.output if shape == "object" else redacted["output"] + assert _web_call_summaries(output) == [ + ("ws_srvtoolu_01Fetch", "completed", {"type": "open_page", "url": REDACTED_BY_LITELLM}), + ("ws_srvtoolu_01Search", "completed", {"type": "search", "query": REDACTED_BY_LITELLM}), + ] + assert SECRET in str(result) + + @pytest.mark.parametrize("shape", ["object", "dict"]) def test_redaction_hides_a_compaction_summary_in_responses_api_output(shape: str) -> None: response: Final = ResponsesAPIResponse( diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index c2ee58b67ed..8affe2e136e 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -594,6 +594,57 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert '"id": "srvtoolu_1"' in raw and '"id": "toolu_1"' in raw assert "persim" not in raw and '"query"' not in raw + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_use_rewrite_lands_on_its_own_block_when_server_tool_use_follows(self): + class MaskFruitLookup(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + for tool_call in inputs.get("tool_calls", []): + if tool_call.function.name == "lookup_fruit": + tool_call.function.arguments = '{"fruit": "[MASKED]"}' + return inputs + + handler: Final = AnthropicMessagesHandler() + tool_use: Final = self._ended_tool_use_sse_chunks() + server_tool_use: Final = [ + ( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}, + }, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 1}), + ] + chunks: Final = ( + tool_use[:-2] + + [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in server_tool_use] + + tool_use[-2:] + ) + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MaskFruitLookup(guardrail_name="test"), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert self._partial_jsons(chunks) == ['{"fruit": "[MASKED]"}', "", "", '{"query": "fruit"}'] + @pytest.mark.asyncio async def test_deliver_ended_stream_tool_use_rewrite_on_an_unindexed_tool_block_fails_closed(self): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index e81ee2a4e02..2fb4eb9931c 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -17,7 +17,9 @@ from litellm._uuid import uuid from litellm.constants import NON_LLM_CONNECTION_TIMEOUT, RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.prompt_templates import image_handling from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt +from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call, make_sync_call +from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params from litellm.llms.base_llm.base_utils import response_format_schema, unwrap_response_format_envelope from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.openai import ( @@ -25,7 +27,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallFunctionChunk, ) from litellm.types.responses.main import OutputCodeInterpreterCall -from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponseStream +from litellm.types.utils import ChatCompletionDeltaToolCall, ChatCompletionMessageToolCall, Choices, ModelResponseStream from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE @@ -1966,19 +1968,113 @@ def test_web_search_tool_result_no_extra_tool_calls(): for tc in parsed.choices[0].delta.tool_calls: tool_calls_emitted.append(tc) - # Should have exactly 2 tool calls: - # 1. From content_block_start (server_tool_use) with id and name - # 2. From content_block_delta with the actual query - assert len(tool_calls_emitted) == 2, f"Expected 2 tool calls, got {len(tool_calls_emitted)}" + assert tool_calls_emitted == [] - # First tool call should have the id and name - assert tool_calls_emitted[0]["id"] == "srvtoolu_01ABC123" - assert tool_calls_emitted[0]["function"]["name"] == "web_search" - # Second tool call should have the query arguments - assert tool_calls_emitted[1]["function"]["arguments"] == '{"query": "test"}' +def _server_tool_input_delta(fragment: str) -> dict[str, object]: + return {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": fragment}} - # The {} chunks from web_search_tool_result should NOT have been emitted as tool calls + +def _server_search_then_answer(*partial_json: str) -> tuple[dict[str, object], ...]: + return ( + { + "type": "message_start", + "message": {"id": "msg_01", "model": "claude-opus-5-5", "usage": {"input_tokens": 9}}, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "server_tool_use", "id": "srvtoolu_01Web", "name": "web_search", "input": {}}, + }, + *(_server_tool_input_delta(fragment) for fragment in partial_json), + {"type": "content_block_stop", "index": 0}, + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_01Web", "content": []}, + }, + {"type": "content_block_stop", "index": 1}, + {"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 2, "delta": {"type": "text_delta", "text": "Paris is in France."}}, + {"type": "content_block_stop", "index": 2}, + ) + + +_SERVER_SEARCH_THEN_ANSWER: Final = _server_search_then_answer('{"query": "Paris"}') +_CLIENT_TOOL_CALL: Final[tuple[dict[str, object], ...]] = ( + { + "type": "content_block_start", + "index": 3, + "content_block": {"type": "tool_use", "id": "toolu_01Weather", "name": "get_weather", "input": {}}, + }, + { + "type": "content_block_delta", + "index": 3, + "delta": {"type": "input_json_delta", "partial_json": '{"city": "Paris"}'}, + }, + {"type": "content_block_stop", "index": 3}, +) + + +def _message_delta(stop_reason: str) -> dict[str, object]: + return {"type": "message_delta", "delta": {"stop_reason": stop_reason}, "usage": {"output_tokens": 20}} + + +_PARIS_SEARCH: Final = ("srvtoolu_01Web", "web_search", '{"query": "Paris"}') + + +@pytest.mark.parametrize( + ("events", "expected_tool_calls", "expected_finish_reason", "expected_server_tool_call"), + [ + pytest.param( + (*_SERVER_SEARCH_THEN_ANSWER, _message_delta("end_turn")), [], "stop", _PARIS_SEARCH, id="server_tool_only" + ), + pytest.param( + (*_SERVER_SEARCH_THEN_ANSWER, *_CLIENT_TOOL_CALL, _message_delta("tool_use")), + [("toolu_01Weather", "get_weather", '{"city": "Paris"}')], + "tool_calls", + _PARIS_SEARCH, + id="server_tool_then_client_tool", + ), + pytest.param( + (*_server_search_then_answer('{"que', 'ry": "Pa', 'ris"}'), _message_delta("end_turn")), + [], + "stop", + _PARIS_SEARCH, + id="server_tool_input_split_across_deltas", + ), + pytest.param( + (*_server_search_then_answer(), _message_delta("end_turn")), + [], + "stop", + ("srvtoolu_01Web", "web_search", "{}"), + id="server_tool_with_empty_input", + ), + ], +) +def test_streamed_server_tool_use_is_not_returned_as_a_client_tool_call( + events: tuple[dict[str, object], ...], + expected_tool_calls: list[tuple[str, str, str]], + expected_finish_reason: str, + expected_server_tool_call: tuple[str, str, str], +): + choice: Final = _assemble_stream(events).choices[0] + assert isinstance(choice, Choices) + + assert (choice.message.tool_calls is None) == (not expected_tool_calls) + assert [ + (call.id, call.function.name, call.function.arguments) + for call in choice.message.tool_calls or () + if isinstance(call, ChatCompletionMessageToolCall) + ] == expected_tool_calls + assert choice.finish_reason == expected_finish_reason + assert choice.message.content == "Paris is in France." + assert [ + (call["id"], call["function"].get("name"), call["function"].get("arguments")) + for call in anthropic_server_tool_call_params(choice.message) + ] == [expected_server_tool_call] + web_search_results: Final = message_field(choice.message.provider_specific_fields, "web_search_results") + assert [message_field(result, "tool_use_id") for result in parts_of(web_search_results)] == ["srvtoolu_01Web"] def test_current_content_block_type_tracking(): diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 95bc8af7270..ae72af200d4 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -2,9 +2,9 @@ import copy import itertools import json import logging -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from itertools import chain -from typing import Final +from typing import Final, cast from unittest.mock import MagicMock, patch import httpx @@ -25,12 +25,18 @@ from litellm.constants import ( RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION, ) from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt +from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of from litellm.litellm_core_utils.prompt_templates.reasoning_utils import ( encrypted_reasoning_signature, ) from litellm.llms.anthropic.chat import ModelResponseIterator +from litellm.llms.anthropic.common_utils import ( + AnthropicError, + anthropic_server_tool_call_params, + process_anthropic_headers, + validated_object_mapping, +) from litellm.llms.anthropic.chat.transformation import AnthropicConfig -from litellm.llms.anthropic.common_utils import AnthropicError, process_anthropic_headers from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, ) @@ -43,10 +49,14 @@ from litellm.llms.bedrock.claude_platform.transformation import BedrockClaudePla from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig, ) -from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_TOOL_SEARCH_TOOL_TYPES +from litellm.types.llms.anthropic import ( + ANTHROPIC_BETA_HEADER_VALUES, + ANTHROPIC_TOOL_SEARCH_TOOL_TYPES, + AllAnthropicPassThroughMessageValues, +) from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse -from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk -from litellm.types.utils import ChatCompletionToolCallChunk, ModelResponse, ServerToolUse, Usage +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallFunctionChunk +from litellm.types.utils import ChatCompletionToolCallChunk, Choices, Message, ModelResponse, ServerToolUse, Usage from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE @@ -3624,21 +3634,17 @@ def test_code_execution_tool_results_extraction(): prefix_prompt=None, ) - # Verify tool calls are present - assert transformed_response.choices[0].message.tool_calls is not None - assert len(transformed_response.choices[0].message.tool_calls) == 2 + assert transformed_response.choices[0].message.tool_calls is None - # Verify first tool call - assert transformed_response.choices[0].message.tool_calls[0].id == "srvtoolu_01ABC" - assert transformed_response.choices[0].message.tool_calls[0].function.name == "bash_code_execution" - - # Verify second tool call - assert transformed_response.choices[0].message.tool_calls[1].id == "srvtoolu_01DEF" - assert transformed_response.choices[0].message.tool_calls[1].function.name == "text_editor_code_execution" - - # Verify tool results are in provider_specific_fields provider_fields = transformed_response.choices[0].message.provider_specific_fields assert provider_fields is not None + server_tool_calls: Final = anthropic_server_tool_call_params(transformed_response.choices[0].message) + assert [(call["id"], call["function"].get("name")) for call in server_tool_calls] == [ + ("srvtoolu_01ABC", "bash_code_execution"), + ("srvtoolu_01DEF", "text_editor_code_execution"), + ] + + # Verify tool results are in provider_specific_fields assert "tool_results" in provider_fields assert provider_fields["tool_results"] is not None assert len(provider_fields["tool_results"]) == 2 @@ -3661,6 +3667,45 @@ def test_code_execution_tool_results_extraction(): assert "Done!" in transformed_response.choices[0].message.content +def test_server_and_client_tool_calls_split_by_block_type_not_id_prefix(): + completion_response: Final = { + "id": "msg_split", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [ + {"type": "server_tool_use", "id": "toolu_server", "name": "web_search", "input": {"query": "persimmon"}}, + {"type": "tool_use", "id": "srvtoolu_client", "name": "get_weather", "input": {"city": "Paris"}}, + ], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + raw_response: Final = MagicMock(spec=httpx.Response) + raw_response.headers = {} + + message: Final = ( + AnthropicConfig() + .transform_parsed_response( + completion_response=completion_response, + raw_response=raw_response, + model_response=ModelResponse(), + json_mode=False, + prefix_prompt=None, + ) + .choices[0] + .message + ) + + assert [ + (call["id"], call["function"].get("name"), call["function"].get("arguments")) + for call in anthropic_server_tool_call_params(message) + ] == [("toolu_server", "web_search", '{"query": "persimmon"}')] + assert [ + (tool_call.id, tool_call.function.name, tool_call.function.arguments) for tool_call in message.tool_calls or () + ] == [("srvtoolu_client", "get_weather", '{"city": "Paris"}')] + + def test_code_execution_tool_results_in_hidden_params(): """ Test that tool_results reaches _hidden_params so the Responses API adapter @@ -3748,7 +3793,7 @@ def _tool_search_turn(tool_references: list[str], last_block: dict[str, object]) ] -def _parsed_assistant_message(content: list[dict[str, object]], stop_reason: str): +def _parsed_choice(content: list[dict[str, object]], stop_reason: str) -> Choices: anthropic_response: Final = { "id": "msg_01XYZ", "type": "message", @@ -3763,7 +3808,7 @@ def _parsed_assistant_message(content: list[dict[str, object]], stop_reason: str raw_response.json.return_value = anthropic_response raw_response.status_code = 200 raw_response.headers = {} - return ( + choice: Final = ( AnthropicConfig() .transform_parsed_response( completion_response=anthropic_response, @@ -3773,14 +3818,19 @@ def _parsed_assistant_message(content: list[dict[str, object]], stop_reason: str prefix_prompt=None, ) .choices[0] - .message ) + assert isinstance(choice, Choices) + return choice -_WEATHER_CALL: Final = {"type": "tool_use", "id": "toolu_01Weather", "name": "get_weather", "input": {"city": "Paris"}} +def _parsed_assistant_message(content: list[dict[str, object]], stop_reason: str) -> Message: + return _parsed_choice(content, stop_reason).message -def _web_search_turn(last_block: dict[str, object]) -> list[dict[str, object]]: +_WEATHER_TOOL_USE: Final = {"type": "tool_use", "id": "toolu_01Weather", "name": "get_weather", "input": {"city": "Paris"}} + + +def _web_search_turn(last_block: Mapping[str, object]) -> list[dict[str, object]]: return [ {"type": "server_tool_use", "id": "srvtoolu_01Web", "name": "web_search", "input": {"query": "Paris facts"}}, { @@ -3796,7 +3846,7 @@ def _web_search_turn(last_block: dict[str, object]) -> list[dict[str, object]]: } ], }, - last_block, + dict(last_block), ] @@ -3811,7 +3861,7 @@ def _web_search_turn(last_block: dict[str, object]) -> list[dict[str, object]]: ) @pytest.mark.parametrize( "content", - [_tool_search_turn(["get_weather"], _WEATHER_CALL), _web_search_turn(_WEATHER_CALL)], + [_tool_search_turn(["get_weather"], _WEATHER_TOOL_USE), _web_search_turn(_WEATHER_TOOL_USE)], ids=["tool_search", "web_search"], ) def test_server_tool_result_round_trips_right_after_its_call( @@ -3843,8 +3893,148 @@ def test_server_tool_result_round_trips_right_after_its_call( assert [block["tool_use_id"] for block in replayed[2]["content"]] == ["toolu_01Weather"] -def _block_ids(message: dict[str, object], block_type: str, id_key: str) -> list[object]: - return sorted(block[id_key] for block in message["content"] if block["type"] == block_type) +_FETCH_STILL_WAITING_ON_ITS_RESULT: Final = { + "type": "server_tool_use", + "id": "srvtoolu_01Fetch", + "name": "web_fetch", + "input": {"url": "https://example.com/"}, +} + + +@pytest.mark.parametrize("modify_params", [False, True]) +def test_server_call_still_waiting_on_its_result_is_replayed_as_a_server_call( + modify_params: bool, monkeypatch: pytest.MonkeyPatch +): + content: Final = [_FETCH_STILL_WAITING_ON_ITS_RESULT, _WEATHER_TOOL_USE] + message: Final = _parsed_assistant_message(content, "tool_use") + monkeypatch.setattr(litellm, "modify_params", modify_params) + + replayed: Final = _replayed( + {"role": "user", "content": "Fetch example.com, then tell me the weather in Paris"}, + message.model_dump(), + {"role": "tool", "tool_call_id": "toolu_01Weather", "content": "18C"}, + ) + + assert replayed[1]["content"] == content + assert [turn["role"] for turn in replayed] == ["user", "assistant", "user"] + assert [block["tool_use_id"] for block in replayed[2]["content"]] == ["toolu_01Weather"] + + +_FETCH_AND_WEATHER_THINKING: Final = { + "type": "thinking", + "thinking": "Fetch the page, then check the weather.", + "signature": "sig_fetch_weather", +} +_CHECKING_BOTH: Final = {"type": "text", "text": "Checking both."} + + +@pytest.mark.parametrize("thinking_inline", [False, True], ids=["thinking_blocks", "inline_thinking"]) +def test_tool_use_turn_waiting_on_a_server_call_keeps_its_thinking_and_both_calls(thinking_inline: bool): + parsed: Final = _parsed_assistant_message( + [_FETCH_AND_WEATHER_THINKING, _CHECKING_BOTH, _FETCH_STILL_WAITING_ON_ITS_RESULT, _WEATHER_TOOL_USE], "tool_use" + ).model_dump() + message: Final = ( + {**parsed, "thinking_blocks": None, "content": [_FETCH_AND_WEATHER_THINKING, _CHECKING_BOTH]} + if thinking_inline + else parsed + ) + + replayed: Final = _replayed( + {"role": "user", "content": "Fetch example.com, then tell me the weather in Paris"}, + message, + {"role": "tool", "tool_call_id": "toolu_01Weather", "content": "18C"}, + ) + + assistant_content: Final = replayed[1]["content"] + assert assistant_content[:2] == [_FETCH_AND_WEATHER_THINKING, _CHECKING_BOTH] + assert sorted(assistant_content[2:], key=lambda block: block["id"]) == [ + _FETCH_STILL_WAITING_ON_ITS_RESULT, + _WEATHER_TOOL_USE, + ] + assert [block["tool_use_id"] for block in replayed[2]["content"]] == ["toolu_01Weather"] + + +_PARIS_TEXT: Final = "Paris is the capital of France." +_PARIS_ANSWER: Final[Mapping[str, str]] = {"type": "text", "text": _PARIS_TEXT} + + +@pytest.mark.parametrize( + ("content", "stop_reason", "expected_tool_call_ids", "expected_finish_reason"), + [ + pytest.param(_web_search_turn(_PARIS_ANSWER), "end_turn", [], "stop", id="web_search_then_answer"), + pytest.param( + _web_search_turn(_WEATHER_TOOL_USE), "tool_use", ["toolu_01Weather"], "tool_calls", id="web_search_then_tool" + ), + pytest.param([_WEATHER_TOOL_USE], "tool_use", ["toolu_01Weather"], "tool_calls", id="client_tool_only"), + ], +) +def test_server_tool_use_is_not_returned_as_a_client_tool_call( + content: list[dict[str, object]], + stop_reason: str, + expected_tool_call_ids: list[str], + expected_finish_reason: str, +): + choice: Final = _parsed_choice(content, stop_reason) + + assert (choice.message.tool_calls is None) == (not expected_tool_call_ids) + assert [call.id for call in choice.message.tool_calls or ()] == expected_tool_call_ids + assert choice.finish_reason == expected_finish_reason + assert [ + (call["id"], call["function"].get("name"), call["function"].get("arguments")) + for call in anthropic_server_tool_call_params(choice.message) + ] == [ + (block["id"], block["name"], json.dumps(block["input"])) + for block in content + if block["type"] == "server_tool_use" + ] + if not expected_tool_call_ids: + assert choice.message.content == _PARIS_TEXT + + +def _replayed(*messages: Mapping[str, object]) -> list[AllAnthropicPassThroughMessageValues]: + return anthropic_messages_pt( + messages=cast(list[AllMessageValues], list(messages)), # cast-ok: replays history the way clients send it back + model="claude-opus-5-5", + llm_provider="anthropic", + ) + + +def test_web_search_answer_round_trips_without_any_client_tool_call(): + message: Final = _parsed_assistant_message(_web_search_turn(_PARIS_ANSWER), "end_turn") + + replayed: Final = _replayed( + {"role": "user", "content": "Tell me about Paris"}, + message.model_dump(), + {"role": "user", "content": "And Lyon?"}, + ) + + assert _block_ids(replayed[1], "server_tool_use", "id") == ["srvtoolu_01Web"] + assert _block_ids(replayed[1], "web_search_tool_result", "tool_use_id") == ["srvtoolu_01Web"] + assert [ + message_field(block, "text") + for block in parts_of(replayed[1]["content"]) + if message_field(block, "type") == "text" + ] == [_PARIS_TEXT] + + +def _block_ids(message: Mapping[str, object], block_type: str, id_key: str) -> list[str]: + return sorted( + block_id + for block in parts_of(message["content"]) + if message_field(block, "type") == block_type and isinstance(block_id := message_field(block, id_key), str) + ) + + +def _as_returned_before_server_tool_calls_moved(message: Message) -> Mapping[str, object]: + dumped: Final[Mapping[str, object]] = message.model_dump() + provider_fields: Final = validated_object_mapping(dumped.get("provider_specific_fields")) or {} + return { + **dumped, + "tool_calls": [*anthropic_server_tool_call_params(message), *parts_of(dumped.get("tool_calls"))], + "provider_specific_fields": { + key: value for key, value in provider_fields.items() if key != "server_tool_calls" + }, + } _SEARCH_THAT_FOUND_NOTHING: Final = _tool_search_turn( @@ -3879,21 +4069,33 @@ def test_replayed_search_that_found_nothing_answers_every_tool_use_and_ends_on_a server_tool_use_ids: list[str], monkeypatch: pytest.MonkeyPatch, ): - from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt - message: Final = _parsed_assistant_message(_SEARCH_THAT_FOUND_NOTHING, "end_turn") monkeypatch.setattr(litellm, "modify_params", modify_params) - replayed: Final = anthropic_messages_pt( - messages=[{"role": "user", "content": "Find a tool and use it"}, message.model_dump(), *follow_ups], - model="claude-opus-5-5", - llm_provider="anthropic", + replayed: Final = _replayed( + {"role": "user", "content": "Find a tool and use it"}, + _as_returned_before_server_tool_calls_moved(message), + *follow_ups, ) assert [turn["role"] for turn in replayed] == ["user", "assistant", "user"] assert _block_ids(replayed[2], "tool_result", "tool_use_id") == _block_ids(replayed[1], "tool_use", "id") assert _block_ids(replayed[1], "server_tool_use", "id") == server_tool_use_ids +def test_replayed_search_that_found_nothing_keeps_its_result_when_the_call_never_reached_the_client( + monkeypatch: pytest.MonkeyPatch, +): + message: Final = _parsed_assistant_message(_SEARCH_THAT_FOUND_NOTHING, "end_turn") + monkeypatch.setattr(litellm, "modify_params", True) + + replayed: Final = _replayed({"role": "user", "content": "Find a tool and use it"}, message.model_dump()) + + assert [turn["role"] for turn in replayed] == ["user", "assistant"] + assert _block_ids(replayed[1], "server_tool_use", "id") == ["srvtoolu_01Search"] + assert _block_ids(replayed[1], "tool_search_tool_result", "tool_use_id") == ["srvtoolu_01Search"] + assert _block_ids(replayed[1], "tool_use", "id") == [] + + _CODE_EXECUTION_BLOCKS: Final = [ { "type": "server_tool_use", diff --git a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py index fb4f19a862c..420de9cb4f5 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py @@ -80,6 +80,59 @@ def test_native_compaction_response_roundtrip(text: str | None, count: int, expe assert result["stop_reason"] == expected_stop +_SERVER_AND_CLIENT_TOOL_TURN: Final[tuple[dict[str, object], ...]] = ( + { + "type": "server_tool_use", + "id": "srvtoolu_01Fetch", + "name": "web_fetch", + "input": {"url": "https://example.com/"}, + }, + { + "type": "web_fetch_tool_result", + "tool_use_id": "srvtoolu_01Fetch", + "content": {"type": "web_fetch_result", "url": "https://example.com/"}, + }, + {"type": "server_tool_use", "id": "srvtoolu_02Search", "name": "web_search", "input": {"query": "Paris"}}, + {"type": "text", "text": "The page is titled Example."}, + {"type": "tool_use", "id": "toolu_01Weather", "name": "get_weather", "input": {"city": "Paris"}}, +) + + +def test_claude_server_tool_blocks_survive_the_round_trip_through_chat_completions() -> None: + from unittest.mock import MagicMock + + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + anthropic_response: Final = { + "id": "msg_01Fetch", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": list(_SERVER_AND_CLIENT_TOOL_TURN), + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 100, "output_tokens": 50}, + } + raw_response: Final = MagicMock(spec=httpx.Response) + raw_response.json.return_value = anthropic_response + raw_response.status_code = 200 + raw_response.headers = {} + chat_response: Final = AnthropicConfig().transform_parsed_response( + completion_response=anthropic_response, + raw_response=raw_response, + model_response=ModelResponse(), + json_mode=False, + prefix_prompt=None, + ) + + result: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(chat_response) + + assert [(block["type"], block.get("id") or block.get("tool_use_id")) for block in result["content"]] == [ + (block["type"], block.get("id") or block.get("tool_use_id")) for block in _SERVER_AND_CLIENT_TOOL_TURN + ] + assert result["stop_reason"] == "tool_use" + + def test_translate_chat_refusal_to_anthropic_response(): response = ModelResponse( id="chatcmpl-refusal", diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index c97585e02c2..14cabb683f1 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -14,6 +14,7 @@ import itertools import json import logging import os +import subprocess import sys import threading from collections.abc import Mapping @@ -4869,3 +4870,40 @@ def test_a_register_model_opt_out_for_one_claude_model_leaves_every_other_routed litellm.register_model({f"anthropic/{opted_out}": {"supports_prompt_caching": False}}) assert supports_anthropic_cache_control(untouched, "anthropic", {"id": "plain-id"}) is True + + +_MASKED_HOSTED_WEB_CALL_PROBE: Final = """ +import json + +from litellm.llms.anthropic.common_utils import write_request_hosted_web_call_arguments + +message = { + "role": "assistant", + "provider_specific_fields": { + "web_search_calls": [ + { + "id": "ws_srvtoolu_1", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "deploy token 12345", "queries": ["deploy token 12345"]}, + } + ] + }, +} +write_request_hosted_web_call_arguments(message, {"srvtoolu_1": json.dumps({"query": "deploy token [REDACTED]"})}) +print(json.dumps(message["provider_specific_fields"]["web_search_calls"])) +""" + + +def test_a_masked_replayed_hosted_web_call_serializes_before_its_openai_model_was_ever_built() -> None: + result: Final = subprocess.run( + (sys.executable, "-I", "-c", _MASKED_HOSTED_WEB_CALL_PROBE), + capture_output=True, + text=True, + check=True, + timeout=120, + ) + hosted: Final = json.loads(result.stdout.strip().splitlines()[-1]) + assert [call["action"] for call in hosted] == [ + {"type": "search", "query": "deploy token [REDACTED]", "queries": ["deploy token [REDACTED]"]} + ] 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 5d71cbe9afc..54832bb35cc 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 @@ -6,17 +6,21 @@ with guardrail transformations, including tool calls. """ import json -from collections.abc import Mapping -from typing import Any, Literal, Optional +from collections.abc import Iterable, Mapping +from typing import Any, Final, Literal, Optional import pytest +from openai.types.responses import ResponseFunctionWebSearch from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.openai.chat.guardrail_translation.handler import ( OpenAIChatCompletionsHandler, ) +from litellm.types.llms.openai import ChatCompletionToolCallChunk +from litellm.types.responses.main import build_web_search_call from litellm.types.utils import ( ChatCompletionMessageToolCall, Choices, @@ -144,6 +148,55 @@ class MockMisalignedToolCallGuardrail(CustomGuardrail): ) +def _server_tool_call(call_id: str, name: str, arguments: Mapping[str, object]) -> dict[str, object]: + return {"id": call_id, "type": "function", "function": {"name": name, "arguments": json.dumps(arguments)}} + + +def _tool_call_ids_names_and_arguments( + tool_calls: Iterable[ChatCompletionToolCallChunk], +) -> list[tuple[str | None, str | None, str]]: + return [ + (tool_call["id"], tool_call["function"]["name"], tool_call["function"]["arguments"]) + for tool_call in tool_calls + ] + + +def _hosted_web_call_inputs(web_search_calls: Iterable[ResponseFunctionWebSearch]) -> list[tuple[str, str, Optional[str]]]: + return [ + (call.id, call.status, call.action.query if call.action.type == "search" else call.action.url) + for call in web_search_calls + ] + + +_SEARCH_INPUT: Final = {"query": "jane@example.com"} +_FETCH_INPUT: Final = {"url": "https://example.com/?email=jane@example.com"} + + +class QueryRedactingGuardrail(CustomGuardrail): + """Returns copied tool calls with any ``query`` or ``url`` argument redacted, leaving other calls as they came.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + def redacted(tool_call: ChatCompletionToolCallChunk) -> ChatCompletionToolCallChunk: + arguments: Final[dict[str, object]] = json.loads(tool_call["function"]["arguments"]) + redacted_arguments: Final = { + key: "[REDACTED]" if key in ("query", "url") else value for key, value in arguments.items() + } + if redacted_arguments == arguments: + return tool_call + return {**tool_call, "function": {**tool_call["function"], "arguments": json.dumps(redacted_arguments)}} + + return GenericGuardrailAPIInputs( + texts=inputs.get("texts", []), + tool_calls=[redacted(tool_call) for tool_call in inputs.get("tool_calls", [])], + ) + + class TestOpenAIChatCompletionsHandlerToolsInput: """Test input processing with tools (function definitions)""" @@ -906,6 +959,192 @@ class TestOpenAIChatCompletionsHandlerToolCallsOutput: assert tool_calls[0].function.arguments == first_args assert tool_calls[1].function.arguments == second_args + @pytest.mark.asyncio + async def test_input_guardrail_rewrite_of_a_replayed_server_tool_call_lands_on_the_request(self): + handler: Final = OpenAIChatCompletionsHandler() + data: Final = { + "messages": [ + {"role": "user", "content": "look up jane and the weather"}, + { + "role": "assistant", + "content": None, + "tool_calls": [_server_tool_call("toolu_01Weather", "get_weather", {"city": "Paris"})], + "provider_specific_fields": { + "server_tool_calls": [_server_tool_call("srvtoolu_01Search", "web_search", _SEARCH_INPUT)] + }, + }, + ] + } + + await handler.process_input_messages(data, QueryRedactingGuardrail(guardrail_name="redact")) + + assistant: Final = data["messages"][1] + assert _tool_call_ids_names_and_arguments(assistant["provider_specific_fields"]["server_tool_calls"]) == [ + ("srvtoolu_01Search", "web_search", '{"query": "[REDACTED]"}') + ] + assert _tool_call_ids_names_and_arguments(assistant["tool_calls"]) == [ + ("toolu_01Weather", "get_weather", '{"city": "Paris"}') + ] + + @pytest.mark.asyncio + async def test_input_guardrail_rewrite_of_a_replayed_server_web_call_reaches_its_hosted_copy(self): + handler: Final = OpenAIChatCompletionsHandler() + data: Final = { + "messages": [ + {"role": "user", "content": "look up jane"}, + { + "role": "assistant", + "content": "Searching.", + "provider_specific_fields": { + "server_tool_calls": [ + _server_tool_call("srvtoolu_01Search", "web_search", _SEARCH_INPUT), + _server_tool_call("srvtoolu_01Fetch", "web_fetch", _FETCH_INPUT), + ], + "web_search_calls": [ + build_web_search_call("srvtoolu_01Search", _SEARCH_INPUT, {"content": []}).model_dump(), + build_web_search_call( + "srvtoolu_01Fetch", _FETCH_INPUT, {"content": []}, tool_name="web_fetch" + ).model_dump(), + ], + }, + }, + ] + } + + await handler.process_input_messages(data, QueryRedactingGuardrail(guardrail_name="redact")) + + hosted: Final = json.loads(json.dumps(data["messages"][1]["provider_specific_fields"]["web_search_calls"])) + assert [(call["id"], call["action"].get("query") or call["action"].get("url")) for call in hosted] == [ + ("ws_srvtoolu_01Search", "[REDACTED]"), + ("ws_srvtoolu_01Fetch", "[REDACTED]"), + ] + + @pytest.mark.asyncio + async def test_output_guardrail_sees_a_response_that_only_carries_server_tool_calls(self): + handler: Final = OpenAIChatCompletionsHandler() + guardrail: Final = TestOpenAIChatCompletionsHandlerStreamingOutput._recording_guardrail() + response: Final = ModelResponse( + model="claude-sonnet-4-5", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content=None, + role="assistant", + provider_specific_fields={ + "server_tool_calls": [_server_tool_call("srvtoolu_01Search", "web_search", {"query": "x"})] + }, + ), + ) + ], + ) + + await handler.process_output_response(response, guardrail) + + assert _tool_call_ids_names_and_arguments(guardrail.seen_inputs[-1]["tool_calls"]) == [ + ("srvtoolu_01Search", "web_search", '{"query": "x"}') + ] + + @pytest.mark.asyncio + async def test_output_guardrail_rewrite_of_a_server_tool_call_lands_in_server_tool_calls_only(self): + handler: Final = OpenAIChatCompletionsHandler() + weather_args: Final = json.dumps({"city": "Paris", "query": "forecast"}) + response: Final = ModelResponse( + model="claude-sonnet-4-5", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content="Looking it up", + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_1", + type="function", + function=Function(name="get_weather", arguments=json.dumps({"city": "Paris"})), + ) + ], + provider_specific_fields={ + "server_tool_calls": [ + _server_tool_call("srvtoolu_01Search", "web_search", {"query": "jane@example.com"}) + ] + }, + ), + ), + Choices( + finish_reason="tool_calls", + index=1, + message=Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_2", + type="function", + function=Function(name="get_weather", arguments=weather_args), + ) + ], + ), + ), + ], + ) + + await handler.process_output_response(response, QueryRedactingGuardrail(guardrail_name="redact")) + + first, second = response.choices + assert _tool_call_ids_names_and_arguments(first.message.provider_specific_fields["server_tool_calls"]) == [ + ("srvtoolu_01Search", "web_search", '{"query": "[REDACTED]"}') + ] + assert [(call.id, call.function.arguments) for call in first.message.tool_calls] == [ + ("call_1", '{"city": "Paris"}') + ] + assert [(call.id, call.function.arguments) for call in second.message.tool_calls] == [ + ("call_2", '{"city": "Paris", "query": "[REDACTED]"}') + ] + assert not (second.message.provider_specific_fields or {}).get("server_tool_calls") + + @pytest.mark.asyncio + async def test_output_guardrail_rewrite_of_a_server_web_call_reaches_its_hosted_copy(self): + handler: Final = OpenAIChatCompletionsHandler() + response: Final = ModelResponse( + model="claude-sonnet-4-5", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content=None, + role="assistant", + provider_specific_fields={ + "server_tool_calls": [ + _server_tool_call("srvtoolu_01Search", "web_search", _SEARCH_INPUT), + _server_tool_call("srvtoolu_01Fetch", "web_fetch", _FETCH_INPUT), + ], + "web_search_calls": [ + build_web_search_call("srvtoolu_01Search", _SEARCH_INPUT, {"content": []}), + build_web_search_call( + "srvtoolu_01Fetch", + _FETCH_INPUT, + {"content": []}, + status="in_progress", + tool_name="web_fetch", + ), + ], + }, + ), + ) + ], + ) + + await handler.process_output_response(response, QueryRedactingGuardrail(guardrail_name="redact")) + + assert _hosted_web_call_inputs(response.choices[0].message.provider_specific_fields["web_search_calls"]) == [ + ("ws_srvtoolu_01Search", "completed", "[REDACTED]"), + ("ws_srvtoolu_01Fetch", "in_progress", "[REDACTED]"), + ] + class MockPassThroughGuardrail(CustomGuardrail): """Mock guardrail that passes through without blocking - for testing streaming fallback behavior""" @@ -1462,6 +1701,136 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert chunks[0].choices[0].delta.content == "hello [MASKED]" assert chunks[1].choices[0].delta.content in (None, "") + @staticmethod + def _ended_server_tool_call_stream_chunks() -> list: + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + def chunk(delta: Delta, finish_reason: Optional[str] = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="claude-sonnet-4-5", + object="chat.completion.chunk", + choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)], + ) + + search: Final = ( + {**_server_tool_call("srvtoolu_01Search", "web_search", {"query": "jane@example.com"}), "index": 0} + ) + return [ + chunk(Delta(content="Searching")), + chunk(Delta(content=None, provider_specific_fields={"server_tool_calls": [search]})), + chunk(Delta(content=None), finish_reason="stop"), + ] + + @staticmethod + def _ended_server_fetch_stream_chunks() -> list: + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + def chunk(delta: Delta, finish_reason: Optional[str] = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="claude-sonnet-4-5", + object="chat.completion.chunk", + choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)], + ) + + def hosted_fetch( + tool_input: Mapping[str, str], + result: Mapping[str, object], + status: Literal["in_progress", "completed"] | None, + ) -> dict[str, ResponseFunctionWebSearch]: + return { + "srvtoolu_01Fetch": build_web_search_call( + "srvtoolu_01Fetch", tool_input, result, status=status, tool_name="web_fetch" + ) + } + + fetch: Final = {**_server_tool_call("srvtoolu_01Fetch", "web_fetch", _FETCH_INPUT), "index": 0} + return [ + chunk( + Delta( + content=None, + provider_specific_fields={"web_search_calls": hosted_fetch({}, {"content": []}, "in_progress")}, + ) + ), + chunk( + Delta( + content=None, + provider_specific_fields={ + "server_tool_calls": [fetch], + "web_search_calls": hosted_fetch(_FETCH_INPUT, {"content": []}, "in_progress"), + }, + ) + ), + chunk( + Delta( + content=None, + provider_specific_fields={ + "web_search_calls": hosted_fetch(_FETCH_INPUT, {"content": {"type": "web_fetch_result"}}, None) + }, + ) + ), + chunk(Delta(content="Fetched"), finish_reason="stop"), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_reach_every_hosted_copy_of_a_server_web_call(self): + handler: Final = OpenAIChatCompletionsHandler() + chunks: Final = self._ended_server_fetch_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=QueryRedactingGuardrail(guardrail_name="redact"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert [ + _hosted_web_call_inputs(chunk.choices[0].delta.provider_specific_fields["web_search_calls"].values()) + for chunk in chunks[:3] + ] == [ + [("ws_srvtoolu_01Fetch", "in_progress", "[REDACTED]")], + [("ws_srvtoolu_01Fetch", "in_progress", "[REDACTED]")], + [("ws_srvtoolu_01Fetch", "completed", "[REDACTED]")], + ] + + @pytest.mark.asyncio + async def test_ended_stream_guardrail_sees_streamed_server_tool_calls(self): + handler: Final = OpenAIChatCompletionsHandler() + guardrail: Final = self._recording_guardrail() + + await handler.process_output_streaming_response( + responses_so_far=self._ended_server_tool_call_stream_chunks(), + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert _tool_call_ids_names_and_arguments(guardrail.seen_inputs[-1]["tool_calls"]) == [ + ("srvtoolu_01Search", "web_search", '{"query": "jane@example.com"}') + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_server_tool_call_arguments_back_into_chunks(self): + handler: Final = OpenAIChatCompletionsHandler() + chunks: Final = self._ended_server_tool_call_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=QueryRedactingGuardrail(guardrail_name="redact"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + delta: Final = chunks[1].choices[0].delta + assert _tool_call_ids_names_and_arguments(delta.provider_specific_fields["server_tool_calls"]) == [ + ("srvtoolu_01Search", "web_search", '{"query": "[REDACTED]"}') + ] + assert delta.tool_calls is None + assert chunks[0].choices[0].delta.content == "Searching" + class TestUndecoratedGuardrailIsRecorded: """LIT-5983 regression: the handler calls apply_guardrail bare, so a custom guardrail @@ -2323,6 +2692,29 @@ class TestStreamingScanKey: assert len(ended_key.tool_calls) == 1 and "run_shell" in ended_key.tool_calls[0] assert ended_key.tool_calls_in_flight is False + def test_server_tool_calls_hold_the_window_and_enter_the_key_once_the_stream_has_ended(self): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler: Final = OpenAIChatCompletionsHandler() + server_chunk: Final = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=None, + provider_specific_fields={ + "server_tool_calls": [_server_tool_call("srvtoolu_01Search", "web_search", {"query": "x"})] + }, + ), + ) + ] + ) + open_key: Final = handler.get_streaming_scan_key([self._chunk("hi"), server_chunk]) + ended_key: Final = handler.get_streaming_scan_key([self._chunk("hi"), server_chunk, self._chunk(None, "stop")]) + assert open_key.tool_calls_in_flight is True + assert len(ended_key.tool_calls) == 1 and "web_search" in ended_key.tool_calls[0] + assert ended_key != open_key + def test_text_after_the_first_choice_finishes_still_changes_the_key(self): handler = OpenAIChatCompletionsHandler() first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)] diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index d8954c4072a..a5bb49f0793 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -20,7 +20,9 @@ import pytest from fastapi import HTTPException from pydantic import BaseModel +from openai.types.responses.response_function_web_search import ActionSearch from openai.types.responses import ( + ResponseFunctionWebSearch, ResponseCustomToolCall, ResponseCustomToolCallInputDeltaEvent, ResponseCustomToolCallInputDoneEvent, @@ -61,6 +63,7 @@ from litellm.types.llms.openai import ResponsesAPIResponse from litellm.responses.litellm_completion_transformation.tool_search import build_tool_search_call from litellm.types.responses.main import ( CustomToolCallOutputItem, + build_web_search_call, GenericResponseOutputItem, OutputText, OutputToolSearchCall, @@ -250,6 +253,30 @@ CUSTOM_TOOL_CALL_ITEM = { } +def _claude_web_calls_and_a_native_search() -> list[ResponseFunctionWebSearch]: + return [ + build_web_search_call("srvtoolu_1", {"query": "persimmon"}, None, "completed"), + build_web_search_call("srvtoolu_2", {"url": "https://persimmon.example"}, None, "in_progress", "web_fetch"), + ResponseFunctionWebSearch( + id="ws_native", + type="web_search_call", + status="completed", + action=ActionSearch(type="search", query="persimmon"), + ), + ] + + +def _web_call_actions(items: list[object]) -> list[dict[str, object]]: + return [ResponseFunctionWebSearch.model_validate(item, from_attributes=True).action.model_dump() for item in items] + + +_MASKED_WEB_CALL_ACTIONS: Final = [ + {"type": "search", "query": "[MASKED]", "queries": ["[MASKED]"], "sources": []}, + {"type": "open_page", "url": "https://[MASKED].example"}, + {"type": "search", "query": "persimmon", "queries": None, "sources": None}, +] + + class TestOpenAIResponsesHandlerDiscovery: """Test that the handler is properly discovered by the guardrail system""" @@ -1823,6 +1850,27 @@ class TestOpenAIResponsesHandlerToolCallExtraction: assert search_call.call_id == "srvtoolu_1" assert "name" not in search_call.model_dump() + @pytest.mark.asyncio + async def test_process_output_response_scans_and_rewrites_claude_web_calls_but_not_native_search(self): + guardrail: Final = RecordingPersimmonMaskingGuardrail(guardrail_name="mask") + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-sonnet-5-5", + object="response", + status="completed", + output=_claude_web_calls_and_a_native_search(), + ) + + result: Final = await OpenAIResponsesHandler().process_output_response(response, guardrail) + + assert guardrail.seen_tool_calls == [ + ("web_search", '{"query": "persimmon"}'), + ("web_fetch", '{"url": "https://persimmon.example"}'), + ] + assert [item.type for item in result.output] == ["web_search_call"] * 3 + assert _web_call_actions(result.output) == _MASKED_WEB_CALL_ACTIONS + @pytest.mark.asyncio async def test_process_output_response_never_returns_original_search_arguments_after_a_non_json_rewrite(self): response = ResponsesAPIResponse( @@ -2758,6 +2806,117 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert completed.response.output[0].arguments == {"query": "[MASKED]"} assert "name" not in done.item.model_dump() + @pytest.mark.asyncio + async def test_deliver_ended_stream_scans_and_rewrites_claude_web_calls_but_not_native_search(self): + from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject + from litellm.types.llms.openai import ( + OutputItemAddedEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ResponsesAPIStreamEvents, + ) + + items: Final = _claude_web_calls_and_a_native_search() + added: Final = [ + OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=index, + item=BaseLiteLLMOpenAIResponseObject(**{**item.model_dump(), "action": None}), + ) + for index, item in enumerate(items) + ] + done: Final = [ + OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=index, + sequence_number=index + 3, + item=BaseLiteLLMOpenAIResponseObject(**item.model_dump()), + ) + for index, item in enumerate(items) + ] + completed: Final = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-sonnet-5-5", + object="response", + status="completed", + output=_claude_web_calls_and_a_native_search(), + ), + ) + guardrail: Final = RecordingPersimmonMaskingGuardrail(guardrail_name="mask") + + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=[*added, *done, completed], + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert sorted(guardrail.seen_tool_calls) == [ + ("web_fetch", '{"url": "https://persimmon.example"}'), + ("web_search", '{"query": "persimmon"}'), + ] + assert [event.item.model_dump()["action"] for event in added] == [None, None, None] + assert _web_call_actions([event.item for event in done]) == _MASKED_WEB_CALL_ACTIONS + assert _web_call_actions(completed.response.output) == _MASKED_WEB_CALL_ACTIONS + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_dict_shaped_claude_web_calls_but_not_native_search(self): + items: Final = [item.model_dump() for item in _claude_web_calls_and_a_native_search()] + done: Final = [ + {"type": "response.output_item.done", "output_index": index, "item": dict(item)} + for index, item in enumerate(items) + ] + completed: Final = { + "type": "response.completed", + "response": { + "id": "resp_1", + "model": "claude-sonnet-5-5", + "status": "completed", + "output": [dict(item) for item in items], + }, + } + guardrail: Final = RecordingPersimmonMaskingGuardrail(guardrail_name="mask") + + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=[*done, completed], + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert sorted(guardrail.seen_tool_calls) == [ + ("web_fetch", '{"url": "https://persimmon.example"}'), + ("web_search", '{"query": "persimmon"}'), + ] + assert _web_call_actions([event["item"] for event in done]) == _MASKED_WEB_CALL_ACTIONS + assert _web_call_actions(completed["response"]["output"]) == _MASKED_WEB_CALL_ACTIONS + + @pytest.mark.asyncio + async def test_streaming_output_item_done_reports_a_claude_web_search_call(self): + from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject + from litellm.types.llms.openai import OutputItemDoneEvent, ResponsesAPIStreamEvents + + done: Final = OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=0, + sequence_number=1, + item=BaseLiteLLMOpenAIResponseObject( + **build_web_search_call("srvtoolu_1", {"query": "persimmon"}, None, "completed").model_dump() + ), + ) + guardrail: Final = RecordingPersimmonMaskingGuardrail(guardrail_name="mask") + + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=[done], + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert guardrail.seen_tool_calls == [("web_search", '{"query": "persimmon"}')] + @pytest.mark.asyncio async def test_deliver_ended_stream_rewrites_a_native_tool_search_call_that_has_no_call_id(self): def search_call(arguments: dict, status: str) -> dict: diff --git a/tests/unit/proxy/db/test_spend_log_tool_index.py b/tests/unit/proxy/db/test_spend_log_tool_index.py index af3268a7cc9..901a9ea8b69 100644 --- a/tests/unit/proxy/db/test_spend_log_tool_index.py +++ b/tests/unit/proxy/db/test_spend_log_tool_index.py @@ -6,7 +6,7 @@ plus the LiteLLM_DailyToolSpend rollup in one transaction. from collections.abc import Awaitable, Callable from types import SimpleNamespace -from typing import Any +from typing import Any, Final from unittest.mock import AsyncMock, MagicMock import httpx @@ -23,6 +23,7 @@ from litellm.proxy.db.spend_log_tool_index import ( ) from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import ModelResponse from tests.unit.proxy.db.fake_prisma_engine import engine_call @@ -204,6 +205,36 @@ class TestResponseToolCallNames: def test_blank_names_are_dropped(self): assert response_tool_call_names(_response_with_tool_calls(" ", "real_tool")) == ("real_tool",) + @pytest.mark.parametrize("as_object", [False, True], ids=["dict", "model_response"]) + def test_server_tool_calls_of_every_choice_are_named_before_client_tool_calls(self, as_object: bool): + def tool_call(call_id: str, name: str) -> dict[str, object]: + return {"id": call_id, "type": "function", "function": {"name": name, "arguments": "{}"}} + + response: dict[str, object] = { + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [tool_call("call_1", "get_weather")], + "provider_specific_fields": {"server_tool_calls": [tool_call("srvtoolu_01", "web_search")]}, + }, + }, + { + "index": 1, + "message": { + "role": "assistant", + "content": "done", + "provider_specific_fields": {"server_tool_calls": [tool_call("srvtoolu_02", "web_fetch")]}, + }, + }, + ] + } + completion_response: Final = ModelResponse(**response) if as_object else response + + assert response_tool_call_names(completion_response) == ("web_search", "web_fetch", "get_weather") + def test_responses_api_output_function_calls(self): # Regression: /v1/responses carries invocations in output[] items of # type function_call, not in choices; they must reach the rollup. diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py index 349030c6c75..a8f68f038c0 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py @@ -1,6 +1,7 @@ import asyncio import json import ssl +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -2237,6 +2238,83 @@ async def test_post_call_success_hook_redacts_both_content_and_tool_arguments(): assert message.tool_calls[0].function.arguments == '{"to": "[NAME_1]"}' +@pytest.mark.asyncio +async def test_post_call_success_hook_inspects_and_redacts_server_tool_call_arguments(): + guard: Final = _make_guardrail() + request_data: Final = {"messages": [{"role": "user", "content": "look up persimmon"}]} + + def side_effect(url: str, *args: object, **kwargs: dict[str, list[dict[str, str]]]) -> Response: + last: Final = kwargs["json"]["messages"][-1]["content"] + return _make_response( + { + "analysis_result": {"policy_drill_down": {"PII": {}}}, + "required_action": {"action_type": "anonymize_action", "policy_name": "PII"}, + "redacted_chat": { + "all_redacted_messages": [ + {"role": "user", "content": "look up persimmon"}, + {"role": "assistant", "content": last.replace("persimmon", "[FRUIT_1]")}, + ] + }, + } + ) + + llm_response: Final = ModelResponse( + choices=[ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "content": None, + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + } + ] + }, + }, + } + ] + ) + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + side_effect=side_effect, + ) as mock_post: + result: Final = await guard.async_post_call_success_hook( + data=request_data, + response=llm_response, + user_api_key_dict=UserAPIKeyAuth(), + ) + + inspected: Final = [call.kwargs["json"]["messages"][-1] for call in mock_post.call_args_list] + assert inspected == [ + {"role": "assistant", "content": '{"query": "persimmon"}'}, + {"role": "assistant", "content": '{"city": "Paris"}'}, + ] + message: Final = result.choices[0].message + assert message.provider_specific_fields == { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "[FRUIT_1]"}'}, + } + ] + } + assert message.content is None + assert message.tool_calls[0].function.arguments == '{"city": "Paris"}' + + def _make_responses_api_response(output: list) -> ResponsesAPIResponse: return ResponsesAPIResponse(id="resp-1", created_at=0, output=output) @@ -2332,6 +2410,53 @@ async def test_post_call_success_hook_redacts_responses_api_function_call_argume assert result.output[0].arguments == '{"recipient": "[NAME_1]"}' +@pytest.mark.asyncio +async def test_post_call_success_hook_redacts_claude_hosted_web_fetch_url(): + from openai.types.responses import ResponseFunctionWebSearch + from openai.types.responses.response_function_web_search import ActionOpenPage + + guard: Final = _make_guardrail() + request_data: Final = {"messages": [{"role": "user", "content": "read the page"}]} + anonymize_response: Final = _make_response( + { + "analysis_result": {"policy_drill_down": {"PII": {}}}, + "required_action": {"action_type": "anonymize_action", "policy_name": "PII"}, + "redacted_chat": { + "all_redacted_messages": [ + {"role": "user", "content": "read the page"}, + {"role": "assistant", "content": '{"url": "https://example.com/?ssn=[SSN_1]"}'}, + ] + }, + } + ) + response: Final = _make_responses_api_response( + [ + ResponseFunctionWebSearch( + id="ws_srvtoolu_01", + type="web_search_call", + status="completed", + action=ActionOpenPage(type="open_page", url="https://example.com/?ssn=123-45-6789"), + ) + ] + ) + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + mock_post.return_value = anonymize_response + result: Final = await guard.async_post_call_success_hook( + data=request_data, + response=response, + user_api_key_dict=UserAPIKeyAuth(), + ) + posted: Final = mock_post.call_args.kwargs["json"]["messages"] + assert posted[-1] == {"role": "assistant", "content": '{"url": "https://example.com/?ssn=123-45-6789"}'} + assert (result.output[0].id, result.output[0].action.url) == ( + "ws_srvtoolu_01", + "https://example.com/?ssn=[SSN_1]", + ) + + @pytest.mark.asyncio async def test_post_call_success_hook_blocks_responses_api_output(): """A ``block_action`` on Responses API output must raise so the blocked text never 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 5ce78fd05d1..4cdabe5c360 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 @@ -40,6 +40,14 @@ from tests.unit.proxy.guardrails.guardrail_hooks._cisco_ai_defense_test_utils im patch, pytest, ) +from typing import Final + +from openai.types.responses import ResponseFunctionWebSearch +from openai.types.responses.response_function_web_search import ( + Action, + ActionOpenPage, + ActionSearch, +) def test_cisco_ai_defense_config_via_init_v2_chat(monkeypatch): @@ -404,6 +412,31 @@ class TestCiscoAIDefenseChatMode: ) +_CLAUDE_HOSTED_WEB_ACTIONS: Final = ( + ActionSearch(type="search", query="ssn 123-45-6789"), + ActionOpenPage(type="open_page", url="https://example.com/?ssn=123-45-6789"), +) + + +def _claude_hosted_web_call_response(action: Action): + 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 +538,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: Action) -> None: + g: Final = _make_guardrail(event_hook="post_call") + post_mock: Final = 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: Final = post_mock.call_args.kwargs["json"] + joined: Final = " ".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 +623,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: Action) -> None: + g: Final = _make_guardrail(event_hook="post_call", on_flagged_action="monitor") + cisco_resp: Final = _redact_response(sanitized_text="[REDACTED]") + + with _patch_inspection_post(g, AsyncMock(return_value=cisco_resp)): + result: Final = 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: Final = 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( @@ -1660,6 +1728,38 @@ class TestCiscoAIDefenseToolCallBypass: f"included in the scanned text. Sent: {sent!r}" ) + @pytest.mark.asyncio + async def test_pre_call_scans_function_tool_call_that_also_carries_a_custom_key(self): + g: Final = _make_guardrail(event_hook="pre_call") + post_mock: Final = AsyncMock(return_value=_safe_response()) + data: Final = { + "messages": [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "custom": {}, + "function": {"name": "send_data", "arguments": '{"data":"SSN 123-45-6789"}'}, + } + ], + } + ] + } + + with _patch_inspection_post(g, post_mock): + await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + sent: Final = post_mock.call_args.kwargs["json"] + assert "123-45-6789" in " ".join(m.get("content", "") for m in (sent.get("messages") or [])) + @pytest.mark.parametrize( "data", [ @@ -1777,6 +1877,78 @@ class TestCiscoAIDefenseToolCallBypass: f"included in the scanned text. Sent: {sent!r}" ) + @staticmethod + def _server_tool_call_response() -> ModelResponse: + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + return ModelResponse( + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_1", + type="function", + function=Function(name="get_weather", arguments='{"city": "Paris"}'), + ) + ], + provider_specific_fields={ + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + } + ] + }, + ), + ) + ] + ) + + @pytest.mark.asyncio + async def test_post_call_scan_sends_server_tool_call_arguments_to_cisco(self): + g: Final = _make_guardrail(event_hook="post_call") + post_mock: Final = AsyncMock(return_value=_safe_response()) + + with _patch_inspection_post(g, post_mock): + await g.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "look up persimmon"}]}, + user_api_key_dict=UserAPIKeyAuth(), + response=self._server_tool_call_response(), + ) + + assert [call.kwargs["json"]["messages"] for call in post_mock.call_args_list] == [ + [ + {"role": "user", "content": "look up persimmon"}, + {"role": "assistant", "content": '{"query": "persimmon"} {"city": "Paris"}'}, + ] + ] + + @pytest.mark.asyncio + async def test_post_call_redact_clears_server_tool_call_arguments(self): + g: Final = _make_guardrail(event_hook="post_call", on_flagged_action="monitor") + cisco_resp: Final = _redact_response(sanitized_text="[REDACTED]") + + with _patch_inspection_post(g, AsyncMock(return_value=cisco_resp)): + result: Final = await g.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "look up persimmon"}]}, + user_api_key_dict=UserAPIKeyAuth(), + response=self._server_tool_call_response(), + ) + + message: Final = result.choices[0].message + assert message.provider_specific_fields == { + "server_tool_calls": [ + {"id": "srvtoolu_01", "type": "function", "function": {"name": "web_search", "arguments": "{}"}} + ] + } + assert message.tool_calls[0].function.arguments == "{}" + class TestCiscoAIDefenseToolDefinitionBypass: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py new file mode 100644 index 00000000000..cb7fb7898ab --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py @@ -0,0 +1,98 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation + + +@pytest.mark.asyncio +async def test_prompt_check_appends_replayed_server_tool_call_arguments_to_system_message( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "add_function_to_prompt", False) + route: Final = respx_mock.post("https://api.lakera.ai/v1/prompt_injection").mock( + return_value=httpx.Response(200, json={"results": [{"flagged": False}]}) + ) + moderation: Final = lakeraAI_Moderation(api_key="test-key", moderation_check="pre_call") + + await moderation._check( + data={ + "messages": [ + {"role": "system", "content": "You are a shopping assistant."}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_01", + "type": "function", + "function": {"name": "lookup_order", "arguments": '{"order_id": "A-17"}'}, + } + ], + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_fetch", "arguments": '{"url": "https://example.com"}'}, + } + ] + }, + }, + ] + }, + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + + assert json.loads(route.calls.last.request.content)["input"][0] == { + "role": "system", + "content": 'You are a shopping assistant. Function Input: {"url": "https://example.com"} {"order_id": "A-17"}', + } + + +@pytest.mark.asyncio +async def test_prompt_check_scans_a_replayed_turn_that_only_has_server_tool_calls( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "add_function_to_prompt", False) + route: Final = respx_mock.post("https://api.lakera.ai/v1/prompt_injection").mock( + return_value=httpx.Response(200, json={"results": [{"flagged": False}]}) + ) + moderation: Final = lakeraAI_Moderation(api_key="test-key", moderation_check="pre_call") + + await moderation._check( + data={ + "messages": [ + {"role": "system", "content": "You are a shopping assistant."}, + { + "role": "assistant", + "content": "Searching.", + "tool_calls": None, + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "order A-17"}'}, + } + ] + }, + }, + ] + }, + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + + assert json.loads(route.calls.last.request.content)["input"][0] == { + "role": "system", + "content": 'You are a shopping assistant. Function Input: {"query": "order A-17"}', + } diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py index b8a219ba863..634e511fb91 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -1,3 +1,4 @@ +import json import os import uuid from typing import Final @@ -1549,3 +1550,254 @@ async def test_api_error_handling_legacy_messages( assert "Failed to verify request safety with Lasso API" in str(excinfo.value) assert error_message in str(excinfo.value) assert route.called + + +def _reply_with_web_search_call( + arguments: str, tool_calls: list[dict[str, object]] | None = None +) -> litellm.ModelResponse: + return litellm.ModelResponse( + choices=[ + litellm.Choices( + index=0, + finish_reason="tool_calls", + message=litellm.Message( + role="assistant", + content=None, + tool_calls=tool_calls, + provider_specific_fields={ + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": arguments}, + } + ] + }, + ), + ) + ] + ) + + +@pytest.mark.asyncio +async def test_post_call_scan_sends_server_tool_calls_as_tool_use_entries( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + lasso_guardrail: Final = LassoGuardrail( + lasso_api_key="test-key", + user_id="test-user", + conversation_id="test-conversation", + guardrail_name="test-guard", + event_hook="post_call", + default_on=True, + ) + route: Final = respx_mock.post("https://server.lasso.security/gateway/v3/classify").mock( + return_value=httpx.Response(200, json={"violations_detected": False, "deputies": {}, "findings": {}}) + ) + response: Final = _reply_with_web_search_call( + '{"query": "persimmon"}', + tool_calls=[ + { + "id": "call_01", + "type": "function", + "function": {"name": "lookup_order", "arguments": '{"order_id": "A-17"}'}, + } + ], + ) + + await lasso_guardrail.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "search persimmons"}]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert json.loads(route.calls.last.request.content) == { + "messages": [ + { + "role": "model", + "content": { + "type": "tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "persimmon"}, + }, + }, + { + "role": "model", + "content": { + "type": "tool_use", + "id": "call_01", + "name": "lookup_order", + "input": {"order_id": "A-17"}, + }, + }, + ], + "messageType": "COMPLETION", + "source": {"type": "litellm"}, + "userId": "test-user", + "sessionId": "test-conversation", + } + + +@pytest.mark.asyncio +async def test_post_call_masking_rewrites_server_tool_call_arguments( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + lasso_guardrail: Final = LassoGuardrail( + lasso_api_key="test-key", + user_id="test-user", + conversation_id="test-conversation", + mask=True, + guardrail_name="test-guard", + event_hook="post_call", + default_on=True, + ) + respx_mock.post("https://server.lasso.security/gateway/v3/classifix").mock( + return_value=httpx.Response( + 200, + json={ + "violations_detected": True, + "deputies": {"pattern-detection": True}, + "findings": {"pattern-detection": [{"action": "AUTO_MASKING", "severity": "HIGH"}]}, + "messages": [ + { + "role": "model", + "content": { + "type": "tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": " persimmon"}, + }, + } + ], + }, + ) + ) + + result: Final = await lasso_guardrail.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "search persimmons"}]}, + user_api_key_dict=UserAPIKeyAuth(), + response=_reply_with_web_search_call('{"query": "jane@example.com persimmon"}'), + ) + + assert result.choices[0].message.provider_specific_fields == { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": " persimmon"}'}, + } + ] + } + + +def _replayed_web_fetch_request(url: str) -> dict[str, object]: + return { + "messages": [ + {"role": "user", "content": "open the order page"}, + { + "role": "assistant", + "content": None, + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_fetch", "arguments": json.dumps({"url": url})}, + } + ] + }, + }, + {"role": "user", "content": "summarize it"}, + ] + } + + +@pytest.mark.asyncio +async def test_pre_call_scan_sends_replayed_server_tool_calls_as_tool_use_entries( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + lasso_guardrail: Final = LassoGuardrail( + lasso_api_key="test-key", + user_id="test-user", + conversation_id="test-conversation", + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + ) + route: Final = respx_mock.post("https://server.lasso.security/gateway/v3/classify").mock( + return_value=httpx.Response(200, json={"violations_detected": False, "deputies": {}, "findings": {}}) + ) + + await lasso_guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=_replayed_web_fetch_request("https://example.com/?email=jane@example.com"), + call_type="completion", + ) + + assert json.loads(route.calls.last.request.content)["messages"] == [ + {"role": "user", "content": "open the order page"}, + { + "role": "model", + "content": { + "type": "tool_use", + "id": "srvtoolu_01", + "name": "web_fetch", + "input": {"url": "https://example.com/?email=jane@example.com"}, + }, + }, + {"role": "user", "content": "summarize it"}, + ] + + +@pytest.mark.asyncio +async def test_pre_call_masking_rewrites_replayed_server_tool_call_arguments( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + lasso_guardrail: Final = LassoGuardrail( + lasso_api_key="test-key", + user_id="test-user", + conversation_id="test-conversation", + mask=True, + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + ) + respx_mock.post("https://server.lasso.security/gateway/v3/classifix").mock( + return_value=httpx.Response( + 200, + json={ + "violations_detected": True, + "deputies": {"pattern-detection": True}, + "findings": {"pattern-detection": [{"action": "AUTO_MASKING", "severity": "HIGH"}]}, + "messages": [ + {"role": "user", "content": "open the order page"}, + { + "role": "model", + "content": { + "type": "tool_use", + "id": "srvtoolu_01", + "name": "web_fetch", + "input": {"url": "https://example.com/?email="}, + }, + }, + {"role": "user", "content": "summarize it"}, + ], + }, + ) + ) + + result: Final = await lasso_guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=_replayed_web_fetch_request("https://example.com/?email=jane@example.com"), + call_type="completion", + ) + + assert isinstance(result, dict) + assert result["messages"] == _replayed_web_fetch_request("https://example.com/?email=")["messages"] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 2603cf92f30..8a1d7f0b4e9 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -893,6 +893,55 @@ class TestRestoration: assert result.choices[0].message.content == "a@b.com" + @pytest.mark.asyncio + async def test_server_tool_call_arguments_are_restored(self): + guardrail, _ = _shielded({"[EMAIL_1]": "a@b.com", "[URL_1]": "https://b.com/a"}) + response: Final = ModelResponse( + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + provider_specific_fields={ + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "[EMAIL_1] persimmon"}'}, + }, + { + "id": "srvtoolu_02", + "type": "function", + "function": {"name": "web_fetch", "arguments": '{"url": "[URL_1]"}'}, + }, + ] + }, + ), + ) + ] + ) + + result: Final = await guardrail.async_post_call_success_hook( + data={"messages": []}, user_api_key_dict=None, response=response + ) + + assert result.choices[0].message.provider_specific_fields == { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "a@b.com persimmon"}'}, + }, + { + "id": "srvtoolu_02", + "type": "function", + "function": {"name": "web_fetch", "arguments": '{"url": "https://b.com/a"}'}, + }, + ] + } + @pytest.mark.asyncio async def test_responses_api_shape_is_restored(self): """The Responses API reply carries output items, not choices. diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py index 25d2556bf26..34221720205 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py @@ -2072,6 +2072,63 @@ class TestCompletionResponseTextPartsToolCalls: sent_text = mock_check.call_args.kwargs["text"] assert '{"credit_card": "4111-1111-1111-1111"}' in sent_text + @pytest.mark.asyncio + async def test_post_call_hook_sends_server_tool_call_args_to_purview(self): + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.types.utils import Choices, Message, ModelResponse + + def graph(request: httpx.Request) -> httpx.Response: + if request.url.host == "login.microsoftonline.com": + return httpx.Response(200, json={"access_token": "graph-token", "expires_in": 3600}) + if not request.url.path.endswith("/processContent"): + return httpx.Response(200, json={"value": []}, headers={"ETag": "scope-etag"}) + return httpx.Response(200, json={"protectionScopeState": "unchanged", "policyActions": []}) + + transport: Final = Mock(side_effect=graph) + guardrail: Final = _make_guardrail() + guardrail.async_handler = AsyncHTTPHandler(transport=httpx.MockTransport(transport)) + response: Final = ModelResponse( + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + provider_specific_fields={ + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + } + ] + }, + ), + ) + ], + ) + + await guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + response=response, + ) + + scanned: Final = [ + json.loads(call.args[0].content)["contentToProcess"]["contentEntries"][0]["content"]["data"] + for call in transport.call_args_list + if call.args[0].url.path.endswith("/processContent") + ] + assert scanned == ['{"query": "persimmon"}\n\n---\n\n{"city": "Paris"}'] + def test_responses_api_function_call_args_included(self): """Function-call arguments in ``ResponsesAPIResponse.output`` must be DLP-scanned.""" from litellm.types.llms.openai import ResponsesAPIResponse @@ -2123,6 +2180,29 @@ class TestCompletionResponseTextPartsToolCalls: parts = guardrail._completion_response_text_parts(response) assert any('{"secret": "hunter2"}' in p for p in parts) + def test_responses_api_claude_hosted_web_fetch_url_included(self): + from openai.types.responses import ResponseFunctionWebSearch + from openai.types.responses.response_function_web_search import ActionOpenPage + + from litellm.types.llms.openai import ResponsesAPIResponse + + guardrail: Final = _make_guardrail() + response: Final = ResponsesAPIResponse( + id="resp-ws-1", + created_at=0, + output=[ + ResponseFunctionWebSearch( + id="ws_srvtoolu_01", + type="web_search_call", + status="completed", + action=ActionOpenPage(type="open_page", url="https://example.com/?ssn=123-45-6789"), + ) + ], + ) + assert guardrail._completion_response_text_parts(response) == [ + '{"url": "https://example.com/?ssn=123-45-6789"}' + ] + # --------------------------------------------------------------- # Graph user id path encoding diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py index 4e31a3c2055..145e03c718b 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py @@ -1,5 +1,9 @@ +from typing import Final from unittest.mock import AsyncMock, patch +from openai.types.responses.response_function_web_search import ResponseFunctionWebSearch +from typing_extensions import ReadOnly, TypedDict + import httpx import pytest from fastapi import HTTPException @@ -11,6 +15,8 @@ from litellm.proxy.guardrails.guardrail_hooks.pangea.pangea import ( ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.types.guardrails import GuardrailEventHooks, LitellmParams +from litellm.types.llms.openai import ChatCompletionAssistantToolCall +from litellm.types.responses.main import build_web_search_call from litellm.types.utils import Choices, Message, ModelResponse @@ -368,3 +374,86 @@ async def test_pangea_ai_guard_response_transformed(pangea_guardrail): ) assert response.choices[0]["message"]["content"] == "Yes, here is an SSN: " + + +class _GuardedMessage(TypedDict): + tool_calls: ReadOnly[list[ChatCompletionAssistantToolCall]] + + +class _GuardedChoice(TypedDict): + message: ReadOnly[_GuardedMessage] + + +class _GuardedOutput(TypedDict): + choices: ReadOnly[list[_GuardedChoice]] + + +class _AiGuardRequest(TypedDict): + input: ReadOnly[_GuardedOutput] + + +def _claude_fetch_call(url: str) -> ChatCompletionAssistantToolCall: + return { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_fetch", "arguments": '{"url": "%s"}' % url}, + } + + +def _with_masked_tool_calls(choice: _GuardedChoice) -> _GuardedChoice: + message: Final = choice["message"] + masked_calls: Final = [_claude_fetch_call("") for _ in message["tool_calls"]] + return {**choice, "message": {**message, "tool_calls": masked_calls}} + + +@pytest.mark.asyncio +async def test_pangea_ai_guard_response_masks_claude_server_tool_call_inputs(pangea_guardrail: PangeaHandler) -> None: + data: Final = {"messages": [{"role": "user", "content": "Read https://example.com/"}]} + fetch: Final = build_web_search_call("srvtoolu_01", {"url": "https://example.com/"}, None, tool_name="web_fetch") + sent_choices: Final[list[_GuardedChoice]] = [] + + async def _mask_tool_call_urls( + _handler: object, url: str, json: _AiGuardRequest, **_kwargs: object + ) -> httpx.Response: + choices: Final = json["input"]["choices"] + sent_choices.extend(choices) + return httpx.Response( + status_code=200, + json={ + "result": { + "blocked": False, + "transformed": True, + "output": {"choices": [_with_masked_tool_calls(choice) for choice in choices]}, + } + }, + request=httpx.Request(method="POST", url=url), + ) + + with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", new=_mask_tool_call_urls): + response: Final = await pangea_guardrail.async_post_call_success_hook( + data=data, + user_api_key_dict=None, + response=ModelResponse( + choices=[ + Choices( + message=Message( + role="assistant", + content="Done", + provider_specific_fields={ + "server_tool_calls": [_claude_fetch_call("https://example.com/")], + "web_search_calls": [fetch], + }, + ) + ) + ] + ), + ) + + assert sent_choices[0]["message"]["tool_calls"] == [_claude_fetch_call("https://example.com/")] + message: Final = response.choices[0]["message"] + assert message["tool_calls"] is None + assert message["provider_specific_fields"]["server_tool_calls"] == [_claude_fetch_call("")] + hosted_copy: Final = ( + ResponseFunctionWebSearch.model_validate(message["provider_specific_fields"]["web_search_calls"][0]) + ) + assert hosted_copy.action.url == "" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index a41331230c2..35cea840df3 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -6513,5 +6513,90 @@ class TestPanwAirsToolCallBlockMaskedDataRouting: assert error["scan_id"] == "scan-tool-2" +class TestPanwAirsServerToolCalls: + SERVER_TOOL_CALL: Final = { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "persimmon"}'}, + } + + @classmethod + def _response(cls, *client_tool_calls: ChatCompletionMessageToolCall) -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=list(client_tool_calls) or None, + provider_specific_fields={"server_tool_calls": [copy.deepcopy(cls.SERVER_TOOL_CALL)]}, + ), + ) + ], + model="claude-opus-4-1", + ) + + @staticmethod + def _airs(scan_result: Mapping[str, object]) -> tuple[PanwPrismaAirsHandler, MagicMock]: + transport: Final = MagicMock(side_effect=lambda request: httpx.Response(200, json=scan_result)) + handler: Final = make_handler( + mask_response_content=True, + http_client=AsyncHTTPHandler(transport=httpx.MockTransport(transport)), + ) + return handler, transport + + @pytest.mark.asyncio + async def test_post_call_scan_sends_server_tool_call_query_to_airs( + self, user_api_key_dict: UserAPIKeyAuth + ) -> None: + handler, transport = self._airs({"action": "allow", "category": "benign", "scan_id": "s1"}) + response: Final = self._response( + ChatCompletionMessageToolCall( + id="call_1", + type="function", + function=Function(name="get_weather", arguments='{"city": "Paris"}'), + ) + ) + + await handler.async_post_call_success_hook( + data={"model": "claude-opus-4-1", "litellm_call_id": "call-srv", "metadata": {}}, + user_api_key_dict=user_api_key_dict, + response=response, + ) + + sent_contents: Final = [json.loads(call.args[0].content)["contents"] for call in transport.call_args_list] + assert sent_contents == [[{"response": '{"query": "persimmon"} {"city": "Paris"}'}]] + + @pytest.mark.asyncio + async def test_post_call_masking_rewrites_server_tool_call_arguments( + self, user_api_key_dict: UserAPIKeyAuth + ) -> None: + masked: Final = '{"query": "*********"}' + handler, _ = self._airs( + { + "action": "block", + "category": "sensitive_data", + "scan_id": "s1", + "response_masked_data": {"data": masked}, + } + ) + + result: Final = await handler.async_post_call_success_hook( + data={"model": "claude-opus-4-1", "litellm_call_id": "call-srv-mask", "metadata": {}}, + user_api_key_dict=user_api_key_dict, + response=self._response(), + ) + + message: Final = result.choices[0].message + assert message.provider_specific_fields == { + "server_tool_calls": [ + {"id": "srvtoolu_01", "type": "function", "function": {"name": "web_search", "arguments": masked}} + ] + } + assert message.tool_calls is None + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 563b2a0f2b0..cc0d2a8008f 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -1783,6 +1783,48 @@ async def test_pii_tokens_in_metadata_used_for_unmasking(): assert response.choices[0].message.content == "Hello John, how can I help you?" +@pytest.mark.asyncio +async def test_output_unmasking_restores_server_tool_call_arguments(mock_user_api_key: UserAPIKeyAuth) -> None: + guardrail: Final = OPTIONAL_PresidioPIIMasking(mock_testing=True, output_parse_pii=True) + response: Final = ModelResponse( + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + provider_specific_fields={ + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": " persimmon"}'}, + } + ] + }, + ), + ) + ] + ) + + result: Final = await guardrail.async_post_call_success_hook( + data={"metadata": {"pii_tokens": {"": "John Smith"}}}, + user_api_key_dict=mock_user_api_key, + response=response, + ) + + assert result.choices[0].message.provider_specific_fields == { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "John Smith persimmon"}'}, + } + ] + } + + @pytest.mark.parametrize( "initial_hook", ["pre_call", "during_call", "pre_mcp_call"], diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_qualifire.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_qualifire.py index dfd54cff730..ef97a335562 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_qualifire.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_qualifire.py @@ -2,6 +2,7 @@ Unit tests for Qualifire guardrail integration. """ +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -735,3 +736,41 @@ class TestQualifireGuardrailRegistry: assert "qualifire" in guardrail_class_registry assert guardrail_class_registry["qualifire"] == QualifireGuardrail + + +def test_convert_messages_sends_replayed_server_tool_calls_with_client_tool_calls() -> None: + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail: Final = QualifireGuardrail(api_key="test_key", guardrail_name="test_guardrail") + + result: Final = guardrail._convert_messages_to_api_format( + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"location": "NYC"}'}, + } + ], + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_fetch", "arguments": '{"url": "https://example.com"}'}, + } + ] + }, + }, + ] + ) + + assert result[0]["tool_calls"] == [ + {"id": "srvtoolu_01", "name": "web_fetch", "arguments": {"url": "https://example.com"}}, + {"id": "call_123", "name": "get_weather", "arguments": {"location": "NYC"}}, + ] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_repelloai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_repelloai.py index 77883e9af0e..ff61af9ee62 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_repelloai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_repelloai.py @@ -1,4 +1,6 @@ import os +from collections.abc import Mapping +from typing import Final import pytest from fastapi import HTTPException @@ -852,6 +854,61 @@ class TestRepelloAIPostCall: in captured["json"]["scan_data"]["response"] ) + @pytest.mark.asyncio + async def test_server_tool_call_arguments_scanned(self, monkeypatch: pytest.MonkeyPatch) -> None: + guardrail: Final = _guardrail(event_hook="post_call") + response: Final = ModelResponse( + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_01", + "type": "function", + "function": { + "name": "lookup_order", + "arguments": '{"order_id": "A-17"}', + }, + } + ], + provider_specific_fields={ + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": { + "name": "web_search", + "arguments": '{"query": "persimmon"}', + }, + } + ] + }, + ), + ) + ] + ) + captured: Final = {} + + async def capture( + url: str, headers: Mapping[str, str], json: Mapping[str, object], **_kwargs: object + ) -> Response: + captured["json"] = json + return _verdict_response("passed", url) + + monkeypatch.setattr(guardrail.async_handler, "post", capture) + await guardrail.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "q"}]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + assert captured["json"]["scan_data"] == { + "response": '{"query": "persimmon"}\n{"order_id": "A-17"}' + } + # ---------------------------------------------------------------------- # verdict handling: unknown / malformed responses must not fail open diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 46e766f753c..a46f1ddeed9 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -23,6 +23,7 @@ from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ( ToolPermissionGuardrail, ) from litellm.types.guardrails import GuardrailEventHooks, LitellmParams +from litellm.types.responses.main import build_web_search_call from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ) @@ -41,6 +42,14 @@ from litellm.types.utils import ( ) +_SERVER_WEB_SEARCH_CALL: Final = { + "id": "srvtoolu_01Search", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "weather in Paris"}'}, +} +_ALLOWED_BASH_CALL: Final = {"id": "toolu_01Bash", "type": "function", "function": {"name": "Bash", "arguments": "{}"}} + + class TestToolPermissionGuardrail: """Test class for Tool Permission Guardrail functionality""" @@ -1376,6 +1385,93 @@ class TestToolPermissionGuardrail: assert "Permission denied" in choice.message.content + def _server_search_response(self) -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={ + "content": "Searching.", + "tool_calls": [_ALLOWED_BASH_CALL], + "provider_specific_fields": {"server_tool_calls": [_SERVER_WEB_SEARCH_CALL]}, + }, + ) + ] + ) + + @pytest.mark.asyncio + async def test_denied_server_tool_call_is_blocked(self): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + await self.guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=UserAPIKeyAuth(), response=self._server_search_response() + ) + + @pytest.mark.asyncio + async def test_denied_server_tool_call_in_a_chat_stream_is_blocked(self): + async def _stream(): + yield ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content="Searching.", + provider_specific_fields={"server_tool_calls": [{**_SERVER_WEB_SEARCH_CALL, "index": 0}]}, + ), + finish_reason="stop", + ) + ] + ) + + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + async for _ in self.guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), response=_stream(), request_data={} + ): + pass + + @pytest.mark.asyncio + async def test_rewrite_mode_removes_the_denied_server_tool_call(self): + rewriting: Final = ToolPermissionGuardrail( + guardrail_name="test-tool-permission", + rules=self.test_rules, + default_action="deny", + on_disallowed_action="rewrite", + ) + + with patch.object(rewriting, "should_run_guardrail", return_value=True): + result: Final = await rewriting.async_post_call_success_hook( + data={}, user_api_key_dict=UserAPIKeyAuth(), response=self._server_search_response() + ) + + message: Final = result.choices[0].message + assert message.provider_specific_fields["server_tool_calls"] == [] + assert [tool_call.id for tool_call in message.tool_calls] == [_ALLOWED_BASH_CALL["id"]] + assert "Permission denied" in message.content + + @pytest.mark.asyncio + async def test_rewrite_mode_removes_the_denied_server_tool_calls_hosted_copy(self): + rewriting: Final = ToolPermissionGuardrail( + guardrail_name="test-tool-permission", + rules=self.test_rules, + default_action="deny", + on_disallowed_action="rewrite", + ) + response: Final = self._server_search_response() + response.choices[0].message.provider_specific_fields = { + **response.choices[0].message.provider_specific_fields, + "web_search_calls": [ + build_web_search_call("srvtoolu_01Search", {"query": "weather in Paris"}, {"content": []}) + ], + } + + with patch.object(rewriting, "should_run_guardrail", return_value=True): + result: Final = await rewriting.async_post_call_success_hook( + data={}, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert result.choices[0].message.provider_specific_fields["web_search_calls"] == [] + def test_modify_response_with_permission_errors_keeps_custom_tool_call_in_another_choice(self): denied_call = ChatCompletionMessageToolCall( id="call_bash", type="function", function=Function(name="Bash", arguments="{}") @@ -1405,6 +1501,7 @@ class TestToolPermissionGuardrail: ] assert kept == [[], ["call_patch"]] + class TestToolPermissionGuardrailIntegration: """Integration tests for Tool Permission Guardrail""" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index de3ced0c583..2229741fce5 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -50,6 +50,7 @@ from litellm.proxy.mcp_server.guardrail_translation.handler import ( ) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.responses.main import build_web_search_call from litellm.types.utils import ( CallTypes, CallTypesLiteral, @@ -1288,6 +1289,38 @@ class TestStreamingTransform: # end of stream (it can block), matching block_only. assert guardrail.received_tool_calls + @pytest.mark.asyncio + async def test_claude_server_tool_call_chunks_keep_their_provider_fields(self): + server_fields: Final = { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_fetch", "arguments": '{"url": "https://example.com/"}'}, + } + ], + "web_search_calls": { + "srvtoolu_01": build_web_search_call( + "srvtoolu_01", {"url": "https://example.com/"}, None, tool_name="web_fetch" + ).model_dump() + }, + } + server_chunk: Final = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=None, provider_specific_fields=server_fields))], + ) + + out: Final = await _drive_stream( + UnifiedLLMGuardrails(), + _StreamingTextGuardrail(), + [server_chunk, _stream_chunk("done", finish_reason="stop")], + ) + + delivered: Final = ( + [getattr(item.choices[0].delta, "provider_specific_fields", None) for item in out if item.choices] + ) + assert server_fields in delivered + assert "".join(_delta_text(item) for item in out) == "DONE" + @pytest.mark.asyncio async def test_per_choice_finish_reason_when_choices_finish_in_different_chunks(self): """n>1: a choice finishing before the stream's last chunk keeps its own diff --git a/tests/unit/proxy/guardrails/test_pillar_guardrails.py b/tests/unit/proxy/guardrails/test_pillar_guardrails.py index c6be433399a..86c1b1de8a1 100644 --- a/tests/unit/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/unit/proxy/guardrails/test_pillar_guardrails.py @@ -7,7 +7,7 @@ and following LiteLLM testing patterns and best practices. # Standard library imports import importlib -from typing import Any, Dict +from typing import Any, Dict, Final from unittest.mock import Mock, patch # Add parent directory to path for imports @@ -18,7 +18,7 @@ from urllib.parse import unquote import pytest from fastapi.exceptions import HTTPException -from httpx import Request, Response +from httpx import Request, Response, Timeout # LiteLLM imports import litellm @@ -34,6 +34,7 @@ from litellm.proxy.guardrails.guardrail_hooks.pillar.pillar import ( build_pillar_response_headers, ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.types.utils import Choices, Message, ModelResponse # ============================================================================ # FIXTURES @@ -669,6 +670,63 @@ async def test_post_call_hook_with_tool_calls( assert result == mock_llm_response_with_tools +@pytest.mark.asyncio +async def test_post_call_hook_blocks_a_flagged_claude_server_tool_call( + env_setup: None, + sample_request_data: dict[str, object], + user_api_key_dict: UserAPIKeyAuth, + pillar_clean_response: Response, + pillar_flagged_response: Response, +) -> None: + server_call: Final = { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_fetch", "arguments": '{"url": "https://example.com/"}'}, + } + response: Final = ModelResponse( + choices=[ + Choices( + message=Message( + role="assistant", + content="Fetched it", + provider_specific_fields={"server_tool_calls": [server_call]}, + ) + ) + ] + ) + + class _PillarFlagsTheServerCall: + async def post( + self, + *, + url: str, + headers: dict[str, str], + json: dict[str, object], + timeout: float | Timeout | None, + ) -> Response: + messages: Final = json["messages"] + assert isinstance(messages, list) + scanned_tool_calls: Final = messages[-1].get("tool_calls") or [] + return pillar_flagged_response if server_call in scanned_tool_calls else pillar_clean_response + + guardrail: Final = PillarGuardrail( + guardrail_name="pillar-test", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + on_flagged_action="block", + async_handler=_PillarFlagsTheServerCall(), + ) + + with pytest.raises(HTTPException) as excinfo: + await guardrail.async_post_call_success_hook( + data=sample_request_data, + user_api_key_dict=user_api_key_dict, + response=response, + ) + + assert excinfo.value.status_code == 400 + + # ========================================================================= # HEADER CONFIGURATION TESTS # ========================================================================= diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 2427e7eacb1..d5860b7519b 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -34,7 +34,7 @@ from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) from litellm.types.llms.base import HiddenParams -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams +from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams, ResponsesAPIResponse from litellm.types.responses.main import OutputToolSearchCall, build_web_search_call from litellm.types.utils import ( ChatCompletionMessageToolCall, @@ -4957,6 +4957,159 @@ def test_function_call_tool_id_falls_back_to_unique_id_for_degenerate_call_id(): assert convert(openai)["id"] == "call_tokyo" +_WEATHER_TOOL_USE: Final = { + "type": "tool_use", + "id": "toolu_01Weather", + "name": "get_weather", + "input": {"city": "Paris"}, +} +_FETCHED_URL: Final = "https://example.com/" +_BLOCKED_URL: Final = "https://blocked.invalid/" +_ANTHROPIC_WEB_FETCH_TURN: Final[tuple[dict[str, object], ...]] = ( + {"type": "server_tool_use", "id": "srvtoolu_01Fetch", "name": "web_fetch", "input": {"url": _FETCHED_URL}}, + { + "type": "web_fetch_tool_result", + "tool_use_id": "srvtoolu_01Fetch", + "content": { + "type": "web_fetch_result", + "url": _FETCHED_URL, + "content": {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "Example"}}, + }, + }, + {"type": "server_tool_use", "id": "srvtoolu_02Fetch", "name": "web_fetch", "input": {"url": _BLOCKED_URL}}, + { + "type": "web_fetch_tool_result", + "tool_use_id": "srvtoolu_02Fetch", + "content": {"type": "web_fetch_tool_result_error", "error_code": "url_not_allowed"}, + }, + {"type": "text", "text": "The page is titled Example."}, + _WEATHER_TOOL_USE, +) + + +def _anthropic_turn_on_responses(content: Sequence[Mapping[str, object]]) -> ResponsesAPIResponse: + from unittest.mock import MagicMock + + import httpx + + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + anthropic_response: Final = { + "id": "msg_01Fetch", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": list(content), + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 100, "output_tokens": 50}, + } + raw_response: Final = MagicMock(spec=httpx.Response) + raw_response.json.return_value = anthropic_response + raw_response.status_code = 200 + raw_response.headers = {} + chat_response: Final = AnthropicConfig().transform_parsed_response( + completion_response=anthropic_response, + raw_response=raw_response, + model_response=ModelResponse(), + json_mode=False, + prefix_prompt=None, + ) + return LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Use the web and check the weather in Paris.", + responses_api_request={ + "tools": [ + {"type": "web_search_20250305", "name": "web_search"}, + {"type": "web_fetch_20250910", "name": "web_fetch"}, + {"type": "function", "name": "get_weather", "parameters": {"type": "object"}}, + ] + }, + chat_completion_response=chat_response, + ) + + +def test_anthropic_web_fetch_is_a_hosted_open_page_call_on_responses() -> None: + response: Final = _anthropic_turn_on_responses(_ANTHROPIC_WEB_FETCH_TURN) + + assert [ + (item.id, item.status, item.action.model_dump()) for item in response.output if item.type == "web_search_call" + ] == [ + ("ws_srvtoolu_01Fetch", "completed", {"type": "open_page", "url": _FETCHED_URL}), + ("ws_srvtoolu_02Fetch", "failed", {"type": "open_page", "url": _BLOCKED_URL}), + ] + assert [item.name for item in response.output if item.type == "function_call"] == ["get_weather"] + + +@pytest.mark.parametrize( + ("server_tool_use", "expected_action"), + [ + ( + { + "type": "server_tool_use", + "id": "srvtoolu_01Pending", + "name": "web_fetch", + "input": {"url": _FETCHED_URL}, + }, + {"type": "open_page", "url": _FETCHED_URL}, + ), + ( + {"type": "server_tool_use", "id": "srvtoolu_01Pending", "name": "web_search", "input": {"query": "python"}}, + {"type": "search", "query": "python"}, + ), + ], +) +def test_anthropic_server_web_call_without_a_result_is_an_in_progress_hosted_call_on_responses( + server_tool_use: Mapping[str, object], expected_action: Mapping[str, object] +) -> None: + response: Final = _anthropic_turn_on_responses((server_tool_use, _WEATHER_TOOL_USE)) + + web_calls: Final = [item for item in response.output if item.type == "web_search_call"] + assert [(item.id, item.status) for item in web_calls] == [("ws_srvtoolu_01Pending", "in_progress")] + assert expected_action.items() <= web_calls[0].action.model_dump().items() + assert [item.name for item in response.output if item.type == "function_call"] == ["get_weather"] + + +def test_anthropic_code_execution_is_never_a_client_function_call_on_responses() -> None: + response: Final = _anthropic_turn_on_responses( + ( + { + "type": "server_tool_use", + "id": "srvtoolu_01Bash", + "name": "bash_code_execution", + "input": {"command": "ls"}, + }, + { + "type": "bash_code_execution_tool_result", + "tool_use_id": "srvtoolu_01Bash", + "content": { + "type": "bash_code_execution_result", + "stdout": "notes.txt\n", + "stderr": "", + "return_code": 0, + }, + }, + { + "type": "server_tool_use", + "id": "srvtoolu_02Edit", + "name": "text_editor_code_execution", + "input": {"command": "view", "path": "notes.txt"}, + }, + { + "type": "server_tool_use", + "id": "srvtoolu_03Bash", + "name": "bash_code_execution", + "input": {"command": "pwd"}, + }, + _WEATHER_TOOL_USE, + ) + ) + + assert [(item.type, getattr(item, "call_id", item.id)) for item in response.output if item.type != "message"] == [ + ("code_interpreter_call", "srvtoolu_01Bash"), + ("function_call", "toolu_01Weather"), + ] + + class TestHostedWebSearchReplay: def test_emitted_hosted_search_output_round_trips_with_client_tool_result(self) -> None: search_result: Final = { @@ -5027,6 +5180,75 @@ class TestHostedWebSearchReplay: assert messages[2]["content"] == "Paris is sunny." assert history == original + @pytest.mark.parametrize( + "leading_text", ((), ({"type": "text", "text": "Let me look that up."},)), ids=("alone", "after_text") + ) + @pytest.mark.parametrize( + ("server_tool_use", "expected_input"), + ( + ( + {"type": "server_tool_use", "id": "srvtoolu_01Pending", "name": "web_fetch", "input": {"url": _FETCHED_URL}}, + {"url": _FETCHED_URL}, + ), + ( + {"type": "server_tool_use", "id": "srvtoolu_01Pending", "name": "web_search", "input": {"query": "python"}}, + {"query": "python"}, + ), + ), + ids=("web_fetch", "web_search"), + ) + def test_pending_claude_web_call_replays_to_anthropic_as_its_own_server_tool_use( + self, + leading_text: tuple[Mapping[str, object], ...], + server_tool_use: Mapping[str, object], + expected_input: Mapping[str, object], + ) -> None: + response: Final = _anthropic_turn_on_responses((*leading_text, server_tool_use, _WEATHER_TOOL_USE)) + history: Final = [ + {"role": "user", "content": "Use the web and check the weather in Paris."}, + *(item.model_dump(exclude_none=True) for item in response.output), + {"type": "function_call_output", "call_id": "toolu_01Weather", "output": "Paris is sunny."}, + ] + + anthropic_turns: Final = anthropic_messages_pt( + messages=LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=history, responses_api_request={} + ), + model="claude-sonnet-5-5", + llm_provider="anthropic", + ) + + assert [turn["role"] for turn in anthropic_turns] == ["user", "assistant", "user"] + assert [block for block in anthropic_turns[1]["content"] if block["type"] != "text"] == [ + {"type": "server_tool_use", "id": "srvtoolu_01Pending", "name": server_tool_use["name"], "input": expected_input}, + {"type": "tool_use", "id": "toolu_01Weather", "name": "get_weather", "input": {"city": "Paris"}}, + ] + assert [(block["type"], block["tool_use_id"]) for block in anthropic_turns[2]["content"]] == [ + ("tool_result", "toolu_01Weather") + ] + + @pytest.mark.parametrize( + ("call_id", "status"), + (("ws_srvtoolu_01Fetch", "completed"), ("ws_68a1b2c3d4", "in_progress"), ("srvtoolu_01Fetch", "in_progress")), + ids=("completed-claude-call", "pending-non-claude-call", "pending-unprefixed-id"), + ) + def test_web_call_that_is_not_a_pending_claude_call_replays_as_text_only( + self, call_id: str, status: Literal["completed", "in_progress"] + ) -> None: + search: Final = ResponseFunctionWebSearch( + id=call_id, type="web_search_call", status=status, action=ActionOpenPage(type="open_page", url=_FETCHED_URL) + ) + + messages: Final = ( + LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=search.model_dump(exclude_none=True) + ) + ) + + assert messages == [ + {"role": "assistant", "content": "Hosted web search: " + search.model_dump_json(exclude_none=True)} + ] + @pytest.mark.parametrize( "action", ( diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 87858741908..1a6e3cc7832 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -356,6 +356,319 @@ async def test_web_search_stream_preserves_hosted_and_client_calls(sync_mode, to ) +_STREAMED_WEB_FETCH_CALL: Final[tuple[dict[str, object], ...]] = ( + { + "type": "message_start", + "message": { + "id": "msg_01Fetch", + "type": "message", + "role": "assistant", + "content": [], + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "server_tool_use", "id": "srvtoolu_01Fetch", "name": "web_fetch"}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"url": "https://example.com/"}'}, + }, + {"type": "content_block_stop", "index": 0}, +) +_STREAMED_WEB_FETCH_RESULT: Final[tuple[dict[str, object], ...]] = ( + { + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "web_fetch_tool_result", + "tool_use_id": "srvtoolu_01Fetch", + "content": {"type": "web_fetch_result", "url": "https://example.com/"}, + }, + }, + {"type": "content_block_stop", "index": 1}, +) + + +def _streamed_text_then_weather_call(first_index: int) -> tuple[dict[str, object], ...]: + return ( + {"type": "content_block_start", "index": first_index, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": first_index, "delta": {"type": "text_delta", "text": "Fetched."}}, + {"type": "content_block_stop", "index": first_index}, + { + "type": "content_block_start", + "index": first_index + 1, + "content_block": {"type": "tool_use", "id": "toolu_01Weather", "name": "get_weather", "input": {}}, + }, + { + "type": "content_block_delta", + "index": first_index + 1, + "delta": {"type": "input_json_delta", "partial_json": '{"city": "Paris"}'}, + }, + {"type": "content_block_stop", "index": first_index + 1}, + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 20}, + }, + ) + + +def _anthropic_web_fetch_stream_events(fetch_finished: bool) -> tuple[dict[str, object], ...]: + if fetch_finished: + return (*_STREAMED_WEB_FETCH_CALL, *_STREAMED_WEB_FETCH_RESULT, *_streamed_text_then_weather_call(2)) + return (*_STREAMED_WEB_FETCH_CALL, *_streamed_text_then_weather_call(1)) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + ("fetch_finished", "expected_status", "expected_search_events"), + [ + pytest.param( + True, + "completed", + [ + ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + ResponsesAPIStreamEvents.WEB_SEARCH_CALL_IN_PROGRESS, + ResponsesAPIStreamEvents.WEB_SEARCH_CALL_SEARCHING, + ResponsesAPIStreamEvents.WEB_SEARCH_CALL_COMPLETED, + ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + ], + id="fetch_finished_in_this_turn", + ), + pytest.param( + False, + "in_progress", + [ + ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + ResponsesAPIStreamEvents.WEB_SEARCH_CALL_IN_PROGRESS, + ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + ], + id="fetch_still_waiting_on_its_result", + ), + ], +) +@pytest.mark.asyncio +async def test_streamed_anthropic_web_fetch_is_a_hosted_open_page_call( + sync_mode: bool, + fetch_finished: bool, + expected_status: str, + expected_search_events: list[ResponsesAPIStreamEvents], +) -> None: + from litellm.llms.anthropic.chat.handler import ModelResponseIterator + + parser: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=sync_mode, json_mode=False) + chunks: Final = [parser.chunk_parser(event) for event in _anthropic_web_fetch_stream_events(fetch_finished)] + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-sonnet-5-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper(chunks), + request_input="Fetch example.com and check the weather in Paris.", + responses_api_request={ + "tools": [ + {"type": "web_fetch_20250910", "name": "web_fetch"}, + {"type": "function", "name": "get_weather", "parameters": {"type": "object"}}, + ] + }, + custom_llm_provider="anthropic", + ) + + events: Final = await _collect_events(iterator, sync_mode) + + completed: Final = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + assert [ + (item.id, item.status, item.action.model_dump()) + for item in completed.response.output + if item.type == "web_search_call" + ] == [("ws_srvtoolu_01Fetch", expected_status, {"type": "open_page", "url": "https://example.com/"})] + assert [ + event.type + for event in events + if getattr(event, "item_id", getattr(getattr(event, "item", None), "id", None)) == "ws_srvtoolu_01Fetch" + ] == expected_search_events + assert [item.name for item in completed.response.output if item.type == "function_call"] == ["get_weather"] + assert not any( + event.item.name == "web_fetch" + for event in events + if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED + and event.item.type == "function_call" + ) + + +_STREAMED_PENDING_CODE_EXECUTION: Final[tuple[dict[str, object], ...]] = ( + _STREAMED_WEB_FETCH_CALL[0], + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "server_tool_use", "id": "srvtoolu_01Bash", "name": "bash_code_execution"}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"command": "ls"}'}, + }, + {"type": "content_block_stop", "index": 0}, + *_streamed_text_then_weather_call(1), +) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_streamed_anthropic_code_execution_is_never_a_client_function_call(sync_mode: bool) -> None: + from litellm.llms.anthropic.chat.handler import ModelResponseIterator + + parser: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=sync_mode, json_mode=False) + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-sonnet-5-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper( + [parser.chunk_parser(event) for event in _STREAMED_PENDING_CODE_EXECUTION] + ), + request_input="List the files, then check the weather in Paris.", + responses_api_request={ + "tools": [ + {"type": "code_execution_20250825", "name": "code_execution"}, + {"type": "function", "name": "get_weather", "parameters": {"type": "object"}}, + ] + }, + custom_llm_provider="anthropic", + ) + + events: Final = await _collect_events(iterator, sync_mode) + + item_events: Final = [event for event in events if getattr(event, "item", None) is not None] + assert {event.item.call_id for event in item_events if event.item.type == "function_call"} == {"toolu_01Weather"} + completed: Final = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + assert [item.call_id for item in completed.response.output if item.type == "function_call"] == ["toolu_01Weather"] + + +def _streamed_anthropic_tool_search(search_arguments: str) -> tuple[dict[str, object], ...]: + half: Final = len(search_arguments) // 2 + return ( + _STREAMED_WEB_FETCH_CALL[0], + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Searching."}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "server_tool_use", "id": "srvtoolu_01Search", "name": "tool_search_tool_regex"}, + }, + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": search_arguments[:half]}, + }, + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": search_arguments[half:]}, + }, + {"type": "content_block_stop", "index": 1}, + *_streamed_text_then_weather_call(2)[3:], + ) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + ("search_arguments", "expected_status", "expected_arguments"), + [ + pytest.param('{"pattern": "weather"}', "completed", {"pattern": "weather"}, id="arguments_parse"), + pytest.param('{"pattern": ', "incomplete", {}, id="arguments_cut_off"), + ], +) +@pytest.mark.asyncio +async def test_streamed_anthropic_tool_search_keeps_its_place_before_the_function_call( + sync_mode: bool, search_arguments: str, expected_status: str, expected_arguments: dict[str, str] +) -> None: + from litellm.llms.anthropic.chat.handler import ModelResponseIterator + + parser: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=sync_mode, json_mode=False) + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-sonnet-5-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper( + [parser.chunk_parser(event) for event in _streamed_anthropic_tool_search(search_arguments)] + ), + request_input="What is the weather in Paris right now?", + responses_api_request={ + "tools": [ + {"type": "tool_search", "execution": "server"}, + {"type": "function", "name": "get_weather", "parameters": {"type": "object"}, "defer_loading": True}, + ] + }, + custom_llm_provider="anthropic", + ) + + events: Final = await _collect_events(iterator, sync_mode) + + completed: Final = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + output: Final = completed.response.output + assert [item.type for item in output] == ["message", "tool_search_call", "function_call"] + assert (output[1].status, output[1].arguments) == (expected_status, expected_arguments) + search_item_events: Final = [ + event for event in events if getattr(getattr(event, "item", None), "type", None) == "tool_search_call" + ] + assert [(event.type, event.output_index) for event in search_item_events] == [ + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 1), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 1), + ] + + +_STREAMED_PAUSED_WEB_SEARCH: Final[tuple[dict[str, object], ...]] = ( + _STREAMED_WEB_FETCH_CALL[0], + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "server_tool_use", "id": "srvtoolu_01Pending", "name": "web_search"}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"query": "next city"}'}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "pause_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, +) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_streamed_paused_anthropic_web_search_lists_only_the_unrun_search(sync_mode: bool) -> None: + from litellm.llms.anthropic.chat.handler import ModelResponseIterator + + parser: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=sync_mode, json_mode=False) + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-sonnet-5-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper( + [parser.chunk_parser(event) for event in _STREAMED_PAUSED_WEB_SEARCH] + ), + request_input="Find the next city.", + responses_api_request={"tools": [{"type": "web_search_20250305", "name": "web_search"}]}, + custom_llm_provider="anthropic", + ) + + events: Final = await _collect_events(iterator, sync_mode) + + completed: Final = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + assert [(item.type, item.status, item.action.query) for item in completed.response.output] == [ + ("web_search_call", "in_progress", "next city") + ] + assert not any(_is_message_item(event) for event in events) + + def test_tool_calls_present_only_in_final_response_are_emitted_before_completed(): iterator = LiteLLMCompletionStreamingIterator( model="test-model", diff --git a/ui/litellm-dashboard/src/components/logs/detail/PrettyMessagesView.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/PrettyMessagesView.test.tsx index 807aaeeae00..6dbca340b4e 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/PrettyMessagesView.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/PrettyMessagesView.test.tsx @@ -17,6 +17,45 @@ describe("PrettyMessagesView", () => { expect(screen.getByText("Hi there!")).toBeInTheDocument(); }); + it("shows Claude's server tool calls next to the client tool calls on chat completions", () => { + const serverToolCalls = [ + { id: "srvtoolu_01", type: "function", function: { name: "web_search", arguments: '{"query": "persimmon"}' } }, + ]; + const weatherCall = { + id: "toolu_01", + type: "function", + function: { name: "get_weather", arguments: '{"city": "Paris"}' }, + }; + const request = { + messages: [ + { role: "user", content: "Search, then get the weather" }, + { + role: "assistant", + content: "Searching.", + tool_calls: [weatherCall], + provider_specific_fields: { server_tool_calls: serverToolCalls }, + }, + { role: "tool", tool_call_id: "toolu_01", content: "Sunny" }, + ], + }; + const response = { + choices: [ + { + message: { + role: "assistant", + content: null, + tool_calls: [weatherCall], + provider_specific_fields: { server_tool_calls: serverToolCalls }, + }, + }, + ], + }; + + render(); + expect(screen.getAllByText("web_search").length).toBeGreaterThanOrEqual(2); + expect(screen.getAllByText("get_weather").length).toBeGreaterThanOrEqual(2); + }); + it("renders input when request is a bare messages array (cold storage payload)", () => { const request = [{ role: "user", content: "Write me a poem" }]; const response = { diff --git a/ui/litellm-dashboard/src/components/logs/detail/prettyMessagesUtils.ts b/ui/litellm-dashboard/src/components/logs/detail/prettyMessagesUtils.ts index 43bbd691490..5482b034419 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/prettyMessagesUtils.ts +++ b/ui/litellm-dashboard/src/components/logs/detail/prettyMessagesUtils.ts @@ -81,7 +81,7 @@ const parseResponseMessage = (payload: ResponsePayload): ParsedMessage | null => return { role: toRole(message.role, "assistant"), content: parseMessageContent(message.content), - toolCalls: parseChatToolCalls(message.tool_calls), + toolCalls: parseChatMessageToolCalls(message), }; } case "responses": { @@ -104,7 +104,7 @@ const parseChatMessage = (message: unknown): ParsedMessage => { return { role: toRole(message.role, "user"), content: parseMessageContent(message.content), - toolCalls: parseChatToolCalls(message.tool_calls), + toolCalls: parseChatMessageToolCalls(message), toolCallId: typeof message.tool_call_id === "string" ? message.tool_call_id : undefined, }; }; @@ -169,6 +169,13 @@ const parseContentPart = (part: unknown): string => { /** * Parse tool calls from response message */ +const parseChatMessageToolCalls = (message: UnknownRecord): ToolCall[] | undefined => { + const fields = message.provider_specific_fields; + const serverToolCalls = isRecord(fields) ? parseChatToolCalls(fields.server_tool_calls) : undefined; + const toolCalls = parseChatToolCalls(message.tool_calls); + return serverToolCalls?.length ? [...serverToolCalls, ...(toolCalls ?? [])] : toolCalls; +}; + const parseChatToolCalls = (toolCalls: unknown): ToolCall[] | undefined => { if (!Array.isArray(toolCalls)) return undefined; return toolCalls.map((toolCall) => { diff --git a/ui/litellm-dashboard/src/components/logs/detail/sections/ToolsSection/utils.test.ts b/ui/litellm-dashboard/src/components/logs/detail/sections/ToolsSection/utils.test.ts index 78372fe56d9..a6693794df6 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/sections/ToolsSection/utils.test.ts +++ b/ui/litellm-dashboard/src/components/logs/detail/sections/ToolsSection/utils.test.ts @@ -131,6 +131,52 @@ describe("ToolsSection utils", () => { expect(result[1].callData).toBeUndefined(); }); + it("should mark Claude server tools as called from provider_specific_fields.server_tool_calls", () => { + const log: Partial = { + request_id: "test-server-tools", + proxy_server_request: { + tools: [ + { type: "web_search_20250305", name: "web_search", max_uses: 3 }, + { type: "function", function: { name: "get_weather", description: "Get weather" } }, + { type: "function", function: { name: "send_email", description: "Send email" } }, + ], + }, + response: { + choices: [ + { + message: { + tool_calls: [ + { + id: "toolu_01Weather", + type: "function", + function: { name: "get_weather", arguments: '{"location": "Paris"}' }, + }, + ], + provider_specific_fields: { + server_tool_calls: [ + { + id: "srvtoolu_01Search", + type: "function", + function: { name: "web_search", arguments: '{"query": "paris weather"}' }, + }, + ], + }, + }, + }, + ], + }, + }; + + const result = parseToolsFromLog(log as LogEntry); + + expect(result.map((tool) => [tool.name, tool.called, tool.callData?.id])).toEqual([ + ["web_search", true, "srvtoolu_01Search"], + ["get_weather", true, "toolu_01Weather"], + ["send_email", false, undefined], + ]); + expect(result[0].callData?.arguments).toEqual({ query: "paris weather" }); + }); + it("should handle string format request and response", () => { const log: Partial = { request_id: "test-5", diff --git a/ui/litellm-dashboard/src/components/logs/detail/sections/ToolsSection/utils.ts b/ui/litellm-dashboard/src/components/logs/detail/sections/ToolsSection/utils.ts index 4fef4a93497..c489f3b2f49 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/sections/ToolsSection/utils.ts +++ b/ui/litellm-dashboard/src/components/logs/detail/sections/ToolsSection/utils.ts @@ -43,6 +43,24 @@ function extractToolsFromRequest(log: LogEntry): ToolDefinition[] { return []; } +interface ChatMessageToolCallFields { + tool_calls?: unknown; + provider_specific_fields?: { server_tool_calls?: unknown } | null; +} + +function toolCallList(value: unknown): ToolCall[] { + return Array.isArray(value) ? value : []; +} + +function extractChatMessageToolCalls(message: ChatMessageToolCallFields | undefined): ToolCall[] | null { + const serverToolCalls = message?.provider_specific_fields?.server_tool_calls; + const clientToolCalls = message?.tool_calls; + if (!Array.isArray(serverToolCalls) && !Array.isArray(clientToolCalls)) { + return null; + } + return [...toolCallList(serverToolCalls), ...toolCallList(clientToolCalls)]; +} + /** * Extract tool calls from response data */ @@ -54,10 +72,9 @@ function extractToolCallsFromResponse(log: LogEntry): ToolCall[] { // OpenAI format: response.choices[0].message.tool_calls const choices = responseData.choices; if (Array.isArray(choices) && choices.length > 0) { - const firstChoice = choices[0]; - const message = firstChoice.message; - if (message && Array.isArray(message.tool_calls)) { - return message.tool_calls; + const chatToolCalls = extractChatMessageToolCalls(choices[0].message); + if (chatToolCalls) { + return chatToolCalls; } }