init test suite for token counter

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}"

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"

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"

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}"

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