From c4458c09fe4b423eec26ce9b5c3bb0b4abe6d821 Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 27 Feb 2026 15:39:35 -0300 Subject: [PATCH] fix(count_tokens): include system and tools in token counting API requests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The /v1/messages/count_tokens proxy endpoint was only passing `messages` to provider token counting APIs, discarding `system` and `tools`. This caused clients like Claude Code to receive artificially low token counts (e.g. 10 instead of 531), preventing proper context window management and leading to context overflow errors. Pass system and tools through the full chain: - TokenCountRequest → proxy_server → provider counters → API handlers - Bedrock: transform tools to toolConfig format, system to text blocks - Anthropic/Azure AI: pass through directly (same API format) --- .../llms/anthropic/count_tokens/handler.py | 4 + .../anthropic/count_tokens/token_counter.py | 4 + .../anthropic/count_tokens/transformation.py | 26 ++-- .../anthropic/count_tokens/handler.py | 4 + .../anthropic/count_tokens/token_counter.py | 4 + litellm/llms/base_llm/base_utils.py | 2 + .../count_tokens/bedrock_token_counter.py | 10 +- .../bedrock/count_tokens/transformation.py | 92 ++++++++++---- litellm/llms/gemini/common_utils.py | 1 + litellm/llms/vertex_ai/common_utils.py | 1 + litellm/proxy/_types.py | 3 + .../proxy/anthropic_endpoints/endpoints.py | 7 +- litellm/proxy/proxy_server.py | 4 + ...t_anthropic_count_tokens_transformation.py | 92 ++++++++++++++ ...est_bedrock_count_tokens_transformation.py | 120 ++++++++++++++++++ 15 files changed, 331 insertions(+), 43 deletions(-) create mode 100644 tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py 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 aeb9950b11c..ea60c1e2bad 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2832,6 +2832,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 f0b1e66818c..53627250fef 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8321,6 +8321,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 @@ -8381,6 +8383,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_"