scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-23 05:34:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 25

fix(websearch_interception): preserve client max_tokens in /v1/messages follow-up call (#44834)

Browse source
This commit is contained in:
İlker Oğulcan Sayın 2026-10-06 13:03:42 +03:00
parent 7a7d27c550
commit f4f6aa4f60
No known key found for this signature in database
2 changed files with 51 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
litellm/integrations/websearch_interception/handler.py
View file

@ -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)

35
tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py
View file

@ -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