mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-11 17:05:43 +00:00
Constantine
3ebec39b74
Fixes #16716. Previously, synchronous Anthropic client was used for token counting, which blocked the event loop. This change switches to AsyncAnthropic and caches the client instance.
43 lines
1.5 KiB
Python
43 lines
1.5 KiB
Python
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()
|