diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index dd2135f4918..21e39c857ef 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -72,8 +72,18 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): verbose_logger.debug("Transformed request: %s", request_body) - # Get endpoint URL - endpoint_url: Final = api_base or self.get_anthropic_count_tokens_endpoint() + # Construct clean endpoint URL handling custom proxy destinations + if api_base: + if api_base.endswith("/messages/count_tokens"): + endpoint_url = api_base + elif api_base.endswith("/v1/messages"): + endpoint_url = f"{api_base.rstrip('/')}/count_tokens" + elif api_base.endswith("/v1"): + endpoint_url = f"{api_base.rstrip('/')}/messages/count_tokens" + else: + endpoint_url = f"{api_base.rstrip('/')}/v1/messages/count_tokens" + else: + endpoint_url = self.get_anthropic_count_tokens_endpoint() verbose_logger.debug("Making request to: %s", endpoint_url) diff --git a/tests/test_anthropic_count_tokens.py b/tests/test_anthropic_count_tokens.py new file mode 100644 index 00000000000..c14a0feb685 --- /dev/null +++ b/tests/test_anthropic_count_tokens.py @@ -0,0 +1,32 @@ +from unittest.mock import AsyncMock, patch + +import pytest + +from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler + + +@pytest.mark.asyncio +async def test_anthropic_count_tokens_custom_api_base(): + handler = AnthropicCountTokensHandler() + + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"input_tokens": 10} + + with patch("litellm.llms.anthropic.count_tokens.handler.get_async_httpx_client") as mock_client_factory: + mock_client = AsyncMock() + mock_client.post.return_value = mock_response + mock_client_factory.return_value = mock_client + + # Test custom api_base URL resolution and formatting + res = await handler.handle_count_tokens_request( + model="claude-3-5-sonnet-20241022", + messages=[{"role": "user", "content": "hello"}], + api_key="test-key", + api_base="https://custom-proxy.internal.com/v1", + ) + + assert res["input_tokens"] == 10 + mock_client.post.assert_called_once() + called_url = mock_client.post.call_args[0][0] + assert called_url == "https://custom-proxy.internal.com/v1/messages/count_tokens"