From a4a140a12230e24d2cee362fcc07d04731ca99a4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 04:06:08 -0700 Subject: [PATCH] fix(anthropic_messages): pass the client's max_tokens to agentic hooks in kwargs Revert the shallow copy of the request params and add the client's max_tokens to the hook kwargs instead, where the built-in hooks already fall back. A hook that passes its own max_tokens while spreading the params dict no longer fails, and the web search hook no longer sees the pre-transform thinking budget on adaptive models. A plan's own max_tokens still wins in the plan runner. --- litellm/integrations/custom_logger.py | 6 +- .../websearch_interception/handler.py | 4 +- litellm/llms/custom_httpx/llm_http_handler.py | 21 +-- .../test_websearch_interception_wire.py | 31 +++- .../test_websearch_thinking_constraint.py | 29 ++++ .../custom_httpx/test_llm_http_handler.py | 132 ++++++++++++++++++ 6 files changed, 204 insertions(+), 19 deletions(-) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a4de1be1842..f9303ce6602 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -703,7 +703,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac anthropic_messages_optional_request_params: Request parameters (tools, etc.) logging_obj: LiteLLM logging object stream: Whether response is streaming - kwargs: Additional request parameters + kwargs: Additional request parameters, including the client's ``max_tokens`` Returns: Final response after executing agentic loop @@ -727,8 +727,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac final_response = await messages.acreate( model=model, messages=messages + [assistant_msg, user_msg], - max_tokens=anthropic_messages_optional_request_params.get("max_tokens"), - **anthropic_messages_optional_request_params + max_tokens=kwargs["max_tokens"], + **{k: v for k, v in anthropic_messages_optional_request_params.items() if k != "max_tokens"} ) return final_response diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 6e1ba3eea08..9f53ddca614 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1377,7 +1377,9 @@ class WebSearchInterceptionLogger(CustomLogger): if max_tokens is None: max_tokens = cast(int, kwargs.get("max_tokens", 1024)) - patch_kwargs: Final = dict[str, object](request_patch.kwargs) + patch_kwargs: Final = { + key: value for key, value in dict[str, object](request_patch.kwargs).items() if key != "max_tokens" + } response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate( max_tokens=max_tokens, messages=request_patch.messages, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 77b9dfe90ee..c5170a17c7c 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2486,7 +2486,12 @@ class BaseLLMHTTPHandler: # Prepare headers kwargs = kwargs or {} - kwargs_for_agentic: Final = self._agentic_hook_kwargs(kwargs=kwargs, api_key=api_key, api_base=api_base) + kwargs_for_agentic: Final = self._agentic_hook_kwargs( + kwargs=kwargs, + api_key=api_key, + api_base=api_base, + max_tokens=_OPTIONAL_PARAMS.validate_python(anthropic_messages_optional_request_params).get("max_tokens"), + ) provider_specific_header: Final = cast( litellm.types.utils.ProviderSpecificHeader | Sequence[litellm.types.utils.ProviderSpecificHeader] | None, kwargs.get("provider_specific_header", None), @@ -2580,7 +2585,7 @@ class BaseLLMHTTPHandler: request_body: Final = anthropic_messages_provider_config.transform_anthropic_messages_request( model=model, messages=messages, - anthropic_messages_optional_request_params=dict(request_params), + anthropic_messages_optional_request_params=request_params, litellm_params=litellm_params, headers=headers, ) @@ -5781,12 +5786,12 @@ class BaseLLMHTTPHandler: @staticmethod def _agentic_hook_kwargs( - kwargs: Mapping[str, object], api_key: str | None, api_base: str | None + kwargs: Mapping[str, object], api_key: str | None, api_base: str | None, max_tokens: object ) -> dict[str, object]: - """``api_key`` and ``api_base`` are named parameters of ``anthropic_messages`` rather than kwargs, so the - follow-up call an agentic hook makes only reaches the same deployment if they are re-added here.""" - deployment_params: Final = {"api_key": api_key, "api_base": api_base} - return {**kwargs, **{key: value for key, value in deployment_params.items() if value}} + """The follow-up call an agentic hook makes needs ``api_key``, ``api_base`` and the client's ``max_tokens``. + None of them are in ``kwargs``, and the request transform pops ``max_tokens`` from the hooks' params dict.""" + named_params: Final = {"api_key": api_key, "api_base": api_base, "max_tokens": max_tokens} + return {**kwargs, **{key: value for key, value in named_params.items() if value}} @staticmethod def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool: @@ -6010,12 +6015,12 @@ class BaseLLMHTTPHandler: response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate( **{ - "max_tokens": max_tokens, "messages": patch.messages, "model": patch.model or full_model_name, "stream": stream, **optional_params, **kwargs_for_followup, + "max_tokens": max_tokens, } ) 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 a38722b54c5..7d8d61b02df 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 @@ -489,14 +489,30 @@ def test_empty_exa_search_reaches_the_model_without_the_exa_response_envelope(ga assert targets[-3:] == (("POST", "/v1/messages"), ("POST", "/exa/search"), ("POST", "/v1/messages")), targets -@pytest.mark.parametrize("stream", [False, True]) -def test_web_search_follow_up_call_keeps_the_client_max_tokens(gateway: Gateway, tmp_path: Path, stream: bool) -> None: +@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 + 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"]) @@ -513,7 +529,7 @@ def test_web_search_follow_up_call_keeps_the_client_max_tokens(gateway: Gateway, "id": "msg_max_tokens", "type": "message", "role": "assistant", - "model": BEDROCK_MODEL, + "model": bedrock_model, "content": content, "stop_reason": "end_turn" if searched else "tool_use", "stop_sequence": None, @@ -547,7 +563,7 @@ def test_web_search_follow_up_call_keeps_the_client_max_tokens(gateway: Gateway, 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}", + model=f"bedrock/{bedrock_model}", api_key="synthetic-bedrock-token", api_base=wire.url, aws_region_name="us-east-1", @@ -562,13 +578,14 @@ def test_web_search_follow_up_call_keeps_the_client_max_tokens(gateway: Gateway, "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] == [ + 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/integrations/websearch_interception/test_websearch_thinking_constraint.py b/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py index b2d5225070c..8c467b051a0 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py @@ -125,6 +125,35 @@ class TestThinkingBudgetTokensConstraint: assert captured_kwargs["max_tokens"] == 5000 + 1024 + @pytest.mark.asyncio + async def test_client_max_tokens_passed_in_kwargs_reaches_the_follow_up(self) -> None: + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + captured_kwargs: Dict[str, Any] = {} + + async def _fake_acreate(**kw): + captured_kwargs.update(kw) + return MagicMock() + + with ( + patch( + "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", + side_effect=_fake_acreate, + ), + patch.object(logger, "_execute_search", return_value=("search result", None)), + ): + await logger._execute_agentic_loop( + model="us.anthropic.claude-opus-4-6-v1", + messages=[{"role": "user", "content": "hi"}], + tool_calls=_make_tool_calls(), + thinking_blocks=[], + anthropic_messages_optional_request_params={}, + logging_obj=_make_logging_obj(), + stream=False, + kwargs={"max_tokens": 4000}, + ) + + assert captured_kwargs["max_tokens"] == 4000 + @pytest.mark.asyncio async def test_max_tokens_unchanged_when_greater_than_budget(self): """max_tokens > thinking.budget_tokens → no adjustment needed.""" 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 f10f5a67b2f..5c07f4da5ee 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -64,6 +64,7 @@ from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator, ) +from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( @@ -2666,6 +2667,137 @@ 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]) -> list[dict]: + return [ + *messages, + {"role": "assistant", "content": [_LOOKUP_CALL]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_lookup", "content": "42"}]}, + ] + + +class _LookupLoop(CustomLogger): + async def async_should_run_agentic_loop( + self, + response: object, + model: str, + messages: list[dict], + tools: list[dict] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict, + ) -> tuple[bool, dict]: + return isinstance(response, dict) and response.get("stop_reason") == "tool_use", {} + + +class _OwnMaxTokensLoop(_LookupLoop): + async def async_run_agentic_loop( + self, + tools: dict, + model: str, + messages: list[dict], + response: object, + anthropic_messages_provider_config: object, + anthropic_messages_optional_request_params: dict, + logging_obj: object, + stream: bool, + kwargs: dict, + ) -> object: + return await litellm.anthropic.messages.acreate( + model="anthropic/claude-haiku-5-5", + messages=_messages_with_the_lookup_result(messages), + max_tokens=300, + api_key="sk-ant-test", + **anthropic_messages_optional_request_params, + ) + + +class _PlannedLoop(_LookupLoop): + def __init__(self, max_tokens: int | None) -> None: + super().__init__() + self.max_tokens: Final = max_tokens + + async def async_build_agentic_loop_plan( + self, + tools: dict, + model: str, + messages: list[dict], + response: object, + anthropic_messages_provider_config: object, + anthropic_messages_optional_request_params: dict, + logging_obj: object, + stream: bool, + kwargs: dict, + ) -> AgenticLoopPlan: + 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={key: value for key, value in kwargs.items() if key != "litellm_logging_obj"}, + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("hook", "follow_up_max_tokens"), + [ + pytest.param(_OwnMaxTokensLoop(), 300, id="hook-passes-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"), + ], +) +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.