Merge d8913d8127 into aeede53112

litellm/integrations/websearch_interception/handler.py

@@ -1315,6 +1315,7 @@ class WebSearchInterceptionLogger(CustomLogger):
     def _resolve_max_tokens(
         optional_params: dict,
         kwargs: dict,
+        client_max_tokens: int | None = None,
     ) -> int:
         """Extract max_tokens and validate against thinking.budget_tokens.
 
@@ -1323,7 +1324,7 @@ class WebSearchInterceptionLogger(CustomLogger):
         """
         max_tokens: int = optional_params.get(
             "max_tokens",
-            kwargs.get("max_tokens", 1024),
+            kwargs.get("max_tokens", 1024 if client_max_tokens is None else client_max_tokens),
         )
         thinking_param: Final = optional_params.get("thinking")
         if thinking_param and isinstance(thinking_param, dict):
@@ -1466,9 +1467,16 @@ class WebSearchInterceptionLogger(CustomLogger):
         # Correlation context for structured logging
         _call_id: Final = getattr(logging_obj, "litellm_call_id", None) or kwargs.get("litellm_call_id", "unknown")
 
-        full_model_name = model  # safe default before try block
+        agentic_view: Final[_AgenticLoopParamsView] = {
+            "agentic_loop_params": {}
+            if logging_obj is None
+            else logging_obj.model_call_details.get("agentic_loop_params", {})
+        }
+        full_model_name: Final = agentic_view["agentic_loop_params"].get("model", model)
 
-        max_tokens: Final = self._resolve_max_tokens(anthropic_messages_optional_request_params, kwargs)
+        max_tokens: Final = self._resolve_max_tokens(
+            anthropic_messages_optional_request_params, kwargs, agentic_view["agentic_loop_params"].get("max_tokens")
+        )
 
         verbose_logger.debug("WebSearchInterception: Using max_tokens=%s for follow-up request", max_tokens)
 
@@ -1477,11 +1485,6 @@ class WebSearchInterceptionLogger(CustomLogger):
         }
         kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs)
 
-        if logging_obj is not None:
-            agentic_view: Final[_AgenticLoopParamsView] = {
-                "agentic_loop_params": logging_obj.model_call_details.get("agentic_loop_params", {})
-            }
-            full_model_name = agentic_view["agentic_loop_params"].get("model", model)
         verbose_logger.debug(
             "WebSearchInterception: Built anthropic request patch [call_id=%s model=%s messages=%d searches=%d]",
             _call_id,

litellm/llms/anthropic/messages/native/handler.py

@@ -696,6 +696,7 @@ def anthropic_messages_handler(
         litellm_logging_obj.model_call_details["agentic_loop_params"] = {
             "model": original_model,
             "custom_llm_provider": custom_llm_provider,
+            "max_tokens": max_tokens,
         }
         litellm_logging_obj.model_call_details.setdefault("litellm_params", {})[CallTypes.aanthropic_messages.value] = (
             is_async

litellm/llms/custom_httpx/llm_http_handler.py

@@ -6127,10 +6127,10 @@ class BaseLLMHTTPHandler:
         if patch.messages is None:
             raise ValueError("Agentic loop plan missing patched messages")
 
-        full_model_name = model
-        if logging_obj is not None:
-            agentic_params: Final[Mapping[str, object]] = logging_obj.model_call_details.get("agentic_loop_params", {})
-            full_model_name = cast(str, agentic_params.get("model", model))
+        agentic_params: Final[Mapping[str, object]] = (
+            {} if logging_obj is None else logging_obj.model_call_details.get("agentic_loop_params", {})
+        )
+        full_model_name: Final = cast(str, agentic_params.get("model", model))
 
         optional_params: Final = dict(anthropic_messages_optional_request_params)
         optional_params.update(patch.optional_params)
@@ -6143,7 +6143,7 @@ class BaseLLMHTTPHandler:
         else:
             optional_params.pop("max_tokens", None)
         if max_tokens is None:
-            max_tokens = cast(int, kwargs.get("max_tokens", 1024))
+            max_tokens = cast(int, kwargs.get("max_tokens", agentic_params.get("max_tokens", 1024)))
 
         internal_keys: Final = {"litellm_logging_obj"}
         kwargs_for_followup: Final = {

litellm/types/utils.py

@@ -283,6 +283,9 @@ class AgenticLoopParams(TypedDict, total=False):
     custom_llm_provider: str
     """The LLM provider name (e.g., 'bedrock', 'anthropic'), when the recorder resolved one"""
 
+    max_tokens: ReadOnly[int]
+    """The client's max_tokens, which the request transform removes from the params the hooks receive"""
+
 
 class OffPeakWindow(TypedDict, total=False):
     """One off-peak rule: UTC time-of-day windows, optionally restricted to weekdays.

tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py

@@ -510,3 +510,105 @@ def test_empty_exa_search_reaches_the_model_without_the_exa_response_envelope(ga
         assert envelope["requestId"] not in response.text, response.text
         targets: Final = tuple((request.method, urlsplit(request.target).path) for request in wire.drain())
         assert targets[-3:] == (("POST", "/v1/messages"), ("POST", "/exa/search"), ("POST", "/v1/messages")), targets
+
+
+@pytest.mark.parametrize(
+    ("stream", "bedrock_model", "thinking"),
+    [
+        pytest.param(False, BEDROCK_MODEL, None, id="non-streaming"),
+        pytest.param(True, BEDROCK_MODEL, None, id="streaming"),
+        pytest.param(
+            False,
+            "us.anthropic.claude-opus-5-5",
+            {"type": "enabled", "budget_tokens": 8000},
+            id="adaptive-model-with-a-thinking-budget-above-max-tokens",
+        ),
+    ],
+)
+def test_web_search_follow_up_call_keeps_the_client_max_tokens(
+    gateway: Gateway, tmp_path: Path, stream: bool, bedrock_model: str, thinking: dict[str, object] | None
+) -> None:
+    invoke_target: Final = f"/model/{bedrock_model}/invoke"
+
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST", request.target
+        body: Final = json.loads(request.body)
+        if request.target == SEARCH_TARGET:
+            return Reply(body=json.dumps({"query": body["query"], "results": [SEARCH_RESULT]}).encode())
+        assert request.target == invoke_target, request.target
+        searched: Final = any(
+            isinstance(message["content"], list)
+            and any(block["type"] == "tool_result" for block in message["content"])
+            for message in body["messages"]
+        )
+        content: Final = (
+            [{"type": "text", "text": "the long answer"}]
+            if searched
+            else [{"type": "tool_use", "id": "toolu_0", "name": "litellm_web_search", "input": {"query": "q"}}]
+        )
+        return Reply(
+            body=json.dumps(
+                {
+                    "id": "msg_max_tokens",
+                    "type": "message",
+                    "role": "assistant",
+                    "model": bedrock_model,
+                    "content": content,
+                    "stop_reason": "end_turn" if searched else "tool_use",
+                    "stop_sequence": None,
+                    "usage": {"input_tokens": 10, "output_tokens": 4},
+                }
+            ).encode()
+        )
+
+    with wire_server(respond) as wire:
+        config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+        config["search_tools"] = [
+            {
+                "search_tool_name": "integration-search",
+                "litellm_params": {
+                    "search_provider": "tavily",
+                    "api_key": "synthetic-tavily-key",
+                    "api_base": wire.url + "/tavily",
+                },
+            }
+        ]
+        config["litellm_settings"].update(
+            {
+                "callbacks": ["websearch_interception"],
+                "websearch_interception_params": {
+                    "enabled_providers": ["bedrock"],
+                    "search_tool_name": "integration-search",
+                },
+            }
+        )
+        path: Final = tmp_path / "websearch-max-tokens.yaml"
+        path.write_text(yaml.safe_dump(config))
+        with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario:
+            model: Final = scenario.model(
+                model=f"bedrock/{bedrock_model}",
+                api_key="synthetic-bedrock-token",
+                api_base=wire.url,
+                aws_region_name="us-east-1",
+                aws_bedrock_runtime_endpoint=wire.url,
+            )
+            response: Final = candidate.request(
+                "POST",
+                "/v1/messages",
+                {
+                    "model": model,
+                    "max_tokens": 4000,
+                    "stream": stream,
+                    "messages": [{"role": "user", "content": "search and answer at length"}],
+                    "tools": [{"type": "web_search_20250305", "name": "web_search"}],
+                    **({"thinking": thinking} if thinking else {}),
+                },
+            )
+        assert response.status_code == 200, response.text
+        assert "the long answer" in response.text, response.text
+        upstream: Final = wire.drain()
+        assert [request.target for request in upstream] == [invoke_target, SEARCH_TARGET, invoke_target]
+        assert [json.loads(request.body)["max_tokens"] for request in upstream if request.target == invoke_target] == [
+            4000,
+            4000,
+        ]

tests/unit/llms/custom_httpx/test_llm_http_handler.py

@@ -7,7 +7,7 @@ import threading
 import time
 from collections.abc import Awaitable, Callable, Mapping, Sequence
 from types import MappingProxyType
-from typing import Final
+from typing import Final, Literal
 from unittest.mock import AsyncMock, Mock, patch
 
 import httpcore
@@ -2671,6 +2671,153 @@ class _FakeWebsocketsModule:
         return outcome
 
 
+_LOOKUP_CALL: Final = {"type": "tool_use", "id": "toolu_lookup", "name": "lookup", "input": {"q": "answer"}}
+_LOOKUP_TOOL: Final = {
+    "name": "lookup",
+    "description": "look a fact up",
+    "input_schema": {"type": "object", "properties": {"q": {"type": "string"}}},
+}
+
+
+def _answer_after_the_lookup(request: httpx.Request) -> httpx.Response:
+    looked_up: Final = b'"tool_result"' in request.content
+    return httpx.Response(
+        200,
+        json={
+            "id": "msg_lookup",
+            "type": "message",
+            "role": "assistant",
+            "model": "claude-haiku-5-5",
+            "content": [{"type": "text", "text": "the answer"}] if looked_up else [_LOOKUP_CALL],
+            "stop_reason": "end_turn" if looked_up else "tool_use",
+            "stop_sequence": None,
+            "usage": {"input_tokens": 5, "output_tokens": 3},
+        },
+    )
+
+
+def _messages_with_the_lookup_result(messages: list[dict[str, object]]) -> list[dict[str, object]]:
+    return [
+        *messages,
+        {"role": "assistant", "content": [_LOOKUP_CALL]},
+        {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_lookup", "content": "42"}]},
+    ]
+
+
+def _without_the_logging_obj(kwargs: Mapping[str, object]) -> dict[str, object]:
+    return {key: value for key, value in kwargs.items() if key != "litellm_logging_obj"}
+
+
+class _LookupLoop(CustomLogger):
+    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 isinstance(response, dict) and response.get("stop_reason") == "tool_use", {}
+
+
+class _OwnMaxTokensLoop(_LookupLoop):
+    def __init__(self, spreads: Literal["params", "kwargs"]) -> None:
+        super().__init__()
+        self.spreads: Final = spreads
+
+    async def async_run_agentic_loop(
+        self,
+        tools: dict[str, object],
+        model: str,
+        messages: list[dict[str, object]],
+        response: object,
+        anthropic_messages_provider_config: object,
+        anthropic_messages_optional_request_params: dict[str, object],
+        logging_obj: object,
+        stream: bool,
+        kwargs: dict[str, object],
+    ) -> object:
+        spread: Final = (
+            {"api_key": "sk-ant-test", **anthropic_messages_optional_request_params}
+            if self.spreads == "params"
+            else _without_the_logging_obj(kwargs)
+        )
+        return await litellm.anthropic.messages.acreate(
+            model="anthropic/claude-haiku-5-5",
+            messages=_messages_with_the_lookup_result(messages),
+            max_tokens=300,
+            **spread,  # pyright: ignore[reportArgumentType]  # a hook forwards the request's params as they arrive
+        )
+
+
+class _PlannedLoop(_LookupLoop):
+    def __init__(self, max_tokens: int | None, kwargs_max_tokens: int | None = None) -> None:
+        super().__init__()
+        self.max_tokens: Final = max_tokens
+        self.kwargs_max_tokens: Final = kwargs_max_tokens
+
+    async def async_build_agentic_loop_plan(
+        self,
+        tools: dict[str, object],
+        model: str,
+        messages: list[dict[str, object]],
+        response: object,
+        anthropic_messages_provider_config: object,
+        anthropic_messages_optional_request_params: dict[str, object],
+        logging_obj: object,
+        stream: bool,
+        kwargs: dict[str, object],
+    ) -> AgenticLoopPlan:
+        own_kwargs: Final = {} if self.kwargs_max_tokens is None else {"max_tokens": self.kwargs_max_tokens}
+        return AgenticLoopPlan(
+            run_agentic_loop=True,
+            request_patch=AgenticLoopRequestPatch(
+                model="anthropic/claude-haiku-5-5",
+                messages=_messages_with_the_lookup_result(messages),
+                max_tokens=self.max_tokens,
+                kwargs={**_without_the_logging_obj(kwargs), **own_kwargs},
+            ),
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("hook", "follow_up_max_tokens"),
+    [
+        pytest.param(_OwnMaxTokensLoop(spreads="params"), 300, id="hook-spreads-the-params-with-its-own-max-tokens"),
+        pytest.param(_OwnMaxTokensLoop(spreads="kwargs"), 300, id="hook-spreads-the-kwargs-with-its-own-max-tokens"),
+        pytest.param(_PlannedLoop(max_tokens=None), 4000, id="plan-leaves-max-tokens-to-the-client"),
+        pytest.param(_PlannedLoop(max_tokens=300), 300, id="plan-sets-its-own-max-tokens"),
+        pytest.param(_PlannedLoop(max_tokens=None, kwargs_max_tokens=300), 300, id="plan-sets-max-tokens-in-kwargs"),
+    ],
+)
+async def test_agentic_follow_up_uses_the_client_max_tokens_unless_the_hook_sets_its_own(
+    hook: CustomLogger,
+    follow_up_max_tokens: int,
+    respx_mock: respx.MockRouter,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setattr(litellm, "callbacks", [hook])
+    upstream: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
+        side_effect=_answer_after_the_lookup
+    )
+
+    result: Final = await litellm.anthropic.messages.acreate(
+        model="anthropic/claude-haiku-5-5",
+        messages=[{"role": "user", "content": "look the answer up"}],
+        max_tokens=4000,
+        api_key="sk-ant-test",
+        tools=[_LOOKUP_TOOL],
+    )
+
+    assert isinstance(result, dict)
+    assert result["content"] == [{"type": "text", "text": "the answer"}]
+    assert [json.loads(call.request.content)["max_tokens"] for call in upstream.calls] == [4000, follow_up_max_tokens]
+
+
 @pytest.mark.asyncio
 async def test_realtime_backend_open_retries_then_succeeds():
     """A hung/slow open handshake is retried; a later fresh attempt connects.