scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

fix: avoid duplicate api_key/api_base kwargs in agentic follow-up calls

Browse source 
When request_patch.kwargs already contains api_key or api_base,
explicitly passing them again caused TypeError (multiple values for
same keyword arg). Pop conflicting keys from followup kwargs before
conditionally injecting credentials from agentic_loop_params.

Also adds tests for:
- websearch interception legacy path credential forwarding
- count_tokens custom api_base URL building
- llm_http_handler _execute_anthropic_agentic_plan credential forwarding

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jesset 2026-05-03 16:10:23 +08:00
parent 394b015f5e
commit e6cdcb0188
5 changed files with 422 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -766,14 +766,22 @@ class WebSearchInterceptionLogger(CustomLogger):
_followup_api_key = agentic_params.get("api_key")
_followup_api_base = agentic_params.get("api_base")
# Avoid duplicate keyword arguments if request_patch.kwargs already
# contains api_key or api_base.
followup_kwargs = dict(request_patch.kwargs)
if _followup_api_key is not None:
followup_kwargs.pop("api_key", None)
if _followup_api_base is not None:
followup_kwargs.pop("api_base", None)
return await anthropic_messages.acreate(
max_tokens=max_tokens,
messages=request_patch.messages,
model=request_patch.model or model,
api_key=_followup_api_key,
api_base=_followup_api_base,
**({"api_key": _followup_api_key} if _followup_api_key else {}),
**({"api_base": _followup_api_base} if _followup_api_base else {}),
**optional_params,
**request_patch.kwargs,
**followup_kwargs,
)
async def _build_anthropic_request_patch(

7
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -4679,6 +4679,13 @@ class BaseLLMHTTPHandler:
kwargs_for_followup["max_agentic_loops"] = max_loops
kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint]
# Avoid duplicate keyword arguments if patch.kwargs already contains
# api_key or api_base, so agentic credentials take precedence.
if agentic_api_key is not None:
kwargs_for_followup.pop("api_key", None)
if agentic_api_base is not None:
kwargs_for_followup.pop("api_base", None)
return await anthropic_messages.acreate(
**{
"max_tokens": max_tokens,

91
tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py
View file

@ -45,6 +45,21 @@ def _make_logging_obj(
return obj
def _make_logging_obj_with_credentials(
model: str = "bedrock/us.anthropic.claude-opus-4-6-v1",
) -> MagicMock:
obj = MagicMock()
obj.model_call_details = {
"agentic_loop_params": {
"model": model,
"custom_llm_provider": "bedrock",
"api_key": "sk-test-agentic-key",
"api_base": "https://custom.provider.host/plan/anthropic",
},
}
return obj
# ---------------------------------------------------------------------------
# M1-I1 / M1-I3: max_tokens validation against thinking.budget_tokens
# ---------------------------------------------------------------------------
@ -473,3 +488,79 @@ 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"
# ---------------------------------------------------------------------------
# Credentials forwarding from agentic_loop_params
# ---------------------------------------------------------------------------
class TestAgenticLoopCredentialsForwarding:
"""Verify api_key and api_base from agentic_loop_params are forwarded
to the follow-up acreate() call so third-party Anthropic-compatible
providers (e.g. Tencent Cloud) receive the correct endpoint."""
@pytest.mark.asyncio
async def test_api_key_and_api_base_forwarded_from_agentic_loop_params(self):
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"),
):
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": 4096},
logging_obj=_make_logging_obj_with_credentials(),
stream=False,
kwargs={},
)
assert captured_kwargs.get("api_key") == "sk-test-agentic-key"
assert (
captured_kwargs.get("api_base")
== "https://custom.provider.host/plan/anthropic"
)
@pytest.mark.asyncio
async def test_no_credentials_forwarded_when_agentic_params_missing(self):
"""When agentic_loop_params lacks api_key/api_base, neither key should appear
in the follow-up call kwargs (avoiding None overrides)."""
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"),
):
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": 4096},
logging_obj=_make_logging_obj(),
stream=False,
kwargs={},
)
assert "api_key" not in captured_kwargs
assert "api_base" not in captured_kwargs

161
tests/test_litellm/llms/anthropic/count_tokens/test_count_tokens_api_base.py Normal file
View file

@ -0,0 +1,161 @@
"""
Tests for Anthropic CountTokens API custom api_base handling.
Verifies that AnthropicTokenCounter correctly builds the count_tokens endpoint
from a custom api_base instead of hardcoding api.anthropic.com.
"""
import os
import sys
from typing import Optional
from unittest.mock import patch
import pytest
sys.path.insert(
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
)
from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter
class TestCountTokensApiBase:
"""Verify api_base is correctly extracted from litellm_params and used to build endpoint URL."""
@pytest.mark.asyncio
async def test_custom_api_base_builds_correct_count_tokens_url(self):
"""When api_base is provided in litellm_params, count_tokens endpoint should use it."""
counter = AnthropicTokenCounter()
captured_api_base: Optional[str] = None
async def _fake_handle_count_tokens_request(
model, messages, api_key, api_base=None, **kwargs
):
nonlocal captured_api_base
captured_api_base = api_base
return {"input_tokens": 42}
with patch(
"litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler.handle_count_tokens_request",
side_effect=_fake_handle_count_tokens_request,
):
result = await counter.count_tokens(
model_to_use="claude-3-5-sonnet",
messages=[{"role": "user", "content": "hello"}],
contents=None,
deployment={
"litellm_params": {
"api_key": "sk-test-key",
"api_base": "https://api.lkeap.cloud.tencent.com/plan/anthropic",
}
},
request_model="claude-3-5-sonnet",
)
assert (
captured_api_base
== "https://api.lkeap.cloud.tencent.com/plan/anthropic/v1/messages/count_tokens"
)
assert result is not None
assert result.total_tokens == 42
@pytest.mark.asyncio
async def test_api_base_without_trailing_slash(self):
"""api_base without trailing slash should still produce correct URL."""
counter = AnthropicTokenCounter()
captured_api_base: Optional[str] = None
async def _fake_handle_count_tokens_request(
model, messages, api_key, api_base=None, **kwargs
):
nonlocal captured_api_base
captured_api_base = api_base
return {"input_tokens": 10}
with patch(
"litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler.handle_count_tokens_request",
side_effect=_fake_handle_count_tokens_request,
):
await counter.count_tokens(
model_to_use="claude-3-5-sonnet",
messages=[{"role": "user", "content": "hi"}],
contents=None,
deployment={
"litellm_params": {
"api_key": "sk-test-key",
"api_base": "https://custom.host/plan/anthropic",
}
},
request_model="claude-3-5-sonnet",
)
assert (
captured_api_base
== "https://custom.host/plan/anthropic/v1/messages/count_tokens"
)
@pytest.mark.asyncio
async def test_no_api_base_uses_default_endpoint(self):
"""When api_base is not provided, api_base should be None (handler will use default)."""
counter = AnthropicTokenCounter()
captured_api_base: Optional[str] = None
async def _fake_handle_count_tokens_request(
model, messages, api_key, api_base=None, **kwargs
):
nonlocal captured_api_base
captured_api_base = api_base
return {"input_tokens": 5}
with patch(
"litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler.handle_count_tokens_request",
side_effect=_fake_handle_count_tokens_request,
):
await counter.count_tokens(
model_to_use="claude-3-5-sonnet",
messages=[{"role": "user", "content": "hi"}],
contents=None,
deployment={
"litellm_params": {
"api_key": "sk-test-key",
}
},
request_model="claude-3-5-sonnet",
)
assert captured_api_base is None
@pytest.mark.asyncio
async def test_api_base_with_trailing_slash_stripped(self):
"""api_base with trailing slash should have it stripped before appending path."""
counter = AnthropicTokenCounter()
captured_api_base: Optional[str] = None
async def _fake_handle_count_tokens_request(
model, messages, api_key, api_base=None, **kwargs
):
nonlocal captured_api_base
captured_api_base = api_base
return {"input_tokens": 7}
with patch(
"litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler.handle_count_tokens_request",
side_effect=_fake_handle_count_tokens_request,
):
await counter.count_tokens(
model_to_use="claude-3-5-sonnet",
messages=[{"role": "user", "content": "hi"}],
contents=None,
deployment={
"litellm_params": {
"api_key": "sk-test-key",
"api_base": "https://custom.host/plan/anthropic/",
}
},
request_model="claude-3-5-sonnet",
)
assert (
captured_api_base
== "https://custom.host/plan/anthropic/v1/messages/count_tokens"
)

153
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -1,7 +1,8 @@
import asyncio
import os
import sys
from unittest.mock import AsyncMock, Mock, patch
from typing import Any, Dict
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import httpx
import pytest
@ -422,3 +423,153 @@ def test_sync_delete_responses_omits_body_for_azure():
assert captured["url"].endswith(
"/openai/responses/resp_xyz?api-version=2025-03-01-preview"
)
class TestExecuteAnthropicAgenticPlan:
"""Verify _execute_anthropic_agentic_plan forwards api_key/api_base
from agentic_loop_params to the follow-up anthropic_messages.acreate call."""
@pytest.mark.asyncio
async def test_api_key_and_api_base_forwarded_to_followup(self):
handler = BaseLLMHTTPHandler()
captured_kwargs: Dict[str, Any] = {}
async def _fake_acreate(**kw):
captured_kwargs.update(kw)
return MagicMock()
from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
AgenticLoopRequestPatch,
)
plan = AgenticLoopPlan(
run_agentic_loop=True,
request_patch=AgenticLoopRequestPatch(
model=None,
messages=[{"role": "user", "content": "test"}],
max_tokens=1024,
),
)
logging_obj = MagicMock()
logging_obj.model_call_details = {
"agentic_loop_params": {
"model": "custom/model-name",
"api_key": "sk-agentic-test-key",
"api_base": "https://custom.provider.host/v1",
}
}
with patch(
"litellm.anthropic_interface.messages.acreate",
side_effect=_fake_acreate,
):
await handler._execute_anthropic_agentic_plan(
plan=plan,
model="claude-3-5-sonnet",
messages=[{"role": "user", "content": "test"}],
anthropic_messages_optional_request_params={},
kwargs={},
logging_obj=logging_obj,
depth=0,
max_loops=1,
fingerprints=[],
fingerprint="fp-test",
)
assert captured_kwargs.get("api_key") == "sk-agentic-test-key"
assert captured_kwargs.get("api_base") == "https://custom.provider.host/v1"
assert captured_kwargs.get("model") == "custom/model-name"
@pytest.mark.asyncio
async def test_no_credentials_when_agentic_loop_params_empty(self):
handler = BaseLLMHTTPHandler()
captured_kwargs: Dict[str, Any] = {}
async def _fake_acreate(**kw):
captured_kwargs.update(kw)
return MagicMock()
from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
AgenticLoopRequestPatch,
)
plan = AgenticLoopPlan(
run_agentic_loop=True,
request_patch=AgenticLoopRequestPatch(
model=None,
messages=[{"role": "user", "content": "test"}],
max_tokens=1024,
),
)
logging_obj = MagicMock()
logging_obj.model_call_details = {
"agentic_loop_params": {"model": "custom/model-name"}
}
with patch(
"litellm.anthropic_interface.messages.acreate",
side_effect=_fake_acreate,
):
await handler._execute_anthropic_agentic_plan(
plan=plan,
model="claude-3-5-sonnet",
messages=[{"role": "user", "content": "test"}],
anthropic_messages_optional_request_params={},
kwargs={},
logging_obj=logging_obj,
depth=0,
max_loops=1,
fingerprints=[],
fingerprint="fp-test",
)
assert "api_key" not in captured_kwargs
assert "api_base" not in captured_kwargs
@pytest.mark.asyncio
async def test_no_logging_obj_uses_original_model(self):
handler = BaseLLMHTTPHandler()
captured_kwargs: Dict[str, Any] = {}
async def _fake_acreate(**kw):
captured_kwargs.update(kw)
return MagicMock()
from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
AgenticLoopRequestPatch,
)
plan = AgenticLoopPlan(
run_agentic_loop=True,
request_patch=AgenticLoopRequestPatch(
model=None,
messages=[{"role": "user", "content": "test"}],
max_tokens=1024,
),
)
with patch(
"litellm.anthropic_interface.messages.acreate",
side_effect=_fake_acreate,
):
await handler._execute_anthropic_agentic_plan(
plan=plan,
model="original-model-name",
messages=[{"role": "user", "content": "test"}],
anthropic_messages_optional_request_params={},
kwargs={},
logging_obj=None,
depth=0,
max_loops=1,
fingerprints=[],
fingerprint="fp-test",
)
assert captured_kwargs.get("model") == "original-model-name"
assert "api_key" not in captured_kwargs
assert "api_base" not in captured_kwargs