diff --git a/docs/my-website/docs/anthropic_count_tokens.md b/docs/my-website/docs/anthropic_count_tokens.md index 963172fec4e..5985516d69c 100644 --- a/docs/my-website/docs/anthropic_count_tokens.md +++ b/docs/my-website/docs/anthropic_count_tokens.md @@ -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 | diff --git a/docs/my-website/docs/count_tokens.md b/docs/my-website/docs/count_tokens.md new file mode 100644 index 00000000000..108e2e650f2 --- /dev/null +++ b/docs/my-website/docs/count_tokens.md @@ -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` + + + + +```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?" + }' +``` + + + + +```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} +``` + + + + +**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 +``` diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index ab102a69670..d19210d31ab 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -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], diff --git a/litellm/llms/openai/responses/count_tokens/__init__.py b/litellm/llms/openai/responses/count_tokens/__init__.py new file mode 100644 index 00000000000..8f129a6ff09 --- /dev/null +++ b/litellm/llms/openai/responses/count_tokens/__init__.py @@ -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", +] diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py new file mode 100644 index 00000000000..721d07796ee --- /dev/null +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -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": } + + 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)}", + ) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py new file mode 100644 index 00000000000..3d3a659075e --- /dev/null +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -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 diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py new file mode 100644 index 00000000000..3893775fc01 --- /dev/null +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -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": } + """ + + 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 diff --git a/litellm/main.py b/litellm/main.py index c3ac4c24ae2..8be923685ab 100644 --- a/litellm/main.py +++ b/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 diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 44e8c42b2c1..b694b31979c 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -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": } + """ + 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, ) + + diff --git a/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py new file mode 100644 index 00000000000..5b97ccf23a6 --- /dev/null +++ b/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -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" diff --git a/tests/test_litellm/test_count_tokens_public_api.py b/tests/test_litellm/test_count_tokens_public_api.py new file mode 100644 index 00000000000..81ba244796d --- /dev/null +++ b/tests/test_litellm/test_count_tokens_public_api.py @@ -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