scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(allowed_tools): stop web search and code interpreter follow-ups forcing another call

Browse source 
Backport of #50467 to lts/1.108.x.
Cherry-picked from 478fb95d6b (litellm_allowed_tools_flat_tool_choice). Resolved: test imports merged with the line's, and the tests that need main's web search fixtures move to the line test commit.
This commit is contained in:
mateo-berri 2026-11-06 19:59:47 -08:00
parent b51f033a61
commit 83b53e48e3
7 changed files with 155 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/integrations/code_interpreter_interception/handler.py
View file

@ -17,7 +17,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,
@ -661,7 +664,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(

5
litellm/integrations/websearch_interception/handler.py
View file

@ -39,6 +39,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,
)
@ -267,7 +268,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):

13
litellm/litellm_core_utils/allowed_tools_tool_choice.py
View file

@ -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

4
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -1,7 +1,7 @@
import re
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
@ -183,7 +183,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:

39
tests/unit/integrations/code_interpreter_interception/test_handler.py
View file

@ -718,6 +718,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

85
tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py
View file

@ -5,9 +5,12 @@ Tests the end-to-end flow of websearch_interception callback with
litellm.acompletion() for transparent server-side web search execution.
"""
import json
from typing import Final
from unittest.mock import MagicMock
import pytest
import respx
import litellm
from litellm.integrations.websearch_interception.handler import (
@ -421,62 +424,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={
@ -495,7 +510,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__":

40
tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py
View file

@ -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