init token counters

litellm/llms/anthropic/count_tokens/__init__.py

@@ -0,0 +1,15 @@
+"""
+Anthropic CountTokens API implementation.
+"""
+
+from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
+from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter
+from litellm.llms.anthropic.count_tokens.transformation import (
+    AnthropicCountTokensConfig,
+)
+
+__all__ = [
+    "AnthropicCountTokensHandler",
+    "AnthropicCountTokensConfig",
+    "AnthropicTokenCounter",
+]

litellm/llms/anthropic/count_tokens/handler.py

@@ -0,0 +1,126 @@
+"""
+Anthropic CountTokens API handler.
+
+Uses httpx for HTTP requests instead of the Anthropic SDK.
+"""
+
+from typing import Any, Dict, List, Optional, Union
+
+import httpx
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.llms.anthropic.common_utils import AnthropicError
+from litellm.llms.anthropic.count_tokens.transformation import (
+    AnthropicCountTokensConfig,
+)
+from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+
+
+class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
+    """
+    Handler for Anthropic CountTokens API requests.
+
+    Uses httpx for HTTP requests, following the same pattern as BedrockCountTokensHandler.
+    """
+
+    async def handle_count_tokens_request(
+        self,
+        model: str,
+        messages: List[Dict[str, Any]],
+        api_key: str,
+        api_base: Optional[str] = None,
+        timeout: Optional[Union[float, httpx.Timeout]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Handle a CountTokens request using httpx.
+
+        Args:
+            model: The model identifier (e.g., "claude-3-5-sonnet-20241022")
+            messages: The messages to count tokens for
+            api_key: The Anthropic API key
+            api_base: Optional custom API base URL
+            timeout: Optional timeout for the request (defaults to litellm.request_timeout)
+
+        Returns:
+            Dictionary containing token count response
+
+        Raises:
+            AnthropicError: If the API request fails
+        """
+        try:
+            # Validate the request
+            self.validate_request(model, messages)
+
+            verbose_logger.debug(
+                f"Processing Anthropic CountTokens request for model: {model}"
+            )
+
+            # Transform request to Anthropic format
+            request_body = self.transform_request_to_count_tokens(
+                model=model,
+                messages=messages,
+            )
+
+            verbose_logger.debug(f"Transformed request: {request_body}")
+
+            # Get endpoint URL
+            endpoint_url = api_base or self.get_anthropic_count_tokens_endpoint()
+
+            verbose_logger.debug(f"Making request to: {endpoint_url}")
+
+            # Get required headers
+            headers = self.get_required_headers(api_key)
+
+            # Use LiteLLM's async httpx client
+            async_client = get_async_httpx_client(
+                llm_provider=litellm.LlmProviders.ANTHROPIC
+            )
+
+            # Use provided timeout or fall back to litellm.request_timeout
+            request_timeout = timeout if timeout is not None else litellm.request_timeout
+
+            response = await async_client.post(
+                endpoint_url,
+                headers=headers,
+                json=request_body,
+                timeout=request_timeout,
+            )
+
+            verbose_logger.debug(f"Response status: {response.status_code}")
+
+            if response.status_code != 200:
+                error_text = response.text
+                verbose_logger.error(f"Anthropic API error: {error_text}")
+                raise AnthropicError(
+                    status_code=response.status_code,
+                    message=error_text,
+                )
+
+            anthropic_response = response.json()
+
+            verbose_logger.debug(f"Anthropic response: {anthropic_response}")
+
+            # Transform response
+            final_response = self.transform_response(anthropic_response)
+
+            verbose_logger.debug(f"Final response: {final_response}")
+
+            return final_response
+
+        except AnthropicError:
+            # Re-raise Anthropic exceptions as-is
+            raise
+        except httpx.HTTPStatusError as e:
+            # HTTP errors - preserve the actual status code
+            verbose_logger.error(f"HTTP error in CountTokens handler: {str(e)}")
+            raise AnthropicError(
+                status_code=e.response.status_code,
+                message=e.response.text,
+            )
+        except Exception as e:
+            verbose_logger.error(f"Error in CountTokens handler: {str(e)}")
+            raise AnthropicError(
+                status_code=500,
+                message=f"CountTokens processing error: {str(e)}",
+            )

litellm/llms/anthropic/count_tokens/token_counter.py

@@ -0,0 +1,104 @@
+"""
+Anthropic Token Counter implementation using the CountTokens API.
+"""
+
+import os
+from typing import Any, Dict, List, Optional
+
+from litellm._logging import verbose_logger
+from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
+from litellm.llms.base_llm.base_utils import BaseTokenCounter
+from litellm.types.utils import LlmProviders, TokenCountResponse
+
+# Global handler instance - reuse across all token counting requests
+anthropic_count_tokens_handler = AnthropicCountTokensHandler()
+
+
+class AnthropicTokenCounter(BaseTokenCounter):
+    """Token counter implementation for Anthropic provider using the CountTokens API."""
+
+    def should_use_token_counting_api(
+        self,
+        custom_llm_provider: Optional[str] = None,
+    ) -> bool:
+        return custom_llm_provider == LlmProviders.ANTHROPIC.value
+
+    async def count_tokens(
+        self,
+        model_to_use: str,
+        messages: Optional[List[Dict[str, Any]]],
+        contents: Optional[List[Dict[str, Any]]],
+        deployment: Optional[Dict[str, Any]] = None,
+        request_model: str = "",
+    ) -> Optional[TokenCountResponse]:
+        """
+        Count tokens using Anthropic's CountTokens API.
+
+        Args:
+            model_to_use: The model identifier
+            messages: The messages to count tokens for
+            contents: Alternative content format (not used for Anthropic)
+            deployment: Deployment configuration containing litellm_params
+            request_model: The original request model name
+
+        Returns:
+            TokenCountResponse with token count, or None if counting fails
+        """
+        from litellm.llms.anthropic.common_utils import AnthropicError
+
+        if not messages:
+            return None
+
+        deployment = deployment or {}
+        litellm_params = deployment.get("litellm_params", {})
+
+        # Get Anthropic API key from deployment config or environment
+        api_key = litellm_params.get("api_key")
+        if not api_key:
+            api_key = os.getenv("ANTHROPIC_API_KEY")
+
+        if not api_key:
+            verbose_logger.warning("No Anthropic API key found for token counting")
+            return None
+
+        try:
+            result = await anthropic_count_tokens_handler.handle_count_tokens_request(
+                model=model_to_use,
+                messages=messages,
+                api_key=api_key,
+            )
+
+            if result is not None:
+                return TokenCountResponse(
+                    total_tokens=result.get("input_tokens", 0),
+                    request_model=request_model,
+                    model_used=model_to_use,
+                    tokenizer_type="anthropic_api",
+                    original_response=result,
+                )
+        except AnthropicError as e:
+            verbose_logger.warning(
+                f"Anthropic CountTokens API error: status={e.status_code}, message={e.message}"
+            )
+            return TokenCountResponse(
+                total_tokens=0,
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type="anthropic_api",
+                error=True,
+                error_message=e.message,
+                status_code=e.status_code,
+            )
+        except Exception as e:
+            verbose_logger.warning(f"Error calling Anthropic CountTokens API: {e}")
+            return TokenCountResponse(
+                total_tokens=0,
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type="anthropic_api",
+                error=True,
+                error_message=str(e),
+                status_code=500,
+            )
+
+        return None

litellm/llms/anthropic/count_tokens/transformation.py

@@ -0,0 +1,121 @@
+"""
+Anthropic CountTokens API transformation logic.
+
+This module handles the transformation of requests to Anthropic's CountTokens API format.
+"""
+
+from typing import Any, Dict, List
+
+from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
+
+
+class AnthropicCountTokensConfig:
+    """
+    Configuration and transformation logic for Anthropic CountTokens API.
+
+    Anthropic CountTokens API Specification:
+    - Endpoint: POST https://api.anthropic.com/v1/messages/count_tokens
+    - Beta header required: anthropic-beta: token-counting-2024-11-01
+    - Response: {"input_tokens": <number>}
+    """
+
+    def get_anthropic_count_tokens_endpoint(self) -> str:
+        """
+        Get the Anthropic CountTokens API endpoint.
+
+        Returns:
+            The endpoint URL for the CountTokens API
+        """
+        return "https://api.anthropic.com/v1/messages/count_tokens"
+
+    def transform_request_to_count_tokens(
+        self,
+        model: str,
+        messages: List[Dict[str, Any]],
+    ) -> Dict[str, Any]:
+        """
+        Transform request to Anthropic CountTokens format.
+
+        Input:
+        {
+            "model": "claude-3-5-sonnet-20241022",
+            "messages": [{"role": "user", "content": "Hello!"}]
+        }
+
+        Output (Anthropic CountTokens format):
+        {
+            "model": "claude-3-5-sonnet-20241022",
+            "messages": [{"role": "user", "content": "Hello!"}]
+        }
+        """
+        return {
+            "model": model,
+            "messages": messages,
+        }
+
+    def transform_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Transform Anthropic CountTokens response.
+
+        Input (Anthropic response):
+        {
+            "input_tokens": 123
+        }
+
+        Output:
+        {
+            "input_tokens": 123
+        }
+        """
+        return {
+            "input_tokens": response.get("input_tokens", 0),
+        }
+
+    def get_required_headers(self, api_key: str) -> Dict[str, str]:
+        """
+        Get the required headers for the CountTokens API.
+
+        Args:
+            api_key: The Anthropic API key
+
+        Returns:
+            Dictionary of required headers
+        """
+        return {
+            "Content-Type": "application/json",
+            "x-api-key": api_key,
+            "anthropic-version": "2023-06-01",
+            "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
+        }
+
+    def validate_request(
+        self, model: str, messages: List[Dict[str, Any]]
+    ) -> None:
+        """
+        Validate the incoming count tokens request.
+
+        Args:
+            model: The model name
+            messages: The messages to count tokens for
+
+        Raises:
+            ValueError: If the request is invalid
+        """
+        if not model:
+            raise ValueError("model parameter is required")
+
+        if not messages:
+            raise ValueError("messages parameter is required")
+
+        if not isinstance(messages, list):
+            raise ValueError("messages must be a list")
+
+        for i, message in enumerate(messages):
+            if not isinstance(message, dict):
+                raise ValueError(f"Message {i} must be a dictionary")
+
+            if "role" not in message:
+                raise ValueError(f"Message {i} must have a 'role' field")
+
+            if "content" not in message:
+                raise ValueError(f"Message {i} must have a 'content' field")

litellm/llms/azure_ai/anthropic/count_tokens/__init__.py

@@ -0,0 +1,19 @@
+"""
+Azure AI Anthropic CountTokens API implementation.
+"""
+
+from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
+    AzureAIAnthropicCountTokensHandler,
+)
+from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import (
+    AzureAIAnthropicTokenCounter,
+)
+from litellm.llms.azure_ai.anthropic.count_tokens.transformation import (
+    AzureAIAnthropicCountTokensConfig,
+)
+
+__all__ = [
+    "AzureAIAnthropicCountTokensHandler",
+    "AzureAIAnthropicCountTokensConfig",
+    "AzureAIAnthropicTokenCounter",
+]

litellm/llms/azure_ai/anthropic/count_tokens/handler.py

@@ -0,0 +1,131 @@
+"""
+Azure AI Anthropic CountTokens API handler.
+
+Uses httpx for HTTP requests with Azure authentication.
+"""
+
+from typing import Any, Dict, List, Optional, Union
+
+import httpx
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.llms.anthropic.common_utils import AnthropicError
+from litellm.llms.azure_ai.anthropic.count_tokens.transformation import (
+    AzureAIAnthropicCountTokensConfig,
+)
+from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+
+
+class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
+    """
+    Handler for Azure AI Anthropic CountTokens API requests.
+
+    Uses httpx for HTTP requests with Azure authentication.
+    """
+
+    async def handle_count_tokens_request(
+        self,
+        model: str,
+        messages: List[Dict[str, Any]],
+        api_key: str,
+        api_base: str,
+        litellm_params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[Union[float, httpx.Timeout]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Handle a CountTokens request using httpx with Azure authentication.
+
+        Args:
+            model: The model identifier (e.g., "claude-3-5-sonnet")
+            messages: The messages to count tokens for
+            api_key: The Azure AI API key
+            api_base: The Azure AI API base URL
+            litellm_params: Optional LiteLLM parameters
+            timeout: Optional timeout for the request (defaults to litellm.request_timeout)
+
+        Returns:
+            Dictionary containing token count response
+
+        Raises:
+            AnthropicError: If the API request fails
+        """
+        try:
+            # Validate the request
+            self.validate_request(model, messages)
+
+            verbose_logger.debug(
+                f"Processing Azure AI Anthropic CountTokens request for model: {model}"
+            )
+
+            # Transform request to Anthropic format
+            request_body = self.transform_request_to_count_tokens(
+                model=model,
+                messages=messages,
+            )
+
+            verbose_logger.debug(f"Transformed request: {request_body}")
+
+            # Get endpoint URL
+            endpoint_url = self.get_count_tokens_endpoint(api_base)
+
+            verbose_logger.debug(f"Making request to: {endpoint_url}")
+
+            # Get required headers with Azure authentication
+            headers = self.get_required_headers(
+                api_key=api_key,
+                litellm_params=litellm_params,
+            )
+
+            # Use LiteLLM's async httpx client
+            async_client = get_async_httpx_client(
+                llm_provider=litellm.LlmProviders.AZURE_AI
+            )
+
+            # Use provided timeout or fall back to litellm.request_timeout
+            request_timeout = timeout if timeout is not None else litellm.request_timeout
+
+            response = await async_client.post(
+                endpoint_url,
+                headers=headers,
+                json=request_body,
+                timeout=request_timeout,
+            )
+
+            verbose_logger.debug(f"Response status: {response.status_code}")
+
+            if response.status_code != 200:
+                error_text = response.text
+                verbose_logger.error(f"Azure AI Anthropic API error: {error_text}")
+                raise AnthropicError(
+                    status_code=response.status_code,
+                    message=error_text,
+                )
+
+            azure_response = response.json()
+
+            verbose_logger.debug(f"Azure AI Anthropic response: {azure_response}")
+
+            # Transform response
+            final_response = self.transform_response(azure_response)
+
+            verbose_logger.debug(f"Final response: {final_response}")
+
+            return final_response
+
+        except AnthropicError:
+            # Re-raise Anthropic exceptions as-is
+            raise
+        except httpx.HTTPStatusError as e:
+            # HTTP errors - preserve the actual status code
+            verbose_logger.error(f"HTTP error in CountTokens handler: {str(e)}")
+            raise AnthropicError(
+                status_code=e.response.status_code,
+                message=e.response.text,
+            )
+        except Exception as e:
+            verbose_logger.error(f"Error in CountTokens handler: {str(e)}")
+            raise AnthropicError(
+                status_code=500,
+                message=f"CountTokens processing error: {str(e)}",
+            )

litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py

@@ -0,0 +1,119 @@
+"""
+Azure AI Anthropic Token Counter implementation using the CountTokens API.
+"""
+
+import os
+from typing import Any, Dict, List, Optional
+
+from litellm._logging import verbose_logger
+from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
+    AzureAIAnthropicCountTokensHandler,
+)
+from litellm.llms.base_llm.base_utils import BaseTokenCounter
+from litellm.types.utils import LlmProviders, TokenCountResponse
+
+# Global handler instance - reuse across all token counting requests
+azure_ai_anthropic_count_tokens_handler = AzureAIAnthropicCountTokensHandler()
+
+
+class AzureAIAnthropicTokenCounter(BaseTokenCounter):
+    """Token counter implementation for Azure AI Anthropic provider using the CountTokens API."""
+
+    def should_use_token_counting_api(
+        self,
+        custom_llm_provider: Optional[str] = None,
+    ) -> bool:
+        return custom_llm_provider == LlmProviders.AZURE_AI.value
+
+    async def count_tokens(
+        self,
+        model_to_use: str,
+        messages: Optional[List[Dict[str, Any]]],
+        contents: Optional[List[Dict[str, Any]]],
+        deployment: Optional[Dict[str, Any]] = None,
+        request_model: str = "",
+    ) -> Optional[TokenCountResponse]:
+        """
+        Count tokens using Azure AI Anthropic's CountTokens API.
+
+        Args:
+            model_to_use: The model identifier
+            messages: The messages to count tokens for
+            contents: Alternative content format (not used for Anthropic)
+            deployment: Deployment configuration containing litellm_params
+            request_model: The original request model name
+
+        Returns:
+            TokenCountResponse with token count, or None if counting fails
+        """
+        from litellm.llms.anthropic.common_utils import AnthropicError
+
+        if not messages:
+            return None
+
+        deployment = deployment or {}
+        litellm_params = deployment.get("litellm_params", {})
+
+        # Get Azure AI API key from deployment config or environment
+        api_key = litellm_params.get("api_key")
+        if not api_key:
+            api_key = os.getenv("AZURE_AI_API_KEY")
+
+        # Get API base from deployment config or environment
+        api_base = litellm_params.get("api_base")
+        if not api_base:
+            api_base = os.getenv("AZURE_AI_API_BASE")
+
+        if not api_key:
+            verbose_logger.warning("No Azure AI API key found for token counting")
+            return None
+
+        if not api_base:
+            verbose_logger.warning("No Azure AI API base found for token counting")
+            return None
+
+        try:
+            result = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request(
+                model=model_to_use,
+                messages=messages,
+                api_key=api_key,
+                api_base=api_base,
+                litellm_params=litellm_params,
+            )
+
+            if result is not None:
+                return TokenCountResponse(
+                    total_tokens=result.get("input_tokens", 0),
+                    request_model=request_model,
+                    model_used=model_to_use,
+                    tokenizer_type="azure_ai_anthropic_api",
+                    original_response=result,
+                )
+        except AnthropicError as e:
+            verbose_logger.warning(
+                f"Azure AI Anthropic CountTokens API error: status={e.status_code}, message={e.message}"
+            )
+            return TokenCountResponse(
+                total_tokens=0,
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type="azure_ai_anthropic_api",
+                error=True,
+                error_message=e.message,
+                status_code=e.status_code,
+            )
+        except Exception as e:
+            verbose_logger.warning(
+                f"Error calling Azure AI Anthropic CountTokens API: {e}"
+            )
+            return TokenCountResponse(
+                total_tokens=0,
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type="azure_ai_anthropic_api",
+                error=True,
+                error_message=str(e),
+                status_code=500,
+            )
+
+        return None

litellm/llms/azure_ai/anthropic/count_tokens/transformation.py

@@ -0,0 +1,88 @@
+"""
+Azure AI Anthropic CountTokens API transformation logic.
+
+Extends the base Anthropic CountTokens transformation with Azure authentication.
+"""
+
+from typing import Any, Dict, Optional
+
+from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
+from litellm.llms.anthropic.count_tokens.transformation import (
+    AnthropicCountTokensConfig,
+)
+from litellm.llms.azure.common_utils import BaseAzureLLM
+from litellm.types.router import GenericLiteLLMParams
+
+
+class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
+    """
+    Configuration and transformation logic for Azure AI Anthropic CountTokens API.
+
+    Extends AnthropicCountTokensConfig with Azure authentication.
+    Azure AI Anthropic uses the same endpoint format but with Azure auth headers.
+    """
+
+    def get_required_headers(
+        self,
+        api_key: str,
+        litellm_params: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, str]:
+        """
+        Get the required headers for the Azure AI Anthropic CountTokens API.
+
+        Uses Azure authentication (api-key header) instead of Anthropic's x-api-key.
+
+        Args:
+            api_key: The Azure AI API key
+            litellm_params: Optional LiteLLM parameters for additional auth config
+
+        Returns:
+            Dictionary of required headers with Azure authentication
+        """
+        # Start with base headers
+        headers = {
+            "Content-Type": "application/json",
+            "anthropic-version": "2023-06-01",
+            "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
+        }
+
+        # Use Azure authentication
+        litellm_params = litellm_params or {}
+        if "api_key" not in litellm_params:
+            litellm_params["api_key"] = api_key
+
+        litellm_params_obj = GenericLiteLLMParams(**litellm_params)
+
+        # Get Azure auth headers
+        azure_headers = BaseAzureLLM._base_validate_azure_environment(
+            headers={}, litellm_params=litellm_params_obj
+        )
+
+        # Merge Azure auth headers
+        headers.update(azure_headers)
+
+        return headers
+
+    def get_count_tokens_endpoint(self, api_base: str) -> str:
+        """
+        Get the Azure AI Anthropic CountTokens API endpoint.
+
+        Args:
+            api_base: The Azure AI API base URL 
+                      (e.g., https://my-resource.services.ai.azure.com or
+                       https://my-resource.services.ai.azure.com/anthropic)
+
+        Returns:
+            The endpoint URL for the CountTokens API
+        """
+        # Azure AI Anthropic endpoint format:
+        # https://<resource>.services.ai.azure.com/anthropic/v1/messages/count_tokens
+        api_base = api_base.rstrip("/")
+
+        # Ensure the URL has /anthropic path
+        if not api_base.endswith("/anthropic"):
+            if "/anthropic" not in api_base:
+                api_base = f"{api_base}/anthropic"
+
+        # Add the count_tokens path
+        return f"{api_base}/v1/messages/count_tokens"

litellm/llms/azure_ai/common_utils.py

@@ -4,123 +4,6 @@ import litellm
 from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import AllMessageValues
-from litellm.types.utils import TokenCountResponse
-
-
-class AzureAIAnthropicTokenCounter(BaseTokenCounter):
-    """Token counter implementation for Azure AI Anthropic provider using the CountTokens API."""
-
-    def should_use_token_counting_api(
-        self,
-        custom_llm_provider: Optional[str] = None,
-    ) -> bool:
-        from litellm.types.utils import LlmProviders
-
-        return custom_llm_provider == LlmProviders.AZURE_AI.value
-
-    async def count_tokens(
-        self,
-        model_to_use: str,
-        messages: Optional[List[Dict[str, Any]]],
-        contents: Optional[List[Dict[str, Any]]],
-        deployment: Optional[Dict[str, Any]] = None,
-        request_model: str = "",
-    ) -> Optional[TokenCountResponse]:
-        """
-        Count tokens using Azure AI Anthropic's CountTokens API.
-
-        Args:
-            model_to_use: The model identifier
-            messages: The messages to count tokens for
-            contents: Alternative content format (not used for Anthropic)
-            deployment: Deployment configuration containing litellm_params
-            request_model: The original request model name
-
-        Returns:
-            TokenCountResponse with token count, or None if counting fails
-        """
-        import os
-
-        from litellm._logging import verbose_logger
-        from litellm.llms.anthropic.common_utils import AnthropicError
-        from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
-            AzureAIAnthropicCountTokensHandler,
-        )
-
-        if not messages:
-            return None
-
-        deployment = deployment or {}
-        litellm_params = deployment.get("litellm_params", {})
-
-        # Get Azure AI API key from deployment config or environment
-        api_key = litellm_params.get("api_key")
-        if not api_key:
-            api_key = os.getenv("AZURE_AI_API_KEY")
-
-        # Get API base from deployment config or environment
-        api_base = litellm_params.get("api_base")
-        if not api_base:
-            api_base = os.getenv("AZURE_AI_API_BASE")
-
-        if not api_key:
-            verbose_logger.warning(
-                "No Azure AI API key found for token counting"
-            )
-            return None
-
-        if not api_base:
-            verbose_logger.warning(
-                "No Azure AI API base found for token counting"
-            )
-            return None
-
-        try:
-            handler = AzureAIAnthropicCountTokensHandler()
-            result = await handler.handle_count_tokens_request(
-                model=model_to_use,
-                messages=messages,
-                api_key=api_key,
-                api_base=api_base,
-                litellm_params=litellm_params,
-            )
-
-            if result is not None:
-                return TokenCountResponse(
-                    total_tokens=result.get("input_tokens", 0),
-                    request_model=request_model,
-                    model_used=model_to_use,
-                    tokenizer_type="azure_ai_anthropic_api",
-                    original_response=result,
-                )
-        except AnthropicError as e:
-            verbose_logger.warning(
-                f"Azure AI Anthropic CountTokens API error: status={e.status_code}, message={e.message}"
-            )
-            return TokenCountResponse(
-                total_tokens=0,
-                request_model=request_model,
-                model_used=model_to_use,
-                tokenizer_type="azure_ai_anthropic_api",
-                error=True,
-                error_message=e.message,
-                status_code=e.status_code,
-            )
-        except Exception as e:
-            verbose_logger.warning(
-                f"Error calling Azure AI Anthropic CountTokens API: {e}"
-            )
-            return TokenCountResponse(
-                total_tokens=0,
-                request_model=request_model,
-                model_used=model_to_use,
-                tokenizer_type="azure_ai_anthropic_api",
-                error=True,
-                error_message=str(e),
-                status_code=500,
-            )
-
-        return None
 
 
 class AzureFoundryModelInfo(BaseLLMModelInfo):
@@ -142,9 +25,7 @@ class AzureFoundryModelInfo(BaseLLMModelInfo):
 
     @staticmethod
     def get_api_base(api_base: Optional[str] = None) -> Optional[str]:
-        return (
-            api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE")
-        )
+        return api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE")
 
     @staticmethod
     def get_api_key(api_key: Optional[str] = None) -> Optional[str]:
@@ -171,6 +52,10 @@ class AzureFoundryModelInfo(BaseLLMModelInfo):
         """
         # Only return token counter for Claude models
         if self._model and "claude" in self._model.lower():
+            from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import (
+                AzureAIAnthropicTokenCounter,
+            )
+
             return AzureAIAnthropicTokenCounter()
         return None