File size: 1,616 Bytes
3370983 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
"""Token counting utilities for context window management."""
from typing import List, Any
import tiktoken
def count_tokens_for_messages(messages: List[Any], model: str = "gpt-4o") -> int:
"""
Calculate token usage for a list of messages using tiktoken.
Handles text content, function calls, and tool outputs with approximate
overhead calculations for the ChatML format.
Args:
messages: List of LangChain message objects.
model: Target model encoding to use.
Returns:
int: Total estimated token count.
"""
try:
encoding = tiktoken.encoding_for_model(model)
except KeyError:
encoding = tiktoken.get_encoding("cl100k_base")
num_tokens = 0
for message in messages:
# Every message follows <im_start>{role/name}\n{content}<im_end>\n
num_tokens += 4
# Handle content which might be a string or list of content blocks
content = getattr(message, "content", "")
if isinstance(content, str):
num_tokens += len(encoding.encode(content))
# If there are additional keys (like name, function_call, etc.) we should add them
if hasattr(message, "name") and message.name:
num_tokens += len(encoding.encode(message.name))
if hasattr(message, "tool_calls") and message.tool_calls:
for tool_call in message.tool_calls:
num_tokens += len(encoding.encode(str(tool_call)))
num_tokens += 2 # every reply is primed with <im_start>assistant
return num_tokens
|