diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index f9303ce6602..a4de1be1842 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, including the client's ``max_tokens`` + kwargs: Additional request parameters 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=kwargs["max_tokens"], - **{k: v for k, v in anthropic_messages_optional_request_params.items() if k != "max_tokens"} + max_tokens=anthropic_messages_optional_request_params.get("max_tokens"), + **anthropic_messages_optional_request_params ) return final_response diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 9f53ddca614..27e16f44afa 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1296,6 +1296,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. @@ -1304,7 +1305,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): @@ -1377,9 +1378,7 @@ class WebSearchInterceptionLogger(CustomLogger): if max_tokens is None: max_tokens = cast(int, kwargs.get("max_tokens", 1024)) - patch_kwargs: Final = { - key: value for key, value in dict[str, object](request_patch.kwargs).items() if key != "max_tokens" - } + patch_kwargs: Final = dict[str, object](request_patch.kwargs) response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate( max_tokens=max_tokens, messages=request_patch.messages, @@ -1449,9 +1448,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) @@ -1460,11 +1466,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 6f97a019a96..02803242f97 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -636,6 +636,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 c5170a17c7c..cee04f71324 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2486,12 +2486,7 @@ 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, - max_tokens=_OPTIONAL_PARAMS.validate_python(anthropic_messages_optional_request_params).get("max_tokens"), - ) + kwargs_for_agentic: Final = self._agentic_hook_kwargs(kwargs=kwargs, api_key=api_key, api_base=api_base) provider_specific_header: Final = cast( litellm.types.utils.ProviderSpecificHeader | Sequence[litellm.types.utils.ProviderSpecificHeader] | None, kwargs.get("provider_specific_header", None), @@ -5786,12 +5781,12 @@ class BaseLLMHTTPHandler: @staticmethod def _agentic_hook_kwargs( - kwargs: Mapping[str, object], api_key: str | None, api_base: str | None, max_tokens: object + kwargs: Mapping[str, object], api_key: str | None, api_base: str | None ) -> dict[str, object]: - """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}} + """``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}} @staticmethod def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool: @@ -5981,10 +5976,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) @@ -5997,7 +5992,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 = { @@ -6015,12 +6010,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/litellm/types/utils.py b/litellm/types/utils.py index 70ca05904ef..aaf0ff990b8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -268,6 +268,9 @@ class AgenticLoopParams(TypedDict, total=False): custom_llm_provider: str """The LLM provider name (e.g., 'bedrock', 'anthropic')""" + 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/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py b/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py index 8c467b051a0..b2d5225070c 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py @@ -125,35 +125,6 @@ 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 5c07f4da5ee..c05c52daf94 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 from types import MappingProxyType -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, Mock, patch import httpx @@ -2692,7 +2692,7 @@ def _answer_after_the_lookup(request: httpx.Request) -> httpx.Response: ) -def _messages_with_the_lookup_result(messages: list[dict]) -> list[dict]: +def _messages_with_the_lookup_result(messages: list[dict[str, object]]) -> list[dict[str, object]]: return [ *messages, {"role": "assistant", "content": [_LOOKUP_CALL]}, @@ -2700,66 +2700,80 @@ def _messages_with_the_lookup_result(messages: list[dict]) -> list[dict]: ] +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], - tools: list[dict] | None, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, stream: bool, custom_llm_provider: str, - kwargs: dict, - ) -> tuple[bool, dict]: + 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, + tools: dict[str, object], model: str, - messages: list[dict], + messages: list[dict[str, object]], response: object, anthropic_messages_provider_config: object, - anthropic_messages_optional_request_params: dict, + anthropic_messages_optional_request_params: dict[str, object], logging_obj: object, stream: bool, - kwargs: dict, + 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, - api_key="sk-ant-test", - **anthropic_messages_optional_request_params, + **spread, # pyright: ignore[reportArgumentType] # a hook forwards the request's params as they arrive ) class _PlannedLoop(_LookupLoop): - def __init__(self, max_tokens: int | None) -> None: + 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, + tools: dict[str, object], model: str, - messages: list[dict], + messages: list[dict[str, object]], response: object, anthropic_messages_provider_config: object, - anthropic_messages_optional_request_params: dict, + anthropic_messages_optional_request_params: dict[str, object], logging_obj: object, stream: bool, - kwargs: dict, + 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={key: value for key, value in kwargs.items() if key != "litellm_logging_obj"}, + kwargs={**_without_the_logging_obj(kwargs), **own_kwargs}, ), ) @@ -2768,9 +2782,11 @@ class _PlannedLoop(_LookupLoop): @pytest.mark.parametrize( ("hook", "follow_up_max_tokens"), [ - pytest.param(_OwnMaxTokensLoop(), 300, id="hook-passes-its-own-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(