diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index d45e46cfab3..f00dda2d516 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -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( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2e6417974cb..f3919573d41 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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, diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py index a939951c430..c8cd405ce7f 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/count_tokens/test_count_tokens_api_base.py b/tests/test_litellm/llms/anthropic/count_tokens/test_count_tokens_api_base.py new file mode 100644 index 00000000000..eda79e12628 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/count_tokens/test_count_tokens_api_base.py @@ -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" + ) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index b846cd600f0..de00e809c13 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -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