diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index fcbe9823ed4..3a8da2a7fd6 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -476,45 +476,11 @@ class AnthropicModelInfo(BaseLLMModelInfo): Returns: AnthropicTokenCounter instance for this provider. """ - return AnthropicTokenCounter() - - -class AnthropicTokenCounter(BaseTokenCounter): - """Token counter implementation for Anthropic provider.""" - - 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.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]: - from litellm.proxy.utils import count_tokens_with_anthropic_api - - result = await count_tokens_with_anthropic_api( - model_to_use=model_to_use, - messages=messages, - deployment=deployment, + from litellm.llms.anthropic.count_tokens.token_counter import ( + AnthropicTokenCounter, ) - - if result is not None: - return TokenCountResponse( - total_tokens=result.get("total_tokens", 0), - request_model=request_model, - model_used=model_to_use, - tokenizer_type=result.get("tokenizer_used", ""), - original_response=result, - ) - - return None + + return AnthropicTokenCounter() def process_anthropic_headers(headers: Union[httpx.Headers, dict]) -> dict: diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 9487c7f83f2..1a88926ae39 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,17 +1,139 @@ -from typing import List, Literal, Optional +from typing import Any, Dict, List, Literal, Optional import litellm -from litellm.llms.base_llm.base_utils import BaseLLMModelInfo +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): + """Model info for Azure AI / Azure Foundry models.""" + + def __init__(self, model: Optional[str] = None): + self._model = model + @staticmethod def get_azure_ai_route(model: str) -> Literal["agents", "default"]: """ Get the Azure AI route for the given model. - + Similar to BedrockModelInfo.get_bedrock_route(). """ if "agents/" in model: @@ -21,33 +143,40 @@ 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") + 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]: return ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("AZURE_AI_API_KEY") - ) - + api_key + or litellm.api_key + or litellm.openai_key + or get_secret_str("AZURE_AI_API_KEY") + ) + @property def api_version(self, api_version: Optional[str] = None) -> Optional[str]: api_version = ( - api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") + api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") ) return api_version - + + def get_token_counter(self) -> Optional[BaseTokenCounter]: + """ + Factory method to create a token counter for Azure AI. + + Returns: + AzureAIAnthropicTokenCounter for Claude models, None otherwise. + """ + # Only return token counter for Claude models + if self._model and "claude" in self._model.lower(): + return AzureAIAnthropicTokenCounter() + return None + ######################################################### # Not implemented methods ######################################################### - @staticmethod def get_base_model(model: str) -> Optional[str]: @@ -64,4 +193,6 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): api_base: Optional[str] = None, ) -> dict: """Azure Foundry sends api key in query params""" - raise NotImplementedError("Azure Foundry does not support environment validation") + raise NotImplementedError( + "Azure Foundry does not support environment validation" + )