diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 0579eb15baf..4804364d09f 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 d41647fe5f1..342d29c70cb 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1396,7 +1396,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 4a10db31606..990642f5ec7 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2586,7 +2586,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), @@ -5958,12 +5963,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: @@ -6160,12 +6165,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 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/integrations/websearch_interception/test_websearch_thinking_constraint.py b/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py index ee4c4d32478..d64858082f6 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 dd0e5762845..77922f9fcf8 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -2671,6 +2671,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.