mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Merge pull request #22809 from Chesars/worktree-count-tokens-api
feat(openai): add litellm.acount_tokens() public API + OpenAI token counting support
This commit is contained in:
Cesar Garcia
GitHub
commit
6693723588
11 changed files with 1203 additions and 0 deletions
|
|
@ -138,6 +138,7 @@ The `/v1/messages/count_tokens` endpoint automatically routes to the appropriate
|
|||
| Provider | Token Counting Method |
|
||||
|----------|----------------------|
|
||||
| Anthropic | [Anthropic Token Counting API](https://docs.anthropic.com/en/docs/build-with-claude/token-counting) |
|
||||
| OpenAI | [OpenAI Responses API `/input_tokens`](https://platform.openai.com/docs/api-reference/responses/input-tokens) — see [Token Counting](./count_tokens.md) |
|
||||
| Vertex AI (Claude) | Vertex AI Partner Models Token Counter |
|
||||
| Bedrock (Claude) | AWS Bedrock CountTokens API |
|
||||
| Gemini | Google AI Studio countTokens API |
|
||||
|
|
|
|||
189
docs/my-website/docs/count_tokens.md
Normal file
189
docs/my-website/docs/count_tokens.md
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
# Token Counting
|
||||
|
||||
## Overview
|
||||
|
||||
LiteLLM provides exact token counting by calling provider-specific token counting APIs. This gives you accurate token counts before sending requests, helping with cost estimation and context window management.
|
||||
|
||||
| Feature | Details |
|
||||
|---------|---------|
|
||||
| SDK Method | `litellm.acount_tokens()` |
|
||||
| Proxy Endpoints | `/v1/messages/count_tokens` (Anthropic format), `/v1/responses/input_tokens` (OpenAI format) |
|
||||
| Fallback | Local tiktoken-based counting for unsupported providers |
|
||||
|
||||
## Supported Providers
|
||||
|
||||
| Provider | Token Counting API | Format |
|
||||
|----------|-------------------|--------|
|
||||
| OpenAI | [Responses API `/input_tokens`](https://platform.openai.com/docs/api-reference/responses/input-tokens) | OpenAI Responses |
|
||||
| Anthropic | [Messages `/count_tokens`](https://docs.anthropic.com/en/docs/build-with-claude/token-counting) | Anthropic Messages |
|
||||
| Vertex AI (Claude) | Vertex AI Partner Models Token Counter | Anthropic Messages |
|
||||
| Bedrock (Claude) | AWS Bedrock CountTokens API | Anthropic Messages |
|
||||
| Gemini | Google AI Studio countTokens API | Anthropic Messages |
|
||||
| Vertex AI (Gemini) | Vertex AI countTokens API | Anthropic Messages |
|
||||
| Other providers | Local tiktoken fallback | N/A |
|
||||
|
||||
## SDK Usage
|
||||
|
||||
### Basic Usage
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
import litellm
|
||||
|
||||
async def main():
|
||||
# OpenAI
|
||||
result = await litellm.acount_tokens(
|
||||
model="openai/gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
)
|
||||
print(f"Token count: {result.total_tokens}")
|
||||
print(f"Tokenizer: {result.tokenizer_type}") # "openai_api"
|
||||
|
||||
# Anthropic
|
||||
result = await litellm.acount_tokens(
|
||||
model="anthropic/claude-3-5-sonnet-20241022",
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
)
|
||||
print(f"Token count: {result.total_tokens}")
|
||||
print(f"Tokenizer: {result.tokenizer_type}") # "anthropic_api"
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
### With Tools and System Message
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
import litellm
|
||||
|
||||
async def main():
|
||||
result = await litellm.acount_tokens(
|
||||
model="openai/gpt-4o",
|
||||
messages=[{"role": "user", "content": "What's the weather in Paris?"}],
|
||||
tools=[{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather for a city",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
}],
|
||||
system="You are a helpful weather assistant.",
|
||||
)
|
||||
print(f"Token count (with tools): {result.total_tokens}")
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
### Response Format
|
||||
|
||||
`litellm.acount_tokens()` returns a `TokenCountResponse`:
|
||||
|
||||
```python
|
||||
TokenCountResponse(
|
||||
total_tokens=15, # Token count
|
||||
request_model="openai/gpt-4o", # Model requested
|
||||
model_used="gpt-4o", # Model used for counting
|
||||
tokenizer_type="openai_api", # "openai_api", "anthropic_api", "local_tokenizer"
|
||||
original_response={"input_tokens": 15}, # Raw API response
|
||||
error=False, # True if counting failed
|
||||
error_message=None, # Error details if failed
|
||||
)
|
||||
```
|
||||
|
||||
### Fallback Behavior
|
||||
|
||||
If a provider doesn't support a token counting API, or if the API key is missing, `acount_tokens()` automatically falls back to local tiktoken-based counting:
|
||||
|
||||
```python
|
||||
# Unsupported provider → automatic fallback
|
||||
result = await litellm.acount_tokens(
|
||||
model="together_ai/meta-llama/Llama-3-8b-chat-hf",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
)
|
||||
print(result.tokenizer_type) # "local_tokenizer"
|
||||
```
|
||||
|
||||
## Proxy Usage
|
||||
|
||||
### OpenAI Format — `/v1/responses/input_tokens`
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="curl" label="curl">
|
||||
|
||||
```bash
|
||||
curl -X POST "http://localhost:4000/v1/responses/input_tokens" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"model": "gpt-4o",
|
||||
"input": "Hello, how are you?"
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="python" label="Python (httpx)">
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
response = httpx.post(
|
||||
"http://localhost:4000/v1/responses/input_tokens",
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": "Bearer sk-1234"
|
||||
},
|
||||
json={
|
||||
"model": "gpt-4o",
|
||||
"input": "Hello, how are you?"
|
||||
}
|
||||
)
|
||||
|
||||
print(response.json())
|
||||
# {"input_tokens": 7}
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
**Response:**
|
||||
```json
|
||||
{"input_tokens": 7}
|
||||
```
|
||||
|
||||
### Anthropic Format — `/v1/messages/count_tokens`
|
||||
|
||||
See [Anthropic Token Counting](./anthropic_count_tokens.md) for full documentation.
|
||||
|
||||
```bash
|
||||
curl -X POST "http://localhost:4000/v1/messages/count_tokens" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"model": "claude-3-5-sonnet-20241022",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello, how are you?"}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
## Proxy Configuration
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-4o
|
||||
litellm_params:
|
||||
model: openai/gpt-4o
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
- model_name: claude-3-5-sonnet
|
||||
litellm_params:
|
||||
model: anthropic/claude-3-5-sonnet-20241022
|
||||
api_key: os.environ/ANTHROPIC_API_KEY
|
||||
```
|
||||
|
|
@ -58,6 +58,7 @@ from ..common_utils import OpenAIError
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.types.llms.openai import ChatCompletionToolParam
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
|
|
@ -758,6 +759,13 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
def get_base_model(model: Optional[str] = None) -> Optional[str]:
|
||||
return model
|
||||
|
||||
def get_token_counter(self) -> Optional["BaseTokenCounter"]:
|
||||
from litellm.llms.openai.responses.count_tokens.token_counter import (
|
||||
OpenAITokenCounter,
|
||||
)
|
||||
|
||||
return OpenAITokenCounter()
|
||||
|
||||
def get_model_response_iterator(
|
||||
self,
|
||||
streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse],
|
||||
|
|
|
|||
19
litellm/llms/openai/responses/count_tokens/__init__.py
Normal file
19
litellm/llms/openai/responses/count_tokens/__init__.py
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
"""
|
||||
OpenAI Responses API token counting implementation.
|
||||
"""
|
||||
|
||||
from litellm.llms.openai.responses.count_tokens.handler import (
|
||||
OpenAICountTokensHandler,
|
||||
)
|
||||
from litellm.llms.openai.responses.count_tokens.token_counter import (
|
||||
OpenAITokenCounter,
|
||||
)
|
||||
from litellm.llms.openai.responses.count_tokens.transformation import (
|
||||
OpenAICountTokensConfig,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"OpenAICountTokensHandler",
|
||||
"OpenAICountTokensConfig",
|
||||
"OpenAITokenCounter",
|
||||
]
|
||||
105
litellm/llms/openai/responses/count_tokens/handler.py
Normal file
105
litellm/llms/openai/responses/count_tokens/handler.py
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
"""
|
||||
OpenAI Responses API token counting handler.
|
||||
|
||||
Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.llms.openai.common_utils import OpenAIError
|
||||
from litellm.llms.openai.responses.count_tokens.transformation import (
|
||||
OpenAICountTokensConfig,
|
||||
)
|
||||
|
||||
|
||||
class OpenAICountTokensHandler(OpenAICountTokensConfig):
|
||||
"""
|
||||
Handler for OpenAI Responses API token counting requests.
|
||||
"""
|
||||
|
||||
async def handle_count_tokens_request(
|
||||
self,
|
||||
model: str,
|
||||
input: Union[str, List[Any]],
|
||||
api_key: str,
|
||||
api_base: Optional[str] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
instructions: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Handle a token counting request to OpenAI's Responses API.
|
||||
|
||||
Returns:
|
||||
Dictionary containing {"input_tokens": <number>}
|
||||
|
||||
Raises:
|
||||
OpenAIError: If the API request fails
|
||||
"""
|
||||
try:
|
||||
self.validate_request(model, input)
|
||||
|
||||
verbose_logger.debug(
|
||||
f"Processing OpenAI CountTokens request for model: {model}"
|
||||
)
|
||||
|
||||
request_body = self.transform_request_to_count_tokens(
|
||||
model=model,
|
||||
input=input,
|
||||
tools=tools,
|
||||
instructions=instructions,
|
||||
)
|
||||
|
||||
endpoint_url = self.get_openai_count_tokens_endpoint(api_base)
|
||||
|
||||
verbose_logger.debug(f"Making request to: {endpoint_url}")
|
||||
|
||||
headers = self.get_required_headers(api_key)
|
||||
|
||||
async_client = get_async_httpx_client(
|
||||
llm_provider=litellm.LlmProviders.OPENAI
|
||||
)
|
||||
|
||||
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"OpenAI API error: {error_text}")
|
||||
raise OpenAIError(
|
||||
status_code=response.status_code,
|
||||
message=error_text,
|
||||
)
|
||||
|
||||
openai_response = response.json()
|
||||
verbose_logger.debug(f"OpenAI response: {openai_response}")
|
||||
return openai_response
|
||||
|
||||
except OpenAIError:
|
||||
raise
|
||||
except httpx.HTTPStatusError as e:
|
||||
verbose_logger.error(f"HTTP error in CountTokens handler: {str(e)}")
|
||||
raise OpenAIError(
|
||||
status_code=e.response.status_code,
|
||||
message=e.response.text,
|
||||
)
|
||||
except (httpx.RequestError, json.JSONDecodeError, ValueError) as e:
|
||||
verbose_logger.error(f"Error in CountTokens handler: {str(e)}")
|
||||
raise OpenAIError(
|
||||
status_code=500,
|
||||
message=f"CountTokens processing error: {str(e)}",
|
||||
)
|
||||
118
litellm/llms/openai/responses/count_tokens/token_counter.py
Normal file
118
litellm/llms/openai/responses/count_tokens/token_counter.py
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
"""
|
||||
OpenAI Token Counter implementation using the Responses API /input_tokens endpoint.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.llms.openai.common_utils import OpenAIError
|
||||
from litellm.llms.openai.responses.count_tokens.handler import (
|
||||
OpenAICountTokensHandler,
|
||||
)
|
||||
from litellm.llms.openai.responses.count_tokens.transformation import (
|
||||
OpenAICountTokensConfig,
|
||||
)
|
||||
from litellm.types.utils import LlmProviders, TokenCountResponse
|
||||
|
||||
# Global handler instance - reuse across all token counting requests
|
||||
openai_count_tokens_handler = OpenAICountTokensHandler()
|
||||
|
||||
|
||||
class OpenAITokenCounter(BaseTokenCounter):
|
||||
"""Token counter implementation for OpenAI provider using the Responses API."""
|
||||
|
||||
def should_use_token_counting_api(
|
||||
self,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
) -> bool:
|
||||
return custom_llm_provider == LlmProviders.OPENAI.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 = "",
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
system: Optional[Any] = None,
|
||||
) -> Optional[TokenCountResponse]:
|
||||
"""
|
||||
Count tokens using OpenAI's Responses API /input_tokens endpoint.
|
||||
"""
|
||||
if not messages:
|
||||
return None
|
||||
|
||||
deployment = deployment or {}
|
||||
litellm_params = deployment.get("litellm_params", {})
|
||||
|
||||
# Get OpenAI API key from deployment config or environment
|
||||
api_key = litellm_params.get("api_key")
|
||||
if not api_key:
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
|
||||
if not api_key:
|
||||
verbose_logger.warning("No OpenAI API key found for token counting")
|
||||
return None
|
||||
|
||||
api_base = litellm_params.get("api_base")
|
||||
|
||||
# Convert chat messages to Responses API input format
|
||||
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
|
||||
messages
|
||||
)
|
||||
|
||||
# Use system param if instructions not extracted from messages
|
||||
if instructions is None and system is not None:
|
||||
instructions = system if isinstance(system, str) else str(system)
|
||||
|
||||
# If no input items were produced (e.g., system-only messages), fall back to local counting
|
||||
if not input_items:
|
||||
return None
|
||||
|
||||
try:
|
||||
result = await openai_count_tokens_handler.handle_count_tokens_request(
|
||||
model=model_to_use,
|
||||
input=input_items if input_items is not None else [],
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
tools=tools,
|
||||
instructions=instructions,
|
||||
)
|
||||
|
||||
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="openai_api",
|
||||
original_response=result,
|
||||
)
|
||||
except OpenAIError as e:
|
||||
verbose_logger.warning(
|
||||
f"OpenAI 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="openai_api",
|
||||
error=True,
|
||||
error_message=e.message,
|
||||
status_code=e.status_code,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error calling OpenAI CountTokens API: {e}")
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="openai_api",
|
||||
error=True,
|
||||
error_message=str(e),
|
||||
status_code=500,
|
||||
)
|
||||
|
||||
return None
|
||||
158
litellm/llms/openai/responses/count_tokens/transformation.py
Normal file
158
litellm/llms/openai/responses/count_tokens/transformation.py
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
"""
|
||||
OpenAI Responses API token counting transformation logic.
|
||||
|
||||
This module handles the transformation of requests to OpenAI's /v1/responses/input_tokens endpoint.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
|
||||
class OpenAICountTokensConfig:
|
||||
"""
|
||||
Configuration and transformation logic for OpenAI Responses API token counting.
|
||||
|
||||
OpenAI Responses API Token Counting Specification:
|
||||
- Endpoint: POST https://api.openai.com/v1/responses/input_tokens
|
||||
- Response: {"input_tokens": <number>}
|
||||
"""
|
||||
|
||||
def get_openai_count_tokens_endpoint(self, api_base: Optional[str] = None) -> str:
|
||||
base = api_base or "https://api.openai.com/v1"
|
||||
base = base.rstrip("/")
|
||||
return f"{base}/responses/input_tokens"
|
||||
|
||||
def transform_request_to_count_tokens(
|
||||
self,
|
||||
model: str,
|
||||
input: Union[str, List[Any]],
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
instructions: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Transform request to OpenAI Responses API token counting format.
|
||||
|
||||
The Responses API uses `input` (not `messages`) and `instructions` (not `system`).
|
||||
"""
|
||||
request: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"input": input,
|
||||
}
|
||||
|
||||
if instructions is not None:
|
||||
request["instructions"] = instructions
|
||||
|
||||
if tools is not None:
|
||||
request["tools"] = self._transform_tools_for_responses_api(tools)
|
||||
|
||||
return request
|
||||
|
||||
def get_required_headers(self, api_key: str) -> Dict[str, str]:
|
||||
return {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
def validate_request(
|
||||
self, model: str, input: Union[str, List[Any]]
|
||||
) -> None:
|
||||
if not model:
|
||||
raise ValueError("model parameter is required")
|
||||
|
||||
if not input:
|
||||
raise ValueError("input parameter is required")
|
||||
|
||||
@staticmethod
|
||||
def _transform_tools_for_responses_api(
|
||||
tools: List[Dict[str, Any]],
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Transform OpenAI chat tools format to Responses API tools format.
|
||||
|
||||
Chat format: {"type": "function", "function": {"name": "...", "parameters": {...}}}
|
||||
Responses format: {"type": "function", "name": "...", "parameters": {...}}
|
||||
"""
|
||||
transformed = []
|
||||
for tool in tools:
|
||||
if tool.get("type") == "function" and "function" in tool:
|
||||
func = tool["function"]
|
||||
item: Dict[str, Any] = {
|
||||
"type": "function",
|
||||
"name": func.get("name", ""),
|
||||
"description": func.get("description", ""),
|
||||
"parameters": func.get("parameters", {}),
|
||||
}
|
||||
if "strict" in func:
|
||||
item["strict"] = func["strict"]
|
||||
transformed.append(item)
|
||||
else:
|
||||
# Pass through non-function tools (e.g., web_search, file_search)
|
||||
transformed.append(tool)
|
||||
return transformed
|
||||
|
||||
@staticmethod
|
||||
def messages_to_responses_input(
|
||||
messages: List[Dict[str, Any]],
|
||||
) -> tuple:
|
||||
"""
|
||||
Convert standard chat messages format to OpenAI Responses API input format.
|
||||
|
||||
Returns:
|
||||
(input_items, instructions) tuple where instructions is extracted
|
||||
from system/developer messages.
|
||||
"""
|
||||
input_items: List[Dict[str, Any]] = []
|
||||
instructions_parts: List[str] = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "")
|
||||
content = msg.get("content") or ""
|
||||
|
||||
if role in ("system", "developer"):
|
||||
# Extract system/developer messages as instructions
|
||||
if isinstance(content, str):
|
||||
instructions_parts.append(content)
|
||||
elif isinstance(content, list):
|
||||
# Handle content blocks - extract text
|
||||
text_parts = []
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
text_parts.append(block.get("text", ""))
|
||||
elif isinstance(block, str):
|
||||
text_parts.append(block)
|
||||
instructions_parts.append("\n".join(text_parts))
|
||||
elif role == "user":
|
||||
if isinstance(content, list):
|
||||
# Extract text from content blocks for Responses API
|
||||
text_parts = []
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
text_parts.append(block.get("text", ""))
|
||||
elif isinstance(block, str):
|
||||
text_parts.append(block)
|
||||
content = "\n".join(text_parts)
|
||||
input_items.append({"role": "user", "content": content})
|
||||
elif role == "assistant":
|
||||
# Map tool_calls to Responses API function_call items
|
||||
tool_calls = msg.get("tool_calls")
|
||||
if content:
|
||||
input_items.append({"role": "assistant", "content": content})
|
||||
if tool_calls:
|
||||
for tc in tool_calls:
|
||||
func = tc.get("function", {})
|
||||
input_items.append({
|
||||
"type": "function_call",
|
||||
"call_id": tc.get("id", ""),
|
||||
"name": func.get("name", ""),
|
||||
"arguments": func.get("arguments", ""),
|
||||
})
|
||||
elif not content:
|
||||
input_items.append({"role": "assistant", "content": content})
|
||||
elif role == "tool":
|
||||
input_items.append({
|
||||
"type": "function_call_output",
|
||||
"call_id": msg.get("tool_call_id", ""),
|
||||
"output": content if isinstance(content, str) else str(content),
|
||||
})
|
||||
|
||||
instructions = "\n".join(instructions_parts) if instructions_parts else None
|
||||
return input_items, instructions
|
||||
105
litellm/main.py
105
litellm/main.py
|
|
@ -7529,6 +7529,111 @@ def stream_chunk_builder( # noqa: PLR0915
|
|||
)
|
||||
|
||||
|
||||
########## Token Counting API ##########
|
||||
|
||||
|
||||
async def acount_tokens(
|
||||
model: str,
|
||||
messages: Optional[List[Dict[str, Any]]] = None,
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
system: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> "TokenCountResponse":
|
||||
"""
|
||||
Count tokens for a given model and messages using provider-specific APIs.
|
||||
|
||||
Routes to the appropriate provider's token counting API (OpenAI, Anthropic, etc.)
|
||||
for exact token counts. Falls back to local tiktoken-based counting for unsupported providers.
|
||||
|
||||
Args:
|
||||
model: The model identifier (e.g., "openai/gpt-4o", "anthropic/claude-3-5-sonnet-20241022")
|
||||
messages: The messages to count tokens for (standard chat format)
|
||||
tools: Optional tools/functions to include in token count
|
||||
system: Optional system message/instructions
|
||||
api_key: Optional API key (falls back to environment variable)
|
||||
api_base: Optional custom API base URL
|
||||
|
||||
Returns:
|
||||
TokenCountResponse with total_tokens and metadata
|
||||
"""
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.types.utils import LlmProviders, TokenCountResponse
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
# Determine provider from model string
|
||||
resolved_model, custom_llm_provider, dynamic_api_key, dynamic_api_base = (
|
||||
get_llm_provider(
|
||||
model=model,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
)
|
||||
)
|
||||
|
||||
# Use dynamic key/base if not explicitly provided
|
||||
if api_key is None:
|
||||
api_key = dynamic_api_key
|
||||
if api_base is None:
|
||||
api_base = dynamic_api_base
|
||||
|
||||
# Build deployment dict for the token counter
|
||||
deployment: Dict[str, Any] = {
|
||||
"litellm_params": {
|
||||
"model": model,
|
||||
"api_key": api_key,
|
||||
"api_base": api_base,
|
||||
}
|
||||
}
|
||||
|
||||
# Try to get provider-specific token counter
|
||||
try:
|
||||
llm_provider_enum = LlmProviders(custom_llm_provider)
|
||||
provider_model_info = ProviderConfigManager.get_provider_model_info(
|
||||
model=model, provider=llm_provider_enum
|
||||
)
|
||||
|
||||
if provider_model_info is not None:
|
||||
token_counter_instance = provider_model_info.get_token_counter()
|
||||
if (
|
||||
token_counter_instance is not None
|
||||
and token_counter_instance.should_use_token_counting_api(
|
||||
custom_llm_provider
|
||||
)
|
||||
):
|
||||
result = await token_counter_instance.count_tokens(
|
||||
model_to_use=resolved_model,
|
||||
messages=messages,
|
||||
contents=None,
|
||||
deployment=deployment,
|
||||
request_model=model,
|
||||
tools=tools,
|
||||
system=system,
|
||||
)
|
||||
if result is not None and not result.error:
|
||||
return result
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
f"Provider token counting failed for model={model}, falling back to local: {e}"
|
||||
)
|
||||
|
||||
# Fallback to local tiktoken-based token counting
|
||||
fallback_messages = messages or []
|
||||
if system and fallback_messages:
|
||||
fallback_messages = [{"role": "system", "content": system}] + fallback_messages
|
||||
local_count = litellm.token_counter(
|
||||
model=model,
|
||||
messages=fallback_messages,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
return TokenCountResponse(
|
||||
total_tokens=local_count,
|
||||
request_model=model,
|
||||
model_used=resolved_model,
|
||||
tokenizer_type="local_tokenizer",
|
||||
)
|
||||
|
||||
|
||||
# Cache for encoding to avoid repeated __getattr__ calls
|
||||
_encoding_cache: Optional[Any] = None
|
||||
|
||||
|
|
|
|||
|
|
@ -400,6 +400,142 @@ async def cursor_chat_completions(
|
|||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/responses/input_tokens",
|
||||
tags=["responses"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
@router.post(
|
||||
"/responses/input_tokens",
|
||||
tags=["responses"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
@router.post(
|
||||
"/openai/v1/responses/input_tokens",
|
||||
tags=["responses"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def count_response_input_tokens(
|
||||
request: Request,
|
||||
):
|
||||
"""
|
||||
Count input tokens for OpenAI Responses API format.
|
||||
|
||||
This endpoint follows the OpenAI Responses API token counting specification.
|
||||
It accepts the same parameters as the /v1/responses endpoint but returns
|
||||
token counts instead of generating a response.
|
||||
|
||||
Example usage:
|
||||
```
|
||||
curl -X POST "http://localhost:4000/v1/responses/input_tokens" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer your-key" \
|
||||
-d '{
|
||||
"model": "gpt-4o",
|
||||
"input": "Hello, how are you?"
|
||||
}'
|
||||
```
|
||||
|
||||
Returns: {"input_tokens": <number>}
|
||||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
_read_request_body,
|
||||
token_counter as internal_token_counter,
|
||||
)
|
||||
|
||||
try:
|
||||
request_data = await _read_request_body(request=request)
|
||||
data: dict = {**request_data}
|
||||
|
||||
model_name = data.get("model")
|
||||
input_data = data.get("input")
|
||||
|
||||
if not model_name:
|
||||
raise HTTPException(
|
||||
status_code=400, detail={"error": "model parameter is required"}
|
||||
)
|
||||
|
||||
if input_data is None:
|
||||
raise HTTPException(
|
||||
status_code=400, detail={"error": "input parameter is required"}
|
||||
)
|
||||
|
||||
# Convert Responses API `input` to chat messages format for the internal token counter
|
||||
messages: list = []
|
||||
instructions = data.get("instructions")
|
||||
|
||||
if isinstance(input_data, str):
|
||||
messages.append({"role": "user", "content": input_data})
|
||||
elif isinstance(input_data, list):
|
||||
for item in input_data:
|
||||
if isinstance(item, dict):
|
||||
role = item.get("role", "user")
|
||||
content = item.get("content", "")
|
||||
if item.get("type") == "function_call_output":
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"content": item.get("output", ""),
|
||||
"tool_call_id": item.get("call_id", ""),
|
||||
})
|
||||
elif item.get("type") == "function_call":
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"tool_calls": [{
|
||||
"id": item.get("call_id", ""),
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": item.get("name", ""),
|
||||
"arguments": item.get("arguments", ""),
|
||||
},
|
||||
}],
|
||||
})
|
||||
else:
|
||||
messages.append({"role": role, "content": content})
|
||||
elif isinstance(item, str):
|
||||
messages.append({"role": "user", "content": item})
|
||||
|
||||
from litellm.proxy._types import TokenCountRequest
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
token_request = TokenCountRequest(
|
||||
model=model_name,
|
||||
messages=messages,
|
||||
tools=data.get("tools"),
|
||||
system=instructions,
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
return {"input_tokens": _token_response_dict.get("total_tokens", 0)}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except ProxyException as e:
|
||||
status_code = int(e.code) if e.code and e.code.isdigit() else 500
|
||||
raise HTTPException(
|
||||
status_code=status_code,
|
||||
detail={"error": e.message},
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.response_api_endpoints.count_response_input_tokens(): Exception occurred - {}".format(
|
||||
str(e)
|
||||
)
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=500, detail={"error": "Internal server error"}
|
||||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/v1/responses/{response_id}",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
|
|
@ -904,3 +1040,5 @@ async def cancel_response(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
version=version,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,202 @@
|
|||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
from litellm.llms.openai.responses.count_tokens.transformation import (
|
||||
OpenAICountTokensConfig,
|
||||
)
|
||||
|
||||
|
||||
def test_transform_basic_request():
|
||||
"""Test basic request with model and input."""
|
||||
config = OpenAICountTokensConfig()
|
||||
|
||||
result = config.transform_request_to_count_tokens(
|
||||
model="gpt-4o",
|
||||
input="Hello, how are you?",
|
||||
)
|
||||
|
||||
assert result == {
|
||||
"model": "gpt-4o",
|
||||
"input": "Hello, how are you?",
|
||||
}
|
||||
|
||||
|
||||
def test_transform_with_list_input():
|
||||
"""Test request with list input format."""
|
||||
config = OpenAICountTokensConfig()
|
||||
|
||||
input_items = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
|
||||
result = config.transform_request_to_count_tokens(
|
||||
model="gpt-4o",
|
||||
input=input_items,
|
||||
)
|
||||
|
||||
assert result["model"] == "gpt-4o"
|
||||
assert result["input"] == input_items
|
||||
|
||||
|
||||
def test_transform_includes_instructions():
|
||||
"""Test that instructions are included when provided."""
|
||||
config = OpenAICountTokensConfig()
|
||||
|
||||
result = config.transform_request_to_count_tokens(
|
||||
model="gpt-4o",
|
||||
input="Hello",
|
||||
instructions="You are a helpful assistant.",
|
||||
)
|
||||
|
||||
assert result["instructions"] == "You are a helpful assistant."
|
||||
assert result["model"] == "gpt-4o"
|
||||
assert result["input"] == "Hello"
|
||||
|
||||
|
||||
def test_transform_includes_tools():
|
||||
"""Test that tools are included when provided."""
|
||||
config = OpenAICountTokensConfig()
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_weather",
|
||||
"description": "Get weather",
|
||||
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}},
|
||||
}
|
||||
]
|
||||
|
||||
result = config.transform_request_to_count_tokens(
|
||||
model="gpt-4o",
|
||||
input="What's the weather?",
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
assert result["tools"] == tools
|
||||
|
||||
|
||||
def test_transform_no_instructions_no_tools():
|
||||
"""Test that None values are not included."""
|
||||
config = OpenAICountTokensConfig()
|
||||
|
||||
result = config.transform_request_to_count_tokens(
|
||||
model="gpt-4o",
|
||||
input="Hello",
|
||||
instructions=None,
|
||||
tools=None,
|
||||
)
|
||||
|
||||
assert "instructions" not in result
|
||||
assert "tools" not in result
|
||||
|
||||
|
||||
def test_messages_to_responses_input_basic():
|
||||
"""Test converting basic chat messages to Responses API input format."""
|
||||
messages = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
{"role": "user", "content": "How are you?"},
|
||||
]
|
||||
|
||||
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
|
||||
|
||||
assert len(input_items) == 3
|
||||
assert input_items[0] == {"role": "user", "content": "Hello"}
|
||||
assert input_items[1] == {"role": "assistant", "content": "Hi there!"}
|
||||
assert input_items[2] == {"role": "user", "content": "How are you?"}
|
||||
assert instructions is None
|
||||
|
||||
|
||||
def test_messages_to_responses_input_with_system():
|
||||
"""Test that system messages are extracted as instructions."""
|
||||
messages = [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "user", "content": "Hello"},
|
||||
]
|
||||
|
||||
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
|
||||
|
||||
assert len(input_items) == 1
|
||||
assert input_items[0] == {"role": "user", "content": "Hello"}
|
||||
assert instructions == "You are helpful."
|
||||
|
||||
|
||||
def test_messages_to_responses_input_with_developer():
|
||||
"""Test that developer messages are extracted as instructions."""
|
||||
messages = [
|
||||
{"role": "developer", "content": "Be concise."},
|
||||
{"role": "user", "content": "Hello"},
|
||||
]
|
||||
|
||||
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
|
||||
|
||||
assert len(input_items) == 1
|
||||
assert instructions == "Be concise."
|
||||
|
||||
|
||||
def test_messages_to_responses_input_with_tool():
|
||||
"""Test that tool messages are converted to function_call_output."""
|
||||
messages = [
|
||||
{"role": "user", "content": "What's the weather?"},
|
||||
{"role": "tool", "content": "72°F", "tool_call_id": "call_123"},
|
||||
]
|
||||
|
||||
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
|
||||
|
||||
assert len(input_items) == 2
|
||||
assert input_items[1] == {
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_123",
|
||||
"output": "72°F",
|
||||
}
|
||||
|
||||
|
||||
def test_validate_request_valid():
|
||||
"""Test that valid requests pass validation."""
|
||||
config = OpenAICountTokensConfig()
|
||||
config.validate_request(model="gpt-4o", input="Hello")
|
||||
|
||||
|
||||
def test_validate_request_missing_model():
|
||||
"""Test that missing model raises ValueError."""
|
||||
config = OpenAICountTokensConfig()
|
||||
try:
|
||||
config.validate_request(model="", input="Hello")
|
||||
assert False, "Should have raised ValueError"
|
||||
except ValueError as e:
|
||||
assert "model" in str(e)
|
||||
|
||||
|
||||
def test_validate_request_missing_input():
|
||||
"""Test that missing input raises ValueError."""
|
||||
config = OpenAICountTokensConfig()
|
||||
try:
|
||||
config.validate_request(model="gpt-4o", input="")
|
||||
assert False, "Should have raised ValueError"
|
||||
except ValueError as e:
|
||||
assert "input" in str(e)
|
||||
|
||||
|
||||
def test_get_endpoint_default():
|
||||
"""Test default endpoint URL."""
|
||||
config = OpenAICountTokensConfig()
|
||||
assert config.get_openai_count_tokens_endpoint() == "https://api.openai.com/v1/responses/input_tokens"
|
||||
|
||||
|
||||
def test_get_endpoint_custom_base():
|
||||
"""Test custom API base URL."""
|
||||
config = OpenAICountTokensConfig()
|
||||
assert config.get_openai_count_tokens_endpoint("https://custom.api.com/v1") == "https://custom.api.com/v1/responses/input_tokens"
|
||||
|
||||
|
||||
def test_get_required_headers():
|
||||
"""Test required headers include Authorization."""
|
||||
config = OpenAICountTokensConfig()
|
||||
headers = config.get_required_headers("sk-test-key")
|
||||
|
||||
assert headers["Authorization"] == "Bearer sk-test-key"
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
160
tests/test_litellm/test_count_tokens_public_api.py
Normal file
160
tests/test_litellm/test_count_tokens_public_api.py
Normal file
|
|
@ -0,0 +1,160 @@
|
|||
"""
|
||||
Tests for litellm.acount_tokens() public API.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
import litellm
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
|
||||
def test_acount_tokens_routes_to_openai():
|
||||
"""Test that acount_tokens routes to OpenAI token counter for openai/ models."""
|
||||
with patch(
|
||||
"litellm.llms.openai.responses.count_tokens.token_counter.openai_count_tokens_handler.handle_count_tokens_request",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"input_tokens": 15},
|
||||
):
|
||||
result = asyncio.run(
|
||||
litellm.acount_tokens(
|
||||
model="openai/gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
api_key="sk-test-key",
|
||||
)
|
||||
)
|
||||
|
||||
assert result.total_tokens == 15
|
||||
assert result.tokenizer_type == "openai_api"
|
||||
assert result.request_model == "openai/gpt-4o"
|
||||
|
||||
|
||||
def test_acount_tokens_routes_to_anthropic():
|
||||
"""Test that acount_tokens routes to Anthropic token counter for anthropic/ models."""
|
||||
with patch(
|
||||
"litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler.handle_count_tokens_request",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"input_tokens": 20},
|
||||
):
|
||||
result = asyncio.run(
|
||||
litellm.acount_tokens(
|
||||
model="anthropic/claude-3-5-sonnet-20241022",
|
||||
messages=[{"role": "user", "content": "Hello Claude!"}],
|
||||
api_key="sk-ant-test-key",
|
||||
)
|
||||
)
|
||||
|
||||
assert result.total_tokens == 20
|
||||
assert result.tokenizer_type == "anthropic_api"
|
||||
assert result.request_model == "anthropic/claude-3-5-sonnet-20241022"
|
||||
|
||||
|
||||
def test_acount_tokens_fallback_to_local():
|
||||
"""Test that unsupported providers fall back to local tiktoken counting."""
|
||||
result = asyncio.run(
|
||||
litellm.acount_tokens(
|
||||
model="together_ai/meta-llama/Llama-3-8b-chat-hf",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
)
|
||||
)
|
||||
|
||||
assert result.total_tokens > 0
|
||||
assert result.tokenizer_type == "local_tokenizer"
|
||||
|
||||
|
||||
def test_acount_tokens_with_tools():
|
||||
"""Test that tools are passed through to the token counter."""
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather info",
|
||||
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
with patch(
|
||||
"litellm.llms.openai.responses.count_tokens.token_counter.openai_count_tokens_handler.handle_count_tokens_request",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"input_tokens": 30},
|
||||
) as mock_handler:
|
||||
result = asyncio.run(
|
||||
litellm.acount_tokens(
|
||||
model="openai/gpt-4o",
|
||||
messages=[{"role": "user", "content": "What's the weather?"}],
|
||||
tools=tools,
|
||||
api_key="sk-test-key",
|
||||
)
|
||||
)
|
||||
|
||||
assert result.total_tokens == 30
|
||||
mock_handler.assert_called_once()
|
||||
call_kwargs = mock_handler.call_args
|
||||
assert call_kwargs.kwargs.get("tools") == tools
|
||||
|
||||
|
||||
def test_acount_tokens_with_system():
|
||||
"""Test that system messages are passed through."""
|
||||
with patch(
|
||||
"litellm.llms.openai.responses.count_tokens.token_counter.openai_count_tokens_handler.handle_count_tokens_request",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"input_tokens": 25},
|
||||
):
|
||||
result = asyncio.run(
|
||||
litellm.acount_tokens(
|
||||
model="openai/gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
system="You are a helpful assistant.",
|
||||
api_key="sk-test-key",
|
||||
)
|
||||
)
|
||||
|
||||
assert result.total_tokens == 25
|
||||
|
||||
|
||||
def test_acount_tokens_api_error_falls_back():
|
||||
"""Test that API errors in token counting return error response."""
|
||||
from litellm.llms.openai.common_utils import OpenAIError
|
||||
|
||||
with patch(
|
||||
"litellm.llms.openai.responses.count_tokens.token_counter.openai_count_tokens_handler.handle_count_tokens_request",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=OpenAIError(status_code=401, message="Invalid API key"),
|
||||
):
|
||||
result = asyncio.run(
|
||||
litellm.acount_tokens(
|
||||
model="openai/gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
api_key="sk-bad-key",
|
||||
)
|
||||
)
|
||||
|
||||
# Should fall back to local tokenizer when provider API errors
|
||||
assert result.error is False
|
||||
assert result.tokenizer_type == "local_tokenizer"
|
||||
assert result.total_tokens > 0
|
||||
|
||||
|
||||
def test_acount_tokens_no_api_key_falls_back():
|
||||
"""Test that missing API key falls back to local counting."""
|
||||
env_backup = os.environ.pop("OPENAI_API_KEY", None)
|
||||
try:
|
||||
result = asyncio.run(
|
||||
litellm.acount_tokens(
|
||||
model="openai/gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
)
|
||||
)
|
||||
|
||||
# Should fall back to local tokenizer since no API key
|
||||
assert result.total_tokens > 0
|
||||
assert result.tokenizer_type == "local_tokenizer"
|
||||
finally:
|
||||
if env_backup:
|
||||
os.environ["OPENAI_API_KEY"] = env_backup
|
||||
Loading…
Add table
Reference in a new issue