tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-20 08:23:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into litellm_mcp_config_fix

Browse Source
This commit is contained in:
Jugal Bhatt
2025-08-12 17:10:18 -07:00
parent f487816b9f afe159bb8b
commit bb32ec991f
14 changed files with 551 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/anthropic/common_utils.py
+18 -24
View File
@@ -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
litellm/llms/base_llm/base_utils.py
+26 -3
View File
@@ -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.
litellm/llms/gemini/common_utils.py
+56 -2
View File
@@ -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
litellm/llms/gemini/count_tokens/handler.py
+125
View File
@@ -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
litellm/proxy/_types.py
+8 -6
View File
@@ -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):
litellm/proxy/anthropic_endpoints/endpoints.py
+12 -3
View File
@@ -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
litellm/proxy/google_endpoints/endpoints.py
+56 -4
View File
@@ -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=[],
)
litellm/proxy/proxy_config.yaml
+2 -9
View File
@@ -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
model: gemini/*
litellm/proxy/proxy_server.py
+64 -38
View File
@@ -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
litellm/types/google_genai/main.py
+3 -2
View File
@@ -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]
litellm/types/llms/vertex_ai.py
+13 -1
View File
@@ -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
litellm/types/utils.py
+11
View File
@@ -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'
tests/proxy_unit_tests/test_proxy_token_counter.py
+82 -41
View File
@@ -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
tests/test_litellm/llms/gemini/test_gemini_common_utils.py
+75 -1
View File
@@ -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
)