litellm/tests/litellm_utils_tests/base_token_counter_test.py

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

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional

import pytest


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