diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index d41647fe5f1..65b813fccb6 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/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, diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 922de386fe4..518c39df560 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index df82891f789..aaea3735259 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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 = { diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 755a0a8cc7a..6bcaf7ee8b7 100644 --- a/litellm/types/utils.py +++ b/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. diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py index 1bd9f18253e..9b3422c160b 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py +++ b/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, + ] 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 a758f4d4250..0bf9ace2667 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/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.