From aec8299d8b3e0664aca9323855cf38db4b738104 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 20 Jan 2026 12:22:02 -0800 Subject: [PATCH] init test suite for token counter --- .../base_token_counter_test.py | 130 ++++++++++++++++++ .../test_anthropic_token_counter.py | 47 +++++++ .../test_azure_ai_anthropic_token_counter.py | 53 +++++++ .../test_bedrock_token_counter.py | 101 ++++++++++++++ .../test_proxy_token_counter.py | 83 ++++++----- 5 files changed, 380 insertions(+), 34 deletions(-) create mode 100644 tests/litellm_utils_tests/base_token_counter_test.py create mode 100644 tests/litellm_utils_tests/test_anthropic_token_counter.py create mode 100644 tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py create mode 100644 tests/litellm_utils_tests/test_bedrock_token_counter.py diff --git a/tests/litellm_utils_tests/base_token_counter_test.py b/tests/litellm_utils_tests/base_token_counter_test.py new file mode 100644 index 00000000000..b5e87021a0b --- /dev/null +++ b/tests/litellm_utils_tests/base_token_counter_test.py @@ -0,0 +1,130 @@ +""" +Base Token Counter Test Suite. + +This module provides an abstract base test class that enforces common tests +across all token counter implementations. Similar to base_llm_unit_tests.py +for LLM chat tests. + +Usage: + Create a test class that inherits from BaseTokenCounterTest and implement + the abstract methods to provide provider-specific configuration. +""" + +import os +import sys +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.types.utils import TokenCountResponse + + +class BaseTokenCounterTest(ABC): + """ + Abstract base test class for token counter implementations. + + Subclasses must implement: + - get_token_counter(): Returns the token counter instance + - get_test_model(): Returns the model name to use for testing + - get_test_messages(): Returns test messages for token counting + - get_deployment_config(): Returns deployment configuration with credentials + - get_custom_llm_provider(): Returns the provider name for should_use_token_counting_api + """ + + @abstractmethod + def get_token_counter(self) -> BaseTokenCounter: + """Must return the token counter instance to test.""" + pass + + @abstractmethod + def get_test_model(self) -> str: + """Must return the model name to use for testing.""" + pass + + @abstractmethod + def get_test_messages(self) -> List[Dict[str, Any]]: + """Must return test messages for token counting.""" + pass + + @abstractmethod + def get_deployment_config(self) -> Dict[str, Any]: + """Must return deployment configuration with credentials.""" + pass + + @abstractmethod + def get_custom_llm_provider(self) -> str: + """Must return the provider name for should_use_token_counting_api check.""" + pass + + @pytest.fixture(autouse=True) + def _handle_missing_credentials(self): + """Fixture to skip tests when credentials are missing.""" + try: + yield + except Exception as e: + error_str = str(e).lower() + if "api key" in error_str or "api_key" in error_str or "unauthorized" in error_str: + pytest.skip(f"Missing or invalid credentials: {e}") + raise + + @pytest.mark.asyncio + async def test_count_tokens_basic(self): + """ + Test basic token counting functionality. + + Verifies that: + - Token counter returns a TokenCountResponse + - total_tokens is greater than 0 + - tokenizer_type is set + - No error occurred + """ + token_counter = self.get_token_counter() + model = self.get_test_model() + messages = self.get_test_messages() + deployment = self.get_deployment_config() + + result = await token_counter.count_tokens( + model_to_use=model, + messages=messages, + contents=None, + deployment=deployment, + request_model=model, + ) + + print(f"Token count result: {result}") + + assert result is not None, "Token counter should return a result" + assert isinstance(result, TokenCountResponse), "Result should be TokenCountResponse" + assert result.total_tokens > 0, f"Token count should be > 0, got {result.total_tokens}" + assert result.tokenizer_type is not None, "tokenizer_type should be set" + assert result.error is not True, f"Token counting should not error: {result.error_message}" + + def test_should_use_token_counting_api(self): + """ + Test that should_use_token_counting_api returns True for the correct provider. + + Verifies that the token counter correctly identifies when it should be used + based on the custom_llm_provider. + """ + token_counter = self.get_token_counter() + provider = self.get_custom_llm_provider() + + result = token_counter.should_use_token_counting_api( + custom_llm_provider=provider + ) + + assert result is True, f"should_use_token_counting_api should return True for {provider}" + + # Also verify it returns False for other providers + other_provider = "some_other_provider_that_doesnt_exist" + result_other = token_counter.should_use_token_counting_api( + custom_llm_provider=other_provider + ) + + assert result_other is False, f"should_use_token_counting_api should return False for {other_provider}" diff --git a/tests/litellm_utils_tests/test_anthropic_token_counter.py b/tests/litellm_utils_tests/test_anthropic_token_counter.py new file mode 100644 index 00000000000..a1fbcecfdd4 --- /dev/null +++ b/tests/litellm_utils_tests/test_anthropic_token_counter.py @@ -0,0 +1,47 @@ +""" +Anthropic Token Counter Tests. + +Tests for the Anthropic token counter implementation using the base test suite. +""" + +import os +import sys +from typing import Any, Dict, List + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from litellm.llms.anthropic.count_tokens import AnthropicTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from tests.litellm_utils_tests.base_token_counter_test import BaseTokenCounterTest + + +class TestAnthropicTokenCounter(BaseTokenCounterTest): + """Test suite for Anthropic token counter.""" + + def get_token_counter(self) -> BaseTokenCounter: + return AnthropicTokenCounter() + + def get_test_model(self) -> str: + return "claude-sonnet-4-20250514" + + def get_test_messages(self) -> List[Dict[str, Any]]: + return [ + {"role": "user", "content": "Hello, how are you today?"} + ] + + def get_deployment_config(self) -> Dict[str, Any]: + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not set") + return { + "litellm_params": { + "api_key": api_key, + } + } + + def get_custom_llm_provider(self) -> str: + return "anthropic" diff --git a/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py b/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py new file mode 100644 index 00000000000..031502cbece --- /dev/null +++ b/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py @@ -0,0 +1,53 @@ +""" +Azure AI Anthropic Token Counter Tests. + +Tests for the Azure AI Anthropic token counter implementation using the base test suite. +""" + +import os +import sys +from typing import Any, Dict, List + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from litellm.llms.azure_ai.anthropic.count_tokens import AzureAIAnthropicTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from tests.litellm_utils_tests.base_token_counter_test import BaseTokenCounterTest + + +class TestAzureAIAnthropicTokenCounter(BaseTokenCounterTest): + """Test suite for Azure AI Anthropic token counter.""" + + def get_token_counter(self) -> BaseTokenCounter: + return AzureAIAnthropicTokenCounter() + + def get_test_model(self) -> str: + return "claude-3-5-sonnet" + + def get_test_messages(self) -> List[Dict[str, Any]]: + return [ + {"role": "user", "content": "Hello, how are you today?"} + ] + + def get_deployment_config(self) -> Dict[str, Any]: + api_key = os.getenv("AZURE_AI_API_KEY") + api_base = os.getenv("AZURE_AI_API_BASE") + + if not api_key: + pytest.skip("AZURE_AI_API_KEY not set") + if not api_base: + pytest.skip("AZURE_AI_API_BASE not set") + + return { + "litellm_params": { + "api_key": api_key, + "api_base": api_base, + } + } + + def get_custom_llm_provider(self) -> str: + return "azure_ai" diff --git a/tests/litellm_utils_tests/test_bedrock_token_counter.py b/tests/litellm_utils_tests/test_bedrock_token_counter.py new file mode 100644 index 00000000000..f7c29918820 --- /dev/null +++ b/tests/litellm_utils_tests/test_bedrock_token_counter.py @@ -0,0 +1,101 @@ +""" +Bedrock Token Counter Tests. + +Tests for the Bedrock token counter implementation using the base test suite. + +Note: Not all Bedrock models support token counting. The CountTokens API +is only available for specific models. If the model doesn't support token +counting, the test will be skipped. +""" + +import os +import sys +from typing import Any, Dict, List + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter +from tests.litellm_utils_tests.base_token_counter_test import BaseTokenCounterTest + + +class TestBedrockTokenCounter(BaseTokenCounterTest): + """Test suite for Bedrock token counter. + + Note: Bedrock CountTokens API support varies by model. Some models + (like older Claude versions) may not support token counting. + Use amazon.nova-* models for reliable token counting support. + """ + + def get_token_counter(self) -> BaseTokenCounter: + return BedrockTokenCounter() + + def get_test_model(self) -> str: + # Use Amazon Nova model which supports token counting + # Alternatively, use environment variable to override + return os.getenv("BEDROCK_TEST_MODEL", "amazon.nova-lite-v1:0") + + def get_test_messages(self) -> List[Dict[str, Any]]: + return [ + {"role": "user", "content": "Hello, how are you today?"} + ] + + def get_deployment_config(self) -> Dict[str, Any]: + # Bedrock uses AWS credentials from environment + # Check for AWS credentials + aws_access_key = os.getenv("AWS_ACCESS_KEY_ID") + aws_secret_key = os.getenv("AWS_SECRET_ACCESS_KEY") + aws_region = os.getenv("AWS_REGION_NAME", "us-east-1") + + if not aws_access_key or not aws_secret_key: + pytest.skip("AWS credentials not set (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)") + + return { + "litellm_params": { + "aws_access_key_id": aws_access_key, + "aws_secret_access_key": aws_secret_key, + "aws_region_name": aws_region, + } + } + + def get_custom_llm_provider(self) -> str: + return "bedrock" + + @pytest.mark.asyncio + async def test_count_tokens_basic(self): + """ + Test basic token counting functionality. + + Override to handle models that don't support token counting. + """ + from litellm.types.utils import TokenCountResponse + + token_counter = self.get_token_counter() + model = self.get_test_model() + messages = self.get_test_messages() + deployment = self.get_deployment_config() + + result = await token_counter.count_tokens( + model_to_use=model, + messages=messages, + contents=None, + deployment=deployment, + request_model=model, + ) + + print(f"Token count result: {result}") + + assert result is not None, "Token counter should return a result" + assert isinstance(result, TokenCountResponse), "Result should be TokenCountResponse" + + # Check if the model doesn't support token counting + if result.error and "doesn't support counting tokens" in str(result.error_message): + pytest.skip(f"Model {model} doesn't support token counting: {result.error_message}") + + assert result.total_tokens > 0, f"Token count should be > 0, got {result.total_tokens}" + assert result.tokenizer_type is not None, "tokenizer_type should be set" + assert result.error is not True, f"Token counting should not error: {result.error_message}" diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index 8e1057bb8e1..4a2ff28cfd3 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -478,18 +478,19 @@ async def test_anthropic_endpoint_error_handling(): @pytest.mark.asyncio async def test_factory_anthropic_endpoint_calls_anthropic_counter(): """Test that /v1/messages/count_tokens with Anthropic model uses Anthropic counter.""" - from unittest.mock import patch, AsyncMock + from unittest.mock import patch, AsyncMock, MagicMock from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app - # Mock the anthropic token counting function + # Mock the Anthropic CountTokens handler with patch( - "litellm.proxy.utils.count_tokens_with_anthropic_api" - ) as mock_anthropic_count: - mock_anthropic_count.return_value = { - "total_tokens": 42, - "tokenizer_used": "anthropic", - } + "litellm.llms.anthropic.common_utils.AnthropicCountTokensHandler" + ) as MockHandler: + mock_handler_instance = MagicMock() + mock_handler_instance.handle_count_tokens_request = AsyncMock( + return_value={"input_tokens": 42} + ) + MockHandler.return_value = mock_handler_instance # Mock router to return Anthropic deployment with patch("litellm.proxy.proxy_server.llm_router") as mock_router: @@ -510,36 +511,44 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): } ) - client = TestClient(app) + # Set ANTHROPIC_API_KEY for the test + with patch.dict("os.environ", {"ANTHROPIC_API_KEY": "test-key"}): + client = TestClient(app) - response = client.post( - "/v1/messages/count_tokens", - json={ - "model": "claude-3-5-sonnet", - "messages": [{"role": "user", "content": "Hello"}], - }, - headers={"Authorization": "Bearer test-key"}, - ) + response = client.post( + "/v1/messages/count_tokens", + json={ + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "Hello"}], + }, + headers={"Authorization": "Bearer test-key"}, + ) - assert response.status_code == 200 - data = response.json() - assert data["input_tokens"] == 42 + assert response.status_code == 200 + data = response.json() + assert data["input_tokens"] == 42 - # Verify that Anthropic API was called - mock_anthropic_count.assert_called_once() + # Verify that Anthropic handler was called + mock_handler_instance.handle_count_tokens_request.assert_called_once() @pytest.mark.asyncio async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): """Test that /v1/messages/count_tokens with GPT-4 does NOT use Anthropic counter.""" - from unittest.mock import patch, AsyncMock + from unittest.mock import patch, AsyncMock, MagicMock from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app - # Mock the anthropic token counting function + # Mock the Anthropic CountTokens handler with patch( - "litellm.proxy.utils.count_tokens_with_anthropic_api" - ) as mock_anthropic_count: + "litellm.llms.anthropic.common_utils.AnthropicCountTokensHandler" + ) as MockHandler: + mock_handler_instance = MagicMock() + mock_handler_instance.handle_count_tokens_request = AsyncMock( + return_value={"input_tokens": 42} + ) + MockHandler.return_value = mock_handler_instance + # Mock litellm token counter with patch("litellm.token_counter") as mock_litellm_counter: mock_litellm_counter.return_value = 50 @@ -578,21 +587,27 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): data = response.json() assert data["input_tokens"] == 50 - # Verify that Anthropic API was NOT called - mock_anthropic_count.assert_not_called() + # Verify that Anthropic handler was NOT called + mock_handler_instance.handle_count_tokens_request.assert_not_called() @pytest.mark.asyncio async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): """Test that /utils/token_counter does NOT use Anthropic counter even with Anthropic model.""" - from unittest.mock import patch, AsyncMock + from unittest.mock import patch, AsyncMock, MagicMock from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app - # Mock the anthropic token counting function + # Mock the Anthropic CountTokens handler with patch( - "litellm.proxy.utils.count_tokens_with_anthropic_api" - ) as mock_anthropic_count: + "litellm.llms.anthropic.common_utils.AnthropicCountTokensHandler" + ) as MockHandler: + mock_handler_instance = MagicMock() + mock_handler_instance.handle_count_tokens_request = AsyncMock( + return_value={"input_tokens": 42} + ) + MockHandler.return_value = mock_handler_instance + # Mock litellm token counter with patch("litellm.token_counter") as mock_litellm_counter: mock_litellm_counter.return_value = 35 @@ -635,8 +650,8 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): data = response.json() assert data["total_tokens"] == 35 - # Verify that Anthropic API was NOT called (since call_endpoint=False) - mock_anthropic_count.assert_not_called() + # Verify that Anthropic handler was NOT called (since call_endpoint=False) + mock_handler_instance.handle_count_tokens_request.assert_not_called() @pytest.mark.asyncio