diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index ec6de4e206d..4a0fd6e26f0 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -35,11 +35,14 @@ from operator import itemgetter from types import MappingProxyType from typing import Final, Literal, TypeAlias, TypeVar +from pydantic import ValidationError + from litellm.types.llms.anthropic import ( ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX, ANTHROPIC_TOOL_SEARCH_TOOL_NAMES, AnthropicMessagesSystemMessageParam, AnthropicSystemMessageContent, + AnthropicToolSearchToolResultBlock, ) from litellm.types.llms.openai import ( AllMessageValues, @@ -299,6 +302,15 @@ def is_anthropic_tool_search_use(tool_id: str, tool_name: str) -> bool: return tool_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX) and tool_name in ANTHROPIC_TOOL_SEARCH_TOOL_NAMES +def anthropic_tool_search_result(block: object) -> AnthropicToolSearchToolResultBlock | None: + if message_field(block, "type") != "tool_search_tool_result": + return None + try: + return AnthropicToolSearchToolResultBlock.model_validate(block) + except ValidationError: + return None + + def find_anthropic_server_tool_result( tool_id: str, web_search_results: Sequence[object] | None, diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 66c560013fb..00047e52f15 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -812,8 +812,9 @@ class ModelResponseIterator: reasoning_content=reasoning_content, speed=self.speed, ) - if usage.speed is not None: - self.speed = usage.speed + reported_speed: Final = self._cumulative_usage.get("speed") + if isinstance(reported_speed, str): + self.speed = reported_speed return usage def _content_block_delta_helper( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 48bccc7c001..c777f1f6ac9 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -31,6 +31,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.asyncify import is_event_loop_running from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, coerce_token_limit, drops_unsupported_params, map_finish_reason, @@ -151,8 +152,10 @@ from ..common_utils import ( requires_native_compaction_beta, stop_reason_provider_specific_fields, strip_advisor_blocks_from_messages, + tool_search_result_renamed, tools_marked_for_a_tool_config_point, validated_object_mapping, + without_stale_tool_search_results, ) if TYPE_CHECKING: @@ -411,6 +414,17 @@ class _AnthropicToolResultBlock(TypedDict, total=False): content: ReadOnly[object] +def _provider_fields_with_tool_names(message: object, forward_map: Mapping[str, str]) -> Mapping[str, object] | None: + provider_fields: Final = as_str_mapping(message_field(message, "provider_specific_fields")) + replayed_results: Final = parts_of(provider_fields.get("tool_results")) if provider_fields is not None else () + if provider_fields is None or not replayed_results: + return None + return { + **provider_fields, + "tool_results": [tool_search_result_renamed(result, forward_map) for result in replayed_results], + } + + _ENUM_TYPE_CHECKS: Final[Mapping[object, Callable[[object], bool]]] = MappingProxyType( { "null": lambda v: v is None, @@ -457,7 +471,7 @@ def _basic_sanitize_anthropic_tool_name(name: str) -> str: return _ANTHROPIC_TOOL_NAME_INVALID_CHARS.sub("_", name)[:_ANTHROPIC_TOOL_NAME_MAX_LEN] -def _build_anthropic_tool_name_maps( +def build_anthropic_tool_name_maps( original_names: list[str], ) -> tuple[dict[str, str], dict[str, str]]: """Build (forward, reverse) tool-name maps for a single request. @@ -1458,8 +1472,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): messages: list[AllMessageValues], name_forward_map: dict[str, str], ) -> list[AllMessageValues]: - """Return a copy of `messages` with tool_call/function_call names - rewritten using the per-request forward map. + """Return a copy of `messages` with tool_call/function_call names, and + the tools a replayed tool search found, rewritten using the per-request + forward map. Only mutates messages whose tool_call/function_call name is *in* the forward map. Names absent from the map (already valid, no collision) @@ -1475,7 +1490,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): continue tool_calls = msg.get("tool_calls") function_call = msg.get("function_call") - if not tool_calls and not function_call: + renamed_provider_fields = _provider_fields_with_tool_names(msg, name_forward_map) + if not tool_calls and not function_call and renamed_provider_fields is None: new_messages.append(msg) continue new_msg = dict(msg) @@ -1501,6 +1517,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): new_fc = dict(function_call) new_fc["name"] = name_forward_map[fc_name] new_msg["function_call"] = new_fc + if renamed_provider_fields is not None: + new_msg["provider_specific_fields"] = renamed_provider_fields new_messages.append(cast(AllMessageValues, new_msg)) return new_messages @@ -1517,7 +1535,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): callers that need to compute the maps from the raw OpenAI shape -- e.g. test setup or future pre-mapping consumers. - See _build_anthropic_tool_name_maps for the collision rules. Pulls + See build_anthropic_tool_name_maps for the collision rules. Pulls the original name out of either ``{"function": {"name": ...}}`` (legacy OpenAI shape) or ``{"name": ...}`` (rare top-level shape). """ @@ -1530,7 +1548,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): original = tool.get("name") if isinstance(original, str) and original: original_names.append(original) - return _build_anthropic_tool_name_maps(original_names) + return build_anthropic_tool_name_maps(original_names) @staticmethod def _sanitize_tool_names_in_request( @@ -1558,7 +1576,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # 1. Collect originals from the Anthropic-shaped custom-tool entries. # Order matters: the first occurrence wins the canonical slot; # later collisions get numeric suffixes (see - # ``_build_anthropic_tool_name_maps``). + # ``build_anthropic_tool_name_maps``). original_names: Final[list[str]] = [] for t in tools: if not isinstance(t, dict): @@ -1572,7 +1590,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if not original_names: return {}, {} - forward, reverse = _build_anthropic_tool_name_maps(original_names) + forward, reverse = build_anthropic_tool_name_maps(original_names) if not forward: # Every name was already valid -- nothing to do. return forward, reverse @@ -2544,7 +2562,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # The reverse map only contains entries for names that were actually # rewritten -- so a tool legitimately named ``foo_bar`` is never # incorrectly retyped to ``foo/bar`` on the response side. - # See _build_anthropic_tool_name_maps for the collision-handling + # See build_anthropic_tool_name_maps for the collision-handling # rules and rationale. _name_forward_map, _name_reverse_map = self._sanitize_tool_names_in_request( optional_params=optional_params, @@ -2600,8 +2618,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ## Prevents Anthropic 400: advisor_tool_result in history requires advisor tool. _all_tools: Final = optional_params.get("tools") or [] _has_advisor = any(isinstance(t, dict) and t.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE for t in _all_tools) - if not _has_advisor: - anthropic_messages = strip_advisor_blocks_from_messages(anthropic_messages) + anthropic_messages = list( + without_stale_tool_search_results( + anthropic_messages if _has_advisor else strip_advisor_blocks_from_messages(anthropic_messages), + _TOOLS_ADAPTER.validate_python(_all_tools), + ) + ) ## Add code_execution tool if container_upload is in messages _tools: Final = ( @@ -3239,7 +3261,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Reverse-map rewritten tool names back to caller's originals so a # downstream OpenAI-style dispatcher can match on the registered name. - # See _build_anthropic_tool_name_maps for why this is keyed on the + # See build_anthropic_tool_name_maps for why this is keyed on the # per-request reverse map (so a tool legitimately named `foo_bar` is # never incorrectly retyped to `foo/bar`). No-op when the map is # empty (the common case). diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 1c42dd15853..937afeae402 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.prompt_templates.file_id_utils import ( get_file_ids_from_messages, ) from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( + anthropic_tool_search_result, find_anthropic_server_tool_result, message_field, parts_of, @@ -64,10 +65,12 @@ from litellm.types.llms.anthropic import ( ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER, ANTHROPIC_TOOL_SEARCH_TOOL_TYPES, LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS, + AllAnthropicPassThroughMessageValues, AllAnthropicToolsValues, AnthropicMcpServerTool, AnthropicMessagesToolChoice, AnthropicThinkingParam, + AnthropicToolSearchToolResultBlock, ) from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -1918,6 +1921,113 @@ def strip_advisor_blocks_from_messages(messages: list[_MessageT], replace_with_t return messages +def _assistant_blocks(message: AllAnthropicPassThroughMessageValues) -> tuple[object, ...]: + return parts_of(message["content"]) if message["role"] == "assistant" else () + + +def _results_with_unloaded_tools( + messages: Sequence[AllAnthropicPassThroughMessageValues], tool_names: frozenset[str] +) -> tuple[AnthropicToolSearchToolResultBlock, ...]: + blocks: Final = chain.from_iterable(map(_assistant_blocks, messages)) + results: Final = (result for result in map(anthropic_tool_search_result, blocks) if result is not None) + return tuple(result for result in results if not tool_names.issuperset(result.tool_names)) + + +def _is_stale_tool_search_block(block: object, stale_ids: frozenset[str]) -> bool: + match message_field(block, "type"): + case "server_tool_use": + return message_field(block, "id") in stale_ids + case "tool_search_tool_result": + return message_field(block, "tool_use_id") in stale_ids + case _: + return False + + +def _with_loaded_tool_references_only(block: object, tool_names: frozenset[str]) -> object: + result: Final = anthropic_tool_search_result(block) + fields: Final = as_str_mapping(block) + if result is None or fields is None or tool_names.issuperset(result.tool_names): + return block + content: Final = as_str_mapping(fields.get("content")) or MappingProxyType({}) + references: Final = parts_of(content.get("tool_references")) + loaded: Final = [reference for reference in references if message_field(reference, "tool_name") in tool_names] + return {**fields, "content": {**content, "tool_references": loaded}} + + +def _with_loaded_tool_searches_only( + message: AllAnthropicPassThroughMessageValues, stale_ids: frozenset[str], tool_names: frozenset[str] +) -> AllAnthropicPassThroughMessageValues | None: + if message["role"] != "assistant": + return message + content: Final = message["content"] + if isinstance(content, str): + return message + blocks: Final = list(content) + kept: Final = [ + _with_loaded_tool_references_only(block, tool_names) + for block in blocks + if not _is_stale_tool_search_block(block, stale_ids) + ] + if kept == blocks: + return message + return {**message, "content": kept} if kept else None + + +def _tool_reference_renamed(reference: object, name_map: Mapping[str, str]) -> object: + fields: Final = as_str_mapping(reference) + name: Final = fields.get("tool_name") if fields is not None else None + if fields is None or not isinstance(name, str) or name not in name_map: + return reference + return {**fields, "tool_name": name_map[name]} + + +def tool_search_content_renamed(content: object, name_map: Mapping[str, str]) -> object: + fields: Final = as_str_mapping(content) + if fields is None or "tool_references" not in fields: + return content + references: Final = parts_of(fields["tool_references"]) + return {**fields, "tool_references": [_tool_reference_renamed(reference, name_map) for reference in references]} + + +def tool_search_result_renamed(block: object, name_map: Mapping[str, str]) -> object: + fields: Final = as_str_mapping(block) + if fields is None or fields.get("type") != "tool_search_tool_result": + return block + return {**fields, "content": tool_search_content_renamed(fields.get("content"), name_map)} + + +def replays_tool_search(messages: object) -> bool: + contents: Final = (message_field(message, "content") for message in parts_of(messages)) + blocks: Final = chain.from_iterable(map(parts_of, contents)) + return any(anthropic_tool_search_result(block) is not None for block in blocks) + + +def without_stale_tool_search_results( + messages: Sequence[AllAnthropicPassThroughMessageValues], tools: Sequence[object] +) -> tuple[AllAnthropicPassThroughMessageValues, ...]: + """ + Leave out the tools a replayed tool search found that are no longer in ``tools``, + and the whole search, its ``server_tool_use`` and its ``tool_search_tool_result``, + when none of them are left. Anthropic rejects such a history with "Tool reference + '' not found in available tools", for example after a follow-up turn removes + a tool or sends none. + """ + tool_names: Final = frozenset(name for tool in tools if isinstance(name := message_field(tool, "name"), str)) + unloaded: Final = _results_with_unloaded_tools(messages, tool_names) + if not unloaded: + return tuple(messages) + litellm.verbose_logger.debug( + "Leaving out the tools a replayed tool search found that this request no longer sends: %s", + sorted(frozenset(chain.from_iterable(result.tool_names for result in unloaded)) - tool_names), + ) + stale_ids: Final = frozenset(result.tool_use_id for result in unloaded if tool_names.isdisjoint(result.tool_names)) + return tuple( + kept + for message in messages + if (kept := _with_loaded_tool_searches_only(message, stale_ids, tool_names)) is not None + ) + + def is_anthropic_invalid_thinking_block_error(error_text: str) -> bool: """ Detect Anthropic 400 errors caused by invalid thinking blocks in replayed diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py index 1c38642e062..6e0cdbcce08 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py @@ -23,6 +23,7 @@ from litellm.llms.anthropic.common_utils import ( OBJECT_MAPPING_ADAPTER, AnthropicModelInfo, carried_thinking_keys, + replays_tool_search, ) from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, @@ -265,6 +266,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): optional_params=optional_params, headers=headers, thinking=_anthropic_request.get("thinking"), + history_holds_tool_search=replays_tool_search( + OBJECT_MAPPING_ADAPTER.validate_python(_anthropic_request).get("messages") + ), ) if beta_list: _anthropic_request["anthropic_beta"] = beta_list @@ -397,9 +401,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): optional_params: dict, headers: dict, thinking: AnthropicThinkingParam | None, + history_holds_tool_search: bool, ) -> list[str]: tools: Final = optional_params.get("tools") - tool_search_used: Final = self.is_tool_search_used(tools) + tool_search_used: Final = self.is_tool_search_used(tools) or history_holds_tool_search programmatic_tool_calling_used: Final = self.is_programmatic_tool_calling_used(tools) input_examples_used: Final = self.is_input_examples_used(tools) is_mid_conversation_output_config_used: Final = self.is_mid_conversation_output_config_used(messages) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index f81957401ff..5b3b8e88834 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -119,6 +119,7 @@ from litellm.types.llms.openai import ( from litellm.types.responses.main import ( GenericResponseOutputItem, OutputText, + OutputToolSearchOutput, ) from litellm.types.utils import GenericGuardrailAPIInputs @@ -276,6 +277,7 @@ _ITEM_ID_KEYED_TOOL_CALL_TYPES: Final = frozenset({_TOOL_SEARCH_CALL_TYPE, _WEB_ _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]) +_RAW_ITEMS_ADAPTER: Final = TypeAdapter(tuple[object, ...]) _STREAM_EVENT_TYPE_ADAPTER: Final = TypeAdapter(str) _TOOL_CALL_PAYLOAD_FIELDS: Final[Mapping[str, str]] = MappingProxyType( {"function_call": "arguments", "custom_tool_call": "input"} @@ -343,6 +345,18 @@ def _is_tool_call_item(item: object) -> bool: return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES +def _is_tool_search_item(item: object) -> bool: + match item: + case {"type": "tool_search_call" | "tool_search_output"}: + return True + case _: + return False + + +def _joins_assistant_turn(item: object) -> bool: + return _is_tool_call_item(item) or _is_tool_search_item(item) + + 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) @@ -490,14 +504,17 @@ def _provenance_unit_bounds( ) -> tuple[tuple[int, int], ...]: trailing_roles: Final = tuple( accumulate( - (_last_message_role(messages) for messages in solo_conversions), + ( + "assistant" if _is_tool_search_item(item) else _last_message_role(messages) + for item, messages in zip(raw_input, solo_conversions) + ), lambda previous, current: current if current is not None else previous, ) ) start_indexes: Final = tuple( index for index in range(len(raw_input)) - if index == 0 or not (_is_tool_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant") + if index == 0 or not (_joins_assistant_turn(raw_input[index]) and trailing_roles[index - 1] == "assistant") ) return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input)))) @@ -512,6 +529,7 @@ def _input_item_provenance( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=cast("ResponseInputParam", [item]), # cast-ok: items checked as Mappings above responses_api_request=_EMPTY_RESPONSES_REQUEST, + replay_tool_search=True, ) for item in raw_input ) @@ -519,6 +537,7 @@ def _input_item_provenance( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=cast("ResponseInputParam", list(raw_input)), # cast-ok: items checked as Mappings above responses_api_request=_EMPTY_RESPONSES_REQUEST, + replay_tool_search=True, ) ) if full_conversion != tuple(expected_messages): @@ -531,6 +550,7 @@ def _input_item_provenance( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=cast("ResponseInputParam", list(raw_input[start:end])), # cast-ok: checked as Mappings above responses_api_request=_EMPTY_RESPONSES_REQUEST, + replay_tool_search=True, ) ) for start, end in units @@ -744,7 +764,42 @@ def _patch_or_convert_request_fields( input_items, converted_instructions = ( LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(structured_messages) ) - return _RequestFields(input=tuple(input_items), instructions=converted_instructions) + return _RequestFields( + input=_with_tool_searches_restored(input_items, raw_input), instructions=converted_instructions + ) + + +def _tool_search_items_by_call_id(raw_input: Sequence[object], item_type: str) -> Mapping[str, object]: + return MappingProxyType( + { + call_id: item + for item in raw_input + if stream_item_field(item, "type") == item_type + and isinstance(call_id := stream_item_field(item, "call_id"), str) + } + ) + + +def _restored_tool_search( + item: object, calls: Mapping[str, object], outputs: Mapping[str, object] +) -> tuple[object, ...]: + call_id: Final = stream_item_field(item, "call_id") + if stream_item_field(item, "type") != "function_call" or not isinstance(call_id, str) or call_id not in calls: + return (item,) + arguments: Final = stream_item_field(item, "arguments") + restored_call: Final = _rewritten_tool_search_call( + calls[call_id], arguments if isinstance(arguments, str) else None + ) + return (restored_call, outputs[call_id]) if call_id in outputs else (restored_call,) + + +def _with_tool_searches_restored(input_items: Sequence[object], raw_input: object) -> tuple[object, ...]: + raw_items: Final = _RAW_ITEMS_ADAPTER.validate_python(raw_input) if isinstance(raw_input, list) else () + if not any(map(_is_tool_search_item, raw_items)): + return tuple(input_items) + calls: Final = _tool_search_items_by_call_id(raw_items, _TOOL_SEARCH_CALL_TYPE) + outputs: Final = _tool_search_items_by_call_id(raw_items, "tool_search_output") + return tuple(chain.from_iterable(_restored_tool_search(item, calls, outputs) for item in input_items)) def _next_stream_sequence_number(responses_so_far: Sequence[object] | None) -> int: @@ -791,6 +846,7 @@ class OpenAIResponsesHandler(BaseTranslation): messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=input_data, responses_api_request=data, + replay_tool_search=True, ) return cast(list[AllMessageValues], messages) if messages else None @@ -1953,7 +2009,7 @@ class OpenAIResponsesHandler(BaseTranslation): if isinstance(output_item, GenericResponseOutputItem): if output_item.content and content_idx < len(output_item.content): self._write_content_part_text(output_item.content[content_idx], guardrail_response) - elif isinstance(output_item, BaseModel): + elif isinstance(output_item, BaseModel) and not isinstance(output_item, OutputToolSearchOutput): # Handle other Pydantic models by converting to GenericResponseOutputItem try: generic_item = GenericResponseOutputItem.model_validate(output_item.model_dump()) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index f57d092d552..2cadfd6c4a3 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -27308,6 +27308,65 @@ "title": "OutputToolSearchCall", "type": "object" }, + "OutputToolSearchOutput": { + "additionalProperties": true, + "properties": { + "call_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Call Id" + }, + "execution": { + "enum": [ + "server", + "client" + ], + "title": "Execution", + "type": "string" + }, + "id": { + "title": "Id", + "type": "string" + }, + "status": { + "enum": [ + "in_progress", + "completed", + "incomplete" + ], + "title": "Status", + "type": "string" + }, + "tools": { + "default": [], + "items": { + "additionalProperties": true, + "type": "object" + }, + "title": "Tools", + "type": "array" + }, + "type": { + "const": "tool_search_output", + "title": "Type", + "type": "string" + } + }, + "required": [ + "id", + "execution", + "status", + "type" + ], + "title": "OutputToolSearchOutput", + "type": "object" + }, "PendingSafetyCheck": { "additionalProperties": true, "description": "A pending safety check for the computer call.", @@ -29678,6 +29737,9 @@ { "$ref": "#/components/schemas/OutputToolSearchCall" }, + { + "$ref": "#/components/schemas/OutputToolSearchOutput" + }, { "$ref": "#/components/schemas/ResponseCompactionItem" } 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 0b3d56a016b..dff5f5a7e6e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py @@ -325,6 +325,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=input_data if input_data is not None else "", responses_api_request=data, + replay_tool_search=True, ) return self.get_prompt_text_for_dlp(cast(list["AllMessageValues"], messages)) except Exception: diff --git a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py index 7282e59d746..fc88d8da1a3 100644 --- a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py @@ -760,7 +760,7 @@ def register_tools_from_openapi(spec: _OpenAPISpec, base_url: str) -> None: # by sanitized characters (e.g. "foo/list" and "foo.list") # would both become "foo_list". Append _2, _3, … to keep # every tool reachable, mirroring the Anthropic-side logic - # in _build_anthropic_tool_name_maps. + # in build_anthropic_tool_name_maps. unique = tool_name n = 1 while unique in used_names: diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index 98e1984550c..2a2cea77efe 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -34,7 +34,7 @@ _MAX_ARGUMENTS_LEN: Final = 1_000_000 _DECODED_JSON: Final = TypeAdapter(object) TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE: Final = MappingProxyType( - {"function_call": "fc", "custom_tool_call": "ctc", "tool_search_call": "tsc"} + {"function_call": "fc", "custom_tool_call": "ctc", "tool_search_call": "tsc", "tool_search_output": "tso"} ) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 900cfa3d597..6dbd7d69d11 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -468,6 +468,7 @@ class ResponsesSessionHandler: item_cache_control=item_cache_control, tool_call_provider_fields=tool_call_provider_fields, only_since_compaction=True, + replay_tool_search=True, ) ) @@ -484,6 +485,7 @@ class ResponsesSessionHandler: item_cache_control=item_cache_control, tool_call_provider_fields=tool_call_provider_fields, only_since_compaction=True, + replay_tool_search=True, ) ) elif isinstance(_response_output, dict) and _response_output and _response_output != {}: diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 8b6c53a272d..040e2af2a08 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -5,10 +5,13 @@ from itertools import filterfalse from types import MappingProxyType from typing import Any, Final, Protocol, cast +from pydantic import TypeAdapter + import litellm from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items, require_attributes 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.litellm_core_utils.prompt_templates.mid_conversation_system import parts_of 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 @@ -25,6 +28,10 @@ from litellm.responses.litellm_completion_transformation.tool_search import ( build_tool_search_call, tool_search_routing, ) +from litellm.responses.litellm_completion_transformation.tool_search_results import ( + build_tool_search_output, + tool_search_results_by_call_id, +) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -76,6 +83,8 @@ from litellm.types.utils import ( TextCompletionResponse, ) +_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) + def _index_of_output_item_type(items: Sequence[object], item_type: str) -> int | None: return next( @@ -237,6 +246,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._tool_search_output_index_by_call_id: Mapping[str, int] = MappingProxyType({}) self._done_tool_search_call_ids: frozenset[str] = frozenset() async def aclose(self) -> None: @@ -275,12 +285,43 @@ 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: + if call_id not in self._done_tool_search_call_ids: + 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 + ) + self._pending_tool_events.append( + OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=output_index, + item=BaseLiteLLMOpenAIResponseObject(**item.model_dump()), + ) + ) + self._queue_tool_search_output(call_id) + + def _queue_arrived_tool_search_outputs(self) -> None: + if "tool_results" not in self._accumulated_provider_specific_fields: 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 + for call_id in tool_search_results_by_call_id((self._accumulated_provider_specific_fields,)): + if self._tool_search_executions.get(call_id) == "server": + self._queue_tool_search_call_done(call_id, self._get_or_assign_tool_output_index(call_id), "", "server") + + def _queue_tool_search_output(self, call_id: str) -> None: + if call_id in self._tool_search_output_index_by_call_id: + return + result: Final = tool_search_results_by_call_id((self._accumulated_provider_specific_fields,)).get(call_id) + if result is None: + return + output_index: Final = self._allocate_output_index() + self._tool_search_output_index_by_call_id = MappingProxyType( + {**self._tool_search_output_index_by_call_id, call_id: output_index} + ) + request_tools: Final = parts_of(_REQUEST_FIELDS.validate_python(self.responses_api_request).get("tools")) + item: Final = build_tool_search_output(result, request_tools) + self._queue_output_item_added( + output_index, + BaseLiteLLMOpenAIResponseObject(**item.model_copy(update={"status": "in_progress"}).model_dump()), ) self._pending_tool_events.append( OutputItemDoneEvent( @@ -333,6 +374,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return self._reasoning_output_index case "compaction": return self._compaction_output_index + case "tool_search_output": + return self._tool_search_output_index_by_call_id.get(str(getattr(item, "call_id", None))) case _: call_id: Final = getattr(item, "call_id", None) or str(getattr(item, "id", "")).removeprefix("ws_") return self._tool_output_index_by_call_id.get(str(call_id)) @@ -1331,6 +1374,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if src and isinstance(src, dict): self._merge_provider_specific_fields(src) self._reserve_web_search_indexes(src) + self._queue_arrived_tool_search_outputs() self._ensure_output_item_for_chunk(chunk) # Proceed to transformation self.collected_chat_completion_chunks.append( @@ -1394,6 +1438,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if src and isinstance(src, dict): self._merge_provider_specific_fields(src) self._reserve_web_search_indexes(src) + self._queue_arrived_tool_search_outputs() self._ensure_output_item_for_chunk(chunk) # Always snapshot before returning any pending events so that # finish_reason (e.g. content_filter) is captured even when diff --git a/litellm/responses/litellm_completion_transformation/tool_search.py b/litellm/responses/litellm_completion_transformation/tool_search.py index 5a1b27f9092..08b683f40a2 100644 --- a/litellm/responses/litellm_completion_transformation/tool_search.py +++ b/litellm/responses/litellm_completion_transformation/tool_search.py @@ -41,6 +41,8 @@ __all__ = [ "is_server_tool_search_tool", "provider_builds_anthropic_request", "provider_chat_config", + "request_runs_tool_search", + "route_runs_tool_search", "runs_tool_search_on_the_server", "tool_search_call_ids", "tool_search_routing", @@ -217,6 +219,20 @@ def provider_builds_anthropic_request(model: str, custom_llm_provider: str | Non return chat_config_builds_anthropic_request(provider_chat_config(model, custom_llm_provider)) +def route_runs_tool_search(model: str, custom_llm_provider: str | None) -> bool: + from litellm.utils import supports_tool_search + + return supports_tool_search( + model=model, custom_llm_provider=custom_llm_provider + ) and provider_builds_anthropic_request(model, custom_llm_provider) + + +def request_runs_tool_search(tools: Sequence[object], model: str, custom_llm_provider: str | None) -> bool: + if any(is_anthropic_tool_search_chat_tool(tool) for tool in tools): + return provider_builds_anthropic_request(model, custom_llm_provider) + return has_server_tool_search(tools) and route_runs_tool_search(model, custom_llm_provider) + + def chat_config_builds_anthropic_request(chat_config: object) -> bool: from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig diff --git a/litellm/responses/litellm_completion_transformation/tool_search_results.py b/litellm/responses/litellm_completion_transformation/tool_search_results.py new file mode 100644 index 00000000000..dd14124e5a6 --- /dev/null +++ b/litellm/responses/litellm_completion_transformation/tool_search_results.py @@ -0,0 +1,299 @@ +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType +from typing import Final, Literal, TypeVar, cast # noqa: TID251 # validating a message union drops its untyped keys + +from openai.types.responses.response_function_tool_call_param import ResponseFunctionToolCallParam +from openai.types.responses.response_input_param import ResponseInputItemParam +from pydantic import ConfigDict, TypeAdapter, ValidationError + +from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( + anthropic_tool_search_result, + is_anthropic_tool_search_use, + message_field, + parts_of, +) +from litellm.types.llms.anthropic import ( + ANTHROPIC_TOOL_SEARCH_RESULT_CONTENT_TYPE, + AnthropicToolReferenceBlock, + AnthropicToolSearchResultContent, + AnthropicToolSearchToolResultBlock, +) +from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.openai import ResponseInputParam +from litellm.types.responses.main import OutputToolSearchCall, OutputToolSearchOutput + +from .custom_tools import openai_shaped_tool_call_item_id, serialize_tool_call_arguments +from .tool_search import ANTHROPIC_TOOL_SEARCH_TOOL_NAME + +_ItemT: Final = TypeVar("_ItemT") +_MessageT: Final = TypeVar("_MessageT") + +_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_TOOL_SEARCH_ITEM_TYPES: Final = frozenset({"tool_search_call", "tool_search_output"}) +_BM25_TOOL_SEARCH_TOOL_NAME: Final = "tool_search_tool_bm25" + + +class _ProviderToolResults(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + tool_results: tuple[object, ...] = () + + +def _provider_tool_results(provider_fields: object) -> tuple[object, ...]: + try: + return _ProviderToolResults.model_validate(provider_fields).tool_results + except ValidationError: + return () + + +def tool_search_results_by_call_id( + provider_fields: Iterable[object], +) -> Mapping[str, AnthropicToolSearchToolResultBlock]: + results: Final = map( + anthropic_tool_search_result, chain.from_iterable(map(_provider_tool_results, provider_fields)) + ) + return MappingProxyType({result.tool_use_id: result for result in results if result is not None}) + + +class _ToolFields(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + name: str = "" + tools: tuple[object, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class _LoadableTool: + chat_name: str + position: int + tool: Mapping[str, object] + member: Mapping[str, object] | None + + +def _tool_entry(tool: object) -> tuple[Mapping[str, object], _ToolFields] | None: + try: + mapping: Final = _MAPPING_ADAPTER.validate_python(tool) + return mapping, _ToolFields.model_validate(mapping) + except ValidationError: + return None + + +def _member_chat_name(namespace: str, member: _ToolFields) -> str | None: + match member.type: + case "function": + return f"{namespace}__{member.name}" + case "custom": + return member.name + case _: + return None + + +def _namespace_members( + position: int, namespace: Mapping[str, object], fields: _ToolFields +) -> tuple[_LoadableTool, ...]: + entries: Final = tuple(entry for raw in fields.tools or () if (entry := _tool_entry(raw)) is not None) + return tuple( + _LoadableTool(chat_name, position, namespace, member) + for member, member_fields in entries + if (chat_name := _member_chat_name(fields.name, member_fields)) is not None + ) + + +def _loadable_tools(position: int, tool: object) -> tuple[_LoadableTool, ...]: + entry: Final = _tool_entry(tool) + if entry is None: + return () + mapping, fields = entry + if fields.type == "namespace" and fields.tools is not None: + return _namespace_members(position, mapping, fields) + return (_LoadableTool(fields.name, position, mapping, None),) if fields.name else () + + +def _loaded_tool(hits: Sequence[_LoadableTool]) -> Mapping[str, object]: + members: Final = [hit.member for hit in hits if hit.member is not None] + return {**hits[0].tool, "tools": members} if members else hits[0].tool + + +def _catalog_entries( + loadable: _LoadableTool, anthropic_names: Mapping[str, str] +) -> tuple[tuple[str, _LoadableTool], ...]: + anthropic_name: Final = anthropic_names.get(loadable.chat_name) + names: Final = (loadable.chat_name,) if anthropic_name is None else (loadable.chat_name, anthropic_name) + return tuple((name, loadable) for name in names) + + +def loaded_tools(tool_names: Iterable[str], request_tools: Sequence[object] | None) -> tuple[Mapping[str, object], ...]: + from litellm.llms.anthropic.chat.transformation import build_anthropic_tool_name_maps + + loadables: Final = tuple( + chain.from_iterable(_loadable_tools(position, tool) for position, tool in enumerate(request_tools or ())) + ) + anthropic_names, _ = build_anthropic_tool_name_maps([loadable.chat_name for loadable in loadables]) + catalog: Final = MappingProxyType( + dict(chain.from_iterable(_catalog_entries(loadable, anthropic_names) for loadable in loadables)) + ) + hits: Final = tuple(catalog[name] for name in dict.fromkeys(tool_names) if name in catalog) + positions: Final = dict.fromkeys(hit.position for hit in hits) + return tuple(_loaded_tool(tuple(hit for hit in hits if hit.position == position)) for position in positions) + + +def build_tool_search_output( + result: AnthropicToolSearchToolResultBlock, request_tools: Sequence[object] | None +) -> OutputToolSearchOutput: + return OutputToolSearchOutput( + id=openai_shaped_tool_call_item_id("tool_search_output", result.tool_use_id), + call_id=result.tool_use_id, + execution="server", + status="completed" if result.found_tools else "incomplete", + tools=loaded_tools(result.tool_names, request_tools) if result.found_tools else (), + type="tool_search_output", + ) + + +def _with_tool_search_output( + item: _ItemT, + results: Mapping[str, AnthropicToolSearchToolResultBlock], + request_tools: Sequence[object] | None, +) -> tuple[_ItemT | OutputToolSearchOutput, ...]: + if not isinstance(item, OutputToolSearchCall) or item.call_id is None: + return (item,) + result: Final = results.get(item.call_id) + return (item,) if result is None else (item, build_tool_search_output(result, request_tools)) + + +def with_tool_search_outputs( + items: Sequence[_ItemT], + results: Mapping[str, AnthropicToolSearchToolResultBlock], + request_tools: Sequence[object] | None, +) -> tuple[_ItemT | OutputToolSearchOutput, ...]: + return tuple(chain.from_iterable(_with_tool_search_output(item, results, request_tools) for item in items)) + + +class _ToolSearchCallItem(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + type: Literal["tool_search_call"] + call_id: str + arguments: object = None + + +class _ToolSearchOutputItem(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + type: Literal["tool_search_output"] + call_id: str + status: Literal["in_progress", "completed", "incomplete"] | None = None + tools: tuple[object, ...] = () + + +def _tool_search_call(item: object) -> _ToolSearchCallItem | None: + if message_field(item, "type") != "tool_search_call": + return None + try: + return _ToolSearchCallItem.model_validate(item) + except ValidationError: + return None + + +def _tool_search_output(item: object) -> _ToolSearchOutputItem | None: + if message_field(item, "type") != "tool_search_output": + return None + try: + return _ToolSearchOutputItem.model_validate(item) + except ValidationError: + return None + + +def holds_tool_search_items(input: str | ResponseInputParam) -> bool: + return not isinstance(input, str) and any(message_field(item, "type") in _TOOL_SEARCH_ITEM_TYPES for item in input) + + +def _anthropic_tool_search_result(output: _ToolSearchOutputItem) -> Mapping[str, object]: + loaded: Final = chain.from_iterable(_loadable_tools(0, tool) for tool in output.tools) + result: Final = AnthropicToolSearchToolResultBlock( + tool_use_id=output.call_id, + content=AnthropicToolSearchResultContent( + type=ANTHROPIC_TOOL_SEARCH_RESULT_CONTENT_TYPE, + tool_references=tuple(AnthropicToolReferenceBlock(tool_name=tool.chat_name) for tool in loaded), + ), + ) + return result.model_dump(mode="json") + + +def _search_arguments(arguments: object) -> Mapping[str, object]: + try: + if isinstance(arguments, str): + return _MAPPING_ADAPTER.validate_json(arguments) + return _MAPPING_ADAPTER.validate_python(arguments) + except ValidationError: + return MappingProxyType({}) + + +def _search_tool_name(arguments: object) -> str: + return _BM25_TOOL_SEARCH_TOOL_NAME if "query" in _search_arguments(arguments) else ANTHROPIC_TOOL_SEARCH_TOOL_NAME + + +def _replayed_item(item: ResponseInputItemParam, replayed_ids: frozenset[str]) -> ResponseInputItemParam | None: + call: Final = _tool_search_call(item) + if call is not None and call.call_id in replayed_ids: + return ResponseFunctionToolCallParam( + type="function_call", + call_id=call.call_id, + name=_search_tool_name(call.arguments), + arguments=serialize_tool_call_arguments(call.arguments, "{}"), + ) + output: Final = _tool_search_output(item) + return None if output is not None and output.call_id in replayed_ids else item + + +def _with_tool_search_results(message: _MessageT, results: Mapping[str, Mapping[str, object]]) -> _MessageT: + tool_call_ids: Final = ( + message_field(tool_call, "id") for tool_call in parts_of(message_field(message, "tool_calls")) + ) + found: Final = tuple( + results[tool_call_id] + for tool_call_id in tool_call_ids + if isinstance(tool_call_id, str) and tool_call_id in results + ) + if not found: + return message + fields: Final = _MAPPING_ADAPTER.validate_python(message) + provider_fields: Final = _MAPPING_ADAPTER.validate_python(fields.get("provider_specific_fields") or {}) + tool_results: Final = (*parts_of(provider_fields.get("tool_results")), *found) + with_results: Final = { + **fields, + "provider_specific_fields": {**provider_fields, "tool_results": list(tool_results)}, + } + return cast(_MessageT, with_results) # cast-ok: a copy of the same message dict with provider fields added + + +def converted_with_tool_searches_replayed( + input: str | ResponseInputParam, + convert: Callable[[str | ResponseInputParam], Sequence[_MessageT]], +) -> Sequence[_MessageT]: + if isinstance(input, str) or not holds_tool_search_items(input): + return convert(input) + outputs: Final = MappingProxyType( + { + output.call_id: output + for output in map(_tool_search_output, input) + if output is not None and output.status in (None, "completed") + } + ) + replayed_ids: Final = frozenset( + call.call_id + for call in map(_tool_search_call, input) + if call is not None + and call.call_id in outputs + and is_anthropic_tool_search_use(call.call_id, ANTHROPIC_TOOL_SEARCH_TOOL_NAME) + ) + if not replayed_ids: + return convert(input) + replayed_input: Final = [kept for item in input if (kept := _replayed_item(item, replayed_ids)) is not None] + results: Final = MappingProxyType( + {call_id: _anthropic_tool_search_result(outputs[call_id]) for call_id in replayed_ids} + ) + return tuple(_with_tool_search_results(message, results) for message in convert(replayed_input)) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 05379ffde3e..3c1abbfe492 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -71,7 +71,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( is_present_non_string, is_unhashable_value_behind, ) -from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of +from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( + is_anthropic_tool_search_use, + message_field, + parts_of, +) from litellm.llms.anthropic.common_utils import ( anthropic_server_tool_call_params, anthropic_server_tool_calls, @@ -126,7 +130,6 @@ from litellm.types.llms.openai import ( ResponsesAPIStatus, ToolChoice, ToolMessageContentPart, - ValidChatCompletionMessageContentTypes, ValidChatCompletionMessageContentTypesLiteral, ) from litellm.types.responses.main import ( @@ -138,6 +141,7 @@ from litellm.types.responses.main import ( OutputImageGenerationCall, OutputText, OutputToolSearchCall, + OutputToolSearchOutput, ) from litellm.types.utils import ( ChatCompletionAnnotation, @@ -186,11 +190,19 @@ from .tool_search import ( is_anthropic_tool_search_chat_tool, is_server_tool_search_tool, provider_builds_anthropic_request, + request_runs_tool_search, + route_runs_tool_search, tool_search_call_ids, tool_search_routing, without_tool_search_calls, without_tool_search_uses, ) +from .tool_search_results import ( + converted_with_tool_searches_replayed, + holds_tool_search_items, + tool_search_results_by_call_id, + with_tool_search_outputs, +) NamespaceNameMap: TypeAlias = Mapping[str, tuple[str, str]] NamespaceTool: TypeAlias = Mapping[str, object] @@ -279,6 +291,20 @@ TOOL_CALLS_CACHE: Final = InMemoryCache() _ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) + + +_CONTENT_TYPE_ADAPTER: Final[TypeAdapter[ValidChatCompletionMessageContentTypesLiteral]] = TypeAdapter( + ValidChatCompletionMessageContentTypesLiteral +) + + +def _valid_content_type(content_type: str) -> ValidChatCompletionMessageContentTypesLiteral | None: + try: + return _CONTENT_TYPE_ADAPTER.validate_python(content_type) + except ValidationError: + return None + + _AUDIO_PART_ADAPTER: Final = TypeAdapter(ChatCompletionAudioObject) _IMAGE_URL_ADAPTER: Final = TypeAdapter(ChatCompletionImageUrlObject) _FILE_PART_ADAPTER: Final = TypeAdapter(_ResponsesFilePart) @@ -595,16 +621,58 @@ def _replayed_assistant_without_tool_search_uses(message: Mapping[str, object]) ) -def _without_stored_tool_search_uses(message: SessionMessage, search_call_ids: frozenset[str]) -> SessionMessage | None: +def _server_tool_search_ids(message: object) -> frozenset[str]: + return frozenset( + tool_id + for call in anthropic_server_tool_call_params(message) + if is_anthropic_tool_search_use(tool_id := call["id"] or "", call["function"].get("name") or "") + ) + + +def _holds_stored_tool_search(message: object) -> bool: + return holds_tool_search_use(message) or bool(_server_tool_search_ids(message)) + + +def _provider_fields_without_tool_searches(message: object, search_ids: frozenset[str]) -> Mapping[object, object]: + fields: Final = validated_object_mapping(message_field(message, "provider_specific_fields")) or {} + return { + **fields, + "server_tool_calls": [ + call for call in anthropic_server_tool_call_params(message) if (call["id"] or "") not in search_ids + ], + "tool_results": [ + result + for result in parts_of(fields.get("tool_results")) + if message_field(result, "tool_use_id") not in search_ids + ], + } + + +def _without_server_tool_searches(message: SessionMessage) -> SessionMessage: + search_ids: Final = _server_tool_search_ids(message) + if not search_ids: + return message + fields: Final = _provider_fields_without_tool_searches(message, search_ids) if isinstance(message, Message): - return without_tool_search_uses(message) - if answers_tool_search(message, search_call_ids): - return None - if isinstance(message, Mapping) and holds_tool_search_use(message): - return _replayed_assistant_without_tool_search_uses(message) + return message.model_copy(update={"provider_specific_fields": fields}) + if isinstance(message, Mapping): + return cast( # cast-ok: the stored assistant message minus its search keeps every other field as stored + SessionMessage, {**message, "provider_specific_fields": fields} + ) return message +def _without_stored_tool_search_uses(message: SessionMessage, search_call_ids: frozenset[str]) -> SessionMessage | None: + kept: Final = _without_server_tool_searches(message) + if isinstance(kept, Message): + return without_tool_search_uses(kept) + if answers_tool_search(kept, search_call_ids): + return None + if isinstance(kept, Mapping) and holds_tool_search_use(kept): + return _replayed_assistant_without_tool_search_uses(kept) + return kept + + def _instruction_messages(messages: Sequence[SessionMessage]) -> tuple[SessionMessage, ...]: return tuple( message @@ -1094,16 +1162,11 @@ class LiteLLMCompletionResponsesConfig: def _without_unsupported_tool_search( tools: ResponseTools, model: str, custom_llm_provider: str | None ) -> ResponseTools: - from litellm.utils import supports_tool_search - if not has_server_tool_search(tools): return tools if any(is_anthropic_tool_search_chat_tool(tool) for tool in tools or ()): return [tool for tool in tools or () if not is_server_tool_search_tool(tool)] - route_runs_tool_search: Final = supports_tool_search( - model=model, custom_llm_provider=custom_llm_provider - ) and provider_builds_anthropic_request(model, custom_llm_provider) - if route_runs_tool_search: + if route_runs_tool_search(model, custom_llm_provider): return tools verbose_logger.warning( "Dropping Responses API tool of type 'tool_search': model '%s' (provider '%s') cannot run tool search " @@ -1281,6 +1344,8 @@ class LiteLLMCompletionResponsesConfig: tool_call_provider_fields=not embeds_thought_signature_in_tool_call_id(request_kwargs), native_compaction=native_compaction, only_since_compaction=not native_compaction, + replay_tool_search=holds_tool_search_items(request_input) + and request_runs_tool_search(request_tools or (), model, custom_llm_provider), caller_input=numbered_input, ) malformed_field_error: Final = ( @@ -1366,6 +1431,7 @@ class LiteLLMCompletionResponsesConfig: tool_call_provider_fields: bool = False, native_compaction: bool = False, only_since_compaction: bool = False, + replay_tool_search: bool = False, caller_input: str | ResponseInputParam | None = None, ) -> list[ AllMessageValues @@ -1403,6 +1469,11 @@ class LiteLLMCompletionResponsesConfig: messages, the way Anthropic ignores content before a compaction block. Guardrail scans leave it off because they write rewrites back by input item position. + + ``replay_tool_search`` belongs to callers building history for a target + that builds an Anthropic request: a Claude ``tool_search_call`` and its + ``tool_search_output`` then go back as the search Claude ran. Everyone + else drops both items. """ messages: list[ AllMessageValues @@ -1429,7 +1500,11 @@ class LiteLLMCompletionResponsesConfig: native_compaction=native_compaction, ) try: - converted: Final = convert_input(sent_input) + converted: Final = ( + converted_with_tool_searches_replayed(sent_input, convert_input) + if replay_tool_search + else convert_input(sent_input) + ) except TypeError as error: call_id_error: Final = _unhashable_tool_output_call_id_error( dumped_input if caller_input is None else caller_input, @@ -1515,7 +1590,7 @@ class LiteLLMCompletionResponsesConfig: ) replayed_messages: Final = ( LiteLLMCompletionResponsesConfig._without_tool_search_uses_off_anthropic(stored_messages, route) - if any(holds_tool_search_use(message) for message in stored_messages) + if any(_holds_stored_tool_search(message) for message in stored_messages) else stored_messages ) session_messages: Final = ( @@ -2831,8 +2906,8 @@ class LiteLLMCompletionResponsesConfig: if content_type.startswith("input_"): stripped: Final = content_type[len("input_") :] # Validate stripped type is valid, otherwise default to "text" - if stripped in ValidChatCompletionMessageContentTypes: - return stripped + if valid_stripped := _valid_content_type(stripped): + return valid_stripped # Handle input_audio -> input_audio (it's already valid) if stripped == "audio": return "input_audio" @@ -2843,10 +2918,7 @@ class LiteLLMCompletionResponsesConfig: return "text" # Return as-is if it's a valid type, otherwise default to "text" - if content_type in ValidChatCompletionMessageContentTypes: - return content_type - - return "text" + return _valid_content_type(content_type) or "text" @staticmethod def transform_instructions_to_system_message( @@ -3120,7 +3192,13 @@ class LiteLLMCompletionResponsesConfig: def transform_chat_completion_tools_to_responses_tools( chat_completion_response: ModelResponse, responses_api_request: ResponsesAPIOptionalRequestParams | None = None, - ) -> list[ResponseFunctionToolCall | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall]: + ) -> list[ + ResponseFunctionToolCall + | ResponseFunctionWebSearch + | CustomToolCallOutputItem + | OutputToolSearchCall + | OutputToolSearchOutput + ]: """ Transform a Chat Completion tools into a Responses API tools. @@ -3128,7 +3206,8 @@ class LiteLLMCompletionResponsesConfig: with ``type="custom_tool_call"``. For regular function tools, returns ``ResponseFunctionToolCall`` with ``type="function_call"``. A provider-side tool search call made for a request's ``tool_search`` tool returns - ``OutputToolSearchCall`` with ``type="tool_search_call"``. + ``OutputToolSearchCall`` with ``type="tool_search_call"``, followed by + an ``OutputToolSearchOutput`` listing the tools it loaded. """ all_chat_completion_tools: Final[list[ChatCompletionMessageToolCall]] = [] for choice in chat_completion_response.choices: @@ -3217,7 +3296,15 @@ class LiteLLMCompletionResponsesConfig: ) responses_tools.append(output_tool_call) - return responses_tools + return list( + with_tool_search_outputs( + responses_tools, + tool_search_results_by_call_id( + choice.message.provider_specific_fields for choice in chat_completion_response.choices + ), + request_tools, + ) + ) @staticmethod def _server_tool_call_item( @@ -3522,6 +3609,7 @@ class LiteLLMCompletionResponsesConfig: | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | OutputToolSearchOutput | ResponseCompactionItem ]: responses_output: list[ @@ -3533,6 +3621,7 @@ class LiteLLMCompletionResponsesConfig: | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | OutputToolSearchOutput | ResponseCompactionItem ] = [] diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index c322eb25b3e..3abe64df3c6 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -156,6 +156,36 @@ class ToolReference(TypedDict, total=False): tool_name: Required[str] +class AnthropicToolReferenceBlock(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["tool_reference"] = "tool_reference" + tool_name: str + + +class AnthropicToolSearchResultContent(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: str + tool_references: tuple[AnthropicToolReferenceBlock, ...] = () + + +class AnthropicToolSearchToolResultBlock(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["tool_search_tool_result"] = "tool_search_tool_result" + tool_use_id: str + content: AnthropicToolSearchResultContent + + @property + def found_tools(self) -> bool: + return self.content.type == ANTHROPIC_TOOL_SEARCH_RESULT_CONTENT_TYPE + + @property + def tool_names(self) -> tuple[str, ...]: + return tuple(reference.tool_name for reference in self.content.tool_references) + + class DirectToolCaller(TypedDict, total=False): """Indicates a tool was called directly by Claude.""" @@ -863,6 +893,7 @@ 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_" +ANTHROPIC_TOOL_SEARCH_RESULT_CONTENT_TYPE: Final = "tool_search_tool_search_result" def claude_hosted_web_tool_call(search: ResponseFunctionWebSearch) -> tuple[str, str, str] | None: diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 378899b667d..268bce1cf6f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -95,6 +95,7 @@ from litellm.types.responses.main import ( OutputFunctionToolCall, OutputImageGenerationCall, OutputToolSearchCall, + OutputToolSearchOutput, ) from .base import CachedTokensDetails @@ -1595,6 +1596,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | OutputToolSearchOutput | ResponseCompactionItem ] ) diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py index 7a2f58e9c16..61cd91be7ec 100644 --- a/litellm/types/responses/main.py +++ b/litellm/types/responses/main.py @@ -225,6 +225,17 @@ class OutputToolSearchCall(BaseLiteLLMOpenAIResponseObject): type: Literal["tool_search_call"] +class OutputToolSearchOutput(BaseLiteLLMOpenAIResponseObject): + model_config = ConfigDict(extra="allow", frozen=True, protected_namespaces=()) + + id: str + call_id: str | None = None + execution: Literal["server", "client"] + status: Literal["in_progress", "completed", "incomplete"] + tools: tuple[Mapping[str, object], ...] = () + type: Literal["tool_search_output"] + + class GenericResponseOutputItem(BaseLiteLLMOpenAIResponseObject): """ Generic response API output item 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 adce80d3c75..eb5220b9b41 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 @@ -424,37 +424,35 @@ def test_odd_string_tool_answer_ids_keep_the_search_rebuilt(gateway: Gateway, to assert SEARCH_ID not in forwarded, received[0].body -_ODD_RESULTS: Final = ( - pytest.param( - { - "type": "tool_search_tool_result", - "tool_use_id": SEARCH_ID, - "content": {"type": "tool_search_tool_result_error", "error_code": "unavailable"}, - }, - id="error-content", - ), - pytest.param( - { - "type": "tool_search_tool_result", - "tool_use_id": SEARCH_ID, - "content": { - "type": "tool_search_tool_search_result", - "tool_references": [{"type": "tool_reference", "tool_name": "w" * 5120}], - }, - }, - id="5kb-reference", - ), -) - - -@pytest.mark.parametrize("odd", _ODD_RESULTS) @_STREAM_MODES -def test_unusual_search_results_from_the_upstream_are_kept_and_replayed_verbatim( - gateway: Gateway, odd: Event, stream: bool +def test_an_errored_search_result_from_the_upstream_is_kept_and_replayed_verbatim( + gateway: Gateway, stream: bool ) -> None: - assistant, _, _, received, _ = _round_trip(gateway, "anthropic", stream, (SEARCH_USE, odd, WEATHER_USE)) - assert kept_results(assistant) == [odd], assistant - assert wire_assistant(received[1]) == [SEARCH_USE, odd, WEATHER_USE], received[1].body + errored: Final[Event] = { + "type": "tool_search_tool_result", + "tool_use_id": SEARCH_ID, + "content": {"type": "tool_search_tool_result_error", "error_code": "unavailable"}, + } + assistant, _, _, received, _ = _round_trip(gateway, "anthropic", stream, (SEARCH_USE, errored, WEATHER_USE)) + assert kept_results(assistant) == [errored], assistant + assert wire_assistant(received[1]) == [SEARCH_USE, errored, WEATHER_USE], received[1].body + + +@_STREAM_MODES +def test_a_search_result_naming_a_tool_the_request_never_sent_is_kept_but_left_out_of_the_replay( + gateway: Gateway, stream: bool +) -> None: + unknown_tool: Final[Event] = { + "type": "tool_search_tool_result", + "tool_use_id": SEARCH_ID, + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [{"type": "tool_reference", "tool_name": "w" * 5120}], + }, + } + assistant, _, _, received, _ = _round_trip(gateway, "anthropic", stream, (SEARCH_USE, unknown_tool, WEATHER_USE)) + assert kept_results(assistant) == [unknown_tool], assistant + assert wire_assistant(received[1]) == [WEATHER_USE], received[1].body @_STREAM_MODES 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 95c05610fba..da7d5fa3508 100644 --- a/tests/integration/providers/test_responses_bridge_tool_search_wire.py +++ b/tests/integration/providers/test_responses_bridge_tool_search_wire.py @@ -153,18 +153,18 @@ def _prompt(marker: str) -> str: return f"Weather in Paris and the ACME price marker-{marker}" -def _search_result(marker: str) -> dict[str, JsonValue]: +def _search_result(marker: str, found_tool: str = "get_weather") -> dict[str, JsonValue]: return { "type": "tool_search_tool_result", "tool_use_id": f"srvtoolu_{marker}", "content": { "type": "tool_search_tool_search_result", - "tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}], + "tool_references": [{"type": "tool_reference", "tool_name": found_tool}], }, } -def _search_blocks(marker: str) -> tuple[dict[str, JsonValue], ...]: +def _search_blocks(marker: str, found_tool: str = "get_weather") -> tuple[dict[str, JsonValue], ...]: return ( { "type": "server_tool_use", @@ -172,12 +172,12 @@ def _search_blocks(marker: str) -> tuple[dict[str, JsonValue], ...]: "name": _REGEX_TOOL_NAME, "input": json.loads(_SEARCH_ARGUMENTS), }, - _search_result(marker), + _search_result(marker, found_tool), ) -def _message(marker: str, searched: bool, tool_use_name: str) -> dict[str, JsonValue]: - search_blocks: Final = _search_blocks(marker) if searched else () +def _message(marker: str, searched: bool, tool_use_name: str, found_tool: str = "get_weather") -> dict[str, JsonValue]: + search_blocks: Final = _search_blocks(marker, found_tool) if searched else () return { "id": f"msg_{marker}", "type": "message", @@ -194,7 +194,7 @@ def _message(marker: str, searched: bool, tool_use_name: str) -> dict[str, JsonV } -def _search_events(marker: str, search_arguments: str) -> tuple[Event, ...]: +def _search_events(marker: str, search_arguments: str, found_tool: str) -> tuple[Event, ...]: half: Final = len(search_arguments) // 2 opened: Final[Event] = { "type": "server_tool_use", @@ -215,14 +215,16 @@ def _search_events(marker: str, search_arguments: str) -> tuple[Event, ...]: "delta": {"type": "input_json_delta", "partial_json": search_arguments[half:]}, }, {"type": "content_block_stop", "index": 1}, - {"type": "content_block_start", "index": 2, "content_block": _search_result(marker)}, + {"type": "content_block_start", "index": 2, "content_block": _search_result(marker, found_tool)}, {"type": "content_block_stop", "index": 2}, ) -def _events(marker: str, search_arguments: str, searched: bool, tool_use_name: str) -> tuple[Event, ...]: +def _events( + marker: str, search_arguments: str, searched: bool, tool_use_name: str, found_tool: str = "get_weather" +) -> tuple[Event, ...]: opened: Final = {**_message(marker, searched, tool_use_name), "content": [], "stop_reason": None} - search_events: Final = _search_events(marker, search_arguments) if searched else () + search_events: Final = _search_events(marker, search_arguments, found_tool) if searched else () tool_index: Final = 3 if searched else 1 tool_use: Final[Event] = {"type": "tool_use", "id": f"toolu_{marker}", "name": tool_use_name, "input": {}} return ( @@ -266,6 +268,7 @@ class _Vendor: pause: float = 0 searched: bool | None = None tool_use_name: str = "get_weather" + found_tool: str = "get_weather" error: tuple[int, str] | None = None def respond(self, request: Request) -> Reply: @@ -284,15 +287,17 @@ class _Vendor: return _anthropic_error(400, f"tools: {_ALL_DEFERRED_ERROR} marker-{marker}") searched: Final = bool(regex_tools) if self.searched is None else self.searched if request.target.endswith("/invoke-with-response-stream"): - events: Final = _events(marker, self.search_arguments, searched, self.tool_use_name) + events: Final = _events(marker, self.search_arguments, searched, self.tool_use_name, self.found_tool) return Reply(content_type=EVENT_STREAM, chunks=aws_chunks(events), pause_between_chunks=self.pause) if _streams(request, body): return Reply( content_type="text/event-stream", - chunks=sse_chunks(_events(marker, self.search_arguments, searched, self.tool_use_name)), + chunks=sse_chunks( + _events(marker, self.search_arguments, searched, self.tool_use_name, self.found_tool) + ), pause_between_chunks=self.pause, ) - return Reply(body=json.dumps(_message(marker, searched, self.tool_use_name)).encode()) + return Reply(body=json.dumps(_message(marker, searched, self.tool_use_name, self.found_tool)).encode()) def _base_url(gateway: Gateway) -> str: @@ -413,8 +418,8 @@ def _search_call_fields(item: Mapping[str, JsonValue]) -> dict[str, JsonValue]: def _assert_search_call_output(payload: Mapping[str, JsonValue], marker: str) -> None: items: Final = _tool_items(payload) - assert len(items) == 2, payload - search, function = items + assert len(items) == 3, payload + search, search_output, function = items assert search["type"] == "tool_search_call", payload assert _search_call_fields(search) == { "id": f"tsc_srvtoolu_{marker}", @@ -423,6 +428,14 @@ def _assert_search_call_output(payload: Mapping[str, JsonValue], marker: str) -> "execution": "server", "status": "completed", }, search + assert {key: search_output[key] for key in ("type", "id", "call_id", "execution", "status")} == { + "type": "tool_search_output", + "id": f"tso_srvtoolu_{marker}", + "call_id": f"srvtoolu_{marker}", + "execution": "server", + "status": "completed", + }, search_output + assert [tool["name"] for tool in _ITEMS.validate_python(search_output["tools"])] == ["get_weather"], search_output assert (function["type"], function["name"], function["call_id"]) == ( "function_call", "get_weather", @@ -445,6 +458,8 @@ def _item_events(events: Sequence[Mapping[str, JsonValue]], item_type: str) -> l def _assert_search_call_events(events: Sequence[Mapping[str, JsonValue]], marker: str, final_status: str) -> None: + added_indexes: Final = [event["output_index"] for event in events if event["type"] == "response.output_item.added"] + assert added_indexes == list(range(len(added_indexes))), [event["type"] for event in events] search_events: Final = _item_events(events, "tool_search_call") assert [event["type"] for event in search_events] == [ "response.output_item.added", @@ -593,6 +608,24 @@ def test_tool_search_reaches_claude_with_deferred_tools_and_returns_a_tool_searc _assert_spend_rows(model, (str(answer.payload["id"]),)) +def test_the_openai_sdk_stream_helper_reads_a_streamed_claude_tool_search(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_Vendor().respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + body: Final = _responses_body(model, marker, (_tool_search(), _function_tool(), _namespace_tool()), True) + with _sdk(gateway).responses.stream( + model=model, + input=str(body["input"]), + tools=_ITEMS.validate_python(body["tools"]), + extra_body=dict(_CACHE_BUST), + ) as stream: + events: Final = [event.model_dump() for event in stream] + final: Final = stream.get_final_response() + _assert_search_call_output(final.model_dump(), marker) + _assert_search_call_events(events, marker, "completed") + _only_request(wire) + + @pytest.mark.parametrize("stream", [False, True], ids=["sync", "stream"]) @pytest.mark.parametrize("provider", ["bedrock_invoke", "vertex_ai"]) def test_bedrock_invoke_and_vertex_claude_keep_the_tool_search_tool( @@ -665,8 +698,10 @@ def test_search_arguments_that_never_parse_stream_an_incomplete_tool_search_call tools: Final = (_tool_search(), _function_tool()) answer: Final = _raw_answer(gateway, _responses_body(model, marker, tools, True)) _assert_search_call_events(answer.events, marker, "incomplete") - search, function = _tool_items(answer.payload) + search, search_output, function = _tool_items(answer.payload) assert (search["type"], search["status"], search["arguments"]) == ("tool_search_call", "incomplete", {}), search + assert search_output["type"] == "tool_search_output", search_output + assert search_output["call_id"] == search["call_id"], (search, search_output) assert (function["type"], function["name"]) == ("function_call", "get_weather"), function request, body = _only_request(wire) _assert_anthropic_wire(request, body, ("get_weather",)) @@ -678,7 +713,11 @@ def test_tool_search_alone_reaches_claude_as_the_only_tool(gateway: Gateway) -> with wire_server(_Vendor().respond) as wire, gateway.scenario() as scenario: model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) answer: Final = _raw_answer(gateway, _responses_body(model, marker, (_tool_search(),), False)) - assert [item["type"] for item in _tool_items(answer.payload)] == ["tool_search_call", "function_call"], answer + assert [item["type"] for item in _tool_items(answer.payload)] == [ + "tool_search_call", + "tool_search_output", + "function_call", + ], answer request, body = _only_request(wire) _assert_anthropic_wire(request, body, ()) @@ -891,9 +930,18 @@ def _assert_only_the_function_call_is_replayed(messages: Sequence[Mapping[str, J assert [block["tool_use_id"] for block in _tool_results(messages)] == [f"toolu_{marker}"], messages -def test_tool_search_items_are_not_replayed_to_claude(gateway: Gateway) -> None: - marker: Final = uuid.uuid4().hex - history: Final[list[JsonValue]] = [ +def _assert_the_search_is_replayed_before_the_function_call( + messages: Sequence[Mapping[str, JsonValue]], marker: str +) -> None: + (assistant,) = [turn for turn in messages if turn["role"] == "assistant"] + blocks: Final = _ITEMS.validate_python(assistant["content"]) + assert blocks[:2] == list(_search_blocks(marker)), blocks + assert [(block["type"], block.get("id")) for block in blocks[2:]] == [("tool_use", f"toolu_{marker}")], blocks + assert [block["tool_use_id"] for block in _tool_results(messages)] == [f"toolu_{marker}"], messages + + +def _replayed_search_history(marker: str) -> list[JsonValue]: + return [ {"role": "user", "content": _prompt(marker)}, { "type": "tool_search_call", @@ -921,11 +969,16 @@ def test_tool_search_items_are_not_replayed_to_claude(gateway: Gateway) -> None: {"type": "function_call_output", "call_id": f"toolu_{marker}", "output": "22C"}, {"role": "user", "content": f"And tomorrow? marker-{marker}"}, ] + + +@pytest.mark.parametrize("provider", ["anthropic", "bedrock_invoke", "vertex_ai"]) +def test_a_replayed_tool_search_reaches_claude_as_its_search_and_result(gateway: Gateway, provider: Provider) -> None: + marker: Final = uuid.uuid4().hex with wire_server(_Vendor().respond) as wire, gateway.scenario() as scenario: - model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + model: Final = _deployment(scenario, provider, wire.url, gateway.upstream_url) body: Final[dict[str, JsonValue]] = { "model": model, - "input": history, + "input": _replayed_search_history(marker), "tools": [_tool_search(), _function_tool()], **_CACHE_BUST, } @@ -933,10 +986,85 @@ def test_tool_search_items_are_not_replayed_to_claude(gateway: Gateway) -> None: assert response.status_code == 200, response.text _assert_search_call_output(_JSON_OBJECT.validate_json(response.content), marker) request, sent = _only_request(wire) - _assert_anthropic_wire(request, sent, ("get_weather",)) + _assert_claude_wire(provider, False, request, sent, ("get_weather",)) + _assert_the_search_is_replayed_before_the_function_call(_ITEMS.validate_python(sent["messages"]), marker) + + +@pytest.mark.parametrize("provider", ["anthropic", "bedrock_invoke", "vertex_ai"]) +def test_a_replayed_tool_search_whose_tool_is_no_longer_sent_is_left_out(gateway: Gateway, provider: Provider) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_Vendor().respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, provider, wire.url, gateway.upstream_url) + body: Final[dict[str, JsonValue]] = { + "model": model, + "input": _replayed_search_history(marker), + "tools": [_tool_search(), _namespace_tool()], + **_CACHE_BUST, + } + response: Final = _raw(gateway, "/v1/responses", body) + assert response.status_code == 200, response.text + request, sent = _only_request(wire) _assert_only_the_function_call_is_replayed(_ITEMS.validate_python(sent["messages"]), marker) +@pytest.mark.parametrize("stream", [False, True], ids=["sync", "stream"]) +def test_a_tool_name_anthropic_rewrites_comes_back_and_is_replayed_under_the_callers_name( + gateway: Gateway, stream: bool +) -> None: + marker: Final = uuid.uuid4().hex + vendor: Final = _Vendor(tool_use_name="weather_get", found_tool="weather_get") + tools: Final = (_tool_search(), {**_function_tool(), "name": "weather.get"}) + with wire_server(vendor.respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + items: Final = _tool_items(_raw_answer(gateway, _responses_body(model, marker, tools, stream)).payload) + _only_request(wire) + assert [(item["type"], item.get("name")) for item in items] == [ + ("tool_search_call", None), + ("tool_search_output", None), + ("function_call", "weather.get"), + ], items + assert [tool["name"] for tool in _ITEMS.validate_python(items[1]["tools"])] == ["weather.get"], items[1] + replay: Final = _raw( + gateway, + "/v1/responses", + { + "model": model, + "store": False, + "input": [ + {"role": "user", "content": _prompt(marker)}, + *items, + {"type": "function_call_output", "call_id": f"toolu_{marker}", "output": "22C"}, + {"role": "user", "content": f"And tomorrow? marker-{marker}"}, + ], + "tools": [dict(tool) for tool in tools], + **_CACHE_BUST, + }, + ) + assert replay.status_code == 200, replay.text + _, sent = _only_request(wire) + (assistant,) = [turn for turn in _ITEMS.validate_python(sent["messages"]) if turn["role"] == "assistant"] + blocks: Final = _ITEMS.validate_python(assistant["content"]) + assert [block for block in blocks if block["type"] != "text"] == [ + *_search_blocks(marker, "weather_get"), + {"type": "tool_use", "id": f"toolu_{marker}", "name": "weather_get", "input": {"city": "Paris"}}, + ], blocks + + +@pytest.mark.parametrize("stream", [False, True], ids=["sync", "stream"]) +@pytest.mark.parametrize("provider", ["anthropic", "bedrock_invoke"]) +def test_the_tool_search_output_lists_the_callers_tool_when_anthropic_rewrites_its_name( + gateway: Gateway, provider: Provider, stream: bool +) -> None: + marker: Final = uuid.uuid4().hex + vendor: Final = _Vendor(tool_use_name="weather_get", found_tool="weather_get") + tools: Final = (_tool_search(), {**_function_tool(), "name": "weather.get"}) + with wire_server(vendor.respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, provider, wire.url, gateway.upstream_url) + items: Final = _tool_items(_raw_answer(gateway, _responses_body(model, marker, tools, stream)).payload) + (search_output,) = [item for item in items if item["type"] == "tool_search_output"] + assert [tool["name"] for tool in _ITEMS.validate_python(search_output["tools"])] == ["weather.get"], items + + def test_chat_completions_drop_a_replayed_search_call_that_has_no_result(gateway: Gateway) -> None: marker: Final = uuid.uuid4().hex with wire_server(_Vendor().respond) as wire, gateway.scenario() as scenario: 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 bf093f1a0d8..54f55aba46a 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -3794,7 +3794,9 @@ def _tool_search_turn(tool_references: list[str], last_block: dict[str, object]) ] -def _parsed_choice(content: list[dict[str, object]], stop_reason: str) -> Choices: +def _parsed_choice( + content: list[dict[str, object]], stop_reason: str, tool_name_reverse_map: dict[str, str] | None = None +) -> Choices: anthropic_response: Final = { "id": "msg_01XYZ", "type": "message", @@ -3817,6 +3819,7 @@ def _parsed_choice(content: list[dict[str, object]], stop_reason: str) -> Choice model_response=ModelResponse(), json_mode=False, prefix_prompt=None, + tool_name_reverse_map=tool_name_reverse_map, ) .choices[0] ) @@ -3824,8 +3827,10 @@ def _parsed_choice(content: list[dict[str, object]], stop_reason: str) -> Choice return choice -def _parsed_assistant_message(content: list[dict[str, object]], stop_reason: str) -> Message: - return _parsed_choice(content, stop_reason).message +def _parsed_assistant_message( + content: list[dict[str, object]], stop_reason: str, tool_name_reverse_map: dict[str, str] | None = None +) -> Message: + return _parsed_choice(content, stop_reason, tool_name_reverse_map).message _WEATHER_TOOL_USE: Final = { @@ -4092,6 +4097,110 @@ def test_replayed_search_that_found_nothing_answers_every_tool_use_and_ends_on_a assert _block_ids(replayed[1], "server_tool_use", "id") == server_tool_use_ids +def _function_tool(name: str) -> dict[str, object]: + return { + "type": "function", + "function": {"name": name, "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}}, + "defer_loading": True, + } + + +def _anthropic_request_after( + assistant_content: list[dict[str, object]], stop_reason: str, tool_names: tuple[str, ...] | None +) -> dict[str, object]: + config: Final = AnthropicConfig() + message: Final = _parsed_assistant_message(assistant_content, stop_reason) + answers: Final = [ + {"role": "tool", "tool_call_id": block["id"], "content": "18C"} + for block in assistant_content + if block["type"] == "tool_use" + ] + optional_params: Final = ( + {} + if tool_names is None + else config.map_openai_params( + non_default_params={ + "tools": [ + {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}, + *map(_function_tool, tool_names), + ] + }, + optional_params={}, + model="claude-opus-5-5", + drop_params=False, + ) + ) + return config.transform_request( + model="claude-opus-5-5", + messages=[ + {"role": "user", "content": "What is the weather in Paris?"}, + message.model_dump(), + *answers, + {"role": "user", "content": "And in Osaka?"}, + ], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + +@pytest.mark.parametrize( + ("tool_names", "search_replayed"), + [ + pytest.param(("get_weather", "get_local_time"), True, id="every_loaded_tool_is_still_sent"), + pytest.param(("get_local_time",), False, id="a_loaded_tool_is_no_longer_sent"), + pytest.param(None, False, id="no_tools_so_litellm_sends_a_dummy_tool"), + ], +) +def test_replayed_tool_search_is_sent_only_while_every_tool_it_loaded_is_still_sent( + tool_names: tuple[str, ...] | None, search_replayed: bool +): + search_turn: Final = _tool_search_turn(["get_weather"], _WEATHER_TOOL_USE) + + request: Final = _anthropic_request_after(search_turn, "tool_use", tool_names) + + assert request["messages"][1]["content"] == (search_turn if search_replayed else [_WEATHER_TOOL_USE]) + assert [block["type"] for block in request["messages"][2]["content"]] == ["tool_result", "text"] + + +def test_replayed_tool_search_turn_that_held_only_the_search_is_dropped_once_its_tool_is_no_longer_sent(): + request: Final = _anthropic_request_after(_tool_search_turn(["get_weather"], _WEATHER_TOOL_USE)[:2], "end_turn", ()) + + assert [message["role"] for message in request["messages"]] == ["user", "user"] + assert "srvtoolu_01Search" not in json.dumps(request["messages"]) + + +def test_replayed_tool_search_result_litellm_cannot_read_is_sent_as_the_client_replayed_it(): + unreadable_result: Final = { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_01Search", + "content": {"tool_references": "get_weather"}, + } + message: Final = _parsed_assistant_message(_tool_search_turn(["get_weather"], _WEATHER_TOOL_USE), "tool_use") + provider_fields: Final = {**(message.provider_specific_fields or {}), "tool_results": [unreadable_result]} + replayed: Final = message.model_copy(update={"provider_specific_fields": provider_fields}) + config: Final = AnthropicConfig() + + request: Final = config.transform_request( + model="claude-opus-5-5", + messages=[ + {"role": "user", "content": "What is the weather in Paris?"}, + replayed.model_dump(), + {"role": "tool", "tool_call_id": "toolu_01Weather", "content": "18C"}, + ], + optional_params=config.map_openai_params( + non_default_params={"tools": [_function_tool("get_local_time")]}, + optional_params={}, + model="claude-opus-5-5", + drop_params=False, + ), + litellm_params={}, + headers={}, + ) + + assert unreadable_result in request["messages"][1]["content"] + + def test_replayed_search_that_found_nothing_keeps_its_result_when_the_call_never_reached_the_client( monkeypatch: pytest.MonkeyPatch, ): @@ -5656,10 +5765,10 @@ def test_basic_sanitize_anthropic_tool_name_replaces_invalid_chars(): def test_build_anthropic_tool_name_maps_no_collisions(): """Names that need rewriting go in the maps; valid names stay out.""" from litellm.llms.anthropic.chat.transformation import ( - _build_anthropic_tool_name_maps, + build_anthropic_tool_name_maps, ) - forward, reverse = _build_anthropic_tool_name_maps( + forward, reverse = build_anthropic_tool_name_maps( [ "fine_name", "actions/download-job-logs-for-workflow-run", @@ -5681,10 +5790,10 @@ def test_build_anthropic_tool_name_maps_disambiguates_collision_with_existing_va the rewritten one must get a unique suffix and only THAT one shows up in the reverse map. The legitimately-named `foo_bar` round-trips identically.""" from litellm.llms.anthropic.chat.transformation import ( - _build_anthropic_tool_name_maps, + build_anthropic_tool_name_maps, ) - forward, reverse = _build_anthropic_tool_name_maps(["foo_bar", "foo/bar"]) + forward, reverse = build_anthropic_tool_name_maps(["foo_bar", "foo/bar"]) # The original valid name keeps its slot. assert "foo_bar" not in forward # untouched # The rewritten one gets a disambiguating suffix. @@ -5700,10 +5809,10 @@ def test_build_anthropic_tool_name_maps_disambiguates_two_rewrites_to_same_targe """Two different invalid names that collapse to the same candidate must both end up with unique sanitized forms.""" from litellm.llms.anthropic.chat.transformation import ( - _build_anthropic_tool_name_maps, + build_anthropic_tool_name_maps, ) - forward, reverse = _build_anthropic_tool_name_maps(["foo/bar", "foo.bar"]) + forward, reverse = build_anthropic_tool_name_maps(["foo/bar", "foo.bar"]) # First wins the canonical slot, second gets a suffix. assert forward["foo/bar"] == "foo_bar" assert forward["foo.bar"] == "foo_bar_2" @@ -5715,10 +5824,10 @@ def test_build_anthropic_tool_name_maps_disambiguates_two_rewrites_to_same_targe def test_build_anthropic_tool_name_maps_three_way_collision(): """`foo/bar`, `foo.bar`, and an existing `foo_bar` must all coexist.""" from litellm.llms.anthropic.chat.transformation import ( - _build_anthropic_tool_name_maps, + build_anthropic_tool_name_maps, ) - forward, reverse = _build_anthropic_tool_name_maps(["foo_bar", "foo/bar", "foo.bar"]) + forward, reverse = build_anthropic_tool_name_maps(["foo_bar", "foo/bar", "foo.bar"]) assert "foo_bar" not in forward # untouched assert forward["foo/bar"] == "foo_bar_2" assert forward["foo.bar"] == "foo_bar_3" @@ -5733,10 +5842,10 @@ def test_build_anthropic_tool_name_maps_reverse_order_collision(): its sanitized form collides with, both must still end up with distinct names on the wire.""" from litellm.llms.anthropic.chat.transformation import ( - _build_anthropic_tool_name_maps, + build_anthropic_tool_name_maps, ) - forward, reverse = _build_anthropic_tool_name_maps(["foo/bar", "foo_bar"]) + forward, reverse = build_anthropic_tool_name_maps(["foo/bar", "foo_bar"]) # The valid name keeps its slot untouched. assert "foo_bar" not in forward # The rewritten one gets a disambiguating suffix. @@ -5756,10 +5865,10 @@ def test_build_anthropic_tool_name_maps_duplicate_originals(): on duplicates. """ from litellm.llms.anthropic.chat.transformation import ( - _build_anthropic_tool_name_maps, + build_anthropic_tool_name_maps, ) - forward, reverse = _build_anthropic_tool_name_maps(["foo/bar", "foo/bar"]) + forward, reverse = build_anthropic_tool_name_maps(["foo/bar", "foo/bar"]) # Same original sanitizes to the same target — no spurious suffix. assert forward == {"foo/bar": "foo_bar"} assert reverse == {"foo_bar": "foo/bar"} @@ -10707,6 +10816,136 @@ def test_transform_request_tool_config_point_without_tools_drops_the_key(): assert litellm_params["metadata"] == {} +def test_a_renamed_tool_leaves_tool_results_without_tool_references_as_claude_sent_them(): + errored_search: Final = { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_01Search", + "content": {"type": "tool_search_tool_result_error", "error_code": "unavailable"}, + } + results: Final = [_CODE_EXECUTION_BLOCKS[1], errored_search] + message: Final = _parsed_assistant_message( + [ + *_CODE_EXECUTION_BLOCKS, + _tool_search_turn([], _WEATHER_TOOL_USE)[0], + errored_search, + {**_WEATHER_TOOL_USE, "name": "weather_get"}, + ], + "tool_use", + {"weather_get": "weather.get"}, + ) + config: Final = AnthropicConfig() + + request: Final = config.transform_request( + model="claude-opus-5-5", + messages=[ + {"role": "user", "content": "What is the weather in Paris?"}, + message.model_dump(), + {"role": "tool", "tool_call_id": "toolu_01Weather", "content": "18C"}, + ], + optional_params=config.map_openai_params( + non_default_params={ + "tools": [ + {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}, + _function_tool("weather.get"), + ] + }, + optional_params={}, + model="claude-opus-5-5", + drop_params=False, + ), + litellm_params={}, + headers={}, + ) + + assert message.provider_specific_fields["tool_results"] == results + assert message.tool_calls[-1].function.name == "weather.get" + sent: Final = request["messages"][1]["content"] + assert [block for block in sent if block["type"].endswith("_tool_result")] == results + assert [block["name"] for block in sent if block["type"] == "tool_use"] == ["weather_get"] + + +@pytest.mark.parametrize( + ("last_block", "stop_reason"), + [ + pytest.param(dict(_PARIS_ANSWER), "end_turn", id="search_then_answer"), + pytest.param({**_WEATHER_TOOL_USE, "name": "weather_get"}, "tool_use", id="search_then_tool"), + ], +) +def test_a_renamed_tool_keeps_the_replayed_search_that_loaded_it(last_block: dict[str, object], stop_reason: str): + search_turn: Final = _tool_search_turn(["weather_get"], last_block) + message: Final = _parsed_assistant_message(search_turn, stop_reason, {"weather_get": "weather.get"}) + answers: Final = ( + [{"role": "tool", "tool_call_id": "toolu_01Weather", "content": "18C"}] if stop_reason == "tool_use" else [] + ) + config: Final = AnthropicConfig() + + request: Final = config.transform_request( + model="claude-opus-5-5", + messages=[ + {"role": "user", "content": "What is the weather in Paris?"}, + message.model_dump(), + *answers, + {"role": "user", "content": "And in Osaka?"}, + ], + optional_params=config.map_openai_params( + non_default_params={ + "tools": [ + {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}, + _function_tool("weather.get"), + ] + }, + optional_params={}, + model="claude-opus-5-5", + drop_params=False, + ), + litellm_params={}, + headers={}, + ) + + sent: Final = request["messages"][1]["content"] + assert [block for block in sent if block["type"] == "tool_search_tool_result"] == [search_turn[1]] + assert [block["id"] for block in sent if block["type"] == "server_tool_use"] == ["srvtoolu_01Search"] + + +def _parsed_tool_search_fields( + content: list[dict[str, object]], tool_name_reverse_map: dict[str, str] +) -> tuple[object, list[str]]: + message: Final = _parsed_assistant_message(content, "tool_use", tool_name_reverse_map) + return message.provider_specific_fields["tool_results"], [call.function.name for call in message.tool_calls] + + +def _streamed_tool_search_fields( + content: list[dict[str, object]], tool_name_reverse_map: dict[str, str] +) -> tuple[object, list[str]]: + iterator: Final = ModelResponseIterator(None, sync_stream=True, tool_name_reverse_map=tool_name_reverse_map) + events: Final = chain.from_iterable( + ( + {"type": "content_block_start", "index": index, "content_block": block}, + {"type": "content_block_stop", "index": index}, + ) + for index, block in enumerate(content) + ) + deltas: Final = [iterator.chunk_parser(event).choices[0].delta for event in events] + fields: Final = [delta.provider_specific_fields for delta in deltas if delta.provider_specific_fields] + calls: Final = chain.from_iterable(delta.tool_calls or () for delta in deltas) + return fields[-1]["tool_results"], [call.function.name for call in calls] + + +@pytest.mark.parametrize("stream", [False, True], ids=["sync", "stream"]) +def test_a_renamed_tool_search_result_keeps_claudes_names_so_an_older_proxy_accepts_the_replay(stream: bool): + search_turn: Final = _tool_search_turn(["weather_get"], {**_WEATHER_TOOL_USE, "name": "weather_get", "input": {}}) + reverse_map: Final = {"weather_get": "weather.get"} + + tool_results, tool_names = ( + _streamed_tool_search_fields(search_turn, reverse_map) + if stream + else _parsed_tool_search_fields(search_turn, reverse_map) + ) + + assert tool_results == [search_turn[1]] + assert tool_names == ["weather.get"] + + _INPUT_TRANSFORMATIONS: Final = [{"type": "thinking_block_binding", "count": 2}] diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index 14cabb683f1..4e8dc69fd01 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -28,7 +28,12 @@ import pytest import respx import litellm -from litellm.llms.anthropic.common_utils import ToolsWithoutStrict, is_claude_tool_target, tools_without_strict +from litellm.llms.anthropic.common_utils import ( + ToolsWithoutStrict, + is_claude_tool_target, + tools_without_strict, + without_stale_tool_search_results, +) from litellm.llms.custom_httpx.http_handler import HTTPHandler sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) @@ -4709,6 +4714,79 @@ def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: b assert is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected +def test_dropping_a_stale_tool_search_keeps_every_other_turn_as_it_was_sent() -> None: + search_use: Final = { + "type": "server_tool_use", + "id": "srvtoolu_01Search", + "name": "tool_search_tool_regex", + "input": {"pattern": "weather"}, + } + search_result: Final = { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_01Search", + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}], + }, + } + weather_call: Final = { + "type": "tool_use", + "id": "toolu_01Weather", + "name": "get_weather", + "input": {"city": "Paris"}, + } + messages: Final = [ + {"role": "user", "content": "What is the weather in Paris?"}, + {"role": "assistant", "content": [search_use, search_result, weather_call]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01Weather", "content": "18C"}]}, + {"role": "assistant", "content": "It is 18C in Paris."}, + {"role": "user", "content": "Thanks"}, + {"role": "assistant", "content": [{"type": "text", "text": "Anything else?"}]}, + {"role": "user", "content": "And in Osaka?"}, + ] + + kept: Final = without_stale_tool_search_results(messages, [{"name": "get_local_time"}]) + + assert kept == (messages[0], {"role": "assistant", "content": [weather_call]}, *messages[2:]) + + +def test_a_replayed_tool_search_keeps_the_tools_it_found_that_the_follow_up_still_sends() -> None: + def reference(name: str) -> dict[str, str]: + return {"type": "tool_reference", "tool_name": name} + + search_use: Final = { + "type": "server_tool_use", + "id": "srvtoolu_01Search", + "name": "tool_search_tool_regex", + "input": {"pattern": "weather|time"}, + } + + def search_result(*names: str) -> dict[str, object]: + return { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_01Search", + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [reference(name) for name in names], + }, + } + + weather_call: Final = {"type": "tool_use", "id": "toolu_01Weather", "name": "get_weather", "input": {"city": "Paris"}} + messages: Final = [ + {"role": "user", "content": "What is the weather and the time in Paris?"}, + {"role": "assistant", "content": [search_use, search_result("get_weather", "get_local_time"), weather_call]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01Weather", "content": "18C"}]}, + ] + + kept: Final = without_stale_tool_search_results(messages, [{"name": "get_weather"}, {"name": "get_news"}]) + + assert kept == ( + messages[0], + {"role": "assistant", "content": [search_use, search_result("get_weather"), weather_call]}, + messages[2], + ) + + @pytest.mark.parametrize( "tools", [ diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py index 4f17f8b10e1..5e056e494a1 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py @@ -1164,6 +1164,92 @@ def test_bedrock_chat_invoke_tool_search_beta_follows_model_map( assert result.get("anthropic_beta") == expected_betas +_REPLAYED_WEATHER_SEARCH: Final = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "toolu_01Weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01Search", + "type": "function", + "function": {"name": "tool_search_tool_regex", "arguments": '{"pattern": "weather"}'}, + } + ], + "tool_results": [ + { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_01Search", + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}], + }, + } + ], + }, +} +_PLAIN_WEATHER_CALL: Final = { + key: value for key, value in _REPLAYED_WEATHER_SEARCH.items() if key != "provider_specific_fields" +} + + +@pytest.mark.parametrize( + ("model", "assistant_turn", "expected_betas"), + [ + pytest.param( + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + _REPLAYED_WEATHER_SEARCH, + ["tool-search-tool-2025-10-19"], + id="replayed_search", + ), + pytest.param("us.anthropic.claude-sonnet-4-5-20250929-v1:0", _PLAIN_WEATHER_CALL, None, id="no_search"), + pytest.param( + "us.anthropic.claude-opus-4-1-20250805-v1:0", _REPLAYED_WEATHER_SEARCH, None, id="opus_4_1_unsupported" + ), + ], +) +def test_bedrock_chat_invoke_sends_the_tool_search_beta_for_a_search_replayed_without_the_search_tool( + local_model_cost_map, local_beta_headers_config, model, assistant_turn, expected_betas +): + config: Final = AmazonAnthropicClaudeConfig() + result: Final = config.transform_request( + model=model, + messages=[ + {"role": "user", "content": "What is the weather in Paris?"}, + assistant_turn, + {"role": "tool", "tool_call_id": "toolu_01Weather", "content": "18C"}, + {"role": "user", "content": "And in Osaka?"}, + ], + optional_params=config.map_openai_params( + non_default_params={ + "max_tokens": 64, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + } + ], + }, + optional_params={}, + model=model, + drop_params=False, + ), + litellm_params={}, + headers={}, + ) + + assert result.get("anthropic_beta") == expected_betas + + def test_bedrock_chat_invoke_adds_thinking_display_updates_beta( local_model_cost_map, local_beta_headers_config ) -> None: 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 7daf773d55b..8902cc59f0f 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 @@ -9,6 +9,7 @@ import copy import json import uuid from collections.abc import Callable +from itertools import chain from typing import Any, Final, List, Literal, Optional, Tuple from unittest.mock import AsyncMock, MagicMock, patch @@ -5245,6 +5246,130 @@ class TestCompactionOutputSummaryReachesOutputGuardrails: assert ended_key(self.SUMMARY) != ended_key("The user lives elsewhere.") +@pytest.mark.asyncio +async def test_output_guardrail_rewrite_reaches_the_text_of_an_openai_sdk_message_item() -> None: + from openai.types.responses import ResponseOutputMessage, ResponseOutputText + + response: Final = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + model="gpt-4", + object="response", + status="completed", + output=[ + ResponseOutputMessage( + type="message", + id="msg_123", + status="completed", + role="assistant", + content=[ResponseOutputText(type="output_text", text="My SSN is 123-45-6789", annotations=[])], + ) + ], + ) + + await OpenAIResponsesHandler().process_output_response( + response, TextsReplacingGuardrail(guardrail_name="masker", texts=("My SSN is [MASKED]",)) + ) + + assert [part.text for part in response.output[0].content] == ["My SSN is [MASKED]"] + + +class _WordMaskingGuardrail(CustomGuardrail): + def __init__(self, word: str) -> None: + super().__init__(guardrail_name="word-masker") + self.word = word + self.seen_tool_call_arguments: list[str] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + messages: Final = inputs.get("structured_messages") or [] + tool_calls: Final = chain.from_iterable(message.get("tool_calls") or [] for message in messages) + self.seen_tool_call_arguments.extend(tool_call["function"]["arguments"] for tool_call in tool_calls) + return {**inputs, "structured_messages": json.loads(json.dumps(messages).replace(self.word, "[MASKED]"))} + + +_REPLAYED_CLAUDE_SEARCH_INPUT: Final = ( + {"role": "user", "content": "What is the weather in Paris?"}, + { + "type": "tool_search_call", + "id": "tsc_srvtoolu_01Search", + "call_id": "srvtoolu_01Search", + "execution": "server", + "status": "completed", + "arguments": {"pattern": "Paris weather"}, + }, + { + "type": "tool_search_output", + "id": "tso_srvtoolu_01Search", + "call_id": "srvtoolu_01Search", + "execution": "server", + "status": "completed", + "tools": [{"type": "function", "name": "get_weather", "parameters": {"type": "object"}, "defer_loading": True}], + }, + {"type": "function_call", "call_id": "toolu_01GuardrailWeather", "name": "get_weather", "arguments": '{"city": "Paris"}'}, + {"type": "function_call_output", "call_id": "toolu_01GuardrailWeather", "output": "14C"}, + {"role": "user", "content": "Thanks"}, +) + + +@pytest.mark.asyncio +async def test_a_guardrail_masks_the_arguments_of_a_replayed_claude_tool_search_and_keeps_the_search() -> None: + guardrail: Final = _WordMaskingGuardrail("Paris") + data: Final = { + "model": "anthropic/claude-sonnet-5-5", + "input": copy.deepcopy(list(_REPLAYED_CLAUDE_SEARCH_INPUT)), + "tools": [{"type": "tool_search"}], + } + + result: Final = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert guardrail.seen_tool_call_arguments == ['{"pattern": "Paris weather"}', '{"city": "Paris"}'] + assert result["input"][1:4] == [ + {**_REPLAYED_CLAUDE_SEARCH_INPUT[1], "arguments": {"pattern": "[MASKED] weather"}}, + _REPLAYED_CLAUDE_SEARCH_INPUT[2], + {**_REPLAYED_CLAUDE_SEARCH_INPUT[3], "arguments": '{"city": "[MASKED]"}'}, + ] + assert "Paris" not in json.dumps(result["input"]) + + +@pytest.mark.asyncio +async def test_a_follow_up_without_a_search_tool_still_scans_the_claude_search_and_keeps_it_through_a_rewrite() -> None: + guardrail: Final = _WordMaskingGuardrail("Paris") + sent_input: Final = [{"role": "user", "content": "What is the weather?"}, *_REPLAYED_CLAUDE_SEARCH_INPUT[1:]] + data: Final = { + "model": "anthropic/claude-sonnet-5-5", + "input": copy.deepcopy(sent_input), + "tools": [{"type": "function", "name": "get_weather", "parameters": {"type": "object"}}], + } + + result: Final = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert guardrail.seen_tool_call_arguments == ['{"pattern": "Paris weather"}', '{"city": "Paris"}'] + assert result["input"][1:4] == [ + {**sent_input[1], "arguments": {"pattern": "[MASKED] weather"}}, + sent_input[2], + {**sent_input[3], "arguments": '{"city": "[MASKED]"}'}, + ] + + +@pytest.mark.asyncio +async def test_a_guardrail_rewriting_another_turn_leaves_a_replayed_claude_tool_search_as_sent() -> None: + data: Final = { + "model": "anthropic/claude-sonnet-5-5", + "input": copy.deepcopy(list(_REPLAYED_CLAUDE_SEARCH_INPUT)), + "tools": [{"type": "tool_search"}], + } + + result: Final = await OpenAIResponsesHandler().process_input_messages(data, _WordMaskingGuardrail("Thanks")) + + assert result["input"] == [*_REPLAYED_CLAUDE_SEARCH_INPUT[:5], {"role": "user", "content": "[MASKED]"}] + + _REFUSAL: Final = "I can't share where Jane lives." _MASKED_REFUSAL: Final = "I can't share where [NAME] lives." 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 ea6892379c2..8500a1083df 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py @@ -2805,6 +2805,46 @@ class TestRegistration: ) +def test_responses_input_scan_includes_the_arguments_of_a_replayed_claude_tool_search(): + guardrail = _make_guardrail() + data = { + "model": "claude-sonnet-5-5", + "input": [ + {"role": "user", "content": "What is the weather?"}, + { + "type": "tool_search_call", + "id": "tsc_srvtoolu_01Search", + "call_id": "srvtoolu_01Search", + "arguments": {"pattern": "ssn 123-45-6789"}, + "execution": "server", + "status": "completed", + }, + { + "type": "tool_search_output", + "id": "tso_01Search", + "call_id": "srvtoolu_01Search", + "tools": [{"type": "function", "name": "get_weather"}], + "status": "completed", + }, + { + "type": "function_call", + "id": "fc_01Weather", + "call_id": "toolu_01Weather", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + {"type": "function_call_output", "call_id": "toolu_01Weather", "output": "14C"}, + ], + } + + scanned = guardrail._responses_api_input_to_str(data) + + assert scanned is not None + assert "123-45-6789" in scanned + assert "Paris" in scanned + + @pytest.mark.parametrize("odata_type", [None, 7], ids=["null", "number"]) def test_should_block_rejects_a_policy_action_whose_odata_type_is_not_text(odata_type: object) -> None: response: Final = {"policyActions": [{"@odata.type": odata_type, "restrictionAction": "block"}]} diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 5ae985ea7e5..cbd26618312 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -18,6 +18,7 @@ from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) from litellm.responses.litellm_completion_transformation.session_handler import StoredTurns +from litellm.responses.litellm_completion_transformation.tool_search import anthropic_tool_search_chat_tool from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -372,6 +373,9 @@ def _search_turn(tool_calls: Sequence[Mapping[str, object]]) -> Mapping[str, obj } +_ANTHROPIC_SEARCH_TOOL: Final = anthropic_tool_search_chat_tool() + + async def _follow_up_messages( tool_calls: Sequence[Mapping[str, object]], model: str, custom_llm_provider: str ) -> Sequence[object]: @@ -379,7 +383,10 @@ async def _follow_up_messages( async def _follow_up_messages_after( - spend_log: Mapping[str, object], model: str, custom_llm_provider: str + spend_log: Mapping[str, object], + model: str, + custom_llm_provider: str, + tools: Sequence[Mapping[str, object]] = (_ANTHROPIC_SEARCH_TOOL,), ) -> Sequence[object]: with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([spend_log])): request: Final = await LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( @@ -388,6 +395,7 @@ async def _follow_up_messages_after( "model": model, "custom_llm_provider": custom_llm_provider, "messages": [{"role": "user", "content": "And in Osaka?"}], + **({"tools": [dict(tool) for tool in tools]} if tools else {}), }, instructions=None, ) @@ -492,7 +500,12 @@ def _history_tool_ids(messages: Sequence[object]) -> tuple[Sequence[str], Sequen @pytest.mark.parametrize("search_answered", [False, True]) @pytest.mark.parametrize( ("model", "custom_llm_provider", "expect_search"), - [("mistral/mistral-small-latest", "mistral", False), (_MODEL, "anthropic", True)], + [ + ("mistral/mistral-small-latest", "mistral", False), + ("anthropic/claude-3-7-sonnet-20250219", "anthropic", True), + ("azure_ai/claude-deployment", "azure_ai", True), + (_MODEL, "anthropic", True), + ], ) async def test_previous_response_id_drops_a_client_replayed_search_call_only_off_anthropic( model: str, custom_llm_provider: str, expect_search: bool, search_answered: bool @@ -1598,6 +1611,369 @@ async def test_previous_response_id_keeps_the_tools_an_earlier_codex_tool_search assert prisma_client.db.queries == 1 +_WEATHER_TOOL: Final = { + "type": "function", + "name": "get_weather", + "description": "Get the current weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + "defer_loading": True, +} +_SEARCHABLE_TOOLS: Final = ({"type": "tool_search"}, _WEATHER_TOOL) +_WEATHER_QUESTION: Final = ({"role": "user", "content": "What is the weather in Paris?"},) +_SEARCH_USE: Final = { + "type": "server_tool_use", + "id": "srvtoolu_01Search", + "name": "tool_search_tool_regex", + "input": {"pattern": "weather"}, +} +_SEARCH_RESULT: Final = { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_01Search", + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}], + }, +} +_WEATHER_USE: Final = {"type": "tool_use", "id": "toolu_01Weather", "name": "get_weather", "input": {"city": "Paris"}} +_WEATHER_ANSWER: Final = {"type": "function_call_output", "call_id": "toolu_01Weather", "output": "18C and sunny"} + + +_TOOL_NAMES_ON_EACH_SIDE: Final = pytest.mark.parametrize( + ("caller_name", "claude_name"), + [("get_weather", "get_weather"), ("weather.get", "weather_get")], + ids=["valid_name", "name_anthropic_rewrites"], +) + + +def _search_result_finding(tool_name: str) -> Mapping[str, object]: + return { + **_SEARCH_RESULT, + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [{"type": "tool_reference", "tool_name": tool_name}], + }, + } + + +@_TOOL_NAMES_ON_EACH_SIDE +def test_responses_over_anthropic_return_the_tool_search_output_and_send_the_search_back_without_a_stored_turn( + respx_mock: respx.MockRouter, caller_name: str, claude_name: str +) -> None: + tools: Final = [{"type": "tool_search"}, {**_WEATHER_TOOL, "name": caller_name}] + search_result: Final = _search_result_finding(claude_name) + weather_use: Final = {**_WEATHER_USE, "name": claude_name} + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message(_SEARCH_USE, search_result, weather_use) + ) + + first: Final = litellm.responses(model=_MODEL, input=list(_WEATHER_QUESTION), tools=tools, api_key="mock_api_key") + litellm.responses( + model=_MODEL, + input=[*_WEATHER_QUESTION, *(item.model_dump(exclude_none=True) for item in first.output), _WEATHER_ANSWER], + tools=tools, + api_key="mock_api_key", + ) + + assert [item.type for item in first.output] == ["tool_search_call", "tool_search_output", "function_call"] + search_output: Final = first.output[1] + assert (search_output.id, search_output.call_id) == ("tso_srvtoolu_01Search", "srvtoolu_01Search") + assert (search_output.execution, search_output.status) == ("server", "completed") + assert [dict(tool) for tool in search_output.tools] == tools[1:] + assert first.output[2].name == caller_name + assert _sent_assistant_content(route) == [_SEARCH_USE, search_result, weather_use] + + + +@pytest.mark.parametrize( + ("search_tools", "search_use"), + [ + pytest.param((), _SEARCH_USE, id="regex"), + pytest.param( + ({"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"},), + {**_SEARCH_USE, "name": "tool_search_tool_bm25", "input": {"query": "weather"}}, + id="bm25", + ), + ], +) +def test_responses_over_anthropic_send_a_search_back_under_the_anthropic_search_tool_that_ran_it( + respx_mock: respx.MockRouter, search_tools: tuple[Mapping[str, object], ...], search_use: Mapping[str, object] +) -> None: + tools: Final = [{"type": "tool_search"}, *search_tools, _WEATHER_TOOL] + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message(search_use, _SEARCH_RESULT, _WEATHER_USE) + ) + + first: Final = litellm.responses(model=_MODEL, input=list(_WEATHER_QUESTION), tools=tools, api_key="mock_api_key") + litellm.responses( + model=_MODEL, + input=[*_WEATHER_QUESTION, *(item.model_dump(exclude_none=True) for item in first.output), _WEATHER_ANSWER], + tools=tools, + api_key="mock_api_key", + ) + + assert _sent_assistant_content(route) == [search_use, _SEARCH_RESULT, _WEATHER_USE] + +def _tool_search_stream(*events: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_tool_search_stream", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + *events, + {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + + +def _streamed_tool_use(index: int, block: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + return ( + {"type": "content_block_start", "index": index, "content_block": {**block, "input": {}}}, + { + "type": "content_block_delta", + "index": index, + "delta": {"type": "input_json_delta", "partial_json": json.dumps(block["input"])}, + }, + {"type": "content_block_stop", "index": index}, + ) + + +@_TOOL_NAMES_ON_EACH_SIDE +def test_streamed_responses_over_anthropic_emit_the_tool_search_output_in_order_at_its_completed_position( + respx_mock: respx.MockRouter, caller_name: str, claude_name: str +) -> None: + tools: Final = [{"type": "tool_search"}, {**_WEATHER_TOOL, "name": caller_name}] + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_tool_search_stream( + *_streamed_tool_use(0, _SEARCH_USE), + {"type": "content_block_start", "index": 1, "content_block": _search_result_finding(claude_name)}, + {"type": "content_block_stop", "index": 1}, + *_streamed_tool_use(2, {**_WEATHER_USE, "name": claude_name}), + ) + ) + + events: Final = tuple( + litellm.responses(model=_MODEL, input=list(_WEATHER_QUESTION), tools=tools, stream=True, api_key="mock") + ) + + completed: Final = events[-1].response.output + added: Final = { + event.item.id: event.output_index + for event in events + if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED + } + done: Final = [event for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE] + assert [item.type for item in completed] == ["tool_search_call", "tool_search_output", "function_call"] + assert [event.output_index for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] == [ + 0, + 1, + 2, + ] + assert [(event.output_index, event.item.id) for event in done] == [ + (index, item.id) for index, item in enumerate(completed) + ] + assert {event.item.id: event.output_index for event in done} == added + search_output: Final = next(event.item for event in done if event.item.type == "tool_search_output") + assert search_output.model_dump() == completed[1].model_dump() + assert (search_output.call_id, search_output.status) == ("srvtoolu_01Search", "completed") + assert [dict(tool) for tool in search_output.tools] == tools[1:] + assert completed[2].name == caller_name + + +_STATELESS_SEARCH_TURN: Final = { + **_search_turn([]), + "proxy_server_request": { + "model": _MODEL, + "input": [ + *_WEATHER_QUESTION, + { + "type": "tool_search_call", + "id": "tsc_srvtoolu_01Search", + "call_id": "srvtoolu_01Search", + "arguments": {"pattern": "weather"}, + "execution": "server", + "status": "completed", + }, + { + "type": "tool_search_output", + "id": "tso_srvtoolu_01Search", + "call_id": "srvtoolu_01Search", + "execution": "server", + "status": "completed", + "tools": [_WEATHER_TOOL], + }, + { + "type": "function_call", + "id": "fc_toolu_01Weather", + "call_id": "toolu_01Weather", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + }, + _WEATHER_ANSWER, + ], + }, + "response": { + "id": "chatcmpl-search-turn", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [ + {"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "It is 18C in Paris."}} + ], + }, +} + + +@pytest.mark.asyncio +async def test_previous_response_id_sends_back_a_search_the_stored_turn_got_as_input_items() -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_STATELESS_SEARCH_TURN])): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-search-turn", + input="And in Osaka?", + tools=list(_SEARCHABLE_TOOLS), + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + messages: Final = json.loads(anthropic.body)["messages"] + assert next(message["content"] for message in messages if message["role"] == "assistant") == [ + _SEARCH_USE, + _SEARCH_RESULT, + _WEATHER_USE, + ] + + +@pytest.mark.asyncio +async def test_previous_response_id_keeps_a_stored_search_when_the_follow_up_sends_no_search_tool() -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + loaded_weather_tool: Final = {key: value for key, value in _WEATHER_TOOL.items() if key != "defer_loading"} + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_STATELESS_SEARCH_TURN])): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-search-turn", + input="And in Osaka?", + tools=[loaded_weather_tool], + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + messages: Final = json.loads(anthropic.body)["messages"] + assert next(message["content"] for message in messages if message["role"] == "assistant") == [ + _SEARCH_USE, + _SEARCH_RESULT, + _WEATHER_USE, + ] + + +_SERVER_HELD_SEARCH_TURN: Final = { + **_search_turn([]), + "proxy_server_request": {"model": _MODEL, "input": list(_WEATHER_QUESTION)}, + "response": { + "id": "chatcmpl-search-turn", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "toolu_01Weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01Search", + "type": "function", + "function": {"name": "tool_search_tool_regex", "arguments": '{"pattern": "weather"}'}, + } + ], + "tool_results": [_SEARCH_RESULT], + }, + }, + } + ], + }, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tools", + [ + list(_SEARCHABLE_TOOLS), + [{key: value for key, value in _WEATHER_TOOL.items() if key != "defer_loading"}], + ], + ids=["follow_up_sends_the_search_tool", "follow_up_sends_no_search_tool"], +) +async def test_previous_response_id_replays_a_search_claude_ran_whatever_search_tool_the_follow_up_sends( + tools: list[Mapping[str, object]], +) -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_SERVER_HELD_SEARCH_TURN])): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-search-turn", + input=[_WEATHER_ANSWER, {"role": "user", "content": "And in Osaka?"}], + tools=tools, + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + messages: Final = json.loads(anthropic.body)["messages"] + assert next(message["content"] for message in messages if message["role"] == "assistant") == [ + _SEARCH_USE, + _SEARCH_RESULT, + _WEATHER_USE, + ] + + +@pytest.mark.asyncio +async def test_previous_response_id_drops_a_search_the_stored_turn_got_as_input_items_off_anthropic() -> None: + messages: Final = await _follow_up_messages_after(_STATELESS_SEARCH_TURN, "mistral/mistral-small-latest", "mistral") + + assert _history_tool_ids(messages) == (["toolu_01Weather"], ["toolu_01Weather"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "custom_llm_provider"), + [ + ("anthropic/claude-3-7-sonnet-20250219", "anthropic"), + ("azure_ai/claude-deployment", "azure_ai"), + ("bedrock/invoke/us.anthropic.claude-sonnet-5-5", "bedrock"), + ], +) +async def test_previous_response_id_keeps_a_stored_search_on_claude_routes_the_cost_map_does_not_flag( + model: str, custom_llm_provider: str +) -> None: + messages: Final = await _follow_up_messages_after(_STATELESS_SEARCH_TURN, model, custom_llm_provider) + + assert _history_tool_ids(messages)[0] == ["srvtoolu_01Search", "toolu_01Weather"] + + @pytest.mark.asyncio async def test_a_codex_additional_tools_item_sends_claude_its_tools_but_not_its_text() -> None: anthropic: Final = _RecordingRawAnthropicMessages() 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 31ab4c6ac8c..8fb066d437f 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 @@ -35,7 +35,7 @@ from litellm.responses.litellm_completion_transformation.transformation import ( ) from litellm.types.llms.base import HiddenParams from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams, ResponsesAPIResponse -from litellm.types.responses.main import OutputToolSearchCall, build_web_search_call +from litellm.types.responses.main import OutputToolSearchCall, OutputToolSearchOutput, build_web_search_call from litellm.types.utils import ( ChatCompletionMessageToolCall, Choices, @@ -1696,6 +1696,77 @@ def _chat_response_with_tool_search_call( ) +_STOCK_PRICE_TOOL: Final = {"type": "function", "name": "get_stock_price", "defer_loading": True} +_FX_RATE_TOOL: Final = {"type": "function", "name": "get_fx_rate", "defer_loading": True} +_CONVERT_TOOL: Final = {"type": "custom", "name": "convert_currency", "defer_loading": True} +_FINANCE_NAMESPACE: Final = { + "type": "namespace", + "name": "finance", + "description": "Market data", + "tools": [_STOCK_PRICE_TOOL, _FX_RATE_TOOL, _CONVERT_TOOL, {"type": "web_search"}], +} +_UNNAMED_TOOL: Final = {"type": "function", "name": None} + + +def _chat_response_with_tool_search_result(result_content: Mapping[str, object]) -> ModelResponse: + response: Final = _chat_response_with_tool_search_call('{"pattern": "weather"}') + message: Final = response.choices[0].message + return response.model_copy( + update={ + "choices": [ + response.choices[0].model_copy( + update={ + "message": Message( + content=None, + tool_calls=message.tool_calls, + provider_specific_fields={ + "tool_results": [ + { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_01Search", + "content": result_content, + } + ] + }, + ) + } + ) + ] + } + ) + + +def _stateless_tool_search_input( + call_id: str | None, output_status: Literal["completed", "incomplete"] +) -> Sequence[Mapping[str, object]]: + return [ + {"type": "message", "role": "user", "content": "What is the weather in Paris right now?"}, + { + "type": "tool_search_call", + "id": "tsc_01", + "call_id": call_id, + "arguments": {"pattern": "weather"}, + "execution": "server", + "status": "completed", + }, + { + "type": "tool_search_output", + "id": "tso_01", + "call_id": call_id, + "execution": "server", + "status": output_status, + "tools": [_deferred_tools_with_server_tool_search()[1]], + }, + { + "type": "function_call", + "call_id": "toolu_01Weather", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + }, + {"type": "function_call_output", "call_id": "toolu_01Weather", "output": "18C and sunny"}, + ] + + class TestToolTransformation: """Test cases for tool transformation from Responses API to Chat Completion format""" @@ -2250,6 +2321,201 @@ class TestToolTransformation: assert [call["id"] for call in messages[1]["tool_calls"]] == ["toolu_01Weather"] assert "tool_search" not in json.dumps(messages) + @pytest.mark.parametrize( + ("result_content", "status", "loaded_tools"), + [ + pytest.param( + { + "type": "tool_search_tool_search_result", + "tool_references": [ + {"type": "tool_reference", "tool_name": name} + for name in ("get_weather", "finance__get_stock_price", "convert_currency", "not_a_tool") + ], + }, + "completed", + [ + _deferred_tools_with_server_tool_search()[1], + {**_FINANCE_NAMESPACE, "tools": [_STOCK_PRICE_TOOL, _CONVERT_TOOL]}, + ], + id="found_tools", + ), + pytest.param( + {"type": "tool_search_tool_result_error", "error_code": "unavailable"}, + "incomplete", + [], + id="search_failed", + ), + ], + ) + def test_anthropic_tool_search_result_becomes_a_tool_search_output_of_the_tools_as_sent( + self, + result_content: Mapping[str, object], + status: str, + loaded_tools: list[Mapping[str, object]], + ) -> None: + items: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_tools_to_responses_tools( + chat_completion_response=_chat_response_with_tool_search_result(result_content), + responses_api_request={ + "tools": [*_deferred_tools_with_server_tool_search()[:2], _FINANCE_NAMESPACE, _UNNAMED_TOOL], + }, + ) + + assert [item.type for item in items] == ["tool_search_call", "tool_search_output", "function_call"] + search_output: Final = items[1] + assert isinstance(search_output, OutputToolSearchOutput) + assert (search_output.call_id, search_output.execution, search_output.status) == ( + "srvtoolu_01Search", + "server", + status, + ) + assert [dict(tool) for tool in search_output.tools] == loaded_tools + + @pytest.mark.parametrize( + ("model", "custom_llm_provider", "call_id", "output_status", "replayed_call_ids"), + [ + pytest.param( + "anthropic/claude-sonnet-5-5", + "anthropic", + "srvtoolu_01Search", + "completed", + ["srvtoolu_01Search", "toolu_01Weather"], + id="claude_search_on_anthropic", + ), + pytest.param( + "anthropic/claude-sonnet-5-5", + "anthropic", + "srvtoolu_01Search", + "incomplete", + ["toolu_01Weather"], + id="failed_claude_search_on_anthropic", + ), + pytest.param( + "anthropic/claude-sonnet-5-5", "anthropic", None, "completed", ["toolu_01Weather"], id="openai_search" + ), + pytest.param( + "mistral/mistral-small-latest", + "mistral", + "srvtoolu_01Search", + "completed", + ["toolu_01Weather"], + id="claude_search_off_anthropic", + ), + pytest.param( + "anthropic/claude-3-7-sonnet-20250219", + "anthropic", + "srvtoolu_01Search", + "completed", + ["toolu_01Weather"], + id="claude_search_sent_to_a_claude_without_tool_search", + ), + ], + ) + def test_stateless_tool_search_items_are_replayed_only_for_a_completed_claude_search_on_anthropic( + self, + model: str, + custom_llm_provider: str, + call_id: str | None, + output_status: Literal["completed", "incomplete"], + replayed_call_ids: Sequence[str], + ) -> None: + request: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input=cast( # cast-ok: tool search items are not in the SDK's input TypedDicts + ResponseInputParam, list(_stateless_tool_search_input(call_id, output_status)) + ), + responses_api_request={"tools": _deferred_tools_with_server_tool_search()}, + custom_llm_provider=custom_llm_provider, + ) + + assert [message["role"] for message in request["messages"]] == ["user", "assistant", "tool"] + assert [call["id"] for call in request["messages"][1]["tool_calls"]] == replayed_call_ids + + @pytest.mark.parametrize( + ("model", "custom_llm_provider", "replayed_call_ids"), + [ + pytest.param( + "anthropic/claude-3-7-sonnet-20250219", + "anthropic", + ["srvtoolu_01Search", "toolu_01Weather"], + id="on_anthropic", + ), + pytest.param("mistral/mistral-small-latest", "mistral", ["toolu_01Weather"], id="off_anthropic"), + ], + ) + def test_stateless_tool_search_items_follow_anthropics_own_search_tool_on_a_claude_the_cost_map_does_not_flag( + self, model: str, custom_llm_provider: str, replayed_call_ids: Sequence[str] + ) -> None: + request: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input=cast( # cast-ok: tool search items are not in the SDK's input TypedDicts + ResponseInputParam, list(_stateless_tool_search_input("srvtoolu_01Search", "completed")) + ), + responses_api_request={ + "tools": [ + {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}, + *_deferred_tools_with_server_tool_search()[1:], + ] + }, + custom_llm_provider=custom_llm_provider, + ) + + assert [call["id"] for call in request["messages"][1]["tool_calls"]] == replayed_call_ids + + @pytest.mark.parametrize( + "tools", + [ + pytest.param(_deferred_tools_with_server_tool_search()[1:], id="only_the_tools_it_found"), + pytest.param( + [{"type": "tool_search", "execution": "client"}, *_deferred_tools_with_server_tool_search()[1:]], + id="a_search_the_client_runs", + ), + pytest.param([], id="no_tools"), + ], + ) + def test_stateless_claude_search_is_dropped_when_the_follow_up_sends_no_search_tool_the_route_runs( + self, tools: list[dict[str, object]] + ) -> None: + request: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( # cast-ok: tool search items are not in the SDK's input TypedDicts + ResponseInputParam, list(_stateless_tool_search_input("srvtoolu_01Search", "completed")) + ), + responses_api_request={"tools": tools}, + custom_llm_provider="anthropic", + ) + + assert [call["id"] for call in request["messages"][1]["tool_calls"]] == ["toolu_01Weather"] + + @pytest.mark.parametrize( + ("arguments", "search_tool_name"), + [ + pytest.param({"pattern": "weather"}, "tool_search_tool_regex", id="regex_object"), + pytest.param({"query": "weather"}, "tool_search_tool_bm25", id="bm25_object"), + pytest.param('{"query": "weather"}', "tool_search_tool_bm25", id="bm25_json"), + pytest.param("weather", "tool_search_tool_regex", id="not_json"), + ], + ) + def test_stateless_claude_search_is_replayed_under_the_search_tool_its_arguments_belong_to( + self, arguments: object, search_tool_name: str + ) -> None: + items: Final = ( + {**item, "arguments": arguments} if item["type"] == "tool_search_call" else item + for item in _stateless_tool_search_input("srvtoolu_01Search", "completed") + ) + request: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( # cast-ok: tool search items are not in the SDK's input TypedDicts + ResponseInputParam, list(items) + ), + responses_api_request={"tools": _deferred_tools_with_server_tool_search()}, + custom_llm_provider="anthropic", + ) + + assert [(call["id"], call["function"]["name"]) for call in request["messages"][1]["tool_calls"]] == [ + ("srvtoolu_01Search", search_tool_name), + ("toolu_01Weather", "get_weather"), + ] + def test_transform_local_shell_tools_dropped(self) -> None: ( result_tools, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 29cc8782aa5..b229461e981 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -41761,6 +41761,37 @@ export interface components { } & { [key: string]: unknown; }; + /** OutputToolSearchOutput */ + OutputToolSearchOutput: { + /** Call Id */ + call_id?: string | null; + /** + * Execution + * @enum {string} + */ + execution: "server" | "client"; + /** Id */ + id: string; + /** + * Status + * @enum {string} + */ + status: "in_progress" | "completed" | "incomplete"; + /** + * Tools + * @default [] + */ + tools: { + [key: string]: unknown; + }[]; + /** + * Type + * @constant + */ + type: "tool_search_output"; + } & { + [key: string]: unknown; + }; /** * PageLinks * @description Hypermedia for a paginated list. No `first`/`last`: without a total count the last page is unknown. @@ -45853,7 +45884,7 @@ export interface components { /** Output */ output: (components["schemas"]["ResponseOutputMessage"] | components["schemas"]["ResponseFileSearchToolCall"] | components["schemas"]["ResponseFunctionToolCall"] | components["schemas"]["ResponseFunctionToolCallOutputItem"] | components["schemas"]["ResponseFunctionWebSearch"] | components["schemas"]["ResponseComputerToolCall"] | components["schemas"]["ResponseComputerToolCallOutputItem"] | components["schemas"]["ResponseReasoningItem"] | components["schemas"]["ResponseToolSearchCall"] | components["schemas"]["ResponseToolSearchOutputItem"] | components["schemas"]["ResponseCompactionItem"] | components["schemas"]["ImageGenerationCall"] | components["schemas"]["ResponseCodeInterpreterToolCall"] | components["schemas"]["LocalShellCall"] | components["schemas"]["LocalShellCallOutput"] | components["schemas"]["ResponseFunctionShellToolCall"] | components["schemas"]["ResponseFunctionShellToolCallOutput"] | components["schemas"]["ResponseApplyPatchToolCall"] | components["schemas"]["ResponseApplyPatchToolCallOutput"] | components["schemas"]["McpCall"] | components["schemas"]["McpListTools"] | components["schemas"]["McpApprovalRequest"] | components["schemas"]["McpApprovalResponse"] | components["schemas"]["ResponseCustomToolCall"] | components["schemas"]["ResponseCustomToolCallOutputItem"] | { [key: string]: unknown; - })[] | (components["schemas"]["GenericResponseOutputItem"] | components["schemas"]["OutputCodeInterpreterCall"] | components["schemas"]["OutputFunctionToolCall"] | components["schemas"]["OutputImageGenerationCall"] | components["schemas"]["ResponseFunctionToolCall"] | components["schemas"]["ResponseFunctionWebSearch"] | components["schemas"]["CustomToolCallOutputItem"] | components["schemas"]["OutputToolSearchCall"] | components["schemas"]["ResponseCompactionItem"])[]; + })[] | (components["schemas"]["GenericResponseOutputItem"] | components["schemas"]["OutputCodeInterpreterCall"] | components["schemas"]["OutputFunctionToolCall"] | components["schemas"]["OutputImageGenerationCall"] | components["schemas"]["ResponseFunctionToolCall"] | components["schemas"]["ResponseFunctionWebSearch"] | components["schemas"]["CustomToolCallOutputItem"] | components["schemas"]["OutputToolSearchCall"] | components["schemas"]["OutputToolSearchOutput"] | components["schemas"]["ResponseCompactionItem"])[]; /** Parallel Tool Calls */ parallel_tool_calls?: boolean | null; /** Previous Response Id */