diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index 7ee53dcddab..fd0ba2c8e4f 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -1,4 +1,5 @@ from collections.abc import Iterable, Mapping +from types import MappingProxyType from typing import Final, Literal from openai.types.chat.chat_completion_allowed_tool_choice_param import ChatCompletionAllowedToolChoiceParam @@ -92,8 +93,21 @@ def _has_name(selector: Mapping[str, object]) -> bool: return isinstance(name, str) and name != "" -def _names_on_other_tool_types_are_strings(tools: Iterable[Mapping[str, object]]) -> bool: - return all(tool.get("type") in ("function", "custom") or "name" not in tool or _has_name(tool) for tool in tools) +def _names_a_tool_or_nothing(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + if not isinstance(tool_type, str) or tool_type not in ("function", "custom"): + return tool.get("name") is None or _has_name(tool) + match tool.get(tool_type): + case {"name": str()}: + return True + case {"name": object()}: + return False + case _: + return True + + +def _names_are_strings(tools: Iterable[Mapping[str, object]]) -> bool: + return all(map(_names_a_tool_or_nothing, tools)) def _chat_allowed_tool(tool: Mapping[str, object]) -> Mapping[str, object] | None: @@ -124,9 +138,9 @@ def parse_allowed_tools_tool_choice(tool_choice: object) -> ChatAllowedToolsTool return None match parsed: case ChatAllowedToolsToolChoice(): - return parsed if _names_on_other_tool_types_are_strings(parsed.allowed_tools.tools) else None + return parsed if _names_are_strings(parsed.allowed_tools.tools) else None case ResponsesAllowedToolsToolChoice(): - return _chat_layout(parsed) if _names_on_other_tool_types_are_strings(parsed.tools) else None + return _chat_layout(parsed) if _names_are_strings(parsed.tools) else None def azure_api_version_takes_tool_choice(api_version: str) -> bool: @@ -177,11 +191,31 @@ def refuse_a_malformed_allowed_tools_tool_choice(tool_choice: object, model: str ) -def tool_calls_the_choice_excludes(tool_choice: object, tool_names: Iterable[str]) -> tuple[str, ...]: +def _whole_server_label(tool: Mapping[str, object]) -> str | None: + label: Final = tool.get("server_label") + return label if tool.get("type") == "mcp" and tool.get("name") is None and isinstance(label, str) else None + + +def _whole_server_labels(tool_choice: object) -> frozenset[str]: + parsed: Final = parse_allowed_tools_tool_choice(tool_choice) + labels: Final = () if parsed is None else tuple(map(_whole_server_label, parsed.allowed_tools.tools)) + return frozenset(label for label in labels if label is not None) + + +def tool_calls_the_choice_excludes( + tool_choice: object, + tool_names: Iterable[str], + server_labels_by_tool: Mapping[str, frozenset[str]] = MappingProxyType({}), +) -> tuple[str, ...]: if not is_allowed_tools_tool_choice(tool_choice): return () allowed: Final = _allowed_tool_names(tool_choice) - return tuple(name for name in tool_names if name not in allowed) + whole_servers: Final = _whole_server_labels(tool_choice) + return tuple( + name + for name in tool_names + if name not in allowed and whole_servers.isdisjoint(server_labels_by_tool.get(name, frozenset())) + ) def refuse_tool_calls_the_choice_excludes(tool_choice: object, tool_names: Iterable[str], model: str) -> None: diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 628486f377f..ba6292eb1fd 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -15,6 +15,11 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, ) +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + is_allowed_tools_tool_choice, + refuse_tool_calls_the_choice_excludes, +) +from litellm.litellm_core_utils.get_litellm_params import AGENTIC_LOOP_TOOL_CHOICE_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.types.integrations.custom_logger import ( @@ -26,7 +31,11 @@ from litellm.types.integrations.custom_logger import ( is_interception_internal_key, stashed_stream_options, ) -from litellm.types.utils import ModelResponse +from litellm.types.utils import ( + ChatCompletionMessageCustomToolCall, + ChatCompletionMessageToolCall, + ModelResponse, +) from litellm.utils import CustomStreamWrapper _FOLLOWUP_INTERNAL_PARAMS: Final = frozenset( @@ -132,6 +141,32 @@ def _converted_stream_replay( ) +def _tool_call_name(call: ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall) -> str: + match call: + case ChatCompletionMessageCustomToolCall(): + return call.custom.name + case ChatCompletionMessageToolCall(): + return call.function.name or "" + + +def _called_tool_names(response: ModelResponse) -> tuple[str, ...]: + calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in response.choices) + return tuple(map(_tool_call_name, calls)) + + +def _refuse_a_call_the_choice_excludes(response: ModelResponse, kwargs: Mapping[str, object], model: str) -> None: + refuse_tool_calls_the_choice_excludes(kwargs.get(AGENTIC_LOOP_TOOL_CHOICE_KEY), _called_tool_names(response), model) + + +def _with_caller_tool_choice( + kwargs: Mapping[str, object], optional_params: Mapping[str, object] +) -> Mapping[str, object]: + tool_choice: Final = optional_params.get("tool_choice") + if AGENTIC_LOOP_TOOL_CHOICE_KEY in kwargs or not is_allowed_tools_tool_choice(tool_choice): + return kwargs + return {**kwargs, AGENTIC_LOOP_TOOL_CHOICE_KEY: tool_choice} + + def _with_agentic_loop_metadata(kwargs_for_followup: Mapping[str, object]) -> Mapping[str, object]: metadata: Final = kwargs_for_followup.get("litellm_metadata") return MappingProxyType( @@ -261,7 +296,8 @@ async def maybe_run_chat_completion_agentic_loop( import litellm callbacks: Final = litellm.callbacks + (getattr(logging_obj, "dynamic_success_callbacks", None) or []) - depth, max_loops, fingerprints = _agentic_loop_settings(kwargs) + request_kwargs: Final = _with_caller_tool_choice(kwargs, optional_params) + depth, max_loops, fingerprints = _agentic_loop_settings(request_kwargs) tools: Final = optional_params.get("tools", []) for callback in callbacks: @@ -272,7 +308,7 @@ async def maybe_run_chat_completion_agentic_loop( continue hook_kwargs = { - **kwargs, + **request_kwargs, "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, "custom_llm_provider": custom_llm_provider, } @@ -296,6 +332,7 @@ async def maybe_run_chat_completion_agentic_loop( if not should_run: continue + _refuse_a_call_the_choice_excludes(response, request_kwargs, model) fingerprint = _check_agentic_loop_safety( tool_calls=tool_calls, fingerprints=fingerprints, @@ -342,7 +379,7 @@ async def maybe_run_chat_completion_agentic_loop( callback=callback, model=model, optional_params=optional_params, - kwargs=kwargs, + kwargs=dict(request_kwargs), logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, depth=depth, @@ -358,7 +395,7 @@ async def maybe_run_chat_completion_agentic_loop( return _converted_stream_replay( response, - kwargs=kwargs, + kwargs=request_kwargs, depth=depth, model=model, custom_llm_provider=custom_llm_provider, diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 04755f9e1e9..c7821a898f1 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -32,6 +32,7 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( ) PROVIDER_AFFINITY_HEADER_KWARG_KEY: Final = "provider_affinity_header" +AGENTIC_LOOP_TOOL_CHOICE_KEY: Final = "_agentic_loop_tool_choice" OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS: Final = frozenset( { "token_exchange_endpoint", diff --git a/litellm/main.py b/litellm/main.py index cbdf2b0f4cb..97a472e884d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -71,6 +71,7 @@ from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.allowed_tools_tool_choice import ( drops_unsupported_allowed_tools_tool_choice, + is_allowed_tools_tool_choice, ) from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( @@ -83,6 +84,7 @@ from litellm.litellm_core_utils.chat_completion_agentic_loop import ( from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_litellm_params import ( + AGENTIC_LOOP_TOOL_CHOICE_KEY, AWS_CREDENTIAL_KWARGS_KEYS, OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS, OPTIONAL_KWARGS_KEYS, @@ -5353,6 +5355,9 @@ def completion( messages = validate_and_fix_openai_messages(messages=messages) tools = validate_and_fix_openai_tools(tools=tools) # validate tool_choice + requested_tool_choice: Final = cast( # cast-ok: widening an untyped dict to object is always sound + object, tool_choice + ) drops_unsupported: Final = drops_unsupported_params( cast("Mapping[str, object]", kwargs).get("drop_params") # cast-ok: kwargs is the untyped request dict ) @@ -5886,7 +5891,18 @@ def completion( if key in kwargs }, ) - litellm_params: Final = with_control_options(requested_litellm_params, control_options) + litellm_params: Final = with_control_options( + { + **requested_litellm_params, + **( + {AGENTIC_LOOP_TOOL_CHOICE_KEY: requested_tool_choice} + if AGENTIC_LOOP_TOOL_CHOICE_KEY not in kwargs + and is_allowed_tools_tool_choice(requested_tool_choice) + else {} + ), + }, + control_options, + ) if litellm_params.get("provider_affinity_header") is not None: try: headers = add_provider_affinity_header( diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 694525b0fca..152fb13ea9e 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -328,7 +328,10 @@ async def aresponses_api_with_mcp( ######################################################### if should_auto_execute and isinstance(response, ResponsesAPIResponse): tool_calls: Final = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( - LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response), tool_choice + LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response), + tool_choice, + tool_server_map, + mcp_tools_with_litellm_proxy, ) if tool_calls: diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index fc79515b0ba..eaf975efb85 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -435,6 +435,8 @@ async def acompletion_with_mcp( self.tool_calls = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response=complete_response), kwargs.get("tool_choice"), + self.tool_server_map, + mcp_tools_with_litellm_proxy, ) if self.tool_calls: @@ -606,6 +608,8 @@ async def acompletion_with_mcp( tool_calls: Final = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response=initial_response), kwargs.get("tool_choice"), + tool_server_map, + mcp_tools_with_litellm_proxy, ) if not tool_calls: diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 2fba8e2b0de..afe42500fd9 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -2,6 +2,7 @@ import re import traceback from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload from openai.types.chat import ChatCompletionToolParam @@ -89,6 +90,36 @@ def _proxy_path_mcp_name(tool: ToolParam) -> str | None: return None if match is None else match.group(1) +def _server_url_of(tool: object) -> str | None: + match tool: + case {"type": "mcp", "server_url": str() as server_url} if isinstance(tool, dict): + return server_url + case _: + return None + + +def _names_whole_gateway(tool: object) -> bool: + server_url: Final = _server_url_of(tool) or "" + return server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL) and not server_url.startswith( + LITELLM_PROXY_MCP_SERVER_URL_PREFIX + ) + + +def _server_label_serving(tool: object, server_name: str) -> str | None: + match tool: + case {"type": "mcp", "server_label": str() as label} if ( + _names_whole_gateway(tool) or _server_url_of(tool) == f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{server_name}" + ): + return label + case _: + return None + + +def server_labels_serving(mcp_tools: Iterable[object], server_name: str | None) -> frozenset[str]: + labels: Final = () if server_name is None else tuple(_server_label_serving(tool, server_name) for tool in mcp_tools) + return frozenset(label for label in labels if label is not None) + + def _registered_mcp_servers() -> Collection[MCPServer]: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -589,11 +620,20 @@ class LiteLLM_Proxy_MCP_Handler: _should_auto_execute_tools = should_auto_execute_tools @staticmethod - def tool_calls_to_run(tool_calls: Sequence[object], tool_choice: object) -> Sequence[object]: + def tool_calls_to_run( + tool_calls: Sequence[object], + tool_choice: object, + tool_server_map: Mapping[str, str], + mcp_tools: Iterable[object], + ) -> Sequence[object]: names: Final = tuple( LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call)[0] or "" for tool_call in tool_calls ) - return () if tool_calls_the_choice_excludes(tool_choice, names) else tool_calls + references: Final = tuple(mcp_tools) + labels_by_tool: Final = MappingProxyType( + {name: server_labels_serving(references, tool_server_map.get(name)) for name in names} + ) + return () if tool_calls_the_choice_excludes(tool_choice, names, labels_by_tool) else tool_calls @staticmethod def extract_tool_calls_from_response(response: ResponsesAPIResponse) -> list[object]: diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index bf4d1148536..b4d8d23fe4c 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -749,6 +749,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(self.collected_response), self.original_request_params.get("tool_choice"), + self.tool_server_map, + self.mcp_tools_with_litellm_proxy, ) if self.collected_response is not None else () diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 1bb7d0bb048..a1e3827612c 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -347,6 +347,7 @@ class AgenticLoopState: api_surface: Literal["chat_completions", "responses"] | None = field( default=None, metadata=wire("_agentic_loop_api_surface") ) + tool_choice: Mapping[str, object] | None = field(default=None, metadata=wire("_agentic_loop_tool_choice")) code_interpreter_active: bool | None = field(default=None, metadata=wire("_code_interpreter_interception_active")) code_interpreter_sandbox_key: str | None = field( default=None, metadata=wire("_code_interpreter_interception_sandbox_key") diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index 8b124818d6a..7c19e0b673c 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -26,8 +26,15 @@ _CHOICES_WITH_A_NAME_THAT_IS_NOT_A_STRING: Final = ( {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ""}]}, {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "mcp", "name": ["get_time"]}]}}, {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "mcp", "name": 7}]}}, + { + "type": "allowed_tools", + "allowed_tools": {"mode": "auto", "tools": [{"type": "function", "function": {"name": ["get_time"]}}]}, + }, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "function": {"name": 7}}]}, ) +_WHOLE_DOCS_SERVER: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "docs"}]} + _CHAT_LAYOUT: Final[dict[str, object]] = { "type": "allowed_tools", "allowed_tools": { @@ -273,3 +280,39 @@ def test_a_name_that_is_not_a_string_lists_no_tool(tool_choice: object): @pytest.mark.parametrize("tool_choice", [_FLAT_CHOICE, _CHAT_LAYOUT, "required", None]) def test_a_valid_choice_or_any_other_choice_is_not_called_malformed(tool_choice: object): assert refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") is None + + +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "docs", "name": None}]}, + { + "type": "allowed_tools", + "allowed_tools": {"mode": "auto", "tools": [{"type": "mcp", "server_label": "docs", "name": None}]}, + }, + ], +) +def test_an_mcp_entry_with_a_null_name_is_a_valid_choice(tool_choice: object): + assert parse_allowed_tools_tool_choice(tool_choice) is not None + assert refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") is None + + +def test_an_mcp_entry_without_a_name_allows_every_call_its_server_serves(): + assert tool_calls_the_choice_excludes(_WHOLE_DOCS_SERVER, ("search",), {"search": frozenset({"docs"})}) == () + + +@pytest.mark.parametrize( + ("tool_choice", "server_labels"), + [ + (_WHOLE_DOCS_SERVER, frozenset({"other"})), + (_WHOLE_DOCS_SERVER, frozenset()), + ( + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "docs", "name": "read"}]}, + frozenset({"docs"}), + ), + ], +) +def test_an_mcp_entry_allows_no_call_outside_its_server_or_its_name( + tool_choice: object, server_labels: frozenset[str] +): + assert tool_calls_the_choice_excludes(tool_choice, ("search",), {"search": server_labels}) == ("search",) diff --git a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py index 3cbc3591009..89e4b13fd0d 100644 --- a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -34,6 +34,7 @@ from litellm.integrations.code_interpreter_interception.handler import ( from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.chat_completion_agentic_loop import ( maybe_run_chat_completion_agentic_loop, + run_chat_completion_agentic_hooks, ) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -61,6 +62,7 @@ _INTERNAL_CONTROL_FIELDS = ( "_code_interpreter_interception_active", "_code_interpreter_interception_sandbox_key", "_code_interpreter_interception_converted_stream", + "_agentic_loop_tool_choice", "litellm_metadata", ) @@ -126,6 +128,27 @@ def _tool_call_model_response() -> ModelResponse: ) +def _custom_tool_call_model_response() -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_abc", + "type": "custom", + "custom": {"name": "litellm_code_execution", "input": "print(6*7)"}, + } + ], + ), + ) + ] + ) + + def _plain_model_response(content: str = "The answer is 42") -> ModelResponse: return ModelResponse( choices=[ diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index b0e0feab73e..528fc190faf 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1496,7 +1496,21 @@ _ALLOWED_ONLY_ANOTHER_TOOL: Final = { } +_DEMO_MCP_SERVER: Final = MCPServer( + server_id="demo", + name="demo", + server_name="demo", + transport="http", + url="https://demo.example/mcp", + auth_type="none", +) + + def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock) -> None: + monkeypatch.setattr( + "litellm.proxy.mcp_server.server._get_tools_from_mcp_servers", + AsyncMock(return_value=AggregateToolListing(tools=[_MCP_ECHO_TOOL], outcomes={})), + ) monkeypatch.setitem( sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None) ) @@ -1508,6 +1522,8 @@ def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool call_tool=call_tool, get_mcp_server_from_tool_name=MagicMock(return_value=None), get_mcp_server_by_name=MagicMock(return_value=None), + get_allowed_mcp_servers=AsyncMock(return_value=[_DEMO_MCP_SERVER.server_id]), + get_mcp_servers_from_ids=MagicMock(return_value=[_DEMO_MCP_SERVER]), ), ) @@ -1521,21 +1537,16 @@ async def test_an_mcp_tool_the_choice_excludes_never_runs_when_the_model_ignored call_tool: Final = AsyncMock() _served_by_a_recording_mcp_server(monkeypatch, call_tool) - with patch.object( - LiteLLM_Proxy_MCP_Handler, - "process_mcp_tools_without_openai_transform", - AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), - ): - response: Final = await litellm.acompletion( - model="openai/gpt-5.6", - messages=[{"role": "user", "content": "Call echo with PING"}], - tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], - tool_choice=_ALLOWED_ONLY_ANOTHER_TOOL, - api_key="sk-mcp-chat-rounds", - api_base="https://mcp-chat-rounds.example.com/v1", - stream=stream, - ) - chunks: Final = await _read_all(response) if stream else () + response: Final = await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + tool_choice=_ALLOWED_ONLY_ANOTHER_TOOL, + api_key="sk-mcp-chat-rounds", + api_base="https://mcp-chat-rounds.example.com/v1", + stream=stream, + ) + chunks: Final = await _read_all(response) if stream else () handed_back: Final = ( tuple(call.function.name for chunk in chunks for call in chunk.choices[0].delta.tool_calls or ()) @@ -1561,52 +1572,8 @@ def _anthropic_call_to_echo() -> Response: ) -async def _call_echo_on_anthropic_under_drop_params( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], tool_choice: object -) -> tuple[ModelResponse, respx.Route]: - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - litellm.in_memory_llm_clients_cache.flush_cache() - route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=list(rounds)) - with patch.object( - LiteLLM_Proxy_MCP_Handler, - "process_mcp_tools_without_openai_transform", - AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), - ): - response: Final = await litellm.acompletion( - model="anthropic/claude-haiku-4-5", - messages=[{"role": "user", "content": "Call echo with PING"}], - tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], - tool_choice=tool_choice, - drop_params=True, - api_key="sk-ant-fake", - ) - return response, route - - -@pytest.mark.asyncio -async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -): - call_tool: Final = AsyncMock() - _served_by_a_recording_mcp_server(monkeypatch, call_tool) - - response, route = await _call_echo_on_anthropic_under_drop_params( - respx_mock, monkeypatch, (_anthropic_call_to_echo(),), _ALLOWED_ONLY_ANOTHER_TOOL - ) - - call_tool.assert_not_awaited() - assert "tool_choice" not in json.loads(route.calls[0].request.content) - assert [call.function.name for call in response.choices[0].message.tool_calls] == ["echo"] - - -@pytest.mark.asyncio -async def test_a_dropped_choice_still_runs_an_mcp_call_it_lists( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -): - call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) - _served_by_a_recording_mcp_server(monkeypatch, call_tool) - answer: Final = Response( +def _anthropic_answer() -> Response: + return Response( 200, json={ "id": "msg_answer", @@ -1619,8 +1586,67 @@ async def test_a_dropped_choice_still_runs_an_mcp_call_it_lists( }, ) + +_WHOLE_DEMO_SERVER: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "demo"}]} +_WHOLE_OTHER_SERVER: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "other"}]} + + +async def _call_echo_on_anthropic_under_drop_params( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + rounds: Sequence[Response], + tool_choice: object, + server_url: str = "litellm_proxy/mcp/demo", +) -> tuple[ModelResponse, respx.Route]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=list(rounds)) + response: Final = await litellm.acompletion( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_label": "demo", "server_url": server_url, "require_approval": "never"}], + tool_choice=tool_choice, + drop_params=True, + api_key="sk-ant-fake", + ) + return response, route + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tool_choice", [_ALLOWED_ONLY_ANOTHER_TOOL, _WHOLE_OTHER_SERVER]) +async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tool_choice: object +): + call_tool: Final = AsyncMock() + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + + response, route = await _call_echo_on_anthropic_under_drop_params( + respx_mock, monkeypatch, (_anthropic_call_to_echo(),), tool_choice + ) + + call_tool.assert_not_awaited() + assert "tool_choice" not in json.loads(route.calls[0].request.content) + assert [call.function.name for call in response.choices[0].message.tool_calls] == ["echo"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tool_choice", "server_url"), + [ + (_FLAT_ALLOWED_ECHO, "litellm_proxy/mcp/demo"), + (_WHOLE_DEMO_SERVER, "litellm_proxy/mcp/demo"), + (_WHOLE_DEMO_SERVER, "litellm_proxy"), + ], +) +async def test_a_dropped_choice_still_runs_an_mcp_call_it_lists( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tool_choice: object, server_url: str +): + call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + response, _route = await _call_echo_on_anthropic_under_drop_params( - respx_mock, monkeypatch, (_anthropic_call_to_echo(), answer), _FLAT_ALLOWED_ECHO + respx_mock, monkeypatch, (_anthropic_call_to_echo(), _anthropic_answer()), tool_choice, server_url ) assert call_tool.await_count == 1 diff --git a/tests/unit/test_tool_loop.py b/tests/unit/test_tool_loop.py index a4cec7ec0bd..78d28f1b4a8 100644 --- a/tests/unit/test_tool_loop.py +++ b/tests/unit/test_tool_loop.py @@ -663,7 +663,6 @@ async def test_arun_tool_loop_two_rounds(respx_mock: respx.MockRouter, monkeypat assert messages == messages_snapshot -ANTHROPIC_MESSAGES_URL: Final = "https://api.anthropic.com/v1/messages" TOOL_SEARCH_CALL: Final = { "type": "server_tool_use", "id": "srvtoolu_01Search", diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index dca62b8a35d..9f130bedb07 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -225,6 +225,7 @@ AGENTIC_LOOP_STATE_NAMES: Final = ( "_agentic_loop_depth", "_agentic_loop_fingerprints", "_agentic_loop_api_surface", + "_agentic_loop_tool_choice", "_code_interpreter_interception_active", "_code_interpreter_interception_sandbox_key", "_code_interpreter_interception_session_scoped",