scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(count_tokens): include system and tools in token counting API requests

Browse source 
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)
This commit is contained in:
Chesars 2026-02-27 15:39:35 -03:00
parent adba088df2
commit c4458c09fe
15 changed files with 331 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/anthropic/count_tokens/handler.py
View file

@ -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}")

4
litellm/llms/anthropic/count_tokens/token_counter.py
View file

@ -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:

26
litellm/llms/anthropic/count_tokens/transformation.py
View file

@ -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.

4
litellm/llms/azure_ai/anthropic/count_tokens/handler.py
View file

@ -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}")

4
litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py
View file

@ -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:

2
litellm/llms/base_llm/base_utils.py
View file

@ -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

10
litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
View file

@ -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)

92
litellm/llms/bedrock/count_tokens/transformation.py
View file

@ -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]:

1
litellm/llms/gemini/common_utils.py
View file

@ -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

1
litellm/llms/vertex_ai/common_utils.py
View file

@ -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

3
litellm/proxy/_types.py
View file

@ -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"""

7
litellm/proxy/anthropic_endpoints/endpoints.py
View file

@ -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(

4
litellm/proxy/proxy_server.py
View file

@ -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

92
tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py Normal file
View file

@ -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

120
tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py
View file

@ -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_"