diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 5b5354228f9..07481917afe 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -31,6 +31,8 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): api_key: str, api_base: Optional[str] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Dict[str, Any]: """ Handle a CountTokens request using httpx. @@ -60,6 +62,8 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): request_body = self.transform_request_to_count_tokens( model=model, messages=messages, + tools=tools, + system=system, ) verbose_logger.debug(f"Transformed request: {request_body}") diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 266b2794fc3..93989c58547 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -30,6 +30,8 @@ class AnthropicTokenCounter(BaseTokenCounter): 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 Anthropic's CountTokens API. @@ -66,6 +68,8 @@ class AnthropicTokenCounter(BaseTokenCounter): model=model_to_use, messages=messages, api_key=api_key, + tools=tools, + system=system, ) if result is not None: diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index c3ad72436b4..ea4d60a60ef 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -4,7 +4,7 @@ Anthropic CountTokens API transformation logic. This module handles the transformation of requests to Anthropic's CountTokens API format. """ -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION @@ -32,27 +32,27 @@ class AnthropicCountTokensConfig: self, model: str, messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Dict[str, Any]: """ Transform request to Anthropic CountTokens format. - Input: - { - "model": "claude-3-5-sonnet-20241022", - "messages": [{"role": "user", "content": "Hello!"}] - } - - Output (Anthropic CountTokens format): - { - "model": "claude-3-5-sonnet-20241022", - "messages": [{"role": "user", "content": "Hello!"}] - } + Includes optional system and tools fields for accurate token counting. """ - return { + request: Dict[str, Any] = { "model": model, "messages": messages, } + if system is not None: + request["system"] = system + + if tools is not None: + request["tools"] = tools + + return request + def get_required_headers(self, api_key: str) -> Dict[str, str]: """ Get the required headers for the CountTokens API. diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 52a0bb8bb09..2cba27925c6 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -32,6 +32,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): api_base: str, litellm_params: Optional[Dict[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Dict[str, Any]: """ Handle a CountTokens request using httpx with Azure authentication. @@ -62,6 +64,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): request_body = self.transform_request_to_count_tokens( model=model, messages=messages, + tools=tools, + system=system, ) verbose_logger.debug(f"Transformed request: {request_body}") diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 14f92800079..afdfe9bdee9 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -32,6 +32,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): 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 Azure AI Anthropic's CountTokens API. @@ -79,6 +81,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): api_key=api_key, api_base=api_base, litellm_params=litellm_params, + tools=tools, + system=system, ) if result is not None: diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 9172a05e385..ecff9053dc5 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -24,6 +24,8 @@ class BaseTokenCounter(ABC): 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]: pass diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 54f8a8dbd65..772eb169689 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -30,6 +30,8 @@ class BedrockTokenCounter(BaseTokenCounter): 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 AWS Bedrock's CountTokens API. @@ -54,11 +56,17 @@ class BedrockTokenCounter(BaseTokenCounter): litellm_params = deployment.get("litellm_params", {}) # Build request data in the format expected by BedrockCountTokensHandler - request_data = { + request_data: Dict[str, Any] = { "model": model_to_use, "messages": messages, } + if tools: + request_data["tools"] = tools + + if system: + request_data["system"] = system + # Get the resolved model (strip prefixes like bedrock/, converse/, etc.) resolved_model = get_bedrock_base_model(model_to_use) diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index b313cc9df3c..64f1098e640 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -5,7 +5,8 @@ This module handles the transformation of requests from Anthropic Messages API f to AWS Bedrock's CountTokens API format and vice versa. """ -from typing import Any, Dict, List +import re +from typing import Any, Dict, List, Optional from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import get_bedrock_base_model @@ -75,46 +76,81 @@ class BedrockCountTokensConfig(BaseAWSLLM): input_type = self._detect_input_type(request_data) if input_type == "converse": - return self._transform_to_converse_format(request_data.get("messages", [])) + return self._transform_to_converse_format(request_data) else: return self._transform_to_invoke_model_format(request_data) def _transform_to_converse_format( - self, messages: List[Dict[str, Any]] + self, request_data: Dict[str, Any] ) -> Dict[str, Any]: - """Transform to Converse input format.""" - # Extract system messages if present - system_messages = [] + """Transform to Converse input format, including system and tools.""" + messages = request_data.get("messages", []) + system = request_data.get("system") + tools = request_data.get("tools") + + # Transform messages user_messages = [] - for message in messages: - if message.get("role") == "system": - system_messages.append({"text": message.get("content", "")}) - else: - # Transform message content to Bedrock format - transformed_message: Dict[str, Any] = {"role": message.get("role"), "content": []} + transformed_message: Dict[str, Any] = {"role": message.get("role"), "content": []} + content = message.get("content", "") + if isinstance(content, str): + transformed_message["content"].append({"text": content}) + elif isinstance(content, list): + transformed_message["content"] = content + user_messages.append(transformed_message) - # Handle content - ensure it's in the correct array format - content = message.get("content", "") - if isinstance(content, str): - # String content -> convert to text block - transformed_message["content"].append({"text": content}) - elif isinstance(content, list): - # Already in blocks format - use as is - transformed_message["content"] = content + converse_input: Dict[str, Any] = {"messages": user_messages} - user_messages.append(transformed_message) + # Transform system prompt (string or list of blocks → Bedrock format) + system_blocks = self._transform_system(system) + if system_blocks: + converse_input["system"] = system_blocks - # Build the converse input format - converse_input = {"messages": user_messages} + # Transform tools (Anthropic format → Bedrock toolConfig) + tool_config = self._transform_tools(tools) + if tool_config: + converse_input["toolConfig"] = tool_config - # Add system messages if present - if system_messages: - converse_input["system"] = system_messages - - # Build the complete request return {"input": {"converse": converse_input}} + def _transform_system(self, system: Optional[Any]) -> List[Dict[str, Any]]: + """Transform Anthropic system prompt to Bedrock system blocks.""" + if system is None: + return [] + if isinstance(system, str): + return [{"text": system}] + if isinstance(system, list): + # Already in blocks format (e.g. [{"type": "text", "text": "..."}]) + return [{"text": block.get("text", "")} for block in system if isinstance(block, dict)] + return [] + + def _transform_tools(self, tools: Optional[List[Dict[str, Any]]]) -> Optional[Dict[str, Any]]: + """Transform Anthropic tools to Bedrock toolConfig format.""" + if not tools: + return None + + bedrock_tools = [] + for tool in tools: + name = tool.get("name", "") + # Bedrock tool names must match [a-zA-Z][a-zA-Z0-9_]* and max 64 chars + name = re.sub(r"[^a-zA-Z0-9_]", "_", name) + if name and not name[0].isalpha(): + name = "t_" + name + name = name[:64] + + description = tool.get("description") or name + input_schema = tool.get("input_schema", {"type": "object", "properties": {}}) + + bedrock_tools.append({ + "toolSpec": { + "name": name, + "description": description, + "inputSchema": {"json": input_schema}, + } + }) + + return {"tools": bedrock_tools} + def _transform_to_invoke_model_format( self, request_data: Dict[str, Any] ) -> Dict[str, Any]: diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index e53829d3329..f99548c2c45 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -166,6 +166,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): contents: Optional[List[Dict[str, Any]]], deployment: Optional[Dict[str, Any]] = None, request_model: str = "", + **kwargs, ) -> Optional[TokenCountResponse]: import copy diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 02b69b94d94..244ea098ccc 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1042,6 +1042,7 @@ class VertexAITokenCounter(BaseTokenCounter): contents: Optional[List[Dict[str, Any]]], deployment: Optional[Dict[str, Any]] = None, request_model: str = "", + **kwargs, ) -> Optional[TokenCountResponse]: import copy diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 28311ab1b3b..7c022a433ae 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2857,6 +2857,9 @@ class TokenCountRequest(LiteLLMPydanticObjectBase): Google /countTokens endpoint expects contents to be a list of dicts with the following structure: """ + tools: Optional[List[dict]] = None + system: Optional[Any] = None + class CallInfo(LiteLLMPydanticObjectBase): """Used for slack budget alerting""" diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 77bb1f53e62..5b23b47923d 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -204,7 +204,12 @@ async def count_tokens( # Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest - token_request = TokenCountRequest(model=model_name, messages=messages) + token_request = TokenCountRequest( + model=model_name, + messages=messages, + tools=data.get("tools"), + system=data.get("system"), + ) # Call the internal token counter function with direct request flag set to False token_response = await internal_token_counter( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index be76c2ac5fb..6fbfbbc4966 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8355,6 +8355,8 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) prompt = request.prompt messages = request.messages contents = request.contents + tools = request.tools + system = request.system ######################################################### # Validate request @@ -8415,6 +8417,8 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) contents=contents, deployment=deployment, request_model=request.model, + tools=tools, + system=system, ) ######################################################### # Transfrom the Response to the well known format diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py new file mode 100644 index 00000000000..e982f735fd0 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -0,0 +1,92 @@ +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path +from litellm.llms.anthropic.count_tokens.transformation import ( + AnthropicCountTokensConfig, +) + + +def test_transform_basic_request(): + """Test basic request with only model and messages.""" + config = AnthropicCountTokensConfig() + + result = config.transform_request_to_count_tokens( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hello"}], + ) + + assert result == { + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "Hello"}], + } + + +def test_transform_includes_system(): + """Test that system prompt is included when provided.""" + config = AnthropicCountTokensConfig() + + result = config.transform_request_to_count_tokens( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hello"}], + system="You are a helpful assistant.", + ) + + assert result["system"] == "You are a helpful assistant." + assert result["model"] == "claude-3-5-sonnet" + assert result["messages"] == [{"role": "user", "content": "Hello"}] + + +def test_transform_includes_tools(): + """Test that tools are included when provided.""" + config = AnthropicCountTokensConfig() + + tools = [ + { + "name": "read_file", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, + } + ] + + result = config.transform_request_to_count_tokens( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hello"}], + tools=tools, + ) + + assert result["tools"] == tools + + +def test_transform_includes_system_and_tools(): + """Test that both system and tools are included together.""" + config = AnthropicCountTokensConfig() + + result = config.transform_request_to_count_tokens( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hello"}], + system="Be helpful", + tools=[{"name": "my_tool", "input_schema": {"type": "object"}}], + ) + + assert "system" in result + assert "tools" in result + assert "messages" in result + assert "model" in result + + +def test_transform_no_system_no_tools(): + """Test that None system/tools are not included.""" + config = AnthropicCountTokensConfig() + + result = config.transform_request_to_count_tokens( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hello"}], + system=None, + tools=None, + ) + + assert "system" not in result + assert "tools" not in result diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index ed8d6e1b359..699b67911dd 100644 --- a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -34,3 +34,123 @@ def test_transform_anthropic_to_bedrock_request(): assert "input" in result assert "converse" in result["input"] assert "messages" in result["input"]["converse"] + + +def test_transform_includes_system_prompt(): + """Test that system prompt is included in Bedrock converse format.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "system": "You are a helpful assistant.", + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + converse = result["input"]["converse"] + assert "system" in converse + assert converse["system"] == [{"text": "You are a helpful assistant."}] + + +def test_transform_includes_system_prompt_as_list(): + """Test that system prompt as list of blocks is handled.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "system": [{"type": "text", "text": "Block 1"}, {"type": "text", "text": "Block 2"}], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + converse = result["input"]["converse"] + assert converse["system"] == [{"text": "Block 1"}, {"text": "Block 2"}] + + +def test_transform_includes_tools(): + """Test that tools are transformed to Bedrock toolConfig format.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "tools": [ + { + "name": "read_file", + "description": "Read a file", + "input_schema": { + "type": "object", + "properties": {"path": {"type": "string"}}, + "required": ["path"], + }, + } + ], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + converse = result["input"]["converse"] + assert "toolConfig" in converse + tools = converse["toolConfig"]["tools"] + assert len(tools) == 1 + assert tools[0]["toolSpec"]["name"] == "read_file" + assert tools[0]["toolSpec"]["description"] == "Read a file" + assert tools[0]["toolSpec"]["inputSchema"]["json"]["type"] == "object" + + +def test_transform_includes_system_and_tools_together(): + """Test that both system and tools are included together.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "system": "Be helpful", + "tools": [ + {"name": "my_tool", "description": "A tool", "input_schema": {"type": "object", "properties": {}}}, + ], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + converse = result["input"]["converse"] + assert "system" in converse + assert "toolConfig" in converse + assert "messages" in converse + + +def test_transform_no_system_no_tools(): + """Test that missing system and tools don't add extra keys.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + converse = result["input"]["converse"] + assert "system" not in converse + assert "toolConfig" not in converse + + +def test_tool_name_sanitization(): + """Test that tool names are sanitized for Bedrock requirements.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "tools": [ + {"name": "my-tool!", "description": "A tool", "input_schema": {"type": "object", "properties": {}}}, + ], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + tool_name = result["input"]["converse"]["toolConfig"]["tools"][0]["toolSpec"]["name"] + # Should be sanitized: only [a-zA-Z0-9_] + assert tool_name == "my_tool_"