diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index dd36eee978d..ce5cbc15902 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -18,7 +18,10 @@ from pydantic import ValidationError import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger -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_in_auto_mode, + allowed_tools_tool_choice_with_builtin_as_function, +) from litellm.llms.base_llm.sandbox.transformation import ( CodeExecutionResult, ContainerHandle, @@ -707,7 +710,9 @@ class CodeInterpreterInterceptionLogger(CustomLogger): def _get_followup_optional_params(self, optional_params: dict[str, object]) -> dict[str, object]: drop_tool_choice: Final = self._tool_choice_targets_code_interpreter(optional_params.get("tool_choice")) return { - k: v for k, v in optional_params.items() if k != "tools" and not (k == "tool_choice" and drop_tool_choice) + k: allowed_tools_tool_choice_in_auto_mode(v) if k == "tool_choice" else v + for k, v in optional_params.items() + if k != "tools" and not (k == "tool_choice" and drop_tool_choice) } async def async_post_agentic_loop_response_hook( diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 4b2773dbaa3..0a4759d403c 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -40,6 +40,7 @@ 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_in_auto_mode, allowed_tools_tool_choice_with_builtin_as_function, is_allowed_tools_tool_choice, ) @@ -284,7 +285,9 @@ 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({}) + if not is_allowed_tools_tool_choice(tool_choice): + return MappingProxyType({}) + return {"tool_choice": allowed_tools_tool_choice_in_auto_mode(tool_choice)} class WebSearchInterceptionLogger(CustomLogger): diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index ee6a4fa5cb8..53b5b3091b4 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -328,6 +328,19 @@ def allowed_tools_tool_choice_with_function(tool_choice: object, function_name: return _with_entries(layout, (*_entries_of(layout), _function_entry(layout, function_name))) +def allowed_tools_tool_choice_in_auto_mode(tool_choice: object) -> object: + layout: Final = _callers_layout(tool_choice) + match layout: + case None: + return tool_choice + case ChatAllowedToolsToolChoice(): + return layout.model_copy( + update={"allowed_tools": layout.allowed_tools.model_copy(update={"mode": "auto"})} + ).model_dump(mode="json") + case ResponsesAllowedToolsToolChoice(): + return layout.model_copy(update={"mode": "auto"}).model_dump(mode="json") + + 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 diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 9863b448fd4..e3b7d4f59af 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,6 +1,6 @@ import traceback from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload @@ -148,7 +148,7 @@ class _LabelScope: label: str whole_gateway: bool = False server_ids: frozenset[str] = frozenset() - toolset_tools: Mapping[str, tuple[str, ...]] = MappingProxyType({}) + toolset_tools: Mapping[str, tuple[str, ...]] = field(default_factory=lambda: MappingProxyType({})) def serves(self, server: MCPServer | None, tool_name: str) -> bool: if self.whole_gateway: diff --git a/tests/unit/integrations/code_interpreter_interception/test_handler.py b/tests/unit/integrations/code_interpreter_interception/test_handler.py index 5fde826ee23..2a0712f1a53 100644 --- a/tests/unit/integrations/code_interpreter_interception/test_handler.py +++ b/tests/unit/integrations/code_interpreter_interception/test_handler.py @@ -728,6 +728,45 @@ async def test_chat_completion_build_plan_runs_code_and_appends_tool_message(): assert plan.metadata["code_interpreter_calls"][0]["code"] == "print(40 + 2)" +@pytest.mark.asyncio +async def test_chat_completion_build_plan_keeps_an_allowed_tools_choice_without_forcing_another_call(): + allowed: Final = [ + {"type": "function", "function": {"name": LITELLM_CODE_EXECUTION_TOOL_NAME}}, + {"type": "function", "function": {"name": "get_time"}}, + ] + logger: Final = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox(stdout="42")) + + plan: Final = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + {"id": "call_1", "name": LITELLM_CODE_EXECUTION_TOOL_NAME, "arguments": '{"code":"print(1)"}'} + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response={"choices": [{"message": {"tool_calls": [_chat_function_call_item()]}}]}, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={ + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "tool_choice": {"type": "allowed_tools", "allowed_tools": {"mode": "required", "tools": allowed}}, + }, + logging_obj=FakeLogging(litellm_call_id="k1"), + stream=False, + kwargs={ + "acompletion": True, + "litellm_call_id": "k1", + _ACTIVE_KEY: True, + _SANDBOX_KEY: "sbxkey1", + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + }, + ) + + assert plan.request_patch is not None + assert plan.request_patch.optional_params == { + "tool_choice": {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": allowed}} + } + + @pytest.mark.asyncio async def test_pre_call_strips_client_forged_marker_on_initial_request(): """A client cannot pre-set the active marker on the original request: with no 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 89ad8b72070..d49da255c67 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py @@ -6,6 +6,7 @@ litellm.acompletion() for transparent server-side web search execution. """ import datetime +import json from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -424,62 +425,74 @@ 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 = { +async def test_execute_chat_completion_agentic_loop_keeps_an_allowed_tools_choice_without_forcing_another_call( + mode: str, + perplexity_search: respx.Route, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "test-openai-key") + litellm.in_memory_llm_clients_cache.flush_cache() + followup: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + json={ + "id": "followup", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}], + } + ) + tool_choice: Final = { "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") + await WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI])._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={}, + ) - 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 + assert perplexity_search.call_count == 1 + assert json.loads(followup.calls.last.request.content)["tool_choice"] == { + "type": "allowed_tools", + "allowed_tools": {"mode": "auto", "tools": [{"type": "function", "function": {"name": "litellm_web_search"}}]}, + } @pytest.mark.asyncio @pytest.mark.parametrize( - ("tool_choice", "kept"), + ("tool_choice", "followup_tool_choice"), [ ( { "type": "allowed_tools", "allowed_tools": {"mode": "required", "tools": [{"type": "function", "function": {"name": "get_time"}}]}, }, - True, + { + "type": "allowed_tools", + "allowed_tools": {"mode": "auto", "tools": [{"type": "function", "function": {"name": "get_time"}}]}, + }, ), - ({"type": "function", "function": {"name": "litellm_web_search"}}, False), - ("required", False), + ({"type": "function", "function": {"name": "litellm_web_search"}}, None), + ("required", None), ], ) async def test_chat_followup_plan_keeps_only_an_allowed_tools_choice( - tool_choice: object, kept: bool, perplexity_search: respx.Route + tool_choice: object, followup_tool_choice: object, perplexity_search: respx.Route ) -> None: plan: Final = await WebSearchInterceptionLogger().async_build_chat_completion_agentic_loop_plan( tools={ @@ -498,7 +511,7 @@ async def test_chat_followup_plan_keeps_only_an_allowed_tools_choice( ) assert plan.request_patch is not None - assert plan.request_patch.optional_params.get("tool_choice") == (tool_choice if kept else None) + assert plan.request_patch.optional_params.get("tool_choice") == followup_tool_choice if __name__ == "__main__": 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 ec398e8731a..a69690106d1 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 @@ -10,6 +10,7 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ChatAllowedToolsToolChoice, GatewayTool, GatewayTools, + allowed_tools_tool_choice_in_auto_mode, allowed_tools_tool_choice_in_effect, allowed_tools_tool_choice_with_builtin_as_function, allowed_tools_tool_choice_with_function, @@ -549,3 +550,42 @@ def test_an_added_function_joins_the_choice_in_the_callers_layout( ) 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 + + +@pytest.mark.parametrize( + ("tool_choice", "expected"), + [ + ( + { + "type": "allowed_tools", + "allowed_tools": {"mode": "required", "tools": [{"type": "function", "function": {"name": "get_time"}}]}, + }, + { + "type": "allowed_tools", + "allowed_tools": {"mode": "auto", "tools": [{"type": "function", "function": {"name": "get_time"}}]}, + }, + ), + ( + { + "type": "allowed_tools", + "mode": "required", + "tools": [{"type": "function", "name": "get_time"}, {"type": "mcp", "server_label": "docs"}], + }, + { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "get_time"}, {"type": "mcp", "server_label": "docs"}], + }, + ), + (_CHAT_LAYOUT, _CHAT_LAYOUT), + ("required", "required"), + ( + {"type": "function", "function": {"name": "get_time"}}, + {"type": "function", "function": {"name": "get_time"}}, + ), + ], +) +def test_auto_mode_lifts_only_the_requirement_and_keeps_the_callers_layout_and_tools( + tool_choice: object, expected: object +): + assert allowed_tools_tool_choice_in_auto_mode(tool_choice) == expected