diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 1db94e82066..8ffdd5f4b4f 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1284,16 +1284,25 @@ class WebSearchInterceptionLogger(CustomLogger): def _resolve_max_tokens( optional_params: dict, kwargs: dict, + logging_obj: "LiteLLMLoggingObj | None" = None, ) -> int: """Extract max_tokens and validate against thinking.budget_tokens. Anthropic API requires ``max_tokens > thinking.budget_tokens``. If the constraint is violated, auto-adjust to ``budget_tokens + 1024``. """ - max_tokens: int = optional_params.get( - "max_tokens", - kwargs.get("max_tokens", 1024), - ) + # Look in optional_params, kwargs, or the original model_call_details / litellm_params + max_tokens_val = optional_params.get("max_tokens") + if max_tokens_val is None: + max_tokens_val = kwargs.get("max_tokens") + if max_tokens_val is None and logging_obj is not None: + model_details = getattr(logging_obj, "model_call_details", {}) or {} + max_tokens_val = model_details.get("max_tokens") + if max_tokens_val is None: + litellm_params = model_details.get("litellm_params", {}) or {} + max_tokens_val = litellm_params.get("max_tokens") + + max_tokens: int = max_tokens_val if max_tokens_val is not None else 1024 thinking_param: Final = optional_params.get("thinking") if thinking_param and isinstance(thinking_param, dict): budget_tokens: Final = thinking_param.get("budget_tokens") @@ -1437,7 +1446,9 @@ class WebSearchInterceptionLogger(CustomLogger): full_model_name = model # safe default before try block - 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, logging_obj=logging_obj + ) verbose_logger.debug("WebSearchInterception: Using max_tokens=%s for follow-up request", max_tokens) 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..47082324a0f 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py @@ -491,3 +491,38 @@ class TestFollowUpErrorScenarios: # But ALL proxy metadata must be preserved assert captured_kwargs.get("metadata") == proxy_metadata assert captured_kwargs.get("litellm_call_id") == "call-abc-123" + + @pytest.mark.asyncio + async def test_max_tokens_resolved_from_logging_obj_when_popped_from_optional_params(self): + """When max_tokens was popped from optional_params (Messages transformation), it is resolved from logging_obj.""" + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + captured_kwargs: Dict[str, Any] = {} + + async def _fake_acreate(**kw): + captured_kwargs.update(kw) + return MagicMock() + + logging_obj = _make_logging_obj() + logging_obj.model_call_details["max_tokens"] = 4000 + + 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={}, # max_tokens already popped by transformation + logging_obj=logging_obj, + stream=False, + kwargs={}, + ) + + assert captured_kwargs.get("max_tokens") == 4000