diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d1a78534da..bd44cef954 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -131,6 +131,7 @@ else: unified_guardrail = UnifiedLLMGuardrails() +_anthropic_async_clients = {} def print_verbose(print_statement): """ @@ -4254,11 +4255,16 @@ async def count_tokens_with_anthropic_api( if anthropic_api_key and messages: # Call Anthropic API directly for more accurate token counting - client = anthropic.Anthropic(api_key=anthropic_api_key) + + # Use cached client if available to avoid socket exhaustion + if anthropic_api_key not in _anthropic_async_clients: + _anthropic_async_clients[anthropic_api_key] = anthropic.AsyncAnthropic(api_key=anthropic_api_key) + + client = _anthropic_async_clients[anthropic_api_key] # Call with explicit parameters to satisfy type checking # Type ignore for now since messages come from generic dict input - response = client.beta.messages.count_tokens( + response = await client.beta.messages.count_tokens( model=model_to_use, messages=messages, # type: ignore betas=["token-counting-2024-11-01"], diff --git a/tests/test_litellm/test_utils_custom.py b/tests/test_litellm/test_utils_custom.py new file mode 100644 index 0000000000..292da4132b --- /dev/null +++ b/tests/test_litellm/test_utils_custom.py @@ -0,0 +1,42 @@ +import pytest +from unittest.mock import MagicMock, patch, AsyncMock +from litellm.proxy.utils import count_tokens_with_anthropic_api, _anthropic_async_clients + +@pytest.mark.asyncio +async def test_count_tokens_caching(): + """ + Test that count_tokens_with_anthropic_api caches the client. + """ + # Clear cache + _anthropic_async_clients.clear() + + api_key = "sk-ant-test-key" + messages = [{"role": "user", "content": "hello"}] + model = "claude-3-opus-20240229" + + # Mock anthropic + with patch("anthropic.AsyncAnthropic") as mock_cls: + mock_client = MagicMock() + mock_cls.return_value = mock_client + + # Mock response + mock_response = MagicMock() + mock_response.input_tokens = 10 + + # Setup async return for count_tokens + mock_client.beta.messages.count_tokens = AsyncMock(return_value=mock_response) + + # First call + with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}): + await count_tokens_with_anthropic_api(model, messages) + + assert api_key in _anthropic_async_clients + assert _anthropic_async_clients[api_key] == mock_client + mock_cls.assert_called_once() # Should be called once + + # Second call + with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}): + await count_tokens_with_anthropic_api(model, messages) + + # Should still be called once (cached) + mock_cls.assert_called_once()