diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 3845dc5a46..68b5341e95 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -10,10 +10,11 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, ) -from litellm.llms.base_llm.base_utils import BaseLLMModelInfo +from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.anthropic import AllAnthropicToolsValues, AnthropicMcpServerTool from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import TokenCountResponse class AnthropicError(BaseLLMException): @@ -229,7 +230,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): litellm_model_names.append(litellm_model_name) return litellm_model_names - def get_token_counter(self) -> Optional["AnthropicTokenCounter"]: + def get_token_counter(self) -> Optional[BaseTokenCounter]: """ Factory method to create an Anthropic token counter. @@ -239,32 +240,24 @@ class AnthropicModelInfo(BaseLLMModelInfo): return AnthropicTokenCounter() -class AnthropicTokenCounter: +class AnthropicTokenCounter(BaseTokenCounter): """Token counter implementation for Anthropic provider.""" - - def supports_provider( + + def should_use_token_counting_api( self, - deployment: Optional[Dict[str, Any]] = None, - from_endpoint: bool = False + custom_llm_provider: Optional[str] = None, ) -> bool: - if not from_endpoint: - return False - - if deployment is None: - return False - - full_model = deployment.get("litellm_params", {}).get("model", "") - is_anthropic_provider = full_model.startswith("anthropic/") or "anthropic" in full_model.lower() - - return is_anthropic_provider + 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[Dict[str, Any]]: + ) -> Optional[TokenCountResponse]: from litellm.proxy.utils import count_tokens_with_anthropic_api result = await count_tokens_with_anthropic_api( @@ -274,12 +267,13 @@ class AnthropicTokenCounter: ) if result is not None: - return { - "total_tokens": result["total_tokens"], - "request_model": request_model, - "model_used": model_to_use, - "tokenizer_type": result["tokenizer_used"], - } + 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 diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 3961ee2b9e..9172a05e38 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -5,14 +5,37 @@ Utility functions for base LLM classes. import copy import json from abc import ABC, abstractmethod -from typing import List, Optional, Type, Union +from typing import Any, Dict, List, Optional, Type, Union from openai.lib import _parsing, _pydantic from pydantic import BaseModel from litellm._logging import verbose_logger from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk -from litellm.types.utils import Message, ProviderSpecificModelInfo +from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountResponse + + +class BaseTokenCounter(ABC): + @abstractmethod + 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]: + pass + + @abstractmethod + def should_use_token_counting_api( + self, + custom_llm_provider: Optional[str] = None, + ) -> bool: + """ + Returns True if we should the this API for token counting for the selected `custom_llm_provider` + """ + return False class BaseLLMModelInfo(ABC): @@ -70,7 +93,7 @@ class BaseLLMModelInfo(ABC): """ pass - def get_token_counter(self): + def get_token_counter(self) -> Optional[BaseTokenCounter]: """ Factory method to create a token counter for this provider. diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 53de6711ca..e53829d332 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -1,15 +1,16 @@ import base64 import datetime -from typing import Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Union import httpx import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH -from litellm.llms.base_llm.base_utils import BaseLLMModelInfo +from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import TokenCountResponse class GeminiError(BaseLLMException): @@ -89,6 +90,16 @@ class GeminiModelInfo(BaseLLMModelInfo): return GeminiError( status_code=status_code, message=error_message, headers=headers ) + + def get_token_counter(self) -> Optional[BaseTokenCounter]: + """ + Factory method to create a token counter for this provider. + + Returns: + Optional TokenCounterInterface implementation for this provider, + or None if token counting is not supported. + """ + return GoogleAIStudioTokenCounter() def encode_unserializable_types( @@ -137,3 +148,46 @@ def encode_unserializable_types( def get_api_key_from_env() -> Optional[str]: return get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY") + + +class GoogleAIStudioTokenCounter(BaseTokenCounter): + """Token counter implementation for Google AI Studio 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.GEMINI.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]: + import copy + + from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter + deployment = deployment or {} + count_tokens_params_request = copy.deepcopy(deployment.get("litellm_params", {})) + count_tokens_params = { + "model": model_to_use, + "contents": contents, + } + count_tokens_params_request.update(count_tokens_params) + result = await GoogleAIStudioTokenCounter().acount_tokens( + **count_tokens_params_request, + ) + + if result is not None: + return TokenCountResponse( + total_tokens=result.get("totalTokens", 0), + request_model=request_model, + model_used=model_to_use, + tokenizer_type=result.get("tokenizer_used", ""), + original_response=result, + ) + + return None \ No newline at end of file diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py new file mode 100644 index 0000000000..1f13b0c314 --- /dev/null +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -0,0 +1,125 @@ +from typing import TYPE_CHECKING, Any, Dict, Optional, Union + +import httpx + +import litellm +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.types.google_genai.main import GenerateContentContentListUnionDict +else: + GenerateContentContentListUnionDict = Any + +class GoogleAIStudioTokenCounter: + def validate_environment( + self, + api_key: Optional[str] = None, + headers: Optional[Dict[str, Any]] = None, + model: str = "", + litellm_params: Optional[Dict[str, Any]] = None, + ): + from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig + return GoogleGenAIConfig().validate_environment( + api_key=api_key, + headers=headers, + model=model, + litellm_params=litellm_params, + ) + + + async def acount_tokens( + self, + contents: Any, + model: str, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + **kwargs, + ) -> Dict[str, Any]: + """ + Count tokens using Google Gen AI Studio countTokens endpoint. + + Args: + contents: The content to count tokens for (Google Gen AI format) + Example: [{"parts": [{"text": "Hello world"}]}] + model: The model name (e.g. "gemini-1.5-flash") + api_key: Optional Google API key (will fall back to environment) + api_base: Optional API base URL (defaults to Google Gen AI Studio) + timeout: Optional timeout for the request + **kwargs: Additional parameters + + Returns: + Dict containing token count information from Google Gen AI Studio API. + Example response: + { + "totalTokens": 31, + "totalBillableCharacters": 96, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 31 + } + ] + } + + Raises: + ValueError: If API key is missing + litellm.APIError: If the API call fails + litellm.APIConnectionError: If the connection fails + Exception: For any other unexpected errors + """ + # Set up API base URL + base_url = api_base or "https://generativelanguage.googleapis.com" + url = f"{base_url}/v1beta/models/{model}:countTokens" + + # Prepare headers + headers = self.validate_environment( + api_key=api_key, + headers={}, + model=model, + litellm_params=kwargs, + ) + + # Prepare request body + request_body = { + "contents": contents + } + + async_httpx_client = get_async_httpx_client( + llm_provider=LlmProviders.GEMINI, + ) + + try: + response = await async_httpx_client.post( + url=url, + headers=headers, + json=request_body + ) + + # Check for HTTP errors + response.raise_for_status() + + # Parse response + result = response.json() + return result + + except httpx.HTTPStatusError as e: + error_msg = f"Google Gen AI Studio API error: {e.response.status_code} - {e.response.text}" + raise litellm.APIError( + message=error_msg, + llm_provider="gemini", + model=model, + status_code=e.response.status_code + ) from e + except httpx.RequestError as e: + error_msg = f"Request to Google Gen AI Studio failed: {str(e)}" + raise litellm.APIConnectionError( + message=error_msg, + llm_provider="gemini", + model=model + ) from e + except Exception as e: + error_msg = f"Unexpected error during token counting: {str(e)}" + raise Exception(error_msg) from e + diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6a7006ea8d..cf8b3d147f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2094,13 +2094,15 @@ class TokenCountRequest(LiteLLMPydanticObjectBase): model: str prompt: Optional[str] = None messages: Optional[List[dict]] = None + """ + Anthropic token counting endpoint uses /messages + """ - -class TokenCountResponse(LiteLLMPydanticObjectBase): - total_tokens: int - request_model: str - model_used: str - tokenizer_type: str + + contents: Optional[List[dict]] = None + """ + Google /countTokens endpoint expects contents to be a list of dicts with the following structure: + """ class CallInfo(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 8965ae7715..a10a39a6a5 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -16,6 +16,7 @@ from litellm.proxy.common_request_processing import ( ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request +from litellm.types.utils import TokenCountResponse router = APIRouter() @@ -262,10 +263,18 @@ async def count_tokens( ) # Call the internal token counter function with direct request flag set to False - token_response = await internal_token_counter(token_request, is_direct_request=False) - + token_response = await internal_token_counter( + request=token_request, + call_endpoint=True, + ) + _token_response_dict: dict = {} + if isinstance(token_response, TokenCountResponse): + _token_response_dict = token_response.model_dump() + elif isinstance(token_response, dict): + _token_response_dict = token_response + # Convert the internal response to Anthropic API format - return {"input_tokens": token_response.total_tokens} + return {"input_tokens": _token_response_dict.get("total_tokens", 0)} except HTTPException: raise diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 226fcf05aa..373232e22d 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -3,6 +3,7 @@ from fastapi import APIRouter, Depends, Request, Response from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.types.llms.vertex_ai import TokenCountDetailsResponse router = APIRouter( tags=["google genai endpoints"], @@ -145,10 +146,61 @@ async def google_stream_generate_content( -@router.post("/v1beta/models/{model_name}:countTokens", dependencies=[Depends(user_api_key_auth)]) -@router.post("/models/{model_name}:countTokens", dependencies=[Depends(user_api_key_auth)]) +@router.post( + "/v1beta/models/{model_name}:countTokens", + dependencies=[Depends(user_api_key_auth)], + response_model=TokenCountDetailsResponse, +) +@router.post( + "/models/{model_name}:countTokens", + dependencies=[Depends(user_api_key_auth)], + response_model=TokenCountDetailsResponse, +) async def google_count_tokens(request: Request, model_name: str): """ - Not Implemented, this is a placeholder for the google genai countTokens endpoint. + ```json + return { + "totalTokens": 31, + "totalBillableCharacters": 96, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 31 + } + ] + } + ``` """ - return {} + from litellm.proxy.common_utils.http_parsing_utils import _read_request_body + from litellm.proxy.proxy_server import token_counter as internal_token_counter + + data = await _read_request_body(request=request) + contents = data.get("contents", []) + #Create TokenCountRequest for the internal endpoint + from litellm.proxy._types import TokenCountRequest + + token_request = TokenCountRequest( + model=model_name, + contents=contents + ) + + # Call the internal token counter function with direct request flag set to False + token_response = await internal_token_counter( + request=token_request, + call_endpoint=True, + ) + if token_response is not None: + # cast the response to the well known format + original_response: dict = token_response.original_response or {} + return TokenCountDetailsResponse( + totalTokens=original_response.get("totalTokens", 0), + promptTokensDetails=original_response.get("promptTokensDetails", []), + ) + + ######################################################### + # Return the response in the well known format + ######################################################### + return TokenCountDetailsResponse( + totalTokens=0, + promptTokensDetails=[], + ) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 6e9c925257..938c0fc49e 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,11 +1,4 @@ model_list: - - model_name: openai/* + - model_name: vertex_ai/* litellm_params: - model: openai/* - - model_name: anthropic/* - litellm_params: - model: anthropic/* - -litellm_settings: - callbacks: - - langfuse_otel \ No newline at end of file + model: gemini/* diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 939084c38d..78e9ae2483 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11,7 +11,6 @@ import time import traceback import uuid import warnings -from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from datetime import datetime, timedelta from typing import ( TYPE_CHECKING, @@ -35,10 +34,12 @@ from litellm.constants import ( LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS, LITELLM_SETTINGS_SAFE_DB_OVERRIDES, ) +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.types.utils import ( ModelResponse, ModelResponseStream, TextCompletionResponse, + TokenCountResponse, ) if TYPE_CHECKING: @@ -2999,7 +3000,9 @@ class ProxyConfig: if should_reload: # Perform the reload - from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + from litellm.litellm_core_utils.get_model_cost_map import ( + get_model_cost_map, + ) model_cost_map_url = litellm.model_cost_map_url new_model_cost_map = get_model_cost_map(url=model_cost_map_url) litellm.model_cost = new_model_cost_map @@ -5742,9 +5745,10 @@ async def run_thread( # dependencies=[Depends(user_api_key_auth)], # ) # async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)): +from litellm.llms.base_llm.base_utils import BaseTokenCounter -def _get_provider_token_counter(deployment: dict, model_to_use: str): +def _get_provider_token_counter(deployment: dict, model_to_use: str) -> Tuple[Optional[BaseTokenCounter], Optional[str], Optional[str]]: """ Auto-route to the correct provider's token counter based on model/deployment. Uses the existing get_provider_model_info infrastructure with switch-case pattern. @@ -5755,10 +5759,12 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str): from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider full_model = deployment.get("litellm_params", {}).get("model", "") + model: Optional[str] = None + custom_llm_provider: Optional[str] = None try: # Use existing LiteLLM logic to determine provider - model, provider, dynamic_api_key, api_base = get_llm_provider( + model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( model=full_model, custom_llm_provider=deployment.get("litellm_params", {}).get( "custom_llm_provider" @@ -5772,7 +5778,7 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str): from litellm.utils import ProviderConfigManager # Convert string provider to LlmProviders enum - llm_provider_enum = LlmProviders(provider) + llm_provider_enum = LlmProviders(custom_llm_provider) # Add more provider mappings as needed if llm_provider_enum: @@ -5780,7 +5786,7 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str): model=full_model, provider=llm_provider_enum ) if provider_model_info is not None: - return provider_model_info.get_token_counter() + return provider_model_info.get_token_counter(), model, custom_llm_provider except Exception: # If provider detection fails, fall back to manual checks @@ -5788,9 +5794,9 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str): from litellm.llms.anthropic.common_utils import AnthropicModelInfo anthropic_model_info = AnthropicModelInfo() - return anthropic_model_info.get_token_counter() + return anthropic_model_info.get_token_counter(), model, custom_llm_provider - return None + return None, None, None @router.post( @@ -5799,62 +5805,82 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str): dependencies=[Depends(user_api_key_auth)], response_model=TokenCountResponse, ) -async def token_counter(request: TokenCountRequest, is_direct_request: bool = True): - """ """ +async def token_counter( + request: TokenCountRequest, + call_endpoint: bool = False +): + """ + Args: + request: TokenCountRequest + call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs. + + Returns: + TokenCountResponse + """ from litellm import token_counter global llm_router prompt = request.prompt messages = request.messages - if prompt is None and messages is None: + contents = request.contents + + ######################################################### + # Validate request + ######################################################### + if prompt is None and messages is None and contents is None: raise HTTPException( - status_code=400, detail="prompt or messages must be provided" + status_code=400, detail="prompt or messages or contents must be provided" ) - deployment = None + deployment: Optional[Dict[str, Any]] = None litellm_model_name = None model_info: Optional[ModelMapInfo] = None if llm_router is not None: # get 1 deployment corresponding to the model - for _model in llm_router.model_list: - if _model["model_name"] == request.model: - deployment = _model - model_info = deployment.get("model_info", {}) - break + try: + deployment = await llm_router.async_get_available_deployment( + model=request.model, + request_kwargs={}, + ) + except Exception: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.token_counter(): Exception occured while getting deployment" + ) + pass if deployment is not None: litellm_model_name = deployment.get("litellm_params", {}).get("model") # remove the custom_llm_provider_prefix in the litellm_model_name if "/" in litellm_model_name: litellm_model_name = litellm_model_name.split("/", 1)[1] - model_to_use = ( + model_to_use: str = ( litellm_model_name or request.model ) # use litellm model name, if it's not avalable then fallback to request.model # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) - provider_counter = None - if deployment is not None and not is_direct_request: + provider_counter: Optional[BaseTokenCounter] = None + custom_llm_provider: Optional[str] = None + if call_endpoint is True and deployment is not None: # Auto-route to the correct provider based on model - provider_counter = _get_provider_token_counter(deployment, model_to_use) + provider_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use) + if _model is not None: + model_to_use = _model - if provider_counter is not None and provider_counter.supports_provider( - deployment=deployment, from_endpoint=not is_direct_request - ): - result = await provider_counter.count_tokens( - model_to_use=model_to_use, - messages=messages, # type: ignore - deployment=deployment, - request_model=request.model, - ) - - if result is not None: - return TokenCountResponse( - total_tokens=result["total_tokens"], - request_model=result["request_model"], - model_used=result["model_used"], - tokenizer_type=result["tokenizer_type"], + if provider_counter is not None: + if provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider) is True: + result = await provider_counter.count_tokens( + model_to_use=model_to_use or "", + messages=messages, # type: ignore + contents=contents, + deployment=deployment, + request_model=request.model, ) + ######################################################### + # Transfrom the Response to the well known format + ######################################################### + if result is not None: + return result # Default LiteLLM token counting custom_tokenizer: Optional[CustomHuggingfaceTokenizer] = None diff --git a/litellm/types/google_genai/main.py b/litellm/types/google_genai/main.py index 96abc5d2ae..b875495bab 100644 --- a/litellm/types/google_genai/main.py +++ b/litellm/types/google_genai/main.py @@ -1,9 +1,10 @@ # Import types from the Google GenAI SDK -from typing import TYPE_CHECKING, Any, Optional, TypeAlias, TypedDict +from typing import TYPE_CHECKING, Any, List, Optional, TypeAlias # During static type-checking we can rely on the real google-genai types. from google.genai import types as _genai_types # type: ignore from pydantic import BaseModel +from typing_extensions import TypedDict from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject @@ -19,7 +20,7 @@ ToolConfigDict = _genai_types.ToolConfigDict class GenerateContentRequestDict(GenerateContentRequestParametersDict): # type: ignore[misc] generationConfig: Optional[Any] - tools: Optional[ToolConfigDict] + tools: Optional[ToolConfigDict] # type: ignore[assignment] class GenerateContentResponse(GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject): # type: ignore[misc] diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 1c3eb49fb0..2931770cd6 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -1,11 +1,12 @@ import json from enum import Enum -from typing import Any, Dict, List, Literal, Optional, Tuple, TypedDict, Union +from typing import Any, Dict, List, Literal, Optional, Tuple, Union from typing_extensions import ( Protocol, Required, Self, + TypedDict, TypeGuard, get_origin, override, @@ -241,6 +242,17 @@ class UsageMetadata(TypedDict, total=False): responseTokensDetails: List[PromptTokensDetails] +class TokenCountDetailsResponse(TypedDict): + """ + Response structure for token count details with modality breakdown. + + Example: + {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]} + """ + totalTokens: int + promptTokensDetails: List[PromptTokensDetails] + + class CachedContent(TypedDict, total=False): ttl: TTL expire_time: str diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ec56edf92a..7ae92185e7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2356,6 +2356,17 @@ class LiteLLMLoggingBaseClass: pass +class TokenCountResponse(LiteLLMPydanticObjectBase): + total_tokens: int + request_model: str + model_used: str + tokenizer_type: str + original_response: Optional[dict] = None + """ + Original Response from upstream API call - if an API call was made for token counting + """ + + class CustomHuggingfaceTokenizer(TypedDict): identifier: str revision: str # usually 'main' diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index eefe5059b0..e771c9330f 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -24,7 +24,8 @@ from litellm._logging import verbose_proxy_logger verbose_proxy_logger.setLevel(level=logging.DEBUG) -from litellm.proxy._types import TokenCountRequest, TokenCountResponse +from litellm.proxy._types import TokenCountRequest +from litellm.types.utils import TokenCountResponse from litellm import Router @@ -169,12 +170,12 @@ async def test_anthropic_messages_count_tokens_endpoint(): anthropic_endpoints._read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, is_direct_request=True): - assert is_direct_request == False, "Should be called with is_direct_request=False for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=False): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" assert request.messages == [{"role": "user", "content": "Hello Claude!"}] - from litellm.proxy._types import TokenCountResponse + from litellm.types.utils import TokenCountResponse return TokenCountResponse( total_tokens=15, request_model="claude-3-sonnet-20240229", @@ -236,12 +237,12 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): anthropic_endpoints._read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, is_direct_request=True): - assert is_direct_request == False, "Should be called with is_direct_request=False for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=True): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" assert request.messages == [{"role": "user", "content": "Hello GPT!"}] - from litellm.proxy._types import TokenCountResponse + from litellm.types.utils import TokenCountResponse return TokenCountResponse( total_tokens=12, request_model="gpt-4", @@ -300,7 +301,7 @@ async def test_internal_token_counter_anthropic_provider_detection(): model="claude-test", messages=[{"role": "user", "content": "hello"}], ), - is_direct_request=False + call_endpoint=True ) print("Anthropic provider test response:", response) @@ -330,7 +331,7 @@ async def test_internal_token_counter_anthropic_provider_detection(): model="gpt-test", messages=[{"role": "user", "content": "hello"}], ), - is_direct_request=False + call_endpoint=True ) print("Non-Anthropic provider test response:", response) @@ -385,7 +386,7 @@ 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 + from unittest.mock import patch, AsyncMock from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app @@ -404,6 +405,13 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): "model_info": {} }] + # Mock the async method properly + mock_router.async_get_available_deployment = AsyncMock(return_value={ + "model_name": "claude-3-5-sonnet", + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, + "model_info": {} + }) + client = TestClient(app) response = client.post( @@ -426,7 +434,7 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): @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 + from unittest.mock import patch, AsyncMock from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app @@ -444,6 +452,13 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): "model_info": {} }] + # Mock the async method properly + mock_router.async_get_available_deployment = AsyncMock(return_value={ + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": {} + }) + client = TestClient(app) response = client.post( @@ -466,7 +481,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): @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 + from unittest.mock import patch, AsyncMock from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app @@ -484,6 +499,13 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): "model_info": {} }] + # Mock the async method properly + mock_router.async_get_available_deployment = AsyncMock(return_value={ + "model_name": "claude-3-5-sonnet", + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, + "model_info": {} + }) + client = TestClient(app) response = client.post( @@ -499,7 +521,7 @@ 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 is_direct_request=True) + # Verify that Anthropic API was NOT called (since call_endpoint=False) mock_anthropic_count.assert_not_called() @@ -523,40 +545,59 @@ async def test_factory_registration(): } # Test Anthropic counter supports provider - assert counter.supports_provider(anthropic_deployment, from_endpoint=True) - assert not counter.supports_provider(anthropic_deployment, from_endpoint=False) + assert counter.should_use_token_counting_api(custom_llm_provider="anthropic") + assert not counter.should_use_token_counting_api(custom_llm_provider="openai") # Test non-Anthropic provider - assert not counter.supports_provider(non_anthropic_deployment, from_endpoint=True) - assert not counter.supports_provider(non_anthropic_deployment, from_endpoint=False) + assert not counter.should_use_token_counting_api(custom_llm_provider="openai") # Test None deployment - assert not counter.supports_provider(None, from_endpoint=True) - assert not counter.supports_provider(None, from_endpoint=False) + assert not counter.should_use_token_counting_api(custom_llm_provider=None) -@pytest.mark.asyncio -async def test_factory_anthropic_counter_supports_provider(): - """Test AnthropicTokenCounter provider detection logic.""" - from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + +@pytest.mark.asyncio +async def test_vertex_ai_gemini_token_counting_with_contents(): + """ + Test token counting for Vertex AI Gemini model using contents format with call_endpoint=True + """ + llm_router = Router( + model_list=[ + { + "model_name": "gemini-2.5-pro", + "litellm_params": { + "model": "gemini/gemini-2.5-pro", + }, + } + ] + ) - anthropic_model_info = AnthropicModelInfo() - counter = anthropic_model_info.get_token_counter() + setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - # Test Anthropic provider detection - anthropic_deployment = { - "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"} - } - assert counter.supports_provider(anthropic_deployment, from_endpoint=True) - assert not counter.supports_provider(anthropic_deployment, from_endpoint=False) + # Test with contents format and call_endpoint=True + response = await token_counter( + request=TokenCountRequest( + model="gemini-2.5-pro", + contents=[ + { + "parts": [ + { + "text": "Hello world, how are you doing today? i am ij" + } + ] + } + ], + ), + call_endpoint=True + ) - # Test non-Anthropic provider - openai_deployment = { - "litellm_params": {"model": "openai/gpt-4"} - } - assert not counter.supports_provider(openai_deployment, from_endpoint=True) - assert not counter.supports_provider(openai_deployment, from_endpoint=False) - - # Test None deployment - assert not counter.supports_provider(None, from_endpoint=True) - assert not counter.supports_provider(None, from_endpoint=False) + print("Vertex AI Gemini token counting response:", response) + + # validate we have orignal response + assert response.original_response is not None + assert response.original_response.get("totalTokens") is not None + assert response.original_response.get("promptTokensDetails") is not None + + prompt_tokens_details = response.original_response.get("promptTokensDetails") + assert prompt_tokens_details is not None diff --git a/tests/test_litellm/llms/gemini/test_gemini_common_utils.py b/tests/test_litellm/llms/gemini/test_gemini_common_utils.py index 6aca755506..34472b3856 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_common_utils.py +++ b/tests/test_litellm/llms/gemini/test_gemini_common_utils.py @@ -1,6 +1,8 @@ +from unittest.mock import AsyncMock, patch + import pytest -from litellm.llms.gemini.common_utils import GeminiModelInfo +from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter class TestGeminiModelInfo: @@ -84,3 +86,75 @@ class TestGeminiModelInfo: ] assert result == expected + + +class TestGoogleAIStudioTokenCounter: + """Test suite for GoogleAIStudioTokenCounter class""" + + def test_should_use_token_counting_api(self): + """Test should_use_token_counting_api method with different provider values""" + from litellm.types.utils import LlmProviders + + token_counter = GoogleAIStudioTokenCounter() + + # Test with gemini provider - should return True + assert token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) is True + + # Test with other providers - should return False + assert token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) is False + assert token_counter.should_use_token_counting_api("anthropic") is False + assert token_counter.should_use_token_counting_api("vertex_ai") is False + + # Test with None - should return False + assert token_counter.should_use_token_counting_api(None) is False + + @pytest.mark.asyncio + async def test_count_tokens(self): + """Test count_tokens method with mocked API response""" + from litellm.types.utils import TokenCountResponse + + token_counter = GoogleAIStudioTokenCounter() + + # Mock the GoogleAIStudioTokenCounter from handler module + mock_response = { + "totalTokens": 31, + "totalBillableCharacters": 96, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 31 + } + ] + } + + with patch('litellm.llms.gemini.count_tokens.handler.GoogleAIStudioTokenCounter.acount_tokens', + new_callable=AsyncMock) as mock_acount_tokens: + mock_acount_tokens.return_value = mock_response + + # Test data + model_to_use = "gemini-1.5-flash" + contents = [{"parts": [{"text": "Hello world"}]}] + request_model = "gemini/gemini-1.5-flash" + + # Call the method + result = await token_counter.count_tokens( + model_to_use=model_to_use, + messages=None, + contents=contents, + deployment=None, + request_model=request_model + ) + + # Verify the result + assert result is not None + assert isinstance(result, TokenCountResponse) + assert result.total_tokens == 31 + assert result.request_model == request_model + assert result.model_used == model_to_use + assert result.original_response == mock_response + + # Verify the mock was called correctly + mock_acount_tokens.assert_called_once_with( + model=model_to_use, + contents=contents + )