Utilities
Token Counting
The TokenCounter utility provides provider-agnostic token counting for any model supported by LiteLLM. This allows you to monitor context usage regardless of whether you’re using OpenAI, Anthropic, Google, or any other supported provider.
TokenCountResult
A named tuple containing token count information:
- class TokenCountResult
Named tuple containing token count information.
- total: int
LiteLLM token count for the complete serialized context. With tools but no messages, this is schema-only overhead and excludes request framing.
- system_prompt: int
System message tokens, including request framing when a system message is present. In JSON modes, this also includes the output schema.
TokenCounter
The main utility class for counting tokens:
- class TokenCounter
Utility class for counting tokens in messages using LiteLLM.
- count_messages(model: str, messages: List[Dict[str, Any]]) int
Count tokens in a list of messages.
- Parameters:
model – The model name (e.g., “gpt-4”, “claude-3-opus-20240229”)
messages – List of message dictionaries with “role” and “content” keys
- Returns:
Number of tokens
- count_text(model: str, text: str) int
Count tokens in a text string.
- Parameters:
model – The model name
text – The text to count tokens for
- Returns:
Number of tokens
- get_max_tokens(model: str) int | None
Get the maximum context window for a model.
- Parameters:
model – The model name
- Returns:
Maximum tokens, or None if unknown
- count_context(model: str, system_messages: List[Dict], history_messages: List[Dict], tools: List[Dict] | None = None) TokenCountResult
Count tokens for a complete context with an additive breakdown.
- Parameters:
model – The model name
system_messages – System prompt messages
history_messages – Conversation history messages
tools – Optional tool definitions
- Returns:
TokenCountResult with detailed breakdown
Usage Example
from atomic_agents.utils import TokenCounter, TokenCountResult
# Direct usage
counter = TokenCounter()
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello!"},
{"role": "assistant", "content": "Hi there! How can I help?"},
]
# Count tokens in messages
token_count = counter.count_messages("gpt-4", messages)
# Get max context window
max_tokens = counter.get_max_tokens("gpt-4")
# Count complete context with breakdown
result = counter.count_context(
model="gpt-4",
system_messages=[{"role": "system", "content": "You are helpful."}],
history_messages=[{"role": "user", "content": "Hello!"}],
)
print(f"Total: {result.total}, System: {result.system_prompt}, History: {result.history}")
if result.utilization:
print(f"Context utilization: {result.utilization:.1%}")
Using with AtomicAgent
The easiest way to get token counts is through the agent’s get_context_token_count() method. The agent computes accurate token counts on-demand by serializing the context exactly as Instructor does, including output schema overhead and multimodal content:
# Get accurate token count at any time - always returns a result
token_info = agent.get_context_token_count()
print(f"Total tokens: {token_info.total}")
print(f"System prompt (with schema): {token_info.system_prompt} tokens")
print(f"History: {token_info.history} tokens")
print(f"Tools: {token_info.tools} tokens")
if token_info.utilization:
print(f"Context utilization: {token_info.utilization:.1%}")
The breakdown is additive: system_prompt + history + tools == total. Request
framing is attributed to system_prompt when a system message is present, or to
history otherwise. When tools are provided without system or history messages,
total and tools report schema-only overhead and exclude request framing. In
JSON modes, the output schema is included in system_prompt; in TOOLS mode, tool
definitions are reported through tools.
For requests with messages, the total matches LiteLLM’s count for the complete serialized request.
Tool Message Formatting
- atomic_agents.utils.format_tool_message.format_tool_message(tool_call: Type[BaseModel], tool_id: str | None = None) Dict[source]
Formats a message for a tool call.
- Parameters:
tool_call (Type[BaseModel]) – The Pydantic model instance representing the tool call.
tool_id (str, optional) – The unique identifier for the tool call. If not provided, a random UUID will be generated.
- Returns:
A formatted message dictionary for the tool call.
- Return type:
Dict