diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index 69b528a29f8..240669fa425 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -17,6 +17,7 @@ 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.llms.base_llm.sandbox.transformation import ( CodeExecutionResult, ContainerHandle, @@ -141,6 +142,10 @@ class _SupportsOutput(Protocol): _CachedContainer: TypeAlias = tuple[ContainerHandle, SandboxToolParams | None, float, str | None] +def _is_code_interpreter_tool(tool: Mapping[str, object]) -> bool: + return tool.get("type") == "code_interpreter" + + def _output_item_type(item: object) -> object: if isinstance(item, dict): item_mapping: Final[dict[str, object]] = item @@ -301,8 +306,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger): (function_tool if isinstance(tool, dict) and tool.get("type") == "code_interpreter" else tool) for tool in tools ] - if self._tool_choice_targets_code_interpreter(kwargs.get("tool_choice")): - kwargs["tool_choice"] = self._get_function_tool_choice(call_type=call_type) + if "tool_choice" in kwargs: + kwargs["tool_choice"] = self._intercepted_tool_choice(kwargs["tool_choice"], call_type) return kwargs @staticmethod @@ -386,6 +391,14 @@ class CodeInterpreterInterceptionLogger(CustomLogger): or (isinstance(function, dict) and function.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME) ) + @classmethod + def _intercepted_tool_choice(cls, tool_choice: object, call_type: CallTypes | None) -> object: + if cls._tool_choice_targets_code_interpreter(tool_choice): + return cls._get_function_tool_choice(call_type=call_type) + return allowed_tools_tool_choice_with_builtin_as_function( + tool_choice, _is_code_interpreter_tool, LITELLM_CODE_EXECUTION_TOOL_NAME + ) + def _resolve_provider(self, kwargs: dict[str, object]) -> str | None: provider: Final = kwargs.get("custom_llm_provider") if isinstance(provider, str) and provider: diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index a38b8bed595..631a0fcd0fb 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -9,7 +9,7 @@ server-side using litellm router's search tools. import asyncio import math import uuid -from collections.abc import AsyncIterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast @@ -37,6 +37,7 @@ 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.llms.base_llm.search.transformation import SearchResponse from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, @@ -246,6 +247,18 @@ class _ParentRequestCorrelation: _NO_ACOMPLETION_NAMED: Final[_AcompletionNamedParams] = {} +def _tool_choice_naming_litellm_web_search( + kwargs: Mapping[str, object], is_search_tool: Callable[[Mapping[str, object]], bool] +) -> Mapping[str, object]: + if "tool_choice" not in kwargs: + return MappingProxyType({}) + return { + "tool_choice": allowed_tools_tool_choice_with_builtin_as_function( + kwargs["tool_choice"], is_search_tool, LITELLM_WEB_SEARCH_TOOL_NAME + ) + } + + class WebSearchInterceptionLogger(CustomLogger): """ CustomLogger that intercepts WebSearch tool calls for models that don't @@ -493,13 +506,17 @@ class WebSearchInterceptionLogger(CustomLogger): # Keep other tools as-is converted_tools.append(tool) - kwargs["tools"] = converted_tools + converted_kwargs: Final[dict[str, object]] = { + **kwargs, + "tools": converted_tools, + **_tool_choice_naming_litellm_web_search(kwargs, lambda tool: is_search_tool(dict(tool))), + } if kwargs.get("stream"): verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False") - return as_converted_stream(kwargs, WEBSEARCH_INTERCEPTION_PREFIX) + return as_converted_stream(converted_kwargs, WEBSEARCH_INTERCEPTION_PREFIX) - return kwargs + return converted_kwargs def _convert_responses_tools( self, kwargs: Mapping[str, object], tools: Sequence[dict[str, object]] @@ -514,7 +531,11 @@ class WebSearchInterceptionLogger(CustomLogger): get_litellm_web_search_tool_responses() if is_web_search_tool_responses(tool) else tool for tool in tools ] - converted_kwargs: Final = {**kwargs, "tools": converted_tools} + converted_kwargs: Final = { + **kwargs, + "tools": converted_tools, + **_tool_choice_naming_litellm_web_search(kwargs, is_web_search_tool_responses), + } if kwargs.get("stream"): verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False") diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index fd0ba2c8e4f..ae256580b93 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -1,4 +1,4 @@ -from collections.abc import Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping from types import MappingProxyType from typing import Final, Literal @@ -176,7 +176,7 @@ def _allowed_tool_names(tool_choice: object) -> frozenset[object]: if parsed is None: return frozenset() flattened: Final = tuple(_responses_allowed_tool(tool) for tool in parsed.allowed_tools.tools) - return frozenset(tool.get("name") for tool in flattened if tool is not None) + 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: @@ -191,15 +191,22 @@ def refuse_a_malformed_allowed_tools_tool_choice(tool_choice: object, model: str ) -def _whole_server_label(tool: Mapping[str, object]) -> str | None: +def _server_entry(tool: Mapping[str, object]) -> tuple[str, str | None] | None: label: Final = tool.get("server_label") - return label if tool.get("type") == "mcp" and tool.get("name") is None and isinstance(label, str) else None + name: Final = tool.get("name") + if tool.get("type") != "mcp" or not isinstance(label, str): + return None + return label, name if isinstance(name, str) else None -def _whole_server_labels(tool_choice: object) -> frozenset[str]: +def _server_entries(tool_choice: object) -> frozenset[tuple[str, str | None]]: parsed: Final = parse_allowed_tools_tool_choice(tool_choice) - labels: Final = () if parsed is None else tuple(map(_whole_server_label, parsed.allowed_tools.tools)) - return frozenset(label for label in labels if label is not None) + entries: Final = () if parsed is None else tuple(map(_server_entry, parsed.allowed_tools.tools)) + 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) def tool_calls_the_choice_excludes( @@ -210,11 +217,12 @@ def tool_calls_the_choice_excludes( if not is_allowed_tools_tool_choice(tool_choice): return () allowed: Final = _allowed_tool_names(tool_choice) - whole_servers: Final = _whole_server_labels(tool_choice) + server_entries: Final = _server_entries(tool_choice) return tuple( name for name in tool_names - if name not in allowed and whole_servers.isdisjoint(server_labels_by_tool.get(name, frozenset())) + if name not in allowed + and not _a_server_entry_allows(server_entries, name, server_labels_by_tool.get(name, frozenset())) ) @@ -232,6 +240,38 @@ def refuse_tool_calls_the_choice_excludes(tool_choice: object, tool_names: Itera ) +def _with_builtin_as( + tool: Mapping[str, object], is_builtin: Callable[[Mapping[str, object]], bool], function_entry: Mapping[str, object] +) -> Mapping[str, object]: + return function_entry if tool.get("type") not in ("function", "custom", "mcp") and is_builtin(tool) else tool + + +def allowed_tools_tool_choice_with_builtin_as_function( + tool_choice: object, is_builtin: Callable[[Mapping[str, object]], bool], function_name: str +) -> object: + try: + parsed: Final = _ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice) + except ValidationError: + return 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], + } + + def chat_allowed_tools_tool_choice(choice: ChatAllowedToolsToolChoice) -> ChatCompletionAllowedToolChoiceParam: return ChatCompletionAllowedToolChoiceParam( type="allowed_tools", diff --git a/tests/unit/integrations/code_interpreter_interception/test_handler.py b/tests/unit/integrations/code_interpreter_interception/test_handler.py index c037b6be80d..e947260a01e 100644 --- a/tests/unit/integrations/code_interpreter_interception/test_handler.py +++ b/tests/unit/integrations/code_interpreter_interception/test_handler.py @@ -19,6 +19,7 @@ from litellm.integrations.code_interpreter_interception.handler import ( LITELLM_CODE_EXECUTION_TOOL_NAME, CodeInterpreterInterceptionLogger, ) +from litellm.litellm_core_utils.allowed_tools_tool_choice import tool_calls_the_choice_excludes from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, @@ -1200,6 +1201,38 @@ async def test_extract_tool_calls_reads_object_attributes(): # --------------------------------------------------------------------------- +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "tool_choice"), + [ + (CallTypes.aresponses, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "code_interpreter"}]}), + ( + CallTypes.acompletion, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "code_interpreter"}]}}, + ), + ], +) +async def test_pre_call_lets_an_allowed_tools_choice_listing_code_interpreter_allow_the_run( + call_type: CallTypes, tool_choice: dict[str, object] +): + logger: Final = CodeInterpreterInterceptionLogger() + + result: Final = TypeAdapter(dict[str, object]).validate_python( + await logger.async_pre_call_deployment_hook( + { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "tool_choice": tool_choice, + "custom_llm_provider": "openai", + }, + call_type, + ) + ) + + assert tool_calls_the_choice_excludes(result["tool_choice"], (LITELLM_CODE_EXECUTION_TOOL_NAME, "get_weather")) == ( + "get_weather", + ) + + @pytest.mark.asyncio async def test_pre_call_uses_session_id_from_metadata_as_sandbox_key(): """When session_id is in request metadata, it becomes the sandbox key so the diff --git a/tests/unit/integrations/websearch_interception/test_websearch_responses.py b/tests/unit/integrations/websearch_interception/test_websearch_responses.py index 00b7299cc31..221e09c3ac6 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_responses.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_responses.py @@ -279,6 +279,39 @@ async def test_deployment_hook_converts_native_responses_web_search_tool(): assert tool["parameters"]["required"] == ["query"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "tool_choice"), + [ + (CallTypes.aresponses, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "web_search"}]}), + (CallTypes.acompletion, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "web_search"}]}), + ( + CallTypes.acompletion, + {"type": "allowed_tools", "allowed_tools": {"mode": "required", "tools": [{"type": "web_search"}]}}, + ), + ], +) +async def test_deployment_hook_lets_an_allowed_tools_choice_listing_web_search_allow_the_search( + call_type: CallTypes, tool_choice: dict[str, object] +): + logger: Final = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI]) + + result: Final = await logger.async_pre_call_deployment_hook( + kwargs={ + "model": "gpt-4o", + "custom_llm_provider": "openai", + "tools": [{"type": "web_search"}, {"type": "function", "name": "get_weather"}], + "tool_choice": tool_choice, + }, + call_type=call_type, + ) + + assert result is not None + assert tool_calls_the_choice_excludes(result["tool_choice"], ("litellm_web_search", "get_weather")) == ( + "get_weather", + ) + + @pytest.mark.asyncio async def test_deployment_hook_responses_returns_none_without_web_search(): """No web search tool in a responses request -> deployment hook makes no change.""" 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 7c19e0b673c..5d7fbbb2228 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,6 +8,7 @@ import litellm from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ChatAllowedToolsToolChoice, + allowed_tools_tool_choice_with_builtin_as_function, drops_unsupported_allowed_tools_tool_choice, is_allowed_tools_tool_choice, chat_allowed_tools_tool_choice, @@ -297,6 +298,17 @@ def test_an_mcp_entry_with_a_null_name_is_a_valid_choice(tool_choice: object): assert refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") is None +_DOCS_SEARCH: Final = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "mcp", "server_label": "docs", "name": "search"}], +} + + +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 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"})}) == () @@ -310,9 +322,59 @@ def test_an_mcp_entry_without_a_name_allows_every_call_its_server_serves(): {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "docs", "name": "read"}]}, frozenset({"docs"}), ), + (_DOCS_SEARCH, frozenset({"other"})), + (_DOCS_SEARCH, frozenset()), ], ) 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",) + + +def _is_web_search(tool: Mapping[str, object]) -> bool: + return tool.get("type") == "web_search" + + +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "web_search"}, {"type": "function", "name": "get_time"}]}, + { + "type": "allowed_tools", + "allowed_tools": { + "mode": "auto", + "tools": [{"type": "web_search"}, {"type": "function", "function": {"name": "get_time"}}], + }, + }, + ], +) +def test_a_builtin_entry_allows_the_function_that_replaces_it_in_the_same_layout(tool_choice: dict[str, object]): + rewritten: Final = allowed_tools_tool_choice_with_builtin_as_function( + tool_choice, _is_web_search, "litellm_web_search" + ) + + assert tool_calls_the_choice_excludes(tool_choice, ("litellm_web_search", "get_time")) == ("litellm_web_search",) + assert tool_calls_the_choice_excludes(rewritten, ("litellm_web_search", "get_time", "get_weather")) == ( + "get_weather", + ) + assert isinstance(rewritten, dict) and rewritten.keys() == tool_choice.keys() + + +def test_a_function_or_mcp_entry_is_never_taken_for_a_builtin(): + tool_choice: Final = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "web_search"}, {"type": "mcp", "server_label": "docs", "name": "web_search"}], + } + + rewritten: Final = allowed_tools_tool_choice_with_builtin_as_function( + tool_choice, lambda tool: tool.get("name") == "web_search", "litellm_web_search" + ) + + assert parse_allowed_tools_tool_choice(rewritten) == parse_allowed_tools_tool_choice(tool_choice) + + +@pytest.mark.parametrize("tool_choice", ["auto", None, {"type": "allowed_tools", "mode": "auto", "tools": []}]) +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