diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 5352ce6b6a0..6c04d16f817 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -15,6 +15,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.compression import compress from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.allowed_tools_tool_choice import allowed_tools_tool_choice_with_function from litellm.litellm_core_utils.asyncify import asyncify from litellm.types.integrations.compression_interception import ( CompressionInterceptionConfig, @@ -180,6 +181,10 @@ class CompressionInterceptionLogger(CustomLogger): existing_tools=cast(list[dict[str, object]] | None, kwargs.get("tools")), compressed_tools=compressed_tools, ) + if "tool_choice" in kwargs: + kwargs["tool_choice"] = allowed_tools_tool_choice_with_function( + kwargs["tool_choice"], LITELLM_CONTENT_RETRIEVE_TOOL_NAME + ) call_id = cast(str | None, kwargs.get("litellm_call_id")) if not call_id: call_id = str(uuid.uuid4()) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 632864d05c9..4b2773dbaa3 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -14,6 +14,7 @@ from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast +from pydantic import TypeAdapter from typing_extensions import Never, ReadOnly import litellm @@ -38,7 +39,10 @@ from litellm.integrations.websearch_interception.transformation import ( from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) -from litellm.litellm_core_utils.allowed_tools_tool_choice import allowed_tools_tool_choice_with_builtin_as_function +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + allowed_tools_tool_choice_with_builtin_as_function, + is_allowed_tools_tool_choice, +) from litellm.litellm_core_utils.core_helpers import qualify_provider_stripped_model from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.types.integrations.custom_logger import ( @@ -101,6 +105,8 @@ WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks" _RESPONSE_CONTENT_FIELD: Final = "content" +_OPTIONAL_PARAMS: Final = TypeAdapter(Mapping[str, object]) + _CLAUDE_CODE_WEB_SEARCH_PROMPT_PREFIX: Final = "Perform a web search for the query:" _ResponseT: Final = TypeVar("_ResponseT") @@ -276,6 +282,11 @@ def _tool_choice_naming_litellm_web_search( } +def _kept_allowed_tools_tool_choice(optional_params: Mapping[str, object]) -> Mapping[str, object]: + tool_choice: Final = optional_params.get("tool_choice") + return {"tool_choice": tool_choice} if is_allowed_tools_tool_choice(tool_choice) else MappingProxyType({}) + + class WebSearchInterceptionLogger(CustomLogger): """ CustomLogger that intercepts WebSearch tool calls for models that don't @@ -1863,9 +1874,11 @@ class WebSearchInterceptionLogger(CustomLogger): ) if request_patch.messages is None: raise ValueError("WebSearchInterception: missing follow-up messages") - params: Final = dict(optional_params) - params.update(request_patch.optional_params) - params.pop("tool_choice", None) + merged_params: Final = {**optional_params, **request_patch.optional_params} + params: Final = { + **{key: value for key, value in merged_params.items() if key != "tool_choice"}, + **_kept_allowed_tools_tool_choice(merged_params), + } patch_kwargs: Final = dict[str, object](request_patch.kwargs) return await litellm.acompletion( model=request_patch.model or model, @@ -1978,17 +1991,20 @@ class WebSearchInterceptionLogger(CustomLogger): tools_param: Final = optional_params.get("tools") optional_params_clean: Final = { - k: v - for k, v in optional_params.items() - if k - not in { - "tools", - "tool_choice", - "extra_body", - "model_alias_map", - "stream_response", - "custom_prompt_dict", - } + **{ + k: v + for k, v in optional_params.items() + if k + not in { + "tools", + "tool_choice", + "extra_body", + "model_alias_map", + "stream_response", + "custom_prompt_dict", + } + }, + **_kept_allowed_tools_tool_choice(_OPTIONAL_PARAMS.validate_python(optional_params)), } if tools_param is not None: optional_params_clean["tools"] = tools_param diff --git a/litellm/litellm_core_utils/agentic_loop_refusal.py b/litellm/litellm_core_utils/agentic_loop_refusal.py index f3ee15dffdd..828d5fa03ad 100644 --- a/litellm/litellm_core_utils/agentic_loop_refusal.py +++ b/litellm/litellm_core_utils/agentic_loop_refusal.py @@ -40,6 +40,16 @@ async def every_claimed_tool_call( return MappingProxyType({"tool_calls": (*_claimed_calls(refused_tool_calls), *chain.from_iterable(claims))}) +_TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call", "tool_use"}) + + +def called_tool_names(response: object) -> frozenset[str]: + items: Final = (*parts_of(message_field(response, "output")), *parts_of(message_field(response, "content"))) + return _strings( + message_field(item, "name") for item in items if message_field(item, "type") in _TOOL_CALL_ITEM_TYPES + ) + + def _call_ids(call: object) -> frozenset[str]: return _strings((message_field(call, "id"), message_field(call, "call_id"))) diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index 1d254aeb559..ee6a4fa5cb8 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 Callable, Iterable, Mapping +from dataclasses import dataclass from types import MappingProxyType from typing import Final, Literal @@ -179,16 +180,26 @@ def _allowed_tool_names(tool_choice: object) -> frozenset[object]: return frozenset(tool.get("name") for tool in flattened if tool is not None and _server_entry(tool) is None) -def refuse_a_malformed_allowed_tools_tool_choice(tool_choice: object, model: str) -> None: - if is_allowed_tools_tool_choice(tool_choice) and parse_allowed_tools_tool_choice(tool_choice) is None: - raise BadRequestError( - message=( - f"Invalid tool choice, tool_choice={tool_choice}. An allowed_tools tool_choice must be " - f"{CHAT_ALLOWED_TOOLS_LAYOUT} or {RESPONSES_ALLOWED_TOOLS_LAYOUT}" - ), - model=model, - llm_provider="", - ) +def parse_allowed_tools_tool_choice_or_refuse( + tool_choice: object, model: str, drop_params: bool +) -> ChatAllowedToolsToolChoice | None: + parsed: Final = parse_allowed_tools_tool_choice(tool_choice) + if parsed is not None or drop_params: + return parsed + raise BadRequestError( + message=( + f"Invalid tool choice, tool_choice={tool_choice}. An allowed_tools tool_choice must be " + f"{CHAT_ALLOWED_TOOLS_LAYOUT} or {RESPONSES_ALLOWED_TOOLS_LAYOUT}" + ), + model=model, + llm_provider="", + ) + + +def allowed_tools_tool_choice_in_effect(tool_choice: object, model: str, drop_params: bool) -> object: + if not is_allowed_tools_tool_choice(tool_choice): + return tool_choice + return None if parse_allowed_tools_tool_choice_or_refuse(tool_choice, model, drop_params) is None else tool_choice def _server_entry(tool: Mapping[str, object]) -> tuple[str, str | None] | None: @@ -205,14 +216,32 @@ def _server_entries(tool_choice: object) -> frozenset[tuple[str, str | None]]: return frozenset(entry for entry in entries if entry is not None) -def _a_server_entry_allows(entries: frozenset[tuple[str, str | None]], name: str, labels: frozenset[str]) -> bool: - return any(label in labels and entry_name in (None, name) for label, entry_name in entries) +@dataclass(frozen=True, slots=True) +class GatewayTool: + names: frozenset[str] + server_labels: frozenset[str] + + +@dataclass(frozen=True, slots=True) +class GatewayTools: + server_labels: frozenset[str] + tools: Mapping[str, GatewayTool] + + +NO_GATEWAY_TOOLS: Final = GatewayTools(server_labels=frozenset(), tools=MappingProxyType({})) + + +def _server_entry_allows(entry: tuple[str, str | None], tool: GatewayTool | None) -> bool: + label, name = entry + return tool is not None and label in tool.server_labels and (name is None or name in tool.names) + + +def _a_server_entry_allows(entries: Iterable[tuple[str, str | None]], tool: GatewayTool | None) -> bool: + return any(_server_entry_allows(entry, tool) for entry in entries) def tool_calls_the_choice_excludes( - tool_choice: object, - tool_names: Iterable[str], - server_labels_by_tool: Mapping[str, frozenset[str]] = MappingProxyType({}), + tool_choice: object, tool_names: Iterable[str], gateway: GatewayTools = NO_GATEWAY_TOOLS ) -> tuple[str, ...]: if not is_allowed_tools_tool_choice(tool_choice): return () @@ -221,8 +250,7 @@ def tool_calls_the_choice_excludes( return tuple( name for name in tool_names - if name not in allowed - and not _a_server_entry_allows(server_entries, name, server_labels_by_tool.get(name, frozenset())) + if name not in allowed and not _a_server_entry_allows(server_entries, gateway.tools.get(name)) ) @@ -246,29 +274,80 @@ def _with_builtin_as( return function_entry if tool.get("type") not in ("function", "custom", "mcp") and is_builtin(tool) else tool +def _function_entry( + layout: ChatAllowedToolsToolChoice | ResponsesAllowedToolsToolChoice, name: str +) -> Mapping[str, object]: + match layout: + case ChatAllowedToolsToolChoice(): + return {"type": "function", "function": {"name": name}} + case ResponsesAllowedToolsToolChoice(): + return {"type": "function", "name": name} + + +def _entries_of( + layout: ChatAllowedToolsToolChoice | ResponsesAllowedToolsToolChoice, +) -> tuple[Mapping[str, object], ...]: + match layout: + case ChatAllowedToolsToolChoice(): + return layout.allowed_tools.tools + case ResponsesAllowedToolsToolChoice(): + return layout.tools + + +def _with_entries( + layout: ChatAllowedToolsToolChoice | ResponsesAllowedToolsToolChoice, tools: Iterable[Mapping[str, object]] +) -> Mapping[str, object]: + entries: Final = [dict(tool) for tool in tools] + match layout: + case ChatAllowedToolsToolChoice(): + return {"type": "allowed_tools", "allowed_tools": {"mode": layout.allowed_tools.mode, "tools": entries}} + case ResponsesAllowedToolsToolChoice(): + return {"type": "allowed_tools", "mode": layout.mode, "tools": entries} + + +def _callers_layout(tool_choice: object) -> ChatAllowedToolsToolChoice | ResponsesAllowedToolsToolChoice | None: + if parse_allowed_tools_tool_choice(tool_choice) is None: + return None + return _ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice) + + def allowed_tools_tool_choice_with_builtin_as_function( tool_choice: object, is_builtin: Callable[[Mapping[str, object]], bool], function_name: str ) -> object: - if parse_allowed_tools_tool_choice(tool_choice) is None: + layout: Final = _callers_layout(tool_choice) + if layout is None: return tool_choice - parsed: Final = _ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice) - match parsed: - case ChatAllowedToolsToolChoice(): - nested: Final = {"type": "function", "function": {"name": function_name}} - return { - "type": "allowed_tools", - "allowed_tools": { - "mode": parsed.allowed_tools.mode, - "tools": [dict(_with_builtin_as(tool, is_builtin, nested)) for tool in parsed.allowed_tools.tools], - }, - } - case ResponsesAllowedToolsToolChoice(): - flat: Final = {"type": "function", "name": function_name} - return { - "type": "allowed_tools", - "mode": parsed.mode, - "tools": [dict(_with_builtin_as(tool, is_builtin, flat)) for tool in parsed.tools], - } + entry: Final = _function_entry(layout, function_name) + return _with_entries(layout, (_with_builtin_as(tool, is_builtin, entry) for tool in _entries_of(layout))) + + +def allowed_tools_tool_choice_with_function(tool_choice: object, function_name: str) -> object: + layout: Final = _callers_layout(tool_choice) + if layout is None or function_name in _allowed_tool_names(tool_choice): + return tool_choice + return _with_entries(layout, (*_entries_of(layout), _function_entry(layout, function_name))) + + +def _gateway_entry(tool: Mapping[str, object], gateway_labels: frozenset[str]) -> tuple[str, str | None] | None: + entry: Final = _server_entry(tool) + return entry if entry is not None and entry[0] in gateway_labels else None + + +def allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice: object, gateway: GatewayTools) -> object: + layout: Final = _callers_layout(tool_choice) + entries: Final = () if layout is None else _entries_of(layout) + gateway_entries: Final = tuple( + entry for entry in (_gateway_entry(tool, gateway.server_labels) for tool in entries) if entry is not None + ) + if layout is None or not gateway_entries: + return tool_choice + kept: Final = tuple(tool for tool in entries if _gateway_entry(tool, gateway.server_labels) is None) + allowed: Final = tuple( + name for name, tool in gateway.tools.items() if _a_server_entry_allows(gateway_entries, tool) + ) + if not kept and not allowed: + return "none" + return _with_entries(layout, (*kept, *(_function_entry(layout, name) for name in allowed))) def chat_allowed_tools_tool_choice(choice: ChatAllowedToolsToolChoice) -> ChatCompletionAllowedToolChoiceParam: diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 42a8b3013ee..53451b04515 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -1,7 +1,7 @@ # this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api import json -from collections.abc import Awaitable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Iterable, Iterator, Mapping, Sequence from itertools import chain from types import MappingProxyType from typing import Final, TypeVar @@ -19,7 +19,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) from litellm.litellm_core_utils.allowed_tools_tool_choice import ( - is_allowed_tools_tool_choice, + parse_allowed_tools_tool_choice, tool_calls_the_choice_excludes, ) from litellm.litellm_core_utils.core_helpers import qualify_agentic_followup_model @@ -203,15 +203,23 @@ def _called_tool_names(response: ModelResponse) -> tuple[str, ...]: return tuple(map(_tool_call_name, calls)) -def _holds_a_call_the_choice_excludes(response: ModelResponse, kwargs: Mapping[str, object]) -> bool: - return bool(tool_calls_the_choice_excludes(kwargs.get(AGENTIC_LOOP_TOOL_CHOICE_KEY), _called_tool_names(response))) +def _server_fulfilled_tool_names(callbacks: Iterable[CustomLogger]) -> frozenset[str]: + return frozenset(chain.from_iterable(callback.server_fulfilled_tool_names for callback in callbacks)) -def _with_caller_tool_choice( +def holds_a_call_the_choice_excludes( + tool_names: Iterable[str], kwargs: Mapping[str, object], callbacks: Iterable[CustomLogger] +) -> bool: + server_fulfilled: Final = _server_fulfilled_tool_names(callbacks) + client_tool_names: Final = (name for name in tool_names if name not in server_fulfilled) + return bool(tool_calls_the_choice_excludes(kwargs.get(AGENTIC_LOOP_TOOL_CHOICE_KEY), client_tool_names)) + + +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): + if AGENTIC_LOOP_TOOL_CHOICE_KEY in kwargs or parse_allowed_tools_tool_choice(tool_choice) is None: return kwargs return {**kwargs, AGENTIC_LOOP_TOOL_CHOICE_KEY: tool_choice} @@ -277,6 +285,7 @@ def _fingerprint_of_a_turn_to_run( *, response: ModelResponse, request_kwargs: Mapping[str, object], + callbacks: Sequence[CustomLogger], tool_calls: object, fingerprints: Sequence[str], depth: int, @@ -284,7 +293,7 @@ def _fingerprint_of_a_turn_to_run( model: str, stream: bool, ) -> str | None: - if _holds_a_call_the_choice_excludes(response, request_kwargs): + if holds_a_call_the_choice_excludes(_called_tool_names(response), request_kwargs, callbacks): return None try: return _check_agentic_loop_safety( @@ -434,7 +443,7 @@ async def maybe_run_chat_completion_agentic_loop( stream: bool, ) -> ModelResponse | CustomStreamWrapper | None: callbacks: Final = _custom_loggers(logging_obj) - request_kwargs: Final = _with_caller_tool_choice(kwargs, optional_params) + request_kwargs: Final = with_caller_tool_choice(kwargs, optional_params) depth, max_loops, fingerprints = _agentic_loop_settings(request_kwargs) tools: Final = _request_tools(optional_params) @@ -470,6 +479,7 @@ async def maybe_run_chat_completion_agentic_loop( fingerprint = _fingerprint_of_a_turn_to_run( response=response, request_kwargs=request_kwargs, + callbacks=callbacks, tool_calls=tool_calls, fingerprints=fingerprints, depth=depth, @@ -592,6 +602,7 @@ async def run_chat_completion_agentic_hooks( fingerprint = _fingerprint_of_a_turn_to_run( response=response, request_kwargs=kwargs, + callbacks=callbacks, tool_calls=tool_calls, fingerprints=fingerprints, depth=depth, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a10db31606..843be4fca62 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -48,6 +48,7 @@ from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAG from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs from litellm.litellm_core_utils.agentic_loop_refusal import ( + called_tool_names, every_claimed_tool_call, finalize_refused_responses_turn, is_refused_item, @@ -64,7 +65,9 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( ) from litellm.litellm_core_utils.chat_completion_agentic_loop import ( bill_first_round_on_failure, + holds_a_call_the_choice_excludes, run_chat_completion_agentic_hooks, + with_caller_tool_choice, ) from litellm.litellm_core_utils.core_helpers import as_str_mapping, set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value @@ -160,6 +163,7 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, AgenticLoopSafetyError, + AgenticLoopToolChoiceError, converted_stream_requested, is_interception_internal_key, ) @@ -6383,7 +6387,13 @@ class BaseLLMHTTPHandler: tools: Final = anthropic_messages_optional_request_params.get("tools", []) depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs) - hook_kwargs: Final = {**kwargs, "_agentic_loop_api_surface": api_surface} + hook_kwargs: Final = { + **with_caller_tool_choice( + _OPTIONAL_PARAMS.validate_python(kwargs), + _OPTIONAL_PARAMS.validate_python(anthropic_messages_optional_request_params), + ), + "_agentic_loop_api_surface": api_surface, + } for callback in callbacks: if not isinstance(callback, CustomLogger): @@ -6422,6 +6432,10 @@ class BaseLLMHTTPHandler: # Safety guards must run OUTSIDE the callback try/except — they are # bounded-loop / cycle-break rails, not callback bugs. try: + if holds_a_call_the_choice_excludes( + called_tool_names(response), hook_kwargs, (cb for cb in callbacks if isinstance(cb, CustomLogger)) + ): + raise AgenticLoopToolChoiceError(model) fingerprint = self._check_agentic_loop_safety( tool_calls=tool_calls, fingerprints=fingerprints, diff --git a/litellm/main.py b/litellm/main.py index 77220aba5ac..d61b9270b63 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -74,7 +74,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, + parse_allowed_tools_tool_choice, ) from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( @@ -5594,8 +5594,7 @@ def completion( model=model, drop_params=drops_unsupported_params( cast("Mapping[str, object]", kwargs).get("drop_params") # cast-ok: kwargs is the untyped request dict - ) - and not may_name_litellm_mcp_gateway(tools), # pyright: ignore[reportUnknownArgumentType] # untyped tools list + ), ) # validate optional params stop = validate_openai_optional_params(stop=stop) @@ -6212,7 +6211,7 @@ def completion( **( {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) + and parse_allowed_tools_tool_choice(requested_tool_choice) is not None else {} ), }, diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f70f9758cb5..f346ac84baa 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -29,6 +29,10 @@ from litellm.integrations.anthropic_cache_control_hook import ( configured_injection_points, routed_model_info, ) +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + allowed_tools_tool_choice_in_effect, + allowed_tools_tool_choice_with_gateway_tools_as_functions, +) from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params @@ -312,6 +316,11 @@ async def aresponses_api_with_mcp( LiteLLM_Proxy_MCP_Handler, ) + request_kwargs: Final = _STR_MAPPING.validate_python(kwargs) + tool_choice_in_effect: Final = allowed_tools_tool_choice_in_effect( + tool_choice, model, drops_unsupported_params(request_kwargs.get("drop_params")) + ) + # Parse MCP tools and separate from other tools ( mcp_tools_with_litellm_proxy, @@ -330,14 +339,13 @@ async def aresponses_api_with_mcp( else (None, None, None, None) ) - request_kwargs: Final = _STR_MAPPING.validate_python(kwargs) logging_obj: Final = request_kwargs.get("litellm_logging_obj") joins_request_trace: Final = isinstance(logging_obj, LiteLLMLoggingObj) bridged_first_turn: Final = ( joins_request_trace and not stream and previous_response_id is None - and not provider_stores_the_conversation(model, custom_llm_provider, kwargs) + and not provider_stores_the_conversation(model, custom_llm_provider, request_kwargs) ) listing_trace_id: Final = ( @@ -361,6 +369,9 @@ async def aresponses_api_with_mcp( raw_headers=discovery_raw_headers, ) openai_tools: Final = LiteLLM_Proxy_MCP_Handler.transform_mcp_tools_to_openai(original_mcp_tools) + gateway: Final = await LiteLLM_Proxy_MCP_Handler.gateway_tools( + tool_choice_in_effect, tool_server_map, mcp_tools_with_litellm_proxy, user_api_key_auth + ) # Combine with other tools all_tools: Final = openai_tools + other_tools if (openai_tools or other_tools) else None @@ -379,7 +390,7 @@ async def aresponses_api_with_mcp( "stream": stream, "temperature": temperature, "text": text, - "tool_choice": tool_choice, + "tool_choice": allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice_in_effect, gateway), "top_p": top_p, "truncation": truncation, "user": user, @@ -418,6 +429,8 @@ async def aresponses_api_with_mcp( previous_response_id=previous_response_id, tool_server_map=tool_server_map, served_tools=original_mcp_tools, + tool_choice=tool_choice_in_effect, + gateway=gateway, **kwargs, ) await mcp_streaming_response.start_initial_response() @@ -458,9 +471,8 @@ async def aresponses_api_with_mcp( with first_round_billed_on_failure(logging_obj, response): tool_calls: Final = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response), - tool_choice, - tool_server_map, - mcp_tools_with_litellm_proxy, + tool_choice_in_effect, + gateway, ) if tool_calls: diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index f622362dce4..f956459cd1b 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -11,6 +11,9 @@ from pydantic import TypeAdapter from typing_extensions import TypedDict, Unpack from litellm._internal_context import internal_sub_call +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + allowed_tools_tool_choice_with_gateway_tools_as_functions, +) from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.responses.mcp.chat_stream_round_billing import ChatStreamRoundBilling, engaged_round_billing @@ -163,6 +166,9 @@ async def acompletion_with_mcp( deduplicated_mcp_tools, target_format="chat", ) + gateway: Final = await LiteLLM_Proxy_MCP_Handler.gateway_tools( + kwargs.get("tool_choice"), tool_server_map, mcp_tools_with_litellm_proxy, user_api_key_auth + ) # Combine with other tools all_tools: Final = openai_tools + other_tools if (openai_tools or other_tools) else None @@ -182,6 +188,7 @@ async def acompletion_with_mcp( "tools": all_tools, "_skip_mcp_handler": True, # Prevent recursion **clean_kwargs, + "tool_choice": allowed_tools_tool_choice_with_gateway_tools_as_functions(kwargs.get("tool_choice"), gateway), } # If not auto-executing, just make the call with transformed tools @@ -494,8 +501,7 @@ 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, + gateway, ) if self.tool_calls: @@ -677,8 +683,7 @@ 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, + gateway, ) 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 d1188e2ecbc..9863b448fd4 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,5 +1,6 @@ import traceback from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload @@ -11,12 +12,20 @@ from typing_extensions import Unpack from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG -from litellm.litellm_core_utils.allowed_tools_tool_choice import tool_calls_the_choice_excludes +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + NO_GATEWAY_TOOLS, + GatewayTool, + GatewayTools, + parse_allowed_tools_tool_choice, + tool_calls_the_choice_excludes, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy.mcp_server.utils import ( iter_known_server_prefixes, + iter_known_tool_name_spellings, logging_safe_mcp_headers, + match_known_tool_name, split_server_prefix_from_name, strip_known_server_prefix, ) @@ -27,7 +36,6 @@ from litellm.responses.mcp.tool_precheck import ( names_gateway_explicitly, names_whole_gateway, proxy_path_mcp_name, - server_labels_serving, ) from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( @@ -120,6 +128,47 @@ async def _served_mcp_path_names( return await served_names(names) if names else frozenset[str]() +def _served_tool(tool_name: str, server_name: str) -> tuple[MCPServer | None, str]: + from litellm.proxy.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.mcp_server.server import ( + _resolve_display_name_to_original, + ) + + mcp_server: Final = global_mcp_server_manager.get_mcp_server_by_name( + server_name + ) or global_mcp_server_manager.get_mcp_server_from_tool_name(tool_name) + resolved_tool_name: Final = _resolve_display_name_to_original(tool_name, [mcp_server]) if mcp_server else tool_name + return mcp_server, strip_known_server_prefix(resolved_tool_name, mcp_server) + + +@dataclass(frozen=True, slots=True) +class _LabelScope: + label: str + whole_gateway: bool = False + server_ids: frozenset[str] = frozenset() + toolset_tools: Mapping[str, tuple[str, ...]] = MappingProxyType({}) + + def serves(self, server: MCPServer | None, tool_name: str) -> bool: + if self.whole_gateway: + return True + if server is None: + return False + toolset_names: Final = self.toolset_tools.get(server.server_id) + return server.server_id in self.server_ids or ( + toolset_names is not None and match_known_tool_name(tool_name, server, toolset_names) is not None + ) + + +def _gateway_tool(name: str, server: MCPServer | None, tool_name: str, scopes: Iterable[_LabelScope]) -> GatewayTool: + spellings: Final = () if server is None else tuple(iter_known_tool_name_spellings(tool_name, server)) + return GatewayTool( + names=frozenset((name, tool_name, *spellings)), + server_labels=frozenset(scope.label for scope in scopes if scope.serves(server, tool_name)), + ) + + class LiteLLM_Proxy_MCP_Handler: """ Helper class with static methods for MCP integration with Responses API. @@ -612,20 +661,80 @@ class LiteLLM_Proxy_MCP_Handler: _should_auto_execute_tools = should_auto_execute_tools @staticmethod - def tool_calls_to_run( - tool_calls: Sequence[object], + async def _label_scope( + tool: Mapping[str, object], + servers: Sequence[MCPServer], + user_api_key_auth: "UserAPIKeyAuth | None", + granted_toolsets: "GrantedToolsetIds | None", + ) -> _LabelScope | None: + from litellm.proxy.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.mcp_server.server import ( + _get_allowed_mcp_servers_from_mcp_server_names, + ) + + label: Final = tool.get("server_label") + server_url: Final = tool.get("server_url") + if not isinstance(label, str) or not isinstance(server_url, str): + return None + if names_whole_gateway(tool): + return _LabelScope(label=label, whole_gateway=True) + kind, value = await LiteLLM_Proxy_MCP_Handler._resolve_proxy_mcp_name( + server_url.split("/")[-1], user_api_key_auth, granted_toolsets + ) + match kind: + case "server": + resolved: Final = await _get_allowed_mcp_servers_from_mcp_server_names([value], list(servers)) + return _LabelScope(label=label, server_ids=frozenset(server.server_id for server in resolved)) + case "toolset": + permissions: Final = await global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=[value] + ) + return _LabelScope( + label=label, + toolset_tools=MappingProxyType({key: tuple(names) for key, names in permissions.items()}), + ) + case "denied_toolset": + return _LabelScope(label=label) + + @staticmethod + async def gateway_tools( tool_choice: object, tool_server_map: Mapping[str, str], - mcp_tools: Iterable[object], - ) -> Sequence[object]: + mcp_tools_with_litellm_proxy: Iterable[Mapping[str, object]], + user_api_key_auth: "UserAPIKeyAuth | None", + granted_toolsets: "GrantedToolsetIds | None" = None, + ) -> GatewayTools: + if parse_allowed_tools_tool_choice(tool_choice) is None: + return NO_GATEWAY_TOOLS + served: Final = {name: _served_tool(name, server_name) for name, server_name in tool_server_map.items()} + servers: Final = {server.server_id: server for server, _ in served.values() if server is not None} + scopes: Final = tuple( + [ + await LiteLLM_Proxy_MCP_Handler._label_scope( + tool, tuple(servers.values()), user_api_key_auth, granted_toolsets + ) + for tool in mcp_tools_with_litellm_proxy + ] + ) + label_scopes: Final = tuple(scope for scope in scopes if scope is not None) + return GatewayTools( + server_labels=frozenset(scope.label for scope in label_scopes), + tools=MappingProxyType( + { + name: _gateway_tool(name, server, tool_name, label_scopes) + for name, (server, tool_name) in served.items() + } + ), + ) + + @staticmethod + def tool_calls_to_run(tool_calls: Sequence[object], tool_choice: object, gateway: GatewayTools) -> Sequence[object]: names: Final = tuple( LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call)[0] or "" for tool_call in 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 + return () if tool_calls_the_choice_excludes(tool_choice, names, gateway) else tool_calls @staticmethod def extract_tool_calls_from_response(response: ResponsesAPIResponse) -> list[object]: @@ -768,9 +877,6 @@ class LiteLLM_Proxy_MCP_Handler: from litellm.proxy.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - from litellm.proxy.mcp_server.server import ( - _resolve_display_name_to_original, - ) from litellm.proxy.proxy_server import proxy_logging_obj tool_results: Final[list[MCPToolResult]] = [] @@ -799,13 +905,7 @@ class LiteLLM_Proxy_MCP_Handler: server_name = tool_server_map[tool_name] - mcp_server = global_mcp_server_manager.get_mcp_server_by_name( - server_name - ) or global_mcp_server_manager.get_mcp_server_from_tool_name(tool_name) - resolved_tool_name = ( - _resolve_display_name_to_original(tool_name, [mcp_server]) if mcp_server else tool_name - ) - sanitized_tool_name = strip_known_server_prefix(resolved_tool_name, mcp_server) + mcp_server, sanitized_tool_name = _served_tool(tool_name, server_name) start_time = datetime.now() logging_input: Sequence[Mapping[str, object]] = [ @@ -1235,6 +1335,8 @@ class LiteLLM_Proxy_MCP_Handler: previous_response_id: str | None, tool_server_map: dict[str, str], served_tools: Sequence[MCPTool] | None = None, + tool_choice: object = None, + gateway: GatewayTools = NO_GATEWAY_TOOLS, **kwargs, ) -> "MCPEnhancedStreamingIterator": """ @@ -1269,6 +1371,8 @@ class LiteLLM_Proxy_MCP_Handler: user_api_key_auth=kwargs.get("user_api_key_auth") or kwargs.get("litellm_metadata", {}).get("user_api_key_auth"), original_request_params=request_params, + tool_choice=tool_choice, + gateway=gateway, ) @staticmethod diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 2f366837b3b..8e8baf7667c 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final, cast from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.agentic_loop_settings import validated_max_agentic_loops +from litellm.litellm_core_utils.allowed_tools_tool_choice import NO_GATEWAY_TOOLS, GatewayTools from litellm.responses.litellm_completion_transformation.streaming_iterator import LiteLLMCompletionStreamingIterator from litellm.responses.mcp.request_context import MCPRequestContext from litellm.responses.mcp.stream_round_billing import StreamedRoundBilling @@ -286,6 +287,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): user_api_key_auth: "UserAPIKeyAuth | None" = None, original_request_params: dict[str, Any] | None = None, served_tools: Sequence[MCPTool] | None = None, + tool_choice: object = None, + gateway: GatewayTools = NO_GATEWAY_TOOLS, ): # MCP setup self.mcp_tools_with_litellm_proxy = mcp_tools_with_litellm_proxy or [] @@ -310,6 +313,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.mcp_discovery_generated = True # Events are already generated self.mcp_events = mcp_events # Store the initial MCP events for backward compatibility self.tool_server_map = tool_server_map + self.tool_choice: Final = tool_choice + self.gateway: Final = gateway self.served_tools = tuple(served_tools) if served_tools is not None else None # Iterator references @@ -841,9 +846,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): tool_calls: Final = ( 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, + self.tool_choice, + self.gateway, ) if self.collected_response is not None else () diff --git a/litellm/responses/mcp/tool_precheck.py b/litellm/responses/mcp/tool_precheck.py index cbd8a36d049..a67d81e70ae 100644 --- a/litellm/responses/mcp/tool_precheck.py +++ b/litellm/responses/mcp/tool_precheck.py @@ -29,21 +29,6 @@ def names_whole_gateway(tool: object) -> bool: ) -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 _mcp_server_url(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 proxy_path_mcp_name(tool: object) -> str | None: server_url: Final = _mcp_server_url(tool) match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) diff --git a/litellm/tool_loop.py b/litellm/tool_loop.py index 9fc9af3a1c3..8b19471055b 100644 --- a/litellm/tool_loop.py +++ b/litellm/tool_loop.py @@ -7,9 +7,10 @@ from typing_extensions import TypedDict, Unpack from litellm.constants import DEFAULT_TOOL_LOOP_MAX_ROUNDS from litellm.litellm_core_utils.allowed_tools_tool_choice import ( - refuse_a_malformed_allowed_tools_tool_choice, + allowed_tools_tool_choice_in_effect, refuse_tool_calls_the_choice_excludes, ) +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, @@ -36,12 +37,14 @@ class ToolLoopMaxRoundsExceeded(RuntimeError): super().__init__(f"model still requested tool calls on round {max_rounds} of {max_rounds}") -def _validate_tool_loop_args(max_rounds: int, model: str, completion_kwargs: _ToolLoopCompletionKwargs) -> None: +def _tool_choice_in_effect(max_rounds: int, model: str, completion_kwargs: _ToolLoopCompletionKwargs) -> object: if max_rounds < 1: raise ValueError(f"max_rounds must be >= 1, got {max_rounds}") if completion_kwargs.get("stream"): raise ValueError("run_tool_loop requires whole responses; stream=True is not supported") - refuse_a_malformed_allowed_tools_tool_choice(completion_kwargs.get("tool_choice"), model) + return allowed_tools_tool_choice_in_effect( + completion_kwargs.get("tool_choice"), model, drops_unsupported_params(completion_kwargs.get("drop_params")) + ) def _expect_model_response(response: object) -> ModelResponse: @@ -120,12 +123,10 @@ def _client_tool_calls( def _refuse_tool_calls_the_choice_excludes( - tool_calls: tuple[ChatCompletionMessageToolCall, ...], model: str, completion_kwargs: _ToolLoopCompletionKwargs + tool_calls: tuple[ChatCompletionMessageToolCall, ...], model: str, tool_choice: object ) -> None: refuse_tool_calls_the_choice_excludes( - completion_kwargs.get("tool_choice"), - tuple(tool_call.function.name or "" for tool_call in tool_calls), - model, + tool_choice, tuple(tool_call.function.name or "" for tool_call in tool_calls), model ) @@ -163,7 +164,7 @@ def run_tool_loop( """ import litellm - _validate_tool_loop_args(max_rounds, model, completion_kwargs) + tool_choice: Final = _tool_choice_in_effect(max_rounds, model, completion_kwargs) history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns for round_number in range(1, max_rounds + 1): message = ( @@ -181,7 +182,7 @@ def run_tool_loop( return message.content if round_number == max_rounds: break - _refuse_tool_calls_the_choice_excludes(client_tool_calls, model, completion_kwargs) + _refuse_tool_calls_the_choice_excludes(client_tool_calls, model, tool_choice) tool_results = tuple(execute_tool(tool_call) for tool_call in client_tool_calls) history = ( *history, @@ -215,7 +216,7 @@ async def arun_tool_loop( """ import litellm - _validate_tool_loop_args(max_rounds, model, completion_kwargs) + tool_choice: Final = _tool_choice_in_effect(max_rounds, model, completion_kwargs) history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns for round_number in range(1, max_rounds + 1): message = ( @@ -233,7 +234,7 @@ async def arun_tool_loop( return message.content if round_number == max_rounds: break - _refuse_tool_calls_the_choice_excludes(client_tool_calls, model, completion_kwargs) + _refuse_tool_calls_the_choice_excludes(client_tool_calls, model, tool_choice) tool_results = tuple([await execute_tool(tool_call) for tool_call in client_tool_calls]) history = ( *history, diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 14ed719a880..127b3d61753 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -81,6 +81,11 @@ class AgenticLoopLimitError(AgenticLoopSafetyError): pass +class AgenticLoopToolChoiceError(AgenticLoopSafetyError): + def __init__(self, model: str) -> None: + super().__init__(f"The allowed_tools tool_choice excludes a tool call the model made, model={model}") + + class StandardCustomLoggerInitParams(LiteLLMBaseModel): """ Params for initializing a CustomLogger. diff --git a/litellm/utils.py b/litellm/utils.py index 6a960a32a84..7b2464f3489 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -85,10 +85,8 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.allowed_tools_tool_choice import ( - CHAT_ALLOWED_TOOLS_LAYOUT, - RESPONSES_ALLOWED_TOOLS_LAYOUT, chat_allowed_tools_tool_choice, - parse_allowed_tools_tool_choice, + parse_allowed_tools_tool_choice_or_refuse, ) from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, @@ -9669,23 +9667,14 @@ def validate_chat_completion_tool_choice( if tool_choice_type in ("auto", "none", "required") and "function" not in tool_choice: return tool_choice_type if tool_choice_type == "allowed_tools": - allowed_tools: Final = parse_allowed_tools_tool_choice( + allowed_tools: Final = parse_allowed_tools_tool_choice_or_refuse( cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic object, tool_choice - ) + ), + model, + drop_params, ) - if allowed_tools is None and drop_params: - return None - if allowed_tools is None: - raise BadRequestError( - message=( - f"Invalid tool choice, tool_choice={tool_choice}. An allowed_tools tool_choice must be " - f"{CHAT_ALLOWED_TOOLS_LAYOUT} or {RESPONSES_ALLOWED_TOOLS_LAYOUT}" - ), - model=model, - llm_provider="", - ) - return dict(chat_allowed_tools_tool_choice(allowed_tools)) + return None if allowed_tools is None else dict(chat_allowed_tools_tool_choice(allowed_tools)) # Standard OpenAI format: {"type": "function", "function": {...}} if tool_choice_type is None or tool_choice.get("function") is None: diff --git a/tests/unit/integrations/compression_interception/test_compression_interception_handler.py b/tests/unit/integrations/compression_interception/test_compression_interception_handler.py index d7a1d6f14e1..549c0363942 100644 --- a/tests/unit/integrations/compression_interception/test_compression_interception_handler.py +++ b/tests/unit/integrations/compression_interception/test_compression_interception_handler.py @@ -548,3 +548,50 @@ async def test_pre_call_hook_counts_tokens_off_the_event_loop(): assert result["messages"] is messages assert "tools" not in result assert_loop_stayed_free(took, lags) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tool_choice", "retrieval_entry"), + [ + ( + {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "existing_tool"}]}, + {"type": "function", "name": "litellm_content_retrieve"}, + ), + ( + { + "type": "allowed_tools", + "allowed_tools": {"mode": "auto", "tools": [{"type": "function", "function": {"name": "existing_tool"}}]}, + }, + {"type": "function", "function": {"name": "litellm_content_retrieve"}}, + ), + ], +) +async def test_pre_call_hook_lets_an_allowed_tools_choice_allow_the_retrieval_tool_it_adds( + monkeypatch, tool_choice, retrieval_entry +): + monkeypatch.setattr( + "litellm.integrations.compression_interception.handler.compress", + lambda **_: { + "messages": [{"role": "user", "content": "stubbed"}], + "original_tokens": 12000, + "compressed_tokens": 5000, + "cache": {"auth.py": "full file content"}, + "tools": [{"type": "function", "function": {"name": "litellm_content_retrieve", "parameters": {}}}], + }, + ) + + result = await CompressionInterceptionLogger().async_pre_call_deployment_hook( + kwargs={ + "model": "anthropic/claude-sonnet-4-5", + "messages": [{"role": "user", "content": "very large context"}], + "tools": [{"type": "function", "function": {"name": "existing_tool", "parameters": {"type": "object"}}}], + "tool_choice": tool_choice, + }, + call_type=CallTypes.anthropic_messages, + ) + + assert result is not None + entries = result["tool_choice"].get("tools") or result["tool_choice"]["allowed_tools"]["tools"] + assert entries[-1] == retrieval_entry + assert entries[:-1] == (tool_choice.get("tools") or tool_choice["allowed_tools"]["tools"]) diff --git a/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py b/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py index fc415374096..89ad8b72070 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py @@ -422,6 +422,85 @@ async def test_execute_chat_completion_agentic_loop_strips_tool_choice(): ) +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["auto", "required"]) +async def test_execute_chat_completion_agentic_loop_keeps_an_allowed_tools_choice(mode): + websearch_logger = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI]) + tool_choice = { + "type": "allowed_tools", + "allowed_tools": { + "mode": mode, + "tools": [{"type": "function", "function": {"name": "litellm_web_search"}}], + }, + } + captured_kwargs: dict = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return ModelResponse(id="followup", model="gpt-4o", object="chat.completion") + + async def fake_search(query, **_): + return ("Bitcoin price is $60,000", None) + + with patch.object(websearch_logger, "_execute_search", side_effect=fake_search): + with patch("litellm.acompletion", side_effect=fake_acompletion): + await websearch_logger._execute_chat_completion_agentic_loop( + model="gpt-4o", + messages=[{"role": "user", "content": "What is Bitcoin price?"}], + tool_calls=[{"id": "call_1", "name": "litellm_web_search", "input": {"query": "bitcoin price"}}], + optional_params={ + "tools": [ + {"type": "function", "function": {"name": "litellm_web_search"}}, + {"type": "function", "function": {"name": "get_time"}}, + ], + "tool_choice": tool_choice, + }, + logging_obj=MagicMock(), + stream=False, + kwargs={}, + ) + + assert captured_kwargs["tool_choice"] == tool_choice + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tool_choice", "kept"), + [ + ( + { + "type": "allowed_tools", + "allowed_tools": {"mode": "required", "tools": [{"type": "function", "function": {"name": "get_time"}}]}, + }, + True, + ), + ({"type": "function", "function": {"name": "litellm_web_search"}}, False), + ("required", False), + ], +) +async def test_chat_followup_plan_keeps_only_an_allowed_tools_choice( + tool_choice: object, kept: bool, perplexity_search: respx.Route +) -> None: + plan: Final = await WebSearchInterceptionLogger().async_build_chat_completion_agentic_loop_plan( + tools={ + "tool_calls": [{"id": "call_abc", "name": "litellm_web_search", "input": {"query": "latest ai news"}}], + "tool_type": "websearch", + "provider": "openai", + "response_format": "openai", + }, + model="gpt-4o", + messages=[{"role": "user", "content": "What's the latest AI news?"}], + response=MagicMock(), + optional_params={"tool_choice": tool_choice, "max_tokens": 512}, + logging_obj=None, + stream=False, + kwargs={"custom_llm_provider": "openai"}, + ) + + assert plan.request_patch is not None + assert plan.request_patch.optional_params.get("tool_choice") == (tool_choice if kept else None) + + if __name__ == "__main__": # Run with: pytest test_websearch_chat_completion.py -v -s pytest.main([__file__, "-v", "-s"]) 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 5a87f0e6434..ec398e8731a 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 @@ -8,12 +8,16 @@ import litellm from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ChatAllowedToolsToolChoice, + GatewayTool, + GatewayTools, + allowed_tools_tool_choice_in_effect, allowed_tools_tool_choice_with_builtin_as_function, + allowed_tools_tool_choice_with_function, + allowed_tools_tool_choice_with_gateway_tools_as_functions, drops_unsupported_allowed_tools_tool_choice, is_allowed_tools_tool_choice, chat_allowed_tools_tool_choice, parse_allowed_tools_tool_choice, - refuse_a_malformed_allowed_tools_tool_choice, refuse_tool_calls_the_choice_excludes, tool_calls_the_choice_excludes, unenforced_allowed_tools_tool_choice, @@ -266,11 +270,12 @@ def test_a_choice_the_provider_cannot_take_answers_400_without_drop_params(): *_CHOICES_WITH_A_NAME_THAT_IS_NOT_A_STRING, ], ) -def test_a_malformed_choice_is_refused(tool_choice: object): +def test_a_malformed_choice_is_refused_without_drop_params_and_dropped_with_it(tool_choice: object): with pytest.raises(BadRequestError, match="Invalid tool choice") as exc: - refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") + allowed_tools_tool_choice_in_effect(tool_choice, "the-model", False) assert exc.value.status_code == 400 + assert allowed_tools_tool_choice_in_effect(tool_choice, "the-model", True) is None @pytest.mark.parametrize("tool_choice", _CHOICES_WITH_A_NAME_THAT_IS_NOT_A_STRING) @@ -279,8 +284,9 @@ 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("drop_params", [True, False]) +def test_a_valid_choice_or_any_other_choice_stays_in_effect(tool_choice: object, drop_params: bool): + assert allowed_tools_tool_choice_in_effect(tool_choice, "the-model", drop_params) is tool_choice @pytest.mark.parametrize( @@ -295,7 +301,7 @@ def test_a_valid_choice_or_any_other_choice_is_not_called_malformed(tool_choice: ) 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 + assert allowed_tools_tool_choice_in_effect(tool_choice, "the-model", False) is tool_choice _DOCS_SEARCH: Final = { @@ -305,12 +311,36 @@ _DOCS_SEARCH: Final = { } -def test_a_named_mcp_entry_allows_its_call_on_a_server_its_label_serves(): - assert tool_calls_the_choice_excludes(_DOCS_SEARCH, ("search",), {"search": frozenset({"docs", "other"})}) == () +def _gateway_serving_search(*server_labels: str) -> GatewayTools: + return GatewayTools( + server_labels=frozenset({"docs", "other"}), + tools={ + "docs_server-search": GatewayTool( + names=frozenset({"docs_server-search", "search"}), server_labels=frozenset(server_labels) + ) + }, + ) + + +@pytest.mark.parametrize("name", ["search", "docs_server-search"]) +def test_a_named_mcp_entry_allows_its_call_on_a_server_its_label_serves(name: str): + tool_choice: Final = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "mcp", "server_label": "docs", "name": name}], + } + + assert ( + tool_calls_the_choice_excludes(tool_choice, ("docs_server-search",), _gateway_serving_search("docs", "other")) + == () + ) 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"})}) == () + assert ( + tool_calls_the_choice_excludes(_WHOLE_DOCS_SERVER, ("docs_server-search",), _gateway_serving_search("docs")) + == () + ) @pytest.mark.parametrize( @@ -329,7 +359,9 @@ def test_an_mcp_entry_without_a_name_allows_every_call_its_server_serves(): 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",) + assert tool_calls_the_choice_excludes( + tool_choice, ("docs_server-search",), _gateway_serving_search(*server_labels) + ) == ("docs_server-search",) def _is_web_search(tool: Mapping[str, object]) -> bool: @@ -388,3 +420,132 @@ def test_a_function_or_mcp_entry_is_never_taken_for_a_builtin(): ) def test_any_other_or_malformed_choice_is_left_as_sent(tool_choice: object): assert allowed_tools_tool_choice_with_builtin_as_function(tool_choice, _is_web_search, "litellm_web_search") is tool_choice + + +_GATEWAY_WITH_TWO_SERVERS: Final = GatewayTools( + server_labels=frozenset({"docs", "weather"}), + tools={ + "docs_server-search": GatewayTool( + names=frozenset({"docs_server-search", "search"}), server_labels=frozenset({"docs"}) + ), + "docs_server-read": GatewayTool(names=frozenset({"docs_server-read", "read"}), server_labels=frozenset({"docs"})), + "weather_server-forecast": GatewayTool( + names=frozenset({"weather_server-forecast", "forecast"}), server_labels=frozenset({"weather"}) + ), + }, +) + + +@pytest.mark.parametrize( + ("tool_choice", "expected"), + [ + ( + { + "type": "allowed_tools", + "mode": "required", + "tools": [{"type": "mcp", "server_label": "docs"}, {"type": "function", "name": "get_time"}], + }, + { + "type": "allowed_tools", + "mode": "required", + "tools": [ + {"type": "function", "name": "get_time"}, + {"type": "function", "name": "docs_server-search"}, + {"type": "function", "name": "docs_server-read"}, + ], + }, + ), + ( + { + "type": "allowed_tools", + "allowed_tools": { + "mode": "auto", + "tools": [{"type": "mcp", "server_label": "weather", "name": "forecast"}], + }, + }, + { + "type": "allowed_tools", + "allowed_tools": { + "mode": "auto", + "tools": [{"type": "function", "function": {"name": "weather_server-forecast"}}], + }, + }, + ), + ], +) +def test_gateway_mcp_entries_become_the_functions_they_allow_in_the_callers_layout( + tool_choice: dict[str, object], expected: dict[str, object] +): + assert allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice, _GATEWAY_WITH_TWO_SERVERS) == expected + + +def test_an_mcp_entry_of_a_label_the_gateway_does_not_serve_is_forwarded_as_sent(): + tool_choice: Final = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "mcp", "server_label": "remote"}, {"type": "mcp", "server_label": "docs", "name": "read"}], + } + + assert allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice, _GATEWAY_WITH_TWO_SERVERS) == { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "mcp", "server_label": "remote"}, {"type": "function", "name": "docs_server-read"}], + } + + +@pytest.mark.parametrize("mode", ["auto", "required"]) +def test_gateway_entries_that_allow_no_converted_tool_send_none_instead_of_an_empty_list(mode: str): + tool_choice: Final = { + "type": "allowed_tools", + "mode": mode, + "tools": [{"type": "mcp", "server_label": "docs", "name": "delete"}], + } + + assert allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice, _GATEWAY_WITH_TWO_SERVERS) == "none" + + +@pytest.mark.parametrize("tool_choice", ["auto", None, _FLAT_CHOICE, {"type": "allowed_tools", "mode": "auto", "tools": []}]) +def test_a_choice_without_gateway_entries_is_forwarded_as_sent(tool_choice: object): + assert allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice, _GATEWAY_WITH_TWO_SERVERS) is tool_choice + + +@pytest.mark.parametrize( + ("tool_choice", "expected"), + [ + ( + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, + { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "get_time"}, {"type": "function", "name": "retrieve"}], + }, + ), + ( + {"type": "allowed_tools", "allowed_tools": {"mode": "required", "tools": [{"type": "mcp", "server_label": "x"}]}}, + { + "type": "allowed_tools", + "allowed_tools": { + "mode": "required", + "tools": [{"type": "mcp", "server_label": "x"}, {"type": "function", "function": {"name": "retrieve"}}], + }, + }, + ), + ], +) +def test_an_added_function_joins_the_choice_in_the_callers_layout( + tool_choice: dict[str, object], expected: dict[str, object] +): + assert allowed_tools_tool_choice_with_function(tool_choice, "retrieve") == expected + + +@pytest.mark.parametrize( + "tool_choice", + [ + "auto", + None, + {"type": "function", "function": {"name": "get_time"}}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "retrieve"}]}, + ], +) +def test_a_choice_that_is_not_allowed_tools_or_already_lists_the_function_is_left_as_sent(tool_choice: object): + assert allowed_tools_tool_choice_with_function(tool_choice, "retrieve") is tool_choice 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 20bd02811fc..7d41cbd6628 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 @@ -23,7 +23,7 @@ they stay out of the body even without it. import datetime import json from collections.abc import Callable, Mapping, Sequence -from typing import Any, Dict, Final, List, Optional, Tuple +from typing import Any, ClassVar, Dict, Final, List, Optional, Tuple from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1284,3 +1284,60 @@ async def test_acompletion_hands_the_callers_choice_to_the_chat_hooks(tool_choic ) assert (hook.runs, mock_client.post.await_count, response.choices[0].message.tool_calls) == (hook_runs, 1, None) + + + +class _RetrievalHookLogger(_ChatHookLogger): + server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({"litellm_code_execution"}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_response", [_tool_call_model_response, _custom_tool_call_model_response]) +async def test_a_chat_hook_runs_a_tool_litellm_serves_itself_that_no_choice_can_list( + model_response: Callable[[], ModelResponse], +): + hook: Final = _RetrievalHookLogger() + + result: Final = await run_chat_completion_agentic_hooks( + response=model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + optional_params={}, + kwargs={"_agentic_loop_tool_choice": _ALLOWS_ONLY_GET_WEATHER}, + logging_obj=_logging_with(hook), + custom_llm_provider="openai", + stream=False, + ) + + assert hook.runs == 1 + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "ran the hook" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "allowed_tools", "mode": "auto", "tools": []}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["get_weather"]}]}, + ], +) +async def test_a_malformed_choice_dropped_by_drop_params_never_holds_back_the_chat_hooks( + tool_choice: dict[str, object], +): + hook: Final = _ChatHookLogger() + mock_client: Final = MagicMock() + mock_client.post = AsyncMock(return_value=_raw_response_for(_tool_call_model_response())) + + response: Final = await litellm.acompletion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + tools=[{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + tool_choice=tool_choice, + drop_params=True, + api_key="sk-test", + client=mock_client, + success_callback=[hook], + ) + + assert (hook.runs, response.choices[0].message.content) == (1, "ran the hook") diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index dd0e5762845..71abfd84496 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -455,6 +455,95 @@ def test_response_api_handler_sends_no_stream_options_when_web_search_forces_a_n assert "stream_options" not in recorder.seen_kwargs +class _LookupLoop(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.runs = 0 + + async def async_should_run_agentic_loop( + self, + response: object, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict[str, object], + ) -> tuple[bool, dict[str, object]]: + return True, {"tool_calls": [{"id": "call_lookup", "name": "lookup"}]} + + async def async_run_agentic_loop( + self, + tools: dict, + model: str, + messages: list[dict], + response: object, + anthropic_messages_provider_config: object, + anthropic_messages_optional_request_params: dict, + logging_obj: object, + stream: bool, + kwargs: dict, + ) -> object: + self.runs += 1 + return response + + +def _a_lookup_call(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "resp_1", + "object": "response", + "created_at": 0, + "status": "completed", + "model": "gpt-4o", + "output": [ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_lookup", + "name": "lookup", + "arguments": "{}", + "status": "completed", + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + ) + + +@pytest.mark.parametrize("allowed_name,runs", [("get_time", 0), ("lookup", 1)]) +def test_response_api_handler_runs_no_hook_call_the_allowed_tools_choice_excludes(allowed_name: str, runs: int): + hook: Final = _LookupLoop() + + response: Final = BaseLLMHTTPHandler().response_api_handler( + model="gpt-4o", + input="what time is it?", + responses_api_provider_config=AzureOpenAIResponsesAPIConfig(), + response_api_optional_request_params={ + "tools": [ + {"type": "function", "name": "get_time", "parameters": {"type": "object", "properties": {}}}, + {"type": "function", "name": "lookup", "parameters": {"type": "object", "properties": {}}}, + ], + "tool_choice": { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": allowed_name}], + }, + }, + custom_llm_provider="azure", + litellm_params=GenericLiteLLMParams( + api_key="sk-test", api_base="https://example.openai.azure.com", api_version="2025-04-01-preview" + ), + logging_obj=Mock(litellm_call_id="call-1", dynamic_success_callbacks=[hook], model_call_details={}), + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_a_lookup_call))), + ) + + assert hook.runs == runs + assert isinstance(response, ResponsesAPIResponse) + assert len(response.output) == runs + + @pytest.mark.asyncio async def test_async_response_api_handler_streams_when_provider_transform_adds_stream(): handler = BaseLLMHTTPHandler() diff --git a/tests/unit/responses/litellm_completion_transformation/test_handler.py b/tests/unit/responses/litellm_completion_transformation/test_handler.py index 4794f3701c5..1611225dfe9 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_handler.py @@ -332,6 +332,40 @@ async def test_bridged_next_turn_replays_a_turn_stored_under_the_id_the_client_h ] +@pytest.mark.asyncio +async def test_bridged_follow_up_after_a_client_tool_search_reports_auto_for_an_allowed_tools_choice_it_dropped(): + previous_response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="azure_ai", model_id="deployment-a", response_id="chatcmpl-turn-one" + ) + provider: Final = _RecordingAnthropicHandler(_ANTHROPIC_MESSAGE_PAYLOAD) + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider)) + + with patch( + "litellm.proxy.proxy_server.prisma_client", + SimpleNamespace(db=_StoredTurns((_stored_mcp_turn(previous_response_id),))), + ): + response = await litellm.aresponses( + model="azure_ai/claude-sonnet-4-6", + api_base="https://fake-foundry-resource.services.ai.azure.com", + api_key="fake-api-key", + input="Double it", + tools=[ + {"type": "tool_search", "execution": "client"}, + {"type": "function", "name": "get_time", "parameters": {"type": "object", "properties": {}}}, + ], + tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, + drop_params=True, + previous_response_id=previous_response_id, + client=client, + ) + + assert provider.request_body is not None, "the bridged turn never reached the provider" + assert "tool_choice" not in provider.request_body + assert isinstance(response, ResponsesAPIResponse) + assert response.tool_choice == "auto" + + @pytest.mark.asyncio @pytest.mark.parametrize("client_field", ["caller_input", "names_non_string_fields"]) async def test_bridged_follow_up_after_a_client_tool_search_with_a_client_field_named_like_a_bridge_keyword_reaches_the_provider( diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 01b223f5e26..c48433989ea 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -36,6 +36,7 @@ from litellm.types.llms.openai import ( ResponsesAPIStreamingResponse, ToolParam, ) +from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.utils import ModelResponse, StandardLoggingPayload, Usage @@ -1704,6 +1705,15 @@ async def _auto_execute_echo( return await _auto_execute(logging_obj, _ECHO_MCP_TOOL, "PING") +_DEMO_MCP_SERVER: Final = MCPServer( + server_id="demo", name="demo", server_name="demo", transport="http", url="https://demo.example/mcp", auth_type="none" +) + + +def _demo_server_named(name: str) -> MCPServer | None: + return _DEMO_MCP_SERVER if name == "demo" else None + + def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock) -> None: monkeypatch.setitem( sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None) @@ -1715,7 +1725,7 @@ def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool get_registry=MagicMock(return_value={}), call_tool=call_tool, get_mcp_server_from_tool_name=MagicMock(return_value=None), - get_mcp_server_by_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(side_effect=_demo_server_named), ), ) @@ -1806,26 +1816,13 @@ def _anthropic_call_to_echo() -> httpx.Response: ) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "tool_choice", - [ - _ALLOWED_ONLY_ANOTHER_TOOL, - {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["echo"]}]}, - ], -) -async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes( - tool_choice: dict[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -) -> None: +async def _call_echo_on_anthropic_under_drop_params( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[httpx.Response], tool_choice: object +) -> tuple[object, 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( - return_value=_anthropic_call_to_echo() - ) - call_tool: Final = AsyncMock() - _served_by_a_recording_mcp_server(monkeypatch, call_tool) - + route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=list(rounds)) with _echo_tool_listed(): response: Final = await litellm.aresponses( model="anthropic/claude-haiku-4-5", @@ -1835,6 +1832,19 @@ async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes( 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 +) -> None: + 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) @@ -1843,6 +1853,127 @@ async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes( assert _output_types(response.output) == ("function_call",) +def _anthropic_answer() -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_answer", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "PING back"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "allowed_tools", "mode": "auto", "tools": []}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["echo"]}]}, + ], +) +async def test_a_malformed_choice_under_drop_params_runs_an_mcp_call_as_if_never_sent( + tool_choice: dict[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + 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(), _anthropic_answer()), tool_choice + ) + + assert call_tool.await_count == 1 + assert all("tool_choice" not in json.loads(call.request.content) for call in route.calls) + assert isinstance(response, ResponsesAPIResponse) + assert "message" in _output_types(response.output) + + +@pytest.mark.asyncio +async def test_a_malformed_choice_without_drop_params_answers_400_before_listing_or_calling_a_model( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + listing: Final = AsyncMock() + _served_by_a_recording_mcp_server(monkeypatch, AsyncMock()) + monkeypatch.setattr("litellm.proxy.mcp_server.server._get_tools_from_mcp_servers", listing) + + with pytest.raises(litellm.BadRequestError, match="Invalid tool choice"): + await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + tool_choice={"type": "allowed_tools", "mode": "auto", "tools": []}, + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + ) + + listing.assert_not_awaited() + assert not respx_mock.calls + + +_DEMO_LABELED_ECHO_TOOL: Final = MappingProxyType({**_ECHO_MCP_TOOL, "server_label": "demo"}) + + +async def _call_echo_with_the_demo_label(tool_choice: object, stream: bool) -> tuple[object, ...]: + with _echo_tool_listed(): + response: Final = await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_DEMO_LABELED_ECHO_TOOL)], + tool_choice=tool_choice, + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + stream=stream, + ) + return await _drain(response) if stream else (response,) + + +@pytest.mark.asyncio +async def test_the_provider_gets_the_converted_function_in_place_of_a_gateway_mcp_entry( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + route: Final = _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY)) + call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + + await _call_echo_with_the_demo_label( + {"type": "allowed_tools", "mode": "required", "tools": [{"type": "mcp", "server_label": "demo", "name": "echo"}]}, + stream=False, + ) + + assert json.loads(route.calls[0].request.content)["tool_choice"] == { + "type": "allowed_tools", + "mode": "required", + "tools": [{"type": "function", "name": "echo"}], + } + assert call_tool.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_a_gateway_mcp_entry_that_allows_no_converted_tool_sends_none_and_runs_nothing( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: + route: Final = _serve_responses( + respx_mock, + monkeypatch, + (_streamed(_ROUND_ONE_BODY) if stream else httpx.Response(200, content=json.dumps(_ROUND_ONE_BODY, default=dict)),), + ) + call_tool: Final = AsyncMock() + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + + await _call_echo_with_the_demo_label( + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "demo", "name": "reverse"}]}, + stream=stream, + ) + + assert json.loads(route.calls[0].request.content)["tool_choice"] == "none" + call_tool.assert_not_awaited() + + @pytest.mark.asyncio async def test_an_mcp_tool_the_choice_lists_still_runs(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY)) diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index ea013184a15..b0e8af97e3e 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -2773,10 +2773,19 @@ _DEMO_MCP_SERVER: Final = MCPServer( ) -def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock) -> None: +def _demo_server_named(name: str) -> MCPServer | None: + return _DEMO_MCP_SERVER if name.lower() == "demo" else None + + +def _served_by_a_recording_mcp_server( + monkeypatch: pytest.MonkeyPatch, + call_tool: AsyncMock, + listed: MCPTool = _MCP_ECHO_TOOL, + catalog: SimpleNamespace = SimpleNamespace(operation=nullcontext), +) -> None: monkeypatch.setattr( "litellm.proxy.mcp_server.server._get_tools_from_mcp_servers", - AsyncMock(return_value=AggregateToolListing(tools=[_MCP_ECHO_TOOL], outcomes={})), + AsyncMock(return_value=AggregateToolListing(tools=[listed], outcomes={})), ) monkeypatch.setitem( sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None) @@ -2784,11 +2793,11 @@ def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool monkeypatch.setattr( "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", SimpleNamespace( - catalog=SimpleNamespace(operation=nullcontext), + catalog=catalog, get_registry=MagicMock(return_value={}), call_tool=call_tool, get_mcp_server_from_tool_name=MagicMock(return_value=None), - get_mcp_server_by_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(side_effect=_demo_server_named), get_allowed_mcp_servers=AsyncMock(return_value=[_DEMO_MCP_SERVER.server_id]), get_mcp_servers_from_ids=MagicMock(return_value=[_DEMO_MCP_SERVER]), ), @@ -2824,7 +2833,7 @@ async def test_an_mcp_tool_the_choice_excludes_never_runs_when_the_model_ignored call_tool.assert_not_awaited() -def _anthropic_call_to_echo() -> Response: +def _anthropic_call_to_echo(name: str = "echo") -> Response: return Response( 200, json={ @@ -2832,7 +2841,7 @@ def _anthropic_call_to_echo() -> Response: "type": "message", "role": "assistant", "model": "claude-haiku-4-5", - "content": [{"type": "tool_use", "id": "toolu_echo", "name": "echo", "input": {"value": "PING"}}], + "content": [{"type": "tool_use", "id": "toolu_echo", "name": name, "input": {"value": "PING"}}], "stop_reason": "tool_use", "usage": {"input_tokens": 10, "output_tokens": 5}, }, @@ -2918,3 +2927,163 @@ async def test_a_dropped_choice_still_runs_an_mcp_call_it_lists( assert call_tool.await_count == 1 assert response.choices[0].message.content == "PING back" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "allowed_tools", "mode": "auto", "tools": []}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["echo"]}]}, + ], +) +async def test_a_malformed_choice_under_drop_params_runs_an_mcp_call_as_if_never_sent( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tool_choice: object +): + 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(), _anthropic_answer()), tool_choice + ) + + assert call_tool.await_count == 1 + assert response.choices[0].message.content == "PING back" + assert all("tool_choice" not in json.loads(call.request.content) for call in route.calls) + + +_DEMO_ECHO_TOOL: Final = MCPTool.model_validate( + {"name": "demo-echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, +) +_DEMO_IN_AN_ACCESS_GROUP: Final = _DEMO_MCP_SERVER.model_copy(update={"access_groups": ["demo_group"]}) + + +def _choice_of(*tools: Mapping[str, object]) -> Mapping[str, object]: + return {"type": "allowed_tools", "mode": "auto", "tools": list(tools)} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tool_choice", "server_url", "listed"), + [ + (_choice_of({"type": "mcp", "server_label": "demo", "name": "echo"}), "litellm_proxy/mcp/demo", _DEMO_ECHO_TOOL), + ( + _choice_of({"type": "mcp", "server_label": "demo", "name": "demo-echo"}), + "litellm_proxy/mcp/demo", + _DEMO_ECHO_TOOL, + ), + (_WHOLE_DEMO_SERVER, "litellm_proxy/mcp/DEMO", _MCP_ECHO_TOOL), + (_WHOLE_DEMO_SERVER, "litellm_proxy/mcp/demo_group", _MCP_ECHO_TOOL), + ], +) +async def test_an_mcp_entry_allows_the_tool_the_gateway_resolves_for_its_label_and_name( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + tool_choice: object, + server_url: str, + listed: MCPTool, +): + call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) + _served_by_a_recording_mcp_server( + monkeypatch, + call_tool, + listed, + SimpleNamespace( + operation=nullcontext, + current=lambda: SimpleNamespace(servers={_DEMO_MCP_SERVER.server_id: _DEMO_IN_AN_ACCESS_GROUP}), + ), + ) + + response, _route = await _call_echo_on_anthropic_under_drop_params( + respx_mock, monkeypatch, (_anthropic_call_to_echo(listed.name), _anthropic_answer()), tool_choice, server_url + ) + + assert call_tool.await_count == 1 + assert response.choices[0].message.content == "PING back" + + +@pytest.mark.asyncio +async def test_a_named_mcp_entry_allows_no_other_tool_on_its_server( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + call_tool: Final = AsyncMock() + _served_by_a_recording_mcp_server(monkeypatch, call_tool, _DEMO_ECHO_TOOL) + + response, _route = await _call_echo_on_anthropic_under_drop_params( + respx_mock, + monkeypatch, + (_anthropic_call_to_echo("demo-echo"),), + _choice_of({"type": "mcp", "server_label": "demo", "name": "reverse"}), + ) + + call_tool.assert_not_awaited() + assert [call.function.name for call in response.choices[0].message.tool_calls] == ["demo-echo"] + + +async def _call_echo_on_openai_with_the_demo_server( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], tool_choice: object +) -> tuple[object, ...]: + _serve_chat_rounds(respx_mock, monkeypatch, rounds) + await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[ + { + "type": "mcp", + "server_label": "demo", + "server_url": "litellm_proxy/mcp/demo", + "require_approval": "never", + }, + {"type": "function", "function": {"name": "get_time", "parameters": {"type": "object"}}}, + ], + tool_choice=tool_choice, + api_key="sk-mcp-chat-rounds", + api_base="https://mcp-chat-rounds.example.com/v1", + ) + return tuple(json.loads(call.request.content).get("tool_choice") for call in respx_mock.calls) + + +@pytest.mark.asyncio +async def test_the_provider_gets_the_converted_function_in_place_of_a_gateway_mcp_entry( + 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) + + sent: Final = await _call_echo_on_openai_with_the_demo_server( + respx_mock, + monkeypatch, + (_ok(_CHAT_ROUND_ONE), _ok(_CHAT_ROUND_TWO)), + _choice_of({"type": "mcp", "server_label": "demo"}, {"type": "function", "name": "get_time"}), + ) + + assert sent[0] == { + "type": "allowed_tools", + "allowed_tools": { + "mode": "auto", + "tools": [ + {"type": "function", "function": {"name": "get_time"}}, + {"type": "function", "function": {"name": "echo"}}, + ], + }, + } + assert call_tool.await_count == 1 + + +@pytest.mark.asyncio +async def test_a_gateway_mcp_entry_that_allows_no_converted_tool_sends_none_and_runs_nothing( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + call_tool: Final = AsyncMock() + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + + sent: Final = await _call_echo_on_openai_with_the_demo_server( + respx_mock, + monkeypatch, + (_ok(_CHAT_ROUND_ONE),), + _choice_of({"type": "mcp", "server_label": "demo", "name": "reverse"}), + ) + + assert sent == ("none",) + call_tool.assert_not_awaited() diff --git a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py index a16661c0174..220d46d90ea 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -1632,14 +1632,14 @@ async def _no_toolset_grants(context: UserAPIKeyAuth) -> frozenset[str]: return frozenset() -async def _tools_listed_through_the_gateway( +async def _listing_through_the_gateway( monkeypatch: pytest.MonkeyPatch, user_api_key_auth: UserAPIKeyAuth, - server_urls: tuple[str, ...], + mcp_tools: Sequence[Mapping[str, object]], granted_toolsets: GrantedToolsetIds, toolset_table: types.SimpleNamespace | None = None, agent_permission: LiteLLM_ObjectPermissionTable | None = None, -) -> list[str]: +) -> tuple[list[MCPTool], list[str]]: async def upstream_listing(_client: object, server_name: str) -> list[MCPTool]: return list(UPSTREAM_TOOLS[server_name]) @@ -1665,17 +1665,99 @@ async def _tools_listed_through_the_gateway( monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"supported_db_objects": []}) try: - tools, _server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( + return await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( user_api_key_auth=user_api_key_auth, - mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": url} for url in server_urls], + mcp_tools_with_litellm_proxy=mcp_tools, granted_toolsets=granted_toolsets, ) finally: for server in GATEWAY_SERVERS: manager._drop_listed_tools(server.server_id) + + +async def _tools_listed_through_the_gateway( + monkeypatch: pytest.MonkeyPatch, + user_api_key_auth: UserAPIKeyAuth, + server_urls: tuple[str, ...], + granted_toolsets: GrantedToolsetIds, + toolset_table: types.SimpleNamespace | None = None, + agent_permission: LiteLLM_ObjectPermissionTable | None = None, +) -> list[str]: + tools, _server_names = await _listing_through_the_gateway( + monkeypatch, + user_api_key_auth, + [{"type": "mcp", "server_url": url} for url in server_urls], + granted_toolsets, + toolset_table, + agent_permission, + ) return sorted(tool.name for tool in tools) +async def _calls_the_gateway_runs( + monkeypatch: pytest.MonkeyPatch, + labeled_server_urls: Mapping[str, str], + allowed_label: str, + called: tuple[str, ...], + unresolvable_tools: Mapping[str, str] = MappingProxyType({}), +) -> tuple[str, ...]: + user: Final = UserAPIKeyAuth(api_key="sk-team", team_id="team-1", user_role=LitellmUserRoles.INTERNAL_USER) + mcp_tools: Final = [ + {"type": "mcp", "server_label": label, "server_url": url} for label, url in labeled_server_urls.items() + ] + tools, server_names = await _listing_through_the_gateway(monkeypatch, user, mcp_tools, _team_grants("ts-granted")) + _deduplicated, tool_server_map = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools(tools, server_names) + tool_choice: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": allowed_label}]} + gateway: Final = await LiteLLM_Proxy_MCP_Handler.gateway_tools( + tool_choice, {**tool_server_map, **unresolvable_tools}, mcp_tools, user, _team_grants("ts-granted") + ) + return tuple( + name + for name in called + if LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( + ({"type": "function_call", "name": name, "call_id": f"call_{name}", "arguments": "{}"},), + tool_choice, + gateway, + ) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("labeled_server_urls", "allowed_label", "runs"), + [ + ({"team": "litellm_proxy/mcp/team-toolset", "docs": "litellm_proxy/mcp/docs"}, "team", ("calc-add",)), + ({"team": "litellm_proxy/mcp/team-toolset", "docs": "litellm_proxy/mcp/docs"}, "docs", ("docs-lookup",)), + ({"team": "litellm_proxy/mcp/team-toolset", "other": "litellm_proxy/mcp/other-toolset"}, "other", ()), + ], + ids=["granted_toolset", "server_beside_a_toolset", "ungranted_toolset"], +) +async def test_an_mcp_entry_allows_only_the_tools_its_labels_toolset_or_server_serves( + monkeypatch: pytest.MonkeyPatch, + labeled_server_urls: Mapping[str, str], + allowed_label: str, + runs: tuple[str, ...], +): + ran: Final = await _calls_the_gateway_runs( + monkeypatch, labeled_server_urls, allowed_label, ("calc-add", "docs-lookup") + ) + + assert ran == runs + + +@pytest.mark.asyncio +async def test_an_mcp_entry_allows_no_call_on_a_server_the_gateway_cannot_resolve(monkeypatch: pytest.MonkeyPatch): + ran: Final = await _calls_the_gateway_runs( + monkeypatch, + {"team": "litellm_proxy/mcp/team-toolset"}, + "team", + ("calc-add", "ghost-add"), + MappingProxyType({"ghost-add": "ghost"}), + ) + + assert ran == ("calc-add",) + + def _key_holding(*server_ids: str, role: LitellmUserRoles = LitellmUserRoles.INTERNAL_USER) -> UserAPIKeyAuth: return UserAPIKeyAuth( api_key="sk-team", diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 7313a1600f9..37c7a4cd422 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -656,7 +656,7 @@ def test_completion_drop_params_drops_an_allowed_tools_tool_choice_it_cannot_sen {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "mcp", "name": ["get_time"]}]}}, ], ) -async def test_a_malformed_allowed_tools_tool_choice_stays_a_400_on_mcp_gateway_tools( +async def test_a_malformed_allowed_tools_tool_choice_is_a_400_on_mcp_gateway_tools_without_drop_params( tool_choice: dict[str, object], respx_mock: respx.MockRouter ): with pytest.raises(litellm.BadRequestError, match="Invalid tool choice") as exc: @@ -665,7 +665,6 @@ async def test_a_malformed_allowed_tools_tool_choice_stays_a_400_on_mcp_gateway_ messages=[{"role": "user", "content": "What time is it?"}], tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], tool_choice=tool_choice, - drop_params=True, api_base="http://localhost:12346/v1", api_key="fake_api_key", ) diff --git a/tests/unit/test_tool_loop.py b/tests/unit/test_tool_loop.py index e679f9817c1..f56a88e8899 100644 --- a/tests/unit/test_tool_loop.py +++ b/tests/unit/test_tool_loop.py @@ -350,13 +350,45 @@ def test_a_malformed_choice_is_refused_before_any_request( tools=WEATHER_TOOLS, execute_tool=_tool_result, tool_choice=tool_choice, - drop_params=True, api_key="sk-test", ) assert not respx_mock.calls +@pytest.mark.parametrize( + "tool_choice", + [_EMPTY_ALLOWED_TOOLS, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["get_time"]}]}], +) +def test_a_malformed_choice_under_drop_params_runs_the_loop_as_if_never_sent( + tool_choice: dict[str, object], respx_mock: respx.MockRouter +) -> None: + respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + side_effect=[ + httpx.Response(200, json=_openai_response(None, [_TIME_CALL])), + httpx.Response(200, json=_openai_response("done")), + ] + ) + executed: Final[list[str | None]] = [] + + def execute(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + executed.append(tool_call.function.name) + return _tool_result(tool_call) + + answer: Final = litellm.run_tool_loop( + model="openai/gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=execute, + tool_choice=tool_choice, + drop_params=True, + api_key="sk-test", + ) + + assert (answer, executed) == ("done", ["get_time"]) + assert all("tool_choice" not in body for body in _request_bodies(respx_mock)) + + def test_a_choice_on_a_model_alias_reaches_the_loop( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: