diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 1562220c3a8..a8e49870483 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -48,6 +48,7 @@ from litellm.types.llms.anthropic import ( ANTHROPIC_ADVISOR_TOOL_TYPE, ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_HOSTED_TOOLS, + ANTHROPIC_TOOL_SEARCH_TOOL_TYPES, AllAnthropicPassThroughMessageValues, AllAnthropicToolsValues, AnthropicCodeExecutionTool, @@ -147,6 +148,78 @@ else: _ANTHROPIC_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]") _ANTHROPIC_TOOL_NAME_MAX_LEN: Final = 128 _DEFER_LOADING_ADAPTER: Final = TypeAdapter(bool) +_TOOL_FIELDS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_TOOLS_ADAPTER: Final = TypeAdapter(list[object]) + + +def _tool_type(tool: object) -> str | None: + match tool: + case {"type": str() as tool_type}: + return tool_type + case _: + return None + + +def _is_openai_tool_search_tool(tool: object) -> bool: + return _tool_type(tool) == "tool_search" + + +def _is_anthropic_tool_search_tool(tool: object) -> bool: + return _tool_type(tool) in ANTHROPIC_TOOL_SEARCH_TOOL_TYPES + + +def _without_defer_loading(fields: object) -> object: + match fields: + case Mapping(): + return { + key: value + for key, value in _TOOL_FIELDS_ADAPTER.validate_python(fields).items() + if key != "defer_loading" + } + case _: + return fields + + +def _member_without_deferral(member: object) -> object: + match member: + case Mapping(): + return { + key: _without_defer_loading(value) if key == "function" else value + for key, value in _TOOL_FIELDS_ADAPTER.validate_python(member).items() + if key != "defer_loading" + } + case _: + return member + + +def _namespace_members_without_deferral(members: object) -> object: + match members: + case list(): + return [_member_without_deferral(member) for member in _TOOLS_ADAPTER.validate_python(members)] + case _: + return members + + +def _field_without_deferral(key: str, value: object) -> object: + match key: + case "function": + return _without_defer_loading(value) + case "tools": + return _namespace_members_without_deferral(value) + case _: + return value + + +def _without_deferral(tool: object) -> object: + match tool: + case Mapping(): + return { + key: _field_without_deferral(key, value) + for key, value in _TOOL_FIELDS_ADAPTER.validate_python(tool).items() + if key != "defer_loading" + } + case _: + return tool def _validated_defer_loading(value: object) -> bool: @@ -321,6 +394,36 @@ DROP_SERVICE_TIER_PROVIDER_WARNING: Final = ( "Dropping `service_tier` for model=%s on provider=%s. Only the Anthropic API routes accept it." ) +DROP_UNSUPPORTED_TOOL_SEARCH_WARNING: Final = ( + "Dropping the OpenAI `tool_search` tool and the `defer_loading` of the other tools for model=%s on %s " + "(drop_params=True) because this route cannot run that tool search." +) + + +TOOL_SEARCH_MODEL_INFO_REMEDY: Final = ( + " If this model does run tool search, set `supports_tool_search: true` in its `model_info`, or for the SDK " + "with `litellm.register_model`." +) + + +def tools_without_unsupported_tool_search( + tools: Sequence[object], model: str, route: str, drop_params: bool, remedy: str +) -> Sequence[object]: + if not any(_is_openai_tool_search_tool(tool) for tool in tools): + return tools + if not (litellm.drop_params or drop_params): + raise litellm.utils.UnsupportedParamsError( + message=( + f"{model} on {route} cannot run the OpenAI `tool_search` tool as sent. Only a server-executed " + f"`tool_search` on a route that supports tool search becomes Anthropic's tool search.{remedy} " + "To drop it and the `defer_loading` of the other tools, set `litellm.drop_params=True` or for " + "proxy:\n\n`litellm_settings:\n drop_params: true`\n" + ), + status_code=400, + ) + litellm.verbose_logger.warning(DROP_UNSUPPORTED_TOOL_SEARCH_WARNING, model, route) + return [_without_deferral(tool) for tool in tools if not _is_openai_tool_search_tool(tool)] + class AnthropicConfig(AnthropicModelInfo, BaseConfig): """ @@ -1068,6 +1171,53 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): initial_tool["authorization_token"] = authorization_token return initial_tool + def _tools_with_openai_tool_search_mapped( + self, tools: Sequence[object], model: str, drop_params: bool + ) -> Sequence[object]: + from litellm.responses.litellm_completion_transformation.tool_search import ( + anthropic_tool_search_chat_tool, + runs_tool_search_on_the_server, + ) + + first_search_tool_index: Final = next( + (index for index, tool in enumerate(tools) if _is_openai_tool_search_tool(tool)), None + ) + if first_search_tool_index is None: + return tools + if any(_is_anthropic_tool_search_tool(tool) for tool in tools): + return [tool for tool in tools if not _is_openai_tool_search_tool(tool)] + model_runs_tool_search: Final = litellm.supports_tool_search(model, self._resolved_provider) + search_runs_on_the_server: Final = all( + runs_tool_search_on_the_server(tool) for tool in tools if _is_openai_tool_search_tool(tool) + ) + if not (model_runs_tool_search and search_runs_on_the_server): + return tools_without_unsupported_tool_search( + tools=tools, + model=model, + route=self._resolved_provider, + drop_params=drop_params, + remedy="" if model_runs_tool_search else TOOL_SEARCH_MODEL_INFO_REMEDY, + ) + return [ + *tools[:first_search_tool_index], + anthropic_tool_search_chat_tool(), + *(tool for tool in tools[first_search_tool_index + 1 :] if not _is_openai_tool_search_tool(tool)), + ] + + def _tools_param_with_openai_tool_search_mapped( + self, non_default_params: Mapping[str, object], model: str, drop_params: bool + ) -> Mapping[str, object]: + tools: Final = non_default_params.get("tools") + if not isinstance(tools, list): + return {} + return { + "tools": list( + self._tools_with_openai_tool_search_mapped( + tools=_TOOLS_ADAPTER.validate_python(tools), model=model, drop_params=drop_params + ) + ) + } + def _map_tools( self, tools: list, @@ -1636,7 +1786,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if param == "max_tokens" or param == "max_completion_tokens": optional_params["max_tokens"] = value if isinstance(value, int) else max(1, int(round(value))) elif param == "tools": - anthropic_tools, mcp_servers = self._map_tools(value) + anthropic_tools, mcp_servers = self._map_tools( + list(self._tools_with_openai_tool_search_mapped(tools=value, model=model, drop_params=drop_params)) + ) optional_params = self._add_tools_to_optional_params( optional_params=optional_params, tools=anthropic_tools ) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index b0fe0930f9d..cfbb9956988 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -50,11 +50,12 @@ from litellm.llms.anthropic.chat.transformation import ( DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT, AnthropicConfig, + tools_without_unsupported_tool_search, ) from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.tool_result_images import place_tool_result_images -from litellm.llms.bedrock.common_utils import bedrock_model_supports_regex_lookaround +from litellm.llms.bedrock.common_utils import bedrock_model_supports_regex_lookaround, bedrock_supports_tool_search from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, bedrock_request_metadata_is_owned, @@ -144,6 +145,8 @@ UNSUPPORTED_BEDROCK_CONVERSE_BETA_PATTERNS: Final = [ _TOOLS_AS_SENT: Final = TypeAdapter(tuple[Mapping[str, object], ...]) +_TOOLS_AS_RECEIVED: Final = TypeAdapter(list[object]) +CLAUDE_ON_BEDROCK_INVOKE_REMEDY: Final = " This model runs tool search on Bedrock through `bedrock/invoke/`." _DEPLOYMENT_PARAMS: Final = TypeAdapter(Mapping[str, object]) @@ -1094,7 +1097,18 @@ class AmazonConverseConfig(BaseConfig): optional_params["topP" if param == "top_p" else param] = value if param == "tools" and isinstance(value, list): self._apply_tool_call_transformation( - tools=cast(list[OpenAIChatCompletionToolParam], value), + tools=cast( + list[OpenAIChatCompletionToolParam], + list( + tools_without_unsupported_tool_search( + tools=_TOOLS_AS_RECEIVED.validate_python(value), + model=model, + route="bedrock", + drop_params=drop_params, + remedy=CLAUDE_ON_BEDROCK_INVOKE_REMEDY if bedrock_supports_tool_search(model) else "", + ) + ), + ), model=model, non_default_params=non_default_params, optional_params=optional_params, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 68264539da0..9ce3279ecd2 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -41,6 +43,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object]) + class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): """ @@ -98,7 +102,12 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): mapped_params: Final = AnthropicConfig.map_openai_params( self, - non_default_params, + { + **non_default_params, + **self._tools_param_with_openai_tool_search_mapped( + _NON_DEFAULT_PARAMS.validate_python(non_default_params), model, drop_params + ), + }, optional_params, model, drop_params, diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index a98cfdfe52d..e834e459a0a 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -58,6 +58,7 @@ from litellm.types.utils import ( from ...anthropic.chat.transformation import ( REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT, AnthropicConfig, + tools_without_unsupported_tool_search, ) from ...openai_like.chat.transformation import OpenAILikeChatConfig from ..common_utils import DatabricksBase, DatabricksException @@ -65,6 +66,7 @@ from ..common_utils import DatabricksBase, DatabricksException _RESPONSE_FORMAT_ADAPTER: ( # mutable-ok: mirrors the dict return contract of get_json_schema_from_pydantic_object Final[TypeAdapter[dict[str, object] | None]] ) = TypeAdapter(dict[str, object] | None) +_TOOLS_ADAPTER: Final = TypeAdapter(list[object]) def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool: @@ -326,13 +328,23 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): function=function_params, ) - def _map_openai_to_dbrx_tool(self, model: str, tools: list) -> list[DatabricksTool]: + def _map_openai_to_dbrx_tool(self, model: str, tools: list, drop_params: bool) -> list[DatabricksTool]: # if not claude, send as is if "claude" not in model: return tools # if claude, convert to anthropic tool and then to databricks tool - anthropic_tools, _ = self._map_tools(tools=tools) # unclear how mcp tool calling on databricks works + anthropic_tools, _ = self._map_tools( + tools=list( + tools_without_unsupported_tool_search( + tools=_TOOLS_ADAPTER.validate_python(tools), + model=model, + route="databricks", + drop_params=drop_params, + remedy="", + ) + ) + ) # unclear how mcp tool calling on databricks works databricks_tools: Final = [ cast(DatabricksTool, self.convert_anthropic_tool_to_databricks_tool(tool)) for tool in anthropic_tools ] @@ -383,7 +395,9 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): model=model, optional_params=mapped_params, custom_llm_provider="databricks" ) if "tools" in mapped_params: - mapped_params["tools"] = self._map_openai_to_dbrx_tool(model=model, tools=mapped_params["tools"]) + mapped_params["tools"] = self._map_openai_to_dbrx_tool( + model=model, tools=mapped_params["tools"], drop_params=drop_params + ) if "max_completion_tokens" in non_default_params and replace_max_completion_tokens_with_max_tokens: mapped_params["max_tokens"] = non_default_params[ "max_completion_tokens" diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index b57081589ee..f7a1e007c96 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -1,8 +1,10 @@ # What is this? ## Handler file for calling claude-3 on vertex ai +from collections.abc import Mapping from typing import TYPE_CHECKING, Final import httpx +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, inline_remote_image_urls @@ -16,6 +18,8 @@ from .output_params_utils import sanitize_vertex_anthropic_output_params if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +_NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object]) + class VertexAIError(Exception): def __init__(self, status_code, message): @@ -169,7 +173,12 @@ class VertexAIAnthropicConfig(AnthropicConfig): from litellm.llms.anthropic.common_utils import AnthropicModelInfo mapped_params: Final = super().map_openai_params( - non_default_params=non_default_params, + non_default_params={ + **non_default_params, + **self._tools_param_with_openai_tool_search_mapped( + _NON_DEFAULT_PARAMS.validate_python(non_default_params), model, drop_params + ), + }, optional_params=optional_params, model=model, drop_params=drop_params, diff --git a/litellm/main.py b/litellm/main.py index 474d82e8ef1..b474c95db1a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1089,6 +1089,23 @@ def _is_claude_tool_target(custom_llm_provider: str | None, model: str) -> bool: return False +def _history_without_unanswered_tool_search_calls( + messages: Sequence[object], custom_llm_provider: str | None, provider_config: BaseConfig +) -> tuple[object, ...] | None: + from litellm.responses.litellm_completion_transformation.tool_search import ( + chat_messages_without_tool_search_calls, + chat_route_rejects_unanswered_tool_calls, + unanswered_tool_search_call_ids, + ) + + if not chat_route_rejects_unanswered_tool_calls(custom_llm_provider, provider_config): + return None + unanswered_search_call_ids: Final = unanswered_tool_search_call_ids(messages) + if not unanswered_search_call_ids: + return None + return chat_messages_without_tool_search_calls(messages, unanswered_search_call_ids) + + def _without_anthropic_only_tool_keys(tool: dict[str, object]) -> dict[str, object]: kept: Final = {key: value for key, value in tool.items() if key not in _ANTHROPIC_ONLY_TOOL_KEYS} function: Final = tool.get("function") @@ -5556,7 +5573,13 @@ def completion( ) if provider_config is not None: - messages = provider_config.translate_developer_role_to_system_role(messages=messages) + translated_messages: Final = provider_config.translate_developer_role_to_system_role(messages=messages) + history_for_route: Final = _history_without_unanswered_tool_search_calls( + translated_messages, custom_llm_provider, provider_config + ) + messages = ( # rebind-ok: completion() rewrites messages step by step for the provider + translated_messages if history_for_route is None else list(history_for_route) + ) if ( supports_system_message is not None diff --git a/litellm/responses/litellm_completion_transformation/tool_search.py b/litellm/responses/litellm_completion_transformation/tool_search.py index fdeaab255e1..904aa86f841 100644 --- a/litellm/responses/litellm_completion_transformation/tool_search.py +++ b/litellm/responses/litellm_completion_transformation/tool_search.py @@ -19,13 +19,18 @@ __all__ = [ "answers_tool_search", "anthropic_tool_search_chat_tool", "build_tool_search_call", + "chat_config_builds_anthropic_request", + "chat_messages_without_tool_search_calls", + "chat_route_rejects_unanswered_tool_calls", "has_server_tool_search", "holds_tool_search_use", "is_anthropic_tool_search_chat_tool", "is_anthropic_tool_search_use", "is_server_tool_search_tool", "provider_builds_anthropic_request", + "runs_tool_search_on_the_server", "tool_search_call_ids", + "unanswered_tool_search_call_ids", "without_tool_search_calls", "without_tool_search_uses", ] @@ -35,6 +40,7 @@ ANTHROPIC_TOOL_SEARCH_TOOL_NAME: Final = "tool_search_tool_regex" ToolSearchCallStatus: TypeAlias = Literal["in_progress", "completed", "incomplete"] _ARGUMENTS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_MESSAGE_FIELDS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) class _ToolSearchFields(LiteLLMBaseModel): @@ -52,6 +58,14 @@ def is_server_tool_search_tool(tool: object) -> bool: return parsed.type == "tool_search" and parsed.execution != "client" +def runs_tool_search_on_the_server(tool: object) -> bool: + try: + parsed: Final = _ToolSearchFields.model_validate(tool) + except ValidationError: + return False + return parsed.type == "tool_search" and parsed.execution in (None, "server") + + def has_server_tool_search(tools: Sequence[object] | None) -> bool: return any(is_server_tool_search_tool(tool) for tool in tools or ()) @@ -66,8 +80,6 @@ def is_anthropic_tool_search_chat_tool(tool: Mapping[str, object]) -> bool: def provider_builds_anthropic_request(model: str, custom_llm_provider: str | None) -> bool: from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - from litellm.llms.anthropic.chat.transformation import AnthropicConfig - from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.utils import ProviderConfigManager _, provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) @@ -75,10 +87,29 @@ def provider_builds_anthropic_request(model: str, custom_llm_provider: str | Non llm_provider: Final = LlmProviders(provider) except ValueError: return False - chat_config: Final = ProviderConfigManager.get_provider_chat_config(model=model, provider=llm_provider) + return chat_config_builds_anthropic_request( + ProviderConfigManager.get_provider_chat_config(model=model, provider=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 + return isinstance(chat_config, AnthropicConfig) and not isinstance(chat_config, OpenAIGPTConfig) +_PROVIDERS_REJECTING_UNANSWERED_TOOL_CALLS: Final = frozenset({LlmProviders.OPENAI.value, LlmProviders.AZURE.value}) + + +def chat_route_rejects_unanswered_tool_calls(custom_llm_provider: str | None, chat_config: object) -> bool: + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + return custom_llm_provider in _PROVIDERS_REJECTING_UNANSWERED_TOOL_CALLS or isinstance( + chat_config, AmazonConverseConfig + ) + + class CompletionRoute(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) @@ -167,6 +198,60 @@ def without_tool_search_uses(message: Message) -> Message | None: return message.model_copy(update={"tool_calls": list(kept) or None}) +def _is_assistant_message(message: object) -> bool: + match message: + case {"role": "assistant"}: + return True + case Mapping(): + return False + case _: + return _session_message_fields(message).role == "assistant" + + +def _tool_result_ids(messages: Iterable[object]) -> frozenset[str]: + return frozenset( + tool_call_id + for message in messages + if (tool_call_id := _session_message_fields(message).tool_call_id) is not None + ) + + +def unanswered_tool_search_call_ids(messages: Sequence[object]) -> frozenset[str]: + search_call_ids: Final = tool_search_call_ids(message for message in messages if _is_assistant_message(message)) + if not search_call_ids: + return search_call_ids + return search_call_ids - _tool_result_ids(messages) + + +def _assistant_mapping_with_tool_calls( + message: Mapping[str, object], tool_calls: tuple[object, ...] +) -> Mapping[str, object] | None: + if not tool_calls and not message.get("content"): + return None + without_calls: Final = {key: value for key, value in message.items() if key != "tool_calls"} + return {**without_calls, "tool_calls": list(tool_calls)} if tool_calls else without_calls + + +def _without_tool_search_calls(message: object, call_ids: Set[str]) -> object | None: + tool_calls: Final = _assistant_tool_calls(message) + kept_calls: Final = tuple(tool_call for tool_call in tool_calls if _tool_search_call_id(tool_call) not in call_ids) + if len(kept_calls) == len(tool_calls): + return message + match message: + case Message(): + if not kept_calls and not message.content: + return None + return message.model_copy(update={"tool_calls": list(kept_calls) or None}) + case Mapping(): + return _assistant_mapping_with_tool_calls(_MESSAGE_FIELDS_ADAPTER.validate_python(message), kept_calls) + case _: + return message + + +def chat_messages_without_tool_search_calls(messages: Sequence[object], call_ids: Set[str]) -> tuple[object, ...]: + return tuple(kept for message in messages if (kept := _without_tool_search_calls(message, call_ids)) is not None) + + def _parsed_arguments(arguments_json: str) -> Mapping[str, object] | None: try: return _ARGUMENTS_ADAPTER.validate_json(arguments_json or "{}") diff --git a/tests/integration/providers/test_chat_completions_tool_search_chaos.py b/tests/integration/providers/test_chat_completions_tool_search_chaos.py new file mode 100644 index 00000000000..ac9cb214069 --- /dev/null +++ b/tests/integration/providers/test_chat_completions_tool_search_chaos.py @@ -0,0 +1,366 @@ +import asyncio +import re +import socket +import threading +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal, TypeAlias +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support import claude_code as cc +from integration._support import responses_vendor as rv +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.openai_wire import chat_reply +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers.test_chat_completions_tool_search_wire import ( + DEFERRED_NAMES, + UNLISTED_CLAUDE, + Answer, + assert_lands_once_as, + assert_no_search_and_no_deferral, + assert_search_call_answered, + chat_body, + chat_frames, + claude_vendor, + deferred_tools, + response_id, + tool_search, +) +from pydantic import JsonValue + +_CLAUDE: Final = cc.OPUS +_REGEX_TOOL_TYPE: Final = "tool_search_tool_regex_20251119" +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_CACHE_BUST: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}}) +_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 +_SEARCH_ARGUMENTS: Final = '{"pattern": "weather"}' +_CONFIG_MODEL: Final = "chat-tool-search-chaos" + +Kind: TypeAlias = Literal["search", "history"] + + +@dataclass(frozen=True, slots=True) +class _Call: + kind: Kind + stream: bool + marker: str + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + answer: Answer + call_id: str + + +def _history(marker: str) -> list[JsonValue]: + return [ + {"role": "user", "content": f"Weather in Paris marker-{marker}"}, + { + "role": "assistant", + "content": f"searching marker-{marker}", + "tool_calls": [ + { + "id": f"srvtoolu_{marker}", + "type": "function", + "function": {"name": "tool_search_tool_regex", "arguments": _SEARCH_ARGUMENTS}, + }, + { + "id": f"toolu_{marker}", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + }, + ], + }, + {"role": "tool", "tool_call_id": f"toolu_{marker}", "content": f"22C marker-{marker}"}, + {"role": "user", "content": f"And tomorrow? marker-{marker}"}, + ] + + +def _strict_openai(request: Request) -> Reply: + if request.method != "POST": + return Reply() + body: Final = rv.JSON_OBJECT.validate_json(request.body) + marker: Final = rv.newest_marker(request.body.decode()) or "health" + if "srvtoolu_" in request.body.decode(): + return rv.error(400, f"No tool output found for function call srvtoolu_{marker}.", "invalid_request_error") + return chat_reply(f"chatcmpl-{marker}", "gpt-5.6", f"answer marker-{marker}", stream=body.get("stream") is True) + + +def _vendor(request: Request) -> Reply: + if request.target.endswith("/v1/chat/completions"): + return _strict_openai(request) + return claude_vendor(request) + + +def _body(models: Mapping[Kind, str], call: _Call) -> dict[str, JsonValue]: + match call.kind: + case "search": + return chat_body(models["search"], call.marker, (tool_search(), *deferred_tools()), call.stream) + case "history": + return { + "model": models["history"], + "messages": _history(call.marker), + "stream": call.stream, + "num_retries": 0, + } + + +def _calls(count: int, kinds: tuple[Kind, ...]) -> tuple[_Call, ...]: + return tuple( + _Call(kind=kinds[index % len(kinds)], stream=index % 2 == 1, marker=uuid.uuid4().hex) for index in range(count) + ) + + +async def _send(client: httpx.AsyncClient, key: str, models: Mapping[Kind, str], call: _Call) -> _Served: + async with client.stream( + "POST", + "/v1/chat/completions", + json={**_body(models, call), **_CACHE_BUST}, + headers={"Authorization": f"Bearer {key}"}, + ) as response: + raw: Final = (await response.aread()).decode() + payload: Final = rv.JSON_OBJECT.validate_json(raw) if raw.startswith("{") else {} + answer: Final = Answer(response.status_code, payload, chat_frames(raw) if call.stream else (), raw) + return _Served(call, answer, response.headers.get("x-litellm-call-id", "")) + + +async def _burst( + gateway: Gateway, models: Mapping[Kind, str], calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, gateway.key, models, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + unexpected: Final = [ + result + for result in results + if isinstance(result, BaseException) and not isinstance(result, httpx.TransportError) + ] + assert unexpected == [], unexpected + return tuple(result for result in results if isinstance(result, _Served)) + + +def _assert_served(item: _Served) -> None: + match item.call.kind: + case "search": + assert_search_call_answered(item.answer, item.call.marker) + case "history": + assert item.answer.status == 200, item.answer.text + assert f"answer marker-{item.call.marker}" in item.answer.text, item.answer.text + + +def _marked(received: tuple[Request, ...]) -> dict[str, Request]: + posted: Final = tuple(request for request in received if request.method == "POST") + marked: Final = {marker: request for request in posted if (marker := rv.newest_marker(request.body.decode()))} + assert len(marked) == len(posted), [request.target for request in posted] + return marked + + +def _assert_forwarded(request: Request, kind: Kind) -> None: + sent: Final = request.body.decode() + body: Final = rv.JSON_OBJECT.validate_json(request.body) + match kind: + case "search": + assert request.target == "/v1/messages", request.target + tools: Final = rv.ITEMS.validate_python(body["tools"]) + assert tools[0]["type"] == _REGEX_TOOL_TYPE, tools + assert [(tool["name"], tool.get("defer_loading")) for tool in tools[1:]] == [ + (name, True) for name in DEFERRED_NAMES + ], tools + case "history": + assert request.target == "/v1/chat/completions", request.target + assert "srvtoolu_" not in sent and "tool_search_tool_regex" not in sent, body + + +def _assert_every_call_forwarded(wire: Wire, calls: tuple[_Call, ...]) -> None: + forwarded: Final = _marked(wire.drain()) + assert set(forwarded) == {call.marker for call in calls}, sorted(forwarded) + kinds: Final = {call.marker: call.kind for call in calls} + for marker, request in forwarded.items(): + _assert_forwarded(request, kinds[marker]) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind(("127.0.0.1", 0)) + return int(probe.getsockname()[1]) + + +def _identities(served: tuple[_Served, ...], kind: Kind) -> tuple[str, ...]: + return tuple(response_id(item.answer) for item in served if item.call.kind == kind) + + +async def test_a_mixed_burst_carries_the_regex_tool_and_drops_the_history_call_and_lands_each_once( + gateway: Gateway, +) -> None: + calls: Final = _calls(24, ("search", "search", "history")) + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + models: Final = MappingProxyType( + { + "search": scenario.model(model=f"anthropic/{_CLAUDE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY), + "history": scenario.model(model="openai/gpt-5.6", api_base=f"{wire.url}/v1", api_key="scripted"), + } + ) + served: Final = await _burst(gateway, models, calls) + assert len(served) == 24 + for item in served: + _assert_served(item) + _assert_every_call_forwarded(wire, calls) + assert_lands_once_as(models["search"], _identities(served, "search"), "success") + assert_lands_once_as(models["history"], _identities(served, "history"), "success") + + +def _spend_status(model: str, expected: int) -> dict[str, str]: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda found: len(found) >= expected, + seconds=70, + ) + by_id: Final = {str(row["request_id"]): str(row["status"]) for row in rows} + assert len(by_id) == len(rows) == expected, rows + return by_id + + +@pytest.mark.timeout(180) +async def test_a_vendor_outage_fails_each_call_cleanly_and_the_recovered_vendor_gets_the_regex_tool( + gateway: Gateway, +) -> None: + port: Final = _free_port() + while_down: Final = _calls(15, ("search",)) + after: Final = _calls(15, ("search",)) + with gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_CLAUDE}", api_base=f"http://127.0.0.1:{port}", api_key=cc.ANTHROPIC_API_KEY + ) + models: Final = MappingProxyType({"search": model, "history": model}) + failed: Final = await _burst(gateway, models, while_down) + assert len(failed) == 15 + for item in failed: + assert item.answer.status == 500 and "Cannot connect to host" in item.answer.text, item.answer.text + assert "srvtoolu_" not in item.answer.text and item.call_id, item + with wire_server(claude_vendor, port=port) as wire: + wire.drain() + served: Final = await _burst(gateway, models, after) + assert len(served) == 15 + for item in served: + _assert_served(item) + _assert_every_call_forwarded(wire, after) + by_id: Final = _spend_status(model, 30) + assert {by_id.get(item.call_id) for item in failed} == {"failure"}, by_id + matched: Final = [ + [logged for logged in by_id if rv.same_response(logged, response_id(item.answer))] for item in served + ] + assert all(len(ids) == 1 for ids in matched), by_id + assert {by_id[ids[0]] for ids in matched} == {"success"}, by_id + + +def _config(tmp_path: Path, name: str, model: Mapping[str, JsonValue], settings: Mapping[str, JsonValue]) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path: Final = tmp_path / f"{name}.yaml" + path.write_text( + yaml.safe_dump( + { + **config, + "model_list": [{"model_name": name, "litellm_params": dict(model)}], + "litellm_settings": {**config["litellm_settings"], **settings}, + } + ) + ) + return path + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _holding( + vendor: Callable[[Request], Reply], release: threading.Event, held: SimpleQueue[str] +) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + marker: Final = rv.newest_marker(request.body.decode()) + if request.method != "POST" or marker is None: + return vendor(request) + held.put(marker) + assert release.wait(timeout=60), "The burst was never released" + return vendor(request) + + return respond + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +async def test_a_worker_sigkill_mid_burst_leaves_the_sibling_returning_the_regex_tool_search( + gateway: Gateway, tmp_path: Path +) -> None: + calls: Final = _calls(20, ("search",)) + release: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + models: Final = MappingProxyType({"search": _CONFIG_MODEL, "history": _CONFIG_MODEL}) + with wire_server(_holding(claude_vendor, release, held)) as wire: + model: Final[Mapping[str, JsonValue]] = { + "model": f"anthropic/{_CLAUDE}", + "api_base": wire.url, + "api_key": cc.ANTHROPIC_API_KEY, + } + path: Final = _config(tmp_path, _CONFIG_MODEL, model, {}) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task(_burst(candidate, models, calls, tolerate_transport_errors=True)) + await asyncio.to_thread(eventually, held.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + await asyncio.to_thread(owned.kill_worker, victim_pid) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_served(item) + follow_up: Final = _Call("search", False, uuid.uuid4().hex) + (answered,) = await _burst(candidate, models, (follow_up,)) + _assert_served(answered) + _assert_every_call_forwarded(wire, (*calls, follow_up)) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_global_drop_params_drops_the_search_and_every_deferral_for_a_claude_without_tool_search( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(claude_vendor) as wire: + model: Final[Mapping[str, JsonValue]] = { + "model": f"anthropic/{UNLISTED_CLAUDE}", + "api_base": wire.url, + "api_key": cc.ANTHROPIC_API_KEY, + } + path: Final = _config(tmp_path, _CONFIG_MODEL, model, {"drop_params": True}) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + body: Final = {**chat_body(_CONFIG_MODEL, marker, (tool_search(), *deferred_tools()), False), **_CACHE_BUST} + response: Final = owned.gateway.request("POST", "/v1/chat/completions", body) + assert response.status_code == 200, response.text + assert f"toolu_{marker}" in response.text and f"srvtoolu_{marker}" not in response.text, response.text + posted: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert len(posted) == 1, [request.target for request in posted] + assert posted[0].target == "/v1/messages", posted[0].target + assert_no_search_and_no_deferral(rv.JSON_OBJECT.validate_json(posted[0].body), DEFERRED_NAMES) diff --git a/tests/integration/providers/test_chat_completions_tool_search_history_wire.py b/tests/integration/providers/test_chat_completions_tool_search_history_wire.py new file mode 100644 index 00000000000..0157093664e --- /dev/null +++ b/tests/integration/providers/test_chat_completions_tool_search_history_wire.py @@ -0,0 +1,356 @@ +import json +import uuid +from collections.abc import Callable, Mapping, Sequence +from itertools import chain +from typing import Final, Literal, TypeAlias +from urllib.parse import unquote + +import pytest +from _pytest.mark.structures import ParameterSet +from integration._support import responses_vendor as rv +from integration._support.anthropic_thinking import BEDROCK_MODEL +from integration._support.client import Gateway, Scenario +from integration._support.openai_wire import chat_reply +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers.test_bedrock_auth_wire import RESPONSE as CONVERSE_RESPONSE +from integration.providers.test_chat_completions_tool_search_wire import ( + assert_lands_once_as, + raw_answer, + response_id, +) +from pydantic import JsonValue + +_REGEX_TOOL_NAME: Final = "tool_search_tool_regex" +_BM25_TOOL_NAME: Final = "tool_search_tool_bm25" +_SEARCH_ARGUMENTS: Final = '{"pattern": "weather"}' +_AZURE_DEPLOYMENT: Final = "scripted-chat-deployment" +_AZURE_API_VERSION: Final = "2025-04-01-preview" +_GEMINI: Final = "gemini-2.5-flash" +_GEMINI_REPLY: Final[dict[str, JsonValue]] = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "gemini control"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 3, "totalTokenCount": 8}, + "modelVersion": _GEMINI, +} +_CITY_PARAMETERS: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} + +Route: TypeAlias = Literal["openai", "azure", "bedrock_converse", "gemini"] + + +def _tool_call(call_id: str, name: str, arguments: str) -> dict[str, JsonValue]: + return {"id": call_id, "type": "function", "function": {"name": name, "arguments": arguments}} + + +def _search_call(marker: str, name: str = _REGEX_TOOL_NAME) -> dict[str, JsonValue]: + return _tool_call(f"srvtoolu_{marker}", name, _SEARCH_ARGUMENTS) + + +def _weather_call(marker: str) -> dict[str, JsonValue]: + return _tool_call(f"toolu_{marker}", "get_weather", '{"city": "Paris"}') + + +def _weather_result(marker: str) -> dict[str, JsonValue]: + return {"role": "tool", "tool_call_id": f"toolu_{marker}", "content": f"22C marker-{marker}"} + + +def _history(marker: str, calls: Sequence[Mapping[str, JsonValue]], *results: JsonValue) -> list[JsonValue]: + return [ + {"role": "user", "content": f"Weather in Paris marker-{marker}"}, + {"role": "assistant", "content": f"searching marker-{marker}", "tool_calls": [dict(call) for call in calls]}, + *results, + {"role": "user", "content": f"And tomorrow? marker-{marker}"}, + ] + + +def _second_turn(marker: str) -> list[JsonValue]: + return _history(marker, (_search_call(marker), _weather_call(marker)), _weather_result(marker)) + + +def _body(model: str, messages: Sequence[JsonValue], stream: bool) -> dict[str, JsonValue]: + return { + "model": model, + "messages": list(messages), + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Weather for a city", + "parameters": _CITY_PARAMETERS, + }, + } + ], + "stream": stream, + "num_retries": 0, + } + + +def _called_ids(message: Mapping[str, JsonValue]) -> list[str]: + return [str(call["id"]) for call in rv.ITEMS.validate_python(message.get("tool_calls") or [])] + + +def _unanswered_chat_calls(messages: Sequence[Mapping[str, JsonValue]]) -> list[str]: + answered: Final = {str(message.get("tool_call_id")) for message in messages if message.get("role") == "tool"} + called: Final = tuple(chain.from_iterable(_called_ids(message) for message in messages)) + return [call_id for call_id in called if call_id not in answered] + + +def _strict_chat_vendor(refusal: Callable[[str], str]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply() + body: Final = rv.JSON_OBJECT.validate_json(request.body) + marker: Final = rv.newest_marker(request.body.decode()) or "health" + unanswered: Final = _unanswered_chat_calls(rv.ITEMS.validate_python(body["messages"])) + if unanswered: + return rv.error(400, refusal(unanswered[0]), "invalid_request_error") + return chat_reply( + f"chatcmpl-{marker}", "scripted", f"answer marker-{marker}", stream=body.get("stream") is True + ) + + return respond + + +def _openai_refusal(call_id: str) -> str: + return f"No tool output found for function call {call_id}." + + +def _azure_refusal(call_id: str) -> str: + return f"An assistant message with 'tool_calls' must be followed by tool messages responding to {call_id}" + + +def _converse_ids(message: Mapping[str, JsonValue], kind: str) -> set[str]: + blocks: Final = [ + rv.JSON_OBJECT.validate_python(block[kind]) + for block in rv.ITEMS.validate_python(message["content"]) + if kind in block + ] + return {str(block["toolUseId"]) for block in blocks} + + +def _missing_results(messages: Sequence[Mapping[str, JsonValue]], index: int) -> list[str]: + following: Final = messages[index + 1] if index + 1 < len(messages) else {"content": []} + return sorted(_converse_ids(messages[index], "toolUse") - _converse_ids(following, "toolResult")) + + +def _strict_converse_vendor(request: Request) -> Reply: + assert unquote(request.target) == f"/model/{BEDROCK_MODEL}/converse", request.target + messages: Final = rv.ITEMS.validate_python(rv.JSON_OBJECT.validate_json(request.body)["messages"]) + gaps: Final = [ + (index, _missing_results(messages, index)) + for index, message in enumerate(messages) + if message["role"] == "assistant" + ] + unanswered: Final = [(index, missing) for index, missing in gaps if missing] + if not unanswered: + return Reply(body=CONVERSE_RESPONSE) + index, missing = unanswered[0] + refusal: Final = { + "message": f"Expected toolResult blocks at messages.{index + 1}.content for the following Ids: {', '.join(missing)}" + } + return Reply(status=400, body=json.dumps(refusal).encode()) + + +def _gemini_vendor(request: Request) -> Reply: + assert request.target.split("?")[0].endswith(f"/models/{_GEMINI}:generateContent"), request.target + return Reply(body=json.dumps(_GEMINI_REPLY).encode()) + + +def _vendor(route: Route) -> Callable[[Request], Reply]: + match route: + case "openai": + return _strict_chat_vendor(_openai_refusal) + case "azure": + return _strict_chat_vendor(_azure_refusal) + case "bedrock_converse": + return _strict_converse_vendor + case "gemini": + return _gemini_vendor + + +def _deployment(scenario: Scenario, route: Route, wire_url: str) -> str: + match route: + case "openai": + return scenario.model(model="openai/gpt-5.6", api_base=f"{wire_url}/v1", api_key="scripted-openai-key") + case "azure": + return scenario.model( + model=f"azure/{_AZURE_DEPLOYMENT}", + api_base=wire_url, + api_key="scripted-azure-key", + api_version=_AZURE_API_VERSION, + ) + case "bedrock_converse": + return scenario.model( + model=f"bedrock/converse/{BEDROCK_MODEL}", + api_base=wire_url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire_url, + ) + case "gemini": + return scenario.model(model=f"gemini/{_GEMINI}", api_base=wire_url, api_key="synthetic-gemini-key") + + +def _target(route: Route) -> str: + match route: + case "openai": + return "/v1/chat/completions" + case "azure": + return f"/openai/deployments/{_AZURE_DEPLOYMENT}/chat/completions?api-version={_AZURE_API_VERSION}" + case "bedrock_converse": + return f"/model/{BEDROCK_MODEL}/converse" + case "gemini": + return f"/models/{_GEMINI}:generateContent" + + +def _only_post(wire: Wire) -> tuple[Request, dict[str, JsonValue]]: + posted: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert len(posted) == 1, [request.target for request in posted] + return posted[0], rv.JSON_OBJECT.validate_json(posted[0].body) + + +def _wire_messages(body: Mapping[str, JsonValue]) -> list[dict[str, JsonValue]]: + return rv.ITEMS.validate_python(body.get("messages") or body.get("contents") or []) + + +def _chat_call_ids(messages: Sequence[Mapping[str, JsonValue]]) -> list[str]: + (assistant,) = [message for message in messages if message.get("role") == "assistant"] + return [str(call["id"]) for call in rv.ITEMS.validate_python(assistant.get("tool_calls") or [])] + + +def _chat_result_ids(messages: Sequence[Mapping[str, JsonValue]]) -> list[str]: + return [str(message["tool_call_id"]) for message in messages if message.get("role") == "tool"] + + +def _assert_only_the_function_call_went_out(route: Route, body: Mapping[str, JsonValue], marker: str) -> None: + messages: Final = _wire_messages(body) + match route: + case "openai" | "azure": + assert "srvtoolu_" not in json.dumps(body) and _REGEX_TOOL_NAME not in json.dumps(body), body + assert _chat_call_ids(messages) == [f"toolu_{marker}"], messages + assert _chat_result_ids(messages) == [f"toolu_{marker}"], messages + case "bedrock_converse": + assert "srvtoolu_" not in json.dumps(body) and _REGEX_TOOL_NAME not in json.dumps(body), body + assert [message["role"] for message in messages] == ["user", "assistant", "user"], messages + assert _converse_ids(messages[1], "toolUse") == {f"toolu_{marker}"}, messages + assert _converse_ids(messages[2], "toolResult") == {f"toolu_{marker}"}, messages + case "gemini": + assert _REGEX_TOOL_NAME in json.dumps(messages), messages + + +def _answer_text(route: Route, marker: str) -> str: + match route: + case "openai" | "azure": + return f"answer marker-{marker}" + case "bedrock_converse": + return "bedrock wire control" + case "gemini": + return "gemini control" + + +def _history_cells() -> tuple[ParameterSet, ...]: + cases: Final[tuple[tuple[Route, bool], ...]] = ( + ("openai", False), + ("openai", True), + ("azure", False), + ("azure", True), + ("bedrock_converse", False), + ("gemini", False), + ) + return tuple(pytest.param(route, stream, id=f"{route}-{'stream' if stream else 'sync'}") for route, stream in cases) + + +@pytest.mark.parametrize(("route", "stream"), _history_cells()) +def test_a_claude_search_call_with_no_result_is_dropped_only_on_the_routes_that_reject_it( + gateway: Gateway, route: Route, stream: bool +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_vendor(route)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, route, wire.url) + answer: Final = raw_answer(gateway, _body(model, _second_turn(marker), stream)) + assert answer.status == 200, answer.text + assert _answer_text(route, marker) in answer.text, answer.text + request, body = _only_post(wire) + assert unquote(request.target).endswith(_target(route)), request.target + _assert_only_the_function_call_went_out(route, body, marker) + assert_lands_once_as(model, (response_id(answer),), "success") + + +@pytest.mark.parametrize("content", [None, ""], ids=["null", "empty"]) +def test_a_search_only_assistant_turn_with_no_content_is_removed(gateway: Gateway, content: JsonValue) -> None: + marker: Final = uuid.uuid4().hex + messages: Final[list[JsonValue]] = [ + {"role": "user", "content": f"Weather in Paris marker-{marker}"}, + {"role": "assistant", "content": content, "tool_calls": [_search_call(marker)]}, + {"role": "user", "content": f"Just tell me marker-{marker}"}, + ] + with wire_server(_vendor("openai")) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai", wire.url) + answer: Final = raw_answer(gateway, _body(model, messages, False)) + assert answer.status == 200 and f"answer marker-{marker}" in answer.text, answer.text + _, body = _only_post(wire) + sent: Final = _wire_messages(body) + assert [message["role"] for message in sent] == ["user", "user"], sent + assert "srvtoolu_" not in json.dumps(body), body + + +@pytest.mark.parametrize("name", [_REGEX_TOOL_NAME, _BM25_TOOL_NAME]) +def test_a_search_call_the_client_answered_keeps_its_result(gateway: Gateway, name: str) -> None: + marker: Final = uuid.uuid4().hex + search_result: Final[dict[str, JsonValue]] = { + "role": "tool", + "tool_call_id": f"srvtoolu_{marker}", + "content": f"get_weather marker-{marker}", + } + messages: Final = _history( + marker, (_search_call(marker, name), _weather_call(marker)), search_result, _weather_result(marker) + ) + with wire_server(_vendor("openai")) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai", wire.url) + answer: Final = raw_answer(gateway, _body(model, messages, False)) + assert answer.status == 200 and f"answer marker-{marker}" in answer.text, answer.text + _, body = _only_post(wire) + sent: Final = _wire_messages(body) + assert _chat_call_ids(sent) == [f"srvtoolu_{marker}", f"toolu_{marker}"], sent + assert _chat_result_ids(sent) == [f"srvtoolu_{marker}", f"toolu_{marker}"], sent + + +def test_a_function_call_whose_id_only_looks_like_a_search_call_is_kept(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + look_alike: Final = _tool_call(f"srvtoolu_{marker}", "get_weather", '{"city": "Lyon"}') + messages: Final = _history(marker, (look_alike, _weather_call(marker)), _weather_result(marker)) + with wire_server(_vendor("openai")) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai", wire.url) + answer: Final = raw_answer(gateway, _body(model, messages, False)) + assert answer.status == 400, answer.text + assert f"No tool output found for function call srvtoolu_{marker}" in answer.text, answer.text + _, body = _only_post(wire) + assert _chat_call_ids(_wire_messages(body)) == [f"srvtoolu_{marker}", f"toolu_{marker}"], body + + +@pytest.mark.parametrize( + "tool_call_id", + [5, ["srvtoolu_x"], "", "srvtoolu_" + "x" * 5000], + ids=["int", "list", "empty", "five-kb"], +) +def test_a_malformed_answer_id_leaves_the_search_call_unanswered_and_the_proxy_serving( + gateway: Gateway, tool_call_id: JsonValue +) -> None: + marker: Final = uuid.uuid4().hex + malformed: Final[dict[str, JsonValue]] = {"role": "tool", "tool_call_id": tool_call_id, "content": "x"} + messages: Final = _history( + marker, (_search_call(marker), _weather_call(marker)), malformed, _weather_result(marker) + ) + with wire_server(_vendor("openai")) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai", wire.url) + answer: Final = raw_answer(gateway, _body(model, messages, False)) + assert answer.status in (200, 400), answer.text + posted: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert all("srvtoolu_" + marker not in request.body.decode() for request in posted), posted + follow_up: Final = uuid.uuid4().hex + served: Final = raw_answer(gateway, _body(model, _second_turn(follow_up), False)) + assert served.status == 200 and f"answer marker-{follow_up}" in served.text, served.text diff --git a/tests/integration/providers/test_chat_completions_tool_search_wire.py b/tests/integration/providers/test_chat_completions_tool_search_wire.py new file mode 100644 index 00000000000..f99db3fe735 --- /dev/null +++ b/tests/integration/providers/test_chat_completions_tool_search_wire.py @@ -0,0 +1,827 @@ +import asyncio +import json +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType +from typing import Final, Literal, TypeAlias +from urllib.parse import unquote + +import httpx +import openai +import pytest +from integration._support import claude_code as cc +from integration._support import responses_vendor as rv +from integration._support.anthropic_thinking import BEDROCK_MODEL, EVENT_STREAM, MODEL, Event, aws_chunks, sse_chunks +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +from integration._support.openai_wire import chat_reply +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers.test_bedrock_auth_wire import RESPONSE as CONVERSE_RESPONSE +from pydantic import JsonValue, TypeAdapter + +_CLAUDE: Final = cc.OPUS +UNLISTED_CLAUDE: Final = "claude-scripted-unlisted" +_NOVA: Final = "amazon.nova-pro-v1:0" +_REGEX_TOOL_TYPE: Final = "tool_search_tool_regex_20251119" +_REGEX_TOOL_NAME: Final = "tool_search_tool_regex" +_BM25_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"} +) +_REGEX_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType({"type": _REGEX_TOOL_TYPE, "name": _REGEX_TOOL_NAME}) +_SEARCH_TOOL_TYPES: Final = frozenset( + {_REGEX_TOOL_TYPE, _REGEX_TOOL_NAME, "tool_search_tool_bm25_20251119", "tool_search_tool_bm25"} +) +_TOOL_SEARCH_BETA_PREFIXES: Final = ("advanced-tool-use", "tool-search-tool") +_SEARCH_ARGUMENTS: Final = '{"pattern": "weather"}' +_ALL_DEFERRED_ERROR: Final = "At least one tool must have defer_loading=false" +_MODEL_INFO_HINT: Final = "supports_tool_search: true" +_INVOKE_HINT: Final = "bedrock/invoke/" +DEFERRED_NAMES: Final = ("get_weather", "get_stock_price", "lookup_customer") +_CACHE_BUST: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}}) +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_VERTEX_MODEL_PATH: Final = ( + f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}/publishers/anthropic/models/{MODEL}" +) +_BETAS: Final = TypeAdapter(list[JsonValue]) +_CITY_PARAMETERS: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} + +Client: TypeAlias = Literal["sdk", "async_sdk", "raw"] +Provider: TypeAlias = Literal["anthropic", "bedrock_invoke", "vertex_ai", "azure_ai", "claude_platform"] +DropSource: TypeAlias = Literal["off", "request", "deployment", "null", "request_string"] + + +@dataclass(frozen=True, slots=True) +class Answer: + status: int + payload: dict[str, JsonValue] + frames: tuple[dict[str, JsonValue], ...] + text: str + + +def _prompt(marker: str) -> str: + return f"Weather in Paris marker-{marker}" + + +def tool_search(**fields: JsonValue) -> dict[str, JsonValue]: + return {"type": "tool_search", **fields} + + +def _top_level_deferred() -> dict[str, JsonValue]: + return { + "type": "function", + "function": {"name": "get_weather", "description": "Weather for a city", "parameters": _CITY_PARAMETERS}, + "defer_loading": True, + } + + +def _function_level_deferred() -> dict[str, JsonValue]: + return { + "type": "function", + "function": { + "name": "get_stock_price", + "description": "Price of a ticker", + "parameters": _CITY_PARAMETERS, + "defer_loading": True, + }, + } + + +def _namespace() -> dict[str, JsonValue]: + return { + "type": "namespace", + "name": "crm", + "description": "CRM tools", + "tools": [ + { + "type": "function", + "name": "lookup_customer", + "description": "Look up a customer", + "parameters": _CITY_PARAMETERS, + "defer_loading": True, + } + ], + } + + +def _loaded_function() -> dict[str, JsonValue]: + return { + "type": "function", + "function": {"name": "echo", "description": "Echo a city", "parameters": _CITY_PARAMETERS}, + } + + +def deferred_tools() -> tuple[dict[str, JsonValue], ...]: + return (_top_level_deferred(), _function_level_deferred(), _namespace()) + + +def _search_result(marker: str) -> 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"}], + }, + } + + +def _message(marker: str, searched: bool) -> dict[str, JsonValue]: + search_blocks: Final[tuple[dict[str, JsonValue], ...]] = ( + ( + { + "type": "server_tool_use", + "id": f"srvtoolu_{marker}", + "name": _REGEX_TOOL_NAME, + "input": json.loads(_SEARCH_ARGUMENTS), + }, + _search_result(marker), + ) + if searched + else () + ) + return { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": _CLAUDE, + "content": [ + {"type": "text", "text": f"searching marker-{marker}"}, + *search_blocks, + {"type": "tool_use", "id": f"toolu_{marker}", "name": "get_weather", "input": {"city": "Paris"}}, + ], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": { + "input_tokens": 20, + "output_tokens": 7, + "server_tool_use": {"tool_search_requests": 1 if searched else 0}, + }, + } + + +def _search_events(marker: str) -> tuple[Event, ...]: + return ( + { + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "server_tool_use", + "id": f"srvtoolu_{marker}", + "name": _REGEX_TOOL_NAME, + "input": {}, + }, + }, + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": _SEARCH_ARGUMENTS}, + }, + {"type": "content_block_stop", "index": 1}, + {"type": "content_block_start", "index": 2, "content_block": _search_result(marker)}, + {"type": "content_block_stop", "index": 2}, + ) + + +def _events(marker: str, searched: bool) -> tuple[Event, ...]: + opened: Final = {**_message(marker, searched), "content": [], "stop_reason": None} + tool_index: Final = 3 if searched else 1 + return ( + {"type": "message_start", "message": opened}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": f"searching marker-{marker}"}, + }, + {"type": "content_block_stop", "index": 0}, + *(_search_events(marker) if searched else ()), + { + "type": "content_block_start", + "index": tool_index, + "content_block": {"type": "tool_use", "id": f"toolu_{marker}", "name": "get_weather", "input": {}}, + }, + { + "type": "content_block_delta", + "index": tool_index, + "delta": {"type": "input_json_delta", "partial_json": '{"city": "Paris"}'}, + }, + {"type": "content_block_stop", "index": tool_index}, + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 7}, + }, + {"type": "message_stop"}, + ) + + +def _anthropic_error(status: int, message: str) -> Reply: + body: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": message}} + return Reply(status=status, body=json.dumps(body).encode()) + + +def claude_vendor(request: Request) -> Reply: + if request.method != "POST": + return Reply() + body: Final = rv.JSON_OBJECT.validate_json(request.body) + marker: Final = rv.newest_marker(request.body.decode()) or "health" + tools: Final = rv.ITEMS.validate_python(body.get("tools") or []) + search_tools: Final = [tool for tool in tools if tool.get("type") in _SEARCH_TOOL_TYPES] + function_tools: Final = [tool for tool in tools if tool.get("type") not in _SEARCH_TOOL_TYPES] + if len(search_tools) > 1: + return _anthropic_error(400, f"tools: Tool names must be unique marker-{marker}") + if not search_tools and function_tools and any(tool.get("defer_loading") is True for tool in function_tools): + return _anthropic_error(400, f"tools: {_ALL_DEFERRED_ERROR} marker-{marker}") + searched: Final = bool(search_tools) + if request.target.endswith("/invoke-with-response-stream"): + return Reply(content_type=EVENT_STREAM, chunks=aws_chunks(_events(marker, searched))) + if body.get("stream") is True or ":streamRawPredict" in request.target: + return Reply(content_type="text/event-stream", chunks=sse_chunks(_events(marker, searched))) + return Reply(body=json.dumps(_message(marker, searched)).encode()) + + +def _base_url(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def _sdk(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{_base_url(gateway)}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +def _async_sdk(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI( + base_url=f"{_base_url(gateway)}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=60), + ) + + +def raw_chat(gateway: Gateway, body: Mapping[str, JsonValue]) -> httpx.Response: + with httpx.Client(base_url=_base_url(gateway), trust_env=False, timeout=60) as client: + with client.stream( + "POST", + "/v1/chat/completions", + json={**body, **_CACHE_BUST}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + response.read() + return response + + +def chat_frames(text: str) -> tuple[dict[str, JsonValue], ...]: + return tuple(rv.JSON_OBJECT.validate_json(line[6:]) for line in text.splitlines() if line.startswith("data: {")) + + +def chat_body( + model: str, marker: str, tools: Sequence[Mapping[str, JsonValue]], stream: bool, **extra: JsonValue +) -> dict[str, JsonValue]: + return { + "model": model, + "messages": [{"role": "user", "content": _prompt(marker)}], + "tools": [dict(tool) for tool in tools], + "stream": stream, + "num_retries": 0, + **extra, + } + + +def _sdk_answer(client: openai.OpenAI, body: Mapping[str, JsonValue]) -> Answer: + tools: Final = rv.ITEMS.validate_python(body["tools"]) + messages: Final = rv.ITEMS.validate_python(body["messages"]) + if body["stream"] is not True: + created: Final = client.chat.completions.create( + model=str(body["model"]), messages=messages, tools=tools, extra_body=dict(_CACHE_BUST) + ) + return Answer(200, created.model_dump(), (), created.model_dump_json()) + frames: Final = tuple( + chunk.model_dump() + for chunk in client.chat.completions.create( + model=str(body["model"]), messages=messages, tools=tools, stream=True, extra_body=dict(_CACHE_BUST) + ) + ) + return Answer(200, {}, frames, json.dumps(frames)) + + +async def _async_sdk_answer(client: openai.AsyncOpenAI, body: Mapping[str, JsonValue]) -> Answer: + tools: Final = rv.ITEMS.validate_python(body["tools"]) + messages: Final = rv.ITEMS.validate_python(body["messages"]) + if body["stream"] is not True: + created: Final = await client.chat.completions.create( + model=str(body["model"]), messages=messages, tools=tools, extra_body=dict(_CACHE_BUST) + ) + return Answer(200, created.model_dump(), (), created.model_dump_json()) + stream: Final = await client.chat.completions.create( + model=str(body["model"]), messages=messages, tools=tools, stream=True, extra_body=dict(_CACHE_BUST) + ) + frames: Final = tuple([chunk.model_dump() async for chunk in stream]) + return Answer(200, {}, frames, json.dumps(frames)) + + +def raw_answer(gateway: Gateway, body: Mapping[str, JsonValue]) -> Answer: + response: Final = raw_chat(gateway, body) + if response.status_code != 200 or body.get("stream") is True: + return Answer(response.status_code, {}, chat_frames(response.text), response.text) + return Answer(200, rv.JSON_OBJECT.validate_json(response.content), (), response.text) + + +def _answer(gateway: Gateway, client: Client, body: Mapping[str, JsonValue]) -> Answer: + match client: + case "sdk": + return _sdk_answer(_sdk(gateway), body) + case "async_sdk": + return asyncio.run(_async_sdk_answer(_async_sdk(gateway), body)) + case "raw": + return raw_answer(gateway, body) + + +def _first_choice(payload: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return rv.ITEMS.validate_python(payload["choices"])[0] + + +def _streamed_tool_calls(frames: Sequence[Mapping[str, JsonValue]]) -> list[tuple[str, str]]: + deltas: Final = [ + rv.JSON_OBJECT.validate_python(_first_choice(frame)["delta"]) for frame in frames if frame.get("choices") + ] + opened: Final = rv.ITEMS.validate_python( + list(chain.from_iterable(rv.ITEMS.validate_python(delta.get("tool_calls") or []) for delta in deltas)) + ) + return [ + (str(call["id"]), str(rv.JSON_OBJECT.validate_python(call["function"])["name"])) + for call in opened + if call.get("id") + ] + + +def answered_tool_calls(answer: Answer) -> list[tuple[str, str]]: + if answer.frames: + return _streamed_tool_calls(answer.frames) + message: Final = rv.JSON_OBJECT.validate_python(_first_choice(answer.payload)["message"]) + return [ + (str(call["id"]), str(rv.JSON_OBJECT.validate_python(call["function"])["name"])) + for call in rv.ITEMS.validate_python(message.get("tool_calls") or []) + ] + + +def response_id(answer: Answer) -> str: + return str((answer.frames[0] if answer.frames else answer.payload)["id"]) + + +def assert_search_call_answered(answer: Answer, marker: str) -> None: + assert answer.status == 200, answer.text + assert answered_tool_calls(answer) == [ + (f"srvtoolu_{marker}", _REGEX_TOOL_NAME), + (f"toolu_{marker}", "get_weather"), + ], answer.text + if answer.frames: + return + usage: Final = rv.JSON_OBJECT.validate_python(answer.payload["usage"]) + assert rv.JSON_OBJECT.validate_python(usage["server_tool_use"])["tool_search_requests"] == 1, usage + + +def only_request(wire: Wire) -> tuple[Request, dict[str, JsonValue]]: + received: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert len(received) == 1, [request.target for request in received] + return received[0], rv.JSON_OBJECT.validate_json(received[0].body) + + +def _target(provider: Provider, stream: bool) -> str: + match provider: + case "anthropic" | "claude_platform": + return "/v1/messages" + case "azure_ai": + return "/anthropic/v1/messages" + case "bedrock_invoke": + return f"/model/{BEDROCK_MODEL}/invoke-with-response-stream" if stream else f"/model/{BEDROCK_MODEL}/invoke" + case "vertex_ai": + return f"{_VERTEX_MODEL_PATH}:streamRawPredict?alt=sse" if stream else f"{_VERTEX_MODEL_PATH}:rawPredict" + + +def _betas(request: Request, body: Mapping[str, JsonValue]) -> tuple[str, ...]: + header: Final = tuple(beta.strip() for beta in request.headers.get("anthropic-beta", "").split(",") if beta.strip()) + in_body: Final = tuple(str(beta) for beta in _BETAS.validate_python(body.get("anthropic_beta") or [])) + return (*header, *in_body) + + +def assert_search_tool_on_the_wire( + provider: Provider, stream: bool, request: Request, body: Mapping[str, JsonValue], tools_before: int = 0 +) -> list[dict[str, JsonValue]]: + assert unquote(request.target) == _target(provider, stream), request.target + assert provider == "azure_ai" or any( + beta.startswith(_TOOL_SEARCH_BETA_PREFIXES) for beta in _betas(request, body) + ), (dict(request.headers), body.get("anthropic_beta")) + assert "no-cache" not in request.body.decode() and '"tool_search"' not in request.body.decode(), body + tools: Final = rv.ITEMS.validate_python(body["tools"]) + search_tools: Final = [index for index, tool in enumerate(tools) if tool.get("type") in _SEARCH_TOOL_TYPES] + assert search_tools == [tools_before], tools + return tools + + +def assert_deferred_tools_kept(tools: Sequence[Mapping[str, JsonValue]]) -> None: + deferred: Final = [ + (tool["name"], tool.get("defer_loading")) for tool in tools if tool.get("type") not in _SEARCH_TOOL_TYPES + ] + assert deferred == [(name, True) for name in DEFERRED_NAMES], tools + + +def assert_no_search_and_no_deferral(body: Mapping[str, JsonValue], names: Sequence[str]) -> None: + sent: Final = json.dumps(body) + assert "defer_loading" not in sent and "tool_search" not in sent, body + assert [tool["name"] for tool in rv.ITEMS.validate_python(body["tools"])] == list(names), body + + +def _spend_rows(model: str, expected: int) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda found: len(found) >= expected, + seconds=70, + ) + + +def assert_lands_once_as(model: str, identities: Sequence[str], status: str) -> None: + rows: Final = _spend_rows(model, len(identities)) + logged: Final = {str(row["request_id"]): str(row["status"]) for row in rows} + assert len(logged) == len(rows) == len(identities), rows + matched: Final = [ + [logged_id for logged_id in logged if rv.same_response(logged_id, identity)] for identity in identities + ] + assert all(len(ids) == 1 for ids in matched), (identities, rows) + assert {logged[ids[0]] for ids in matched} == {status}, rows + + +def claude_deployment(scenario: Scenario, provider: Provider, wire_url: str, upstream_url: str) -> str: + match provider: + case "anthropic": + return scenario.model(model=f"anthropic/{_CLAUDE}", api_base=wire_url, api_key=cc.ANTHROPIC_API_KEY) + case "bedrock_invoke": + return scenario.model( + model=f"bedrock/invoke/{BEDROCK_MODEL}", + api_base=wire_url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire_url, + ) + case "vertex_ai": + return scenario.model( + model=f"vertex_ai/{MODEL}", + api_base=f"{wire_url}{_VERTEX_MODEL_PATH}", + api_key=None, + vertex_project=_VERTEX_PROJECT, + vertex_location=_VERTEX_LOCATION, + vertex_credentials=service_account_json(_VERTEX_PROJECT, upstream_url.rstrip("/")), + ) + case "azure_ai": + return scenario.model(model=f"azure_ai/{MODEL}", api_base=wire_url, api_key="scripted-azure-key") + case "claude_platform": + return scenario.model( + model=f"bedrock/claude_platform/{MODEL}", + api_base=wire_url, + api_key="scripted-platform-key", + aws_region_name="us-east-1", + workspace_id="scripted-workspace", + ) + + +def _unlisted(scenario: Scenario, wire_url: str, **extra: JsonValue) -> str: + return scenario.model( + model=f"anthropic/{UNLISTED_CLAUDE}", api_base=wire_url, api_key=cc.ANTHROPIC_API_KEY, **extra + ) + + +def _assert_refused_before_the_wire(answer: Answer, wire: Wire, *, hint: str | None, model: str) -> None: + assert answer.status == 400, answer.text + assert f"{model} on " in answer.text and "cannot run the OpenAI `tool_search` tool as sent" in answer.text, ( + answer.text + ) + assert "drop_params" in answer.text, answer.text + assert (hint is not None and hint in answer.text) or ( + hint is None and _MODEL_INFO_HINT not in answer.text and _INVOKE_HINT not in answer.text + ), answer.text + assert [request for request in wire.drain() if request.method == "POST"] == [], answer.text + + +@pytest.mark.parametrize("stream", [False, True], ids=["sync", "stream"]) +@pytest.mark.parametrize("client", ["sdk", "async_sdk", "raw"]) +def test_a_server_tool_search_reaches_claude_as_its_regex_tool_with_every_deferral_kept( + gateway: Gateway, client: Client, stream: bool +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(claude_vendor) as wire, gateway.scenario() as scenario: + model: Final = claude_deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + tools: Final = (tool_search(execution="server"), *deferred_tools()) + answer: Final = _answer(gateway, client, chat_body(model, marker, tools, stream)) + assert_search_call_answered(answer, marker) + request, body = only_request(wire) + assert_deferred_tools_kept(assert_search_tool_on_the_wire("anthropic", stream, request, body)) + assert body.get("stream", False) is stream, body + assert_lands_once_as(model, (response_id(answer),), "success") + + +@pytest.mark.parametrize("stream", [False, True], ids=["sync", "stream"]) +@pytest.mark.parametrize("provider", ["bedrock_invoke", "vertex_ai", "azure_ai", "claude_platform"]) +def test_every_claude_route_that_runs_tool_search_gets_the_regex_tool( + gateway: Gateway, provider: Provider, stream: bool +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(claude_vendor) as wire, gateway.scenario() as scenario: + model: Final = claude_deployment(scenario, provider, wire.url, gateway.upstream_url) + answer: Final = raw_answer(gateway, chat_body(model, marker, (tool_search(), *deferred_tools()), stream)) + assert_search_call_answered(answer, marker) + request, body = only_request(wire) + assert_deferred_tools_kept(assert_search_tool_on_the_wire(provider, stream, request, body)) + assert_lands_once_as(model, (response_id(answer),), "success") + + +def test_several_tool_search_items_go_out_as_one_regex_tool_at_the_first_items_position(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(claude_vendor) as wire, gateway.scenario() as scenario: + model: Final = claude_deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + tools: Final = ( + _loaded_function(), + tool_search(), + _top_level_deferred(), + tool_search(execution="server"), + _function_level_deferred(), + ) + answer: Final = raw_answer(gateway, chat_body(model, marker, tools, False)) + assert_search_call_answered(answer, marker) + request, body = only_request(wire) + sent: Final = assert_search_tool_on_the_wire("anthropic", False, request, body, tools_before=1) + assert sent[1] == dict(_REGEX_TOOL), sent + assert [(tool.get("name"), tool.get("defer_loading")) for tool in sent] == [ + ("echo", None), + (_REGEX_TOOL_NAME, None), + ("get_weather", True), + ("get_stock_price", True), + ], sent + + +@pytest.mark.parametrize("native", [_REGEX_TOOL, _BM25_TOOL], ids=["regex", "bm25"]) +def test_a_native_search_tool_the_caller_sent_goes_out_alone(gateway: Gateway, native: Mapping[str, JsonValue]) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(claude_vendor) as wire, gateway.scenario() as scenario: + model: Final = claude_deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + tools: Final = (dict(native), tool_search(execution="client"), tool_search(), _top_level_deferred()) + answer: Final = raw_answer(gateway, chat_body(model, marker, tools, False)) + assert_search_call_answered(answer, marker) + request, body = only_request(wire) + sent: Final = assert_search_tool_on_the_wire("anthropic", False, request, body) + assert sent[0] == dict(native), sent + assert [(tool["name"], tool.get("defer_loading")) for tool in sent[1:]] == [("get_weather", True)], sent + + +@pytest.mark.parametrize("stream", [False, True], ids=["sync", "stream"]) +def test_a_claude_the_cost_map_does_not_list_is_refused_with_400_naming_model_info( + gateway: Gateway, stream: bool +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(claude_vendor) as wire, gateway.scenario() as scenario: + model: Final = _unlisted(scenario, wire.url) + refused: Final = raw_answer(gateway, chat_body(model, marker, (tool_search(), *deferred_tools()), stream)) + _assert_refused_before_the_wire(refused, wire, hint=_MODEL_INFO_HINT, model=UNLISTED_CLAUDE) + follow_up: Final = uuid.uuid4().hex + served: Final = raw_answer(gateway, chat_body(model, follow_up, (_loaded_function(),), stream)) + assert served.status == 200 and f"marker-{follow_up}" in served.text, served.text + _, body = only_request(wire) + assert [tool["name"] for tool in rv.ITEMS.validate_python(body["tools"])] == ["echo"], body + + +def _drop_params_fields(source: DropSource) -> Mapping[str, JsonValue]: + match source: + case "off" | "deployment": + return {} + case "request": + return {"drop_params": True} + case "null": + return {"drop_params": None} + case "request_string": + return {"drop_params": "true"} + + +@pytest.mark.parametrize("source", ["off", "request", "deployment", "null", "request_string"]) +def test_drop_params_drops_the_search_and_every_deferral_and_without_it_the_call_is_refused( + gateway: Gateway, source: DropSource +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(claude_vendor) as wire, gateway.scenario() as scenario: + model: Final = _unlisted(scenario, wire.url, **({"drop_params": True} if source == "deployment" else {})) + tools: Final = (tool_search(), *deferred_tools()) + answer: Final = raw_answer(gateway, chat_body(model, marker, tools, False, **_drop_params_fields(source))) + if source in ("off", "null"): + _assert_refused_before_the_wire(answer, wire, hint=_MODEL_INFO_HINT, model=UNLISTED_CLAUDE) + return + assert answer.status == 200, answer.text + assert answered_tool_calls(answer) == [(f"toolu_{marker}", "get_weather")], answer.text + request, body = only_request(wire) + assert request.target == "/v1/messages", request.target + assert_no_search_and_no_deferral(body, DEFERRED_NAMES) + assert_lands_once_as(model, (response_id(answer),), "success") + + +@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"]) +def test_model_info_supports_tool_search_turns_the_search_on_for_a_name_the_cost_map_does_not_know( + gateway: Gateway, provider: Literal["anthropic", "azure_ai"] +) -> None: + unknown_name: Final = f"claude-scripted-flagged-{uuid.uuid4().hex[:12]}" + markers: Final = tuple(uuid.uuid4().hex for _ in range(6)) + with wire_server(claude_vendor) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model_info={"supports_tool_search": True}, + model=f"{provider}/{unknown_name}", + api_base=wire.url, + api_key="scripted-flagged-key", + ) + answers: Final = tuple( + raw_answer(gateway, chat_body(model, marker, (tool_search(), *deferred_tools()), False)) + for marker in markers + ) + for answer, marker in zip(answers, markers, strict=True): + assert_search_call_answered(answer, marker) + received: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert len(received) == len(markers), [request.target for request in received] + for request in received: + assert_deferred_tools_kept( + assert_search_tool_on_the_wire(provider, False, request, rv.JSON_OBJECT.validate_json(request.body)) + ) + assert_lands_once_as(model, tuple(response_id(answer) for answer in answers), "success") + + +@pytest.mark.parametrize( + "search_tools", + [ + (tool_search(execution="client"),), + (tool_search(execution="server"), tool_search(execution="client")), + (tool_search(execution=5),), + ], + ids=["client", "server-and-client", "int"], +) +def test_a_client_executed_or_malformed_search_is_refused_with_400_without_the_model_info_hint( + gateway: Gateway, search_tools: tuple[dict[str, JsonValue], ...] +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(claude_vendor) as wire, gateway.scenario() as scenario: + model: Final = claude_deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + refused: Final = raw_answer(gateway, chat_body(model, marker, (*search_tools, *deferred_tools()), False)) + _assert_refused_before_the_wire(refused, wire, hint=None, model=_CLAUDE) + + +def test_a_client_executed_search_with_drop_params_is_dropped_with_every_deferral(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(claude_vendor) as wire, gateway.scenario() as scenario: + model: Final = claude_deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + tools: Final = (tool_search(execution="client"), *deferred_tools()) + answer: Final = raw_answer(gateway, chat_body(model, marker, tools, False, drop_params=True)) + assert answer.status == 200, answer.text + _, body = only_request(wire) + assert_no_search_and_no_deferral(body, DEFERRED_NAMES) + + +def _converse_vendor(request: Request) -> Reply: + assert request.method == "POST" and unquote(request.target).endswith("/converse"), request.target + return Reply(body=CONVERSE_RESPONSE) + + +def _converse(scenario: Scenario, wire_url: str, model: str, **extra: JsonValue) -> str: + return scenario.model( + model=model, + api_base=wire_url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire_url, + **extra, + ) + + +def _converse_tools() -> tuple[dict[str, JsonValue], ...]: + return (tool_search(execution="server"), _top_level_deferred(), _function_level_deferred()) + + +@pytest.mark.parametrize( + ("route_model", "hint"), + [ + (f"bedrock/converse/{BEDROCK_MODEL}", _INVOKE_HINT), + (f"bedrock/us.{BEDROCK_MODEL}", _INVOKE_HINT), + (f"bedrock/converse/{_NOVA}", None), + ], + ids=["claude-converse-prefix", "claude-no-route-prefix", "nova"], +) +def test_bedrock_converse_refuses_tool_search_with_400(gateway: Gateway, route_model: str, hint: str | None) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_converse_vendor) as wire, gateway.scenario() as scenario: + model: Final = _converse(scenario, wire.url, route_model) + refused: Final = raw_answer(gateway, chat_body(model, marker, _converse_tools(), False)) + _assert_refused_before_the_wire(refused, wire, hint=hint, model=route_model.split("/")[-1]) + assert " on bedrock " in refused.text, refused.text + + +@pytest.mark.parametrize("source", ["request", "deployment"]) +def test_bedrock_converse_with_drop_params_drops_the_search_and_every_deferral( + gateway: Gateway, source: Literal["request", "deployment"] +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_converse_vendor) as wire, gateway.scenario() as scenario: + model: Final = _converse( + scenario, + wire.url, + f"bedrock/converse/{BEDROCK_MODEL}", + **({"drop_params": True} if source == "deployment" else {}), + ) + extra: Final[Mapping[str, JsonValue]] = {"drop_params": True} if source == "request" else {} + answer: Final = raw_answer(gateway, chat_body(model, marker, _converse_tools(), False, **extra)) + assert answer.status == 200 and "bedrock wire control" in answer.text, answer.text + request, body = only_request(wire) + assert unquote(request.target) == f"/model/{BEDROCK_MODEL}/converse", request.target + sent: Final = json.dumps(body) + assert "tool_search" not in sent and "defer_loading" not in sent, body + specs: Final = rv.ITEMS.validate_python(rv.JSON_OBJECT.validate_python(body["toolConfig"])["tools"]) + assert [rv.JSON_OBJECT.validate_python(spec["toolSpec"])["name"] for spec in specs] == [ + "get_weather", + "get_stock_price", + ], body + assert_lands_once_as(model, (response_id(answer),), "success") + + +def _databricks_vendor(request: Request) -> Reply: + assert request.method == "POST" and request.target.endswith("/chat/completions"), request.target + marker: Final = rv.newest_marker(request.body.decode()) or "health" + return chat_reply(f"chatcmpl-{marker}", "databricks-claude", f"answer marker-{marker}", stream=False) + + +@pytest.mark.parametrize("drop_params", [False, True], ids=["refused", "dropped"]) +def test_databricks_claude_answers_400_or_drops_the_search_under_drop_params( + gateway: Gateway, drop_params: bool +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_databricks_vendor) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"databricks/databricks-{_CLAUDE}", api_base=f"{wire.url}/serving-endpoints", api_key="scripted-dbx" + ) + answer: Final = raw_answer( + gateway, + chat_body(model, marker, _converse_tools(), False, **({"drop_params": True} if drop_params else {})), + ) + if not drop_params: + _assert_refused_before_the_wire(answer, wire, hint=None, model=f"databricks-{_CLAUDE}") + assert " on databricks " in answer.text, answer.text + return + assert answer.status == 200 and f"answer marker-{marker}" in answer.text, answer.text + _, body = only_request(wire) + sent: Final = json.dumps(body) + assert "tool_search" not in sent and "defer_loading" not in sent, body + names: Final = [ + rv.JSON_OBJECT.validate_python(tool["function"])["name"] for tool in rv.ITEMS.validate_python(body["tools"]) + ] + assert names == ["get_weather", "get_stock_price"], body + + +def test_databricks_non_claude_gets_the_tools_as_sent(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_databricks_vendor) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="databricks/databricks-meta-llama-4-maverick", + api_base=f"{wire.url}/serving-endpoints", + api_key="scripted-dbx", + ) + tools: Final = (tool_search(), _top_level_deferred()) + answer: Final = raw_answer(gateway, chat_body(model, marker, tools, False)) + assert answer.status == 200 and f"answer marker-{marker}" in answer.text, answer.text + _, body = only_request(wire) + assert rv.ITEMS.validate_python(body["tools"]) == [dict(tool) for tool in tools], body + + +def test_a_fallback_from_openai_to_claude_carries_the_tool_search_to_claude(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + + def overloaded_openai(request: Request) -> Reply: + return Reply(status=429, body=json.dumps({"error": {"message": "scripted 429", "type": "rate_limit"}}).encode()) + + with ( + wire_server(overloaded_openai) as openai_wire, + wire_server(claude_vendor) as claude_wire, + gateway.scenario() as scenario, + ): + primary: Final = scenario.model(model="openai/gpt-5.6", api_base=f"{openai_wire.url}/v1", api_key="scripted") + claude: Final = claude_deployment(scenario, "anthropic", claude_wire.url, gateway.upstream_url) + answer: Final = raw_answer( + gateway, chat_body(primary, marker, (tool_search(), *deferred_tools()), False, fallbacks=[claude]) + ) + assert_search_call_answered(answer, marker) + posted: Final = [request for request in openai_wire.drain() if request.method == "POST"] + assert len(posted) == 1 and f"marker-{marker}" in posted[0].body.decode(), posted + request, body = only_request(claude_wire) + assert_deferred_tools_kept(assert_search_tool_on_the_wire("anthropic", False, request, body)) 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 f37e2e87c91..0c093700822 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -30,13 +30,14 @@ from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig +from litellm.llms.bedrock.claude_platform.transformation import BedrockClaudePlatformConfig from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig, ) -from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_TOOL_SEARCH_TOOL_TYPES from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk from litellm.types.utils import ChatCompletionToolCallChunk, ModelResponse, ServerToolUse, Usage from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE @@ -9250,6 +9251,215 @@ def test_completion_surfaces_response_citations_in_provider_specific_fields(resp ] +def _openai_tool_search_with_deferred_functions( + search_tools: tuple[dict[str, object], ...] = ({"type": "tool_search", "execution": "server"},), +) -> list[dict[str, object]]: + return [ + { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, + "defer_loading": True, + }, + *search_tools, + { + "type": "function", + "function": { + "name": "get_stock_price", + "parameters": {"type": "object", "properties": {}}, + "defer_loading": True, + }, + }, + { + "type": "namespace", + "name": "crm", + "tools": [ + { + "type": "function", + "name": "lookup_customer", + "parameters": {"type": "object", "properties": {}}, + "defer_loading": True, + } + ], + }, + ] + + +_DEFERRED_FUNCTION_NAMES: Final = ["get_weather", "get_stock_price", "lookup_customer"] + + +def _search_tools_sent(tools: list[dict[str, object]]) -> list[dict[str, object]]: + return [ + {"type": tool["type"], "name": tool["name"]} + for tool in tools + if tool.get("type") in ANTHROPIC_TOOL_SEARCH_TOOL_TYPES + ] + + +def _deferral_by_function_name(tools: list[dict[str, object]]) -> dict[object, object]: + return { + tool["name"]: tool.get("defer_loading") + for tool in tools + if tool.get("type") not in ANTHROPIC_TOOL_SEARCH_TOOL_TYPES + } + + +@pytest.mark.parametrize( + ("config", "model", "provider"), + [ + pytest.param(AnthropicConfig(), "claude-sonnet-5-5", "anthropic", id="anthropic"), + pytest.param(AmazonAnthropicClaudeConfig(), "us.anthropic.claude-sonnet-5-5", "bedrock", id="bedrock-invoke"), + pytest.param(VertexAIAnthropicConfig(), "claude-sonnet-5-5", "vertex_ai", id="vertex_ai"), + pytest.param(AzureAnthropicConfig(), "claude-sonnet-5-5", "azure_ai", id="azure_ai"), + pytest.param( + BedrockClaudePlatformConfig(), "claude_platform/claude-opus-5-5", "bedrock", id="bedrock-claude-platform" + ), + ], +) +@pytest.mark.parametrize( + "search_tools", + [ + pytest.param(({"type": "tool_search", "execution": "server"},), id="one-search-item"), + pytest.param(({"type": "tool_search"}, {"type": "tool_search", "execution": "server"}), id="two-search-items"), + ], +) +def test_map_openai_params_sends_one_anthropic_search_tool_and_keeps_deferral_on_a_route_that_runs_it( + config, model, provider, search_tools +): + from litellm.responses.litellm_completion_transformation.tool_search import anthropic_tool_search_chat_tool + + assert litellm.supports_tool_search(model, provider) + mapped: Final = config.map_openai_params( + non_default_params={"tools": _openai_tool_search_with_deferred_functions(search_tools)}, + optional_params={}, + model=model, + drop_params=False, + ) + assert _search_tools_sent(mapped["tools"]) == [anthropic_tool_search_chat_tool()] + assert _deferral_by_function_name(mapped["tools"]) == dict.fromkeys(_DEFERRED_FUNCTION_NAMES, True) + + +@pytest.mark.parametrize("native_type", sorted(ANTHROPIC_TOOL_SEARCH_TOOL_TYPES)) +@pytest.mark.parametrize( + ("config", "model"), + [ + pytest.param(AnthropicConfig(), "claude-sonnet-5-5", id="model-with-tool-search"), + pytest.param(AnthropicConfig(), "claude-3-haiku-20240307", id="model-without-tool-search"), + ], +) +def test_map_openai_params_lets_a_native_anthropic_search_tool_replace_the_openai_one( + monkeypatch, config, model, native_type +): + monkeypatch.setattr(litellm, "drop_params", False) + native_name: Final = native_type.removesuffix("_20251119") + mapped: Final = config.map_openai_params( + non_default_params={ + "tools": _openai_tool_search_with_deferred_functions( + ({"type": "tool_search", "execution": "server"}, {"type": native_type, "name": native_name}) + ) + }, + optional_params={}, + model=model, + drop_params=False, + ) + assert _search_tools_sent(mapped["tools"]) == [{"type": native_type, "name": native_name}] + assert _deferral_by_function_name(mapped["tools"]) == dict.fromkeys(_DEFERRED_FUNCTION_NAMES, True) + + +@pytest.mark.parametrize( + ("config", "model", "provider"), + [ + pytest.param(AmazonAnthropicClaudeConfig(), "us.anthropic.claude-sonnet-5-5", "bedrock", id="bedrock-invoke"), + pytest.param(VertexAIAnthropicConfig(), "claude-sonnet-5-5", "vertex_ai", id="vertex_ai"), + ], +) +def test_map_openai_params_judges_tool_search_on_the_real_model_when_tool_based_response_format_swaps_it( + config, model, provider +): + from litellm.responses.litellm_completion_transformation.tool_search import anthropic_tool_search_chat_tool + + assert litellm.supports_tool_search(model, provider) + mapped: Final = config.map_openai_params( + non_default_params={ + "tools": _openai_tool_search_with_deferred_functions(), + "response_format": { + "type": "json_schema", + "json_schema": {"name": "weather", "schema": {"type": "object", "properties": {}}}, + }, + }, + optional_params={}, + model=model, + drop_params=False, + ) + deferral: Final = _deferral_by_function_name(mapped["tools"]) + assert RESPONSE_FORMAT_TOOL_NAME in deferral + assert _search_tools_sent(mapped["tools"]) == [anthropic_tool_search_chat_tool()] + assert {name: deferral[name] for name in _DEFERRED_FUNCTION_NAMES} == dict.fromkeys(_DEFERRED_FUNCTION_NAMES, True) + + +_UNRUNNABLE_TOOL_SEARCH_CASES: Final = [ + pytest.param(AnthropicConfig(), "claude-3-haiku-20240307", "anthropic", "server", id="model-without-tool-search"), + pytest.param(AnthropicConfig(), "claude-sonnet-5-5", None, "client", id="client-executed-search"), + pytest.param(AnthropicConfig(), "claude-sonnet-5-5", None, 123, id="malformed-execution-search"), +] + + +@pytest.mark.parametrize(("config", "model", "provider", "execution"), _UNRUNNABLE_TOOL_SEARCH_CASES) +def test_map_openai_params_answers_400_for_an_openai_tool_search_tool_it_cannot_run( + monkeypatch, config, model, provider, execution +): + monkeypatch.setattr(litellm, "drop_params", False) + assert provider is None or not litellm.supports_tool_search(model, provider) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + config.map_openai_params( + non_default_params={ + "tools": _openai_tool_search_with_deferred_functions(({"type": "tool_search", "execution": execution},)) + }, + optional_params={}, + model=model, + drop_params=False, + ) + assert raised.value.status_code == 400 + assert "tool_search" in str(raised.value) + assert "drop_params: true" in str(raised.value) + assert ("supports_tool_search: true" in str(raised.value)) is (execution == "server") + + +def test_map_openai_params_runs_tool_search_on_a_custom_named_deployment_whose_model_info_says_it_can(monkeypatch): + from litellm.responses.litellm_completion_transformation.tool_search import anthropic_tool_search_chat_tool + + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + tools: Final = _openai_tool_search_with_deferred_functions() + with pytest.raises(litellm.utils.UnsupportedParamsError): + AzureAnthropicConfig().map_openai_params( + non_default_params={"tools": tools}, optional_params={}, model="my-claude-deployment", drop_params=False + ) + litellm.register_model({"azure_ai/my-claude-deployment": {"supports_tool_search": True, "litellm_provider": "azure_ai"}}) + mapped: Final = AzureAnthropicConfig().map_openai_params( + non_default_params={"tools": tools}, optional_params={}, model="my-claude-deployment", drop_params=False + ) + assert _search_tools_sent(mapped["tools"]) == [anthropic_tool_search_chat_tool()] + assert _deferral_by_function_name(mapped["tools"]) == dict.fromkeys(_DEFERRED_FUNCTION_NAMES, True) + + +@pytest.mark.parametrize(("config", "model", "provider", "execution"), _UNRUNNABLE_TOOL_SEARCH_CASES) +def test_map_openai_params_drops_the_openai_tool_search_tool_and_every_deferral_under_drop_params( + monkeypatch, config, model, provider, execution +): + monkeypatch.setattr(litellm, "drop_params", False) + assert provider is None or not litellm.supports_tool_search(model, provider) + mapped: Final = config.map_openai_params( + non_default_params={ + "tools": _openai_tool_search_with_deferred_functions(({"type": "tool_search", "execution": execution},)) + }, + optional_params={}, + model=model, + drop_params=True, + ) + assert [tool["name"] for tool in mapped["tools"]] == _DEFERRED_FUNCTION_NAMES + assert all("defer_loading" not in tool for tool in mapped["tools"]) + + def _refusal_message_body(stop_reason: str, stop_details: dict | None) -> dict: return {**ANTHROPIC_DOCS_REFUSAL_MESSAGE, "stop_reason": stop_reason, "stop_details": stop_details} diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 419a25ff9e3..6c996a7ce75 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -10091,3 +10091,59 @@ async def test_bedrock_user_file_element_becomes_document_block(sync_mode: bool) assert content[0] == {"text": "Summarize this report"} assert content[1]["document"]["source"] == {"bytes": base64.b64encode(pdf_bytes).decode()} assert content[1]["document"]["format"] == "pdf" + + +def _openai_tool_search_with_deferred_functions() -> list[dict[str, object]]: + return [ + {"type": "tool_search", "execution": "server"}, + { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, + "defer_loading": True, + }, + { + "type": "function", + "function": { + "name": "get_stock_price", + "parameters": {"type": "object", "properties": {}}, + "defer_loading": True, + }, + }, + ] + + +@pytest.mark.parametrize( + ("model", "names_invoke_route"), + [ + pytest.param("us.anthropic.claude-sonnet-5-5", True, id="claude-that-runs-tool-search-on-invoke"), + pytest.param("anthropic.claude-3-haiku-20240307-v1:0", False, id="claude-without-tool-search"), + pytest.param("us.amazon.nova-pro-v1:0", False, id="nova"), + ], +) +def test_map_openai_params_answers_400_for_the_openai_tool_search_tool_without_drop_params( + monkeypatch, model, names_invoke_route +): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + AmazonConverseConfig().map_openai_params( + non_default_params={"tools": _openai_tool_search_with_deferred_functions()}, + optional_params={}, + model=model, + drop_params=False, + ) + assert raised.value.status_code == 400 + assert "tool_search" in str(raised.value) + assert "drop_params: true" in str(raised.value) + assert ("bedrock/invoke/" in str(raised.value)) is names_invoke_route + + +def test_map_openai_params_drops_the_openai_tool_search_tool_and_every_deferral_under_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = AmazonConverseConfig().map_openai_params( + non_default_params={"tools": _openai_tool_search_with_deferred_functions()}, + optional_params={}, + model="us.anthropic.claude-sonnet-5-5", + drop_params=True, + ) + assert [tool["function"]["name"] for tool in mapped["tools"]] == ["get_weather", "get_stock_price"] + assert all("defer_loading" not in tool and "defer_loading" not in tool["function"] for tool in mapped["tools"]) diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index d25787d7028..1293b3d8f2c 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -2371,3 +2371,39 @@ def test_transform_response_carries_served_service_tier(served_tier: str | None) ) assert result.to_dict().get("service_tier") == served_tier + + +_DATABRICKS_TOOLS_WITH_OPENAI_TOOL_SEARCH: Final = [ + {"type": "tool_search", "execution": "server"}, + { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, + "defer_loading": True, + }, +] + + +def test_databricks_claude_answers_400_for_the_openai_tool_search_tool_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + DatabricksConfig().map_openai_params( + non_default_params={"tools": _DATABRICKS_TOOLS_WITH_OPENAI_TOOL_SEARCH}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert "tool_search" in str(raised.value) + + +def test_databricks_claude_drops_the_openai_tool_search_tool_under_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = DatabricksConfig().map_openai_params( + non_default_params={"tools": _DATABRICKS_TOOLS_WITH_OPENAI_TOOL_SEARCH}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=True, + ) + assert mapped["tools"] == [ + {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}} + ] diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index de06518e3e1..ae5ac18b3d4 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -415,6 +415,233 @@ def test_completion_strips_eager_input_streaming_before_openai(respx_mock: respx assert sent_tool["function"]["name"] == "write_file" +_CLAUDE_TOOL_SEARCH_CALL: Final = { + "id": "srvtoolu_01search", + "type": "function", + "function": {"name": "tool_search_tool_regex", "arguments": '{"pattern": "weather"}'}, +} +_LOADED_FUNCTION_CALL: Final = { + "id": "toolu_01weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, +} +_USER_TURN: Final = {"role": "user", "content": "What is the weather in Paris?"} +_WEATHER_RESULT: Final = {"role": "tool", "tool_call_id": "toolu_01weather", "content": "18C"} +_SEARCH_RESULT: Final = {"role": "tool", "tool_call_id": "srvtoolu_01search", "content": "searched"} + + +@pytest.mark.parametrize( + ("claude_turn", "sent_after_user"), + [ + pytest.param( + [ + {"role": "assistant", "content": None, "tool_calls": [_CLAUDE_TOOL_SEARCH_CALL, _LOADED_FUNCTION_CALL]}, + _WEATHER_RESULT, + ], + [{"role": "assistant", "tool_calls": [_LOADED_FUNCTION_CALL]}, _WEATHER_RESULT], + id="search-beside-a-function-call", + ), + pytest.param( + [ + litellm.Message(content="Looking", tool_calls=[_CLAUDE_TOOL_SEARCH_CALL, _LOADED_FUNCTION_CALL]), + _WEATHER_RESULT, + ], + [{"role": "assistant", "content": "Looking", "tool_calls": [_LOADED_FUNCTION_CALL]}, _WEATHER_RESULT], + id="sdk-message-object", + ), + pytest.param( + [{"role": "assistant", "content": None, "tool_calls": [_CLAUDE_TOOL_SEARCH_CALL]}], + [], + id="unanswered-search-alone", + ), + pytest.param( + [ + {"role": "assistant", "content": None, "tool_calls": [_CLAUDE_TOOL_SEARCH_CALL, _LOADED_FUNCTION_CALL]}, + _SEARCH_RESULT, + _WEATHER_RESULT, + ], + [ + {"role": "assistant", "tool_calls": [_CLAUDE_TOOL_SEARCH_CALL, _LOADED_FUNCTION_CALL]}, + _SEARCH_RESULT, + _WEATHER_RESULT, + ], + id="answered-search-kept", + ), + ], +) +def test_completion_drops_unanswered_claude_tool_search_calls_from_history_sent_to_openai( + respx_mock: respx.MockRouter, openai_api_response, claude_turn, sent_after_user +): + api_base: Final = "http://localhost:12346/v1" + respx_mock.post(url__regex=rf"{api_base}/chat/completions.*").mock( + return_value=httpx.Response(status_code=200, json=openai_api_response) + ) + + litellm.completion( + model="openai/gpt-5.6", + messages=[_USER_TURN, *claude_turn, {"role": "user", "content": "And tomorrow?"}], + api_base=api_base, + api_key="fake_openai_api_key", + ) + + sent: Final = json.loads(respx_mock.calls[0].request.content)["messages"] + assert [ + {key: message[key] for key in ("role", "content", "tool_calls", "tool_call_id") if key in message} + for message in sent + ] == [_USER_TURN, *sent_after_user, {"role": "user", "content": "And tomorrow?"}] + + +def test_completion_drops_unanswered_claude_tool_search_calls_before_the_responses_bridge_sends_history_to_openai( + respx_mock: respx.MockRouter, +): + route: Final = respx_mock.post("https://api.openai.com/v1/responses").respond( + json={ + "id": "resp_turn2", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_turn2", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "18C and clear", "annotations": []}], + } + ], + "usage": {"input_tokens": 40, "output_tokens": 5, "total_tokens": 45}, + } + ) + + litellm.completion( + model="openai/gpt-5.6", + api_key="sk-test", + api_base="https://api.openai.com/v1", + messages=[ + _USER_TURN, + {"role": "assistant", "content": None, "tool_calls": [_CLAUDE_TOOL_SEARCH_CALL, _LOADED_FUNCTION_CALL]}, + _WEATHER_RESULT, + ], + tools=[{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + ) + + sent_input: Final = json.loads(route.calls.last.request.content)["input"] + assert [(item.get("type"), item.get("call_id")) for item in sent_input if "call_id" in item] == [ + ("function_call", "toolu_01weather"), + ("function_call_output", "toolu_01weather"), + ] + + +def _openai_called_names(body: Mapping[str, object]) -> tuple[str, ...]: + messages: Final = body["messages"] + assert isinstance(messages, list) + return tuple(call["function"]["name"] for message in messages for call in message.get("tool_calls") or ()) + + +def _converse_called_names(body: Mapping[str, object]) -> tuple[str, ...]: + messages: Final = body["messages"] + assert isinstance(messages, list) + blocks: Final = tuple(block for message in messages for block in message["content"]) + return tuple(block["toolUse"]["name"] for block in blocks if "toolUse" in block) + + +def _gemini_called_names(body: Mapping[str, object]) -> tuple[str, ...]: + contents: Final = body["contents"] + assert isinstance(contents, list) + parts: Final = tuple(part for content in contents for part in content["parts"]) + return tuple(part["function_call"]["name"] for part in parts if "function_call" in part) + + +_OPENAI_CHAT_REPLY: Final = { + "id": "chatcmpl-turn2", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "18C"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +@pytest.mark.parametrize( + ("model", "route_kwargs", "url_pattern", "reply", "called_names", "sent_names"), + [ + pytest.param( + "azure/gpt-4o", + {"api_base": "https://turn2.openai.azure.com", "api_key": "fake", "api_version": "2024-10-21"}, + r"https://turn2\.openai\.azure\.com/openai/deployments/gpt-4o/chat/completions.*", + _OPENAI_CHAT_REPLY, + _openai_called_names, + ("get_weather",), + id="azure-openai-strips", + ), + pytest.param( + "bedrock/us.anthropic.claude-opus-5-5", + {"aws_access_key_id": "fake", "aws_secret_access_key": "fake", "aws_region_name": "us-east-1"}, + r"https://bedrock-runtime\.us-east-1\.amazonaws\.com/model/.*/converse", + { + "output": {"message": {"role": "assistant", "content": [{"text": "18C"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + }, + _converse_called_names, + ("get_weather",), + id="bedrock-converse-strips", + ), + pytest.param( + "gemini/gemini-3.8-flash", + {"api_key": "fake"}, + r"https://generativelanguage\.googleapis\.com/.*:generateContent.*", + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "18C"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + _gemini_called_names, + ("tool_search_tool_regex", "get_weather"), + id="gemini-keeps", + ), + ], +) +def test_completion_strips_an_unanswered_claude_tool_search_call_only_for_routes_that_reject_it( + respx_mock: respx.MockRouter, model, route_kwargs, url_pattern, reply, called_names, sent_names +): + route: Final = respx_mock.post(url__regex=url_pattern).respond(json=reply) + + litellm.completion( + model=model, + messages=[ + _USER_TURN, + {"role": "assistant", "content": None, "tool_calls": [_CLAUDE_TOOL_SEARCH_CALL, _LOADED_FUNCTION_CALL]}, + _WEATHER_RESULT, + ], + tools=[{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + **route_kwargs, + ) + + assert called_names(json.loads(route.calls.last.request.content)) == sent_names + + +def test_completion_keeps_a_function_call_whose_id_only_looks_like_a_claude_server_tool( + respx_mock: respx.MockRouter, openai_api_response +): + api_base: Final = "http://localhost:12346/v1" + respx_mock.post(url__regex=rf"{api_base}/chat/completions.*").mock( + return_value=httpx.Response(status_code=200, json=openai_api_response) + ) + lookalike_call: Final = {**_LOADED_FUNCTION_CALL, "id": "srvtoolu_01weather"} + history: Final = [ + _USER_TURN, + {"role": "assistant", "content": None, "tool_calls": [lookalike_call]}, + {"role": "tool", "tool_call_id": "srvtoolu_01weather", "content": "18C"}, + ] + + litellm.completion(model="openai/gpt-5.6", messages=history, api_base=api_base, api_key="fake_openai_api_key") + + sent: Final = json.loads(respx_mock.calls[0].request.content)["messages"] + assert [message["role"] for message in sent] == ["user", "assistant", "tool"] + assert sent[1]["tool_calls"][0]["id"] == "srvtoolu_01weather" + + def test_embedding_keeps_an_internal_prefixed_kwarg_out_of_the_provider_request(respx_mock: respx.MockRouter) -> None: api_base: Final = "http://localhost:12346/v1" mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/embeddings.*").mock(