jpmschweitzerandClaude Opus 4.5
2f97041aa9
- litellm had dependency conflicts with pydantic-ai - tiktoken is lighter and already required by pydantic-ai - Updated documentation (README.md, architecture.md) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
86 lines
2.2 KiB
Python
86 lines
2.2 KiB
Python
"""
|
|
Token counting utilities for context management.
|
|
|
|
Uses tiktoken for token counting. While tiktoken is OpenAI's tokenizer,
|
|
cl100k_base encoding provides reasonable estimates for most LLMs.
|
|
"""
|
|
from functools import lru_cache
|
|
|
|
from src.shared.logging import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def _get_encoding():
|
|
"""Get tiktoken encoding (cached)."""
|
|
import tiktoken
|
|
# cl100k_base is used by GPT-4 and provides reasonable estimates for most models
|
|
return tiktoken.get_encoding("cl100k_base")
|
|
|
|
|
|
def count_tokens(text: str) -> int:
|
|
"""
|
|
Count tokens in a text string.
|
|
|
|
Args:
|
|
text: Text to count tokens for
|
|
|
|
Returns:
|
|
Token count
|
|
"""
|
|
try:
|
|
encoding = _get_encoding()
|
|
return len(encoding.encode(text))
|
|
except Exception as e:
|
|
# Fallback to rough estimate if tiktoken fails
|
|
logger.warning(f"Token counting failed, using estimate: {e}")
|
|
return len(text) // 4
|
|
|
|
|
|
def count_message_tokens(messages: list[dict[str, str]]) -> int:
|
|
"""
|
|
Count tokens for a list of chat messages.
|
|
|
|
Args:
|
|
messages: List of message dicts with 'role' and 'content' keys
|
|
|
|
Returns:
|
|
Total token count including message overhead
|
|
"""
|
|
try:
|
|
encoding = _get_encoding()
|
|
total = 0
|
|
for msg in messages:
|
|
# Each message has ~4 tokens overhead for role/formatting
|
|
total += 4
|
|
total += len(encoding.encode(msg.get("content", "")))
|
|
total += len(encoding.encode(msg.get("role", "")))
|
|
# Add 2 tokens for assistant response priming
|
|
total += 2
|
|
return total
|
|
except Exception as e:
|
|
# Fallback to rough estimate
|
|
logger.warning(f"Token counting failed, using estimate: {e}")
|
|
total = 0
|
|
for msg in messages:
|
|
total += len(msg.get("content", "")) // 4
|
|
total += 4 # Overhead per message
|
|
return total
|
|
|
|
|
|
def estimate_tokens(text: str) -> int:
|
|
"""
|
|
Quick token estimate without external library.
|
|
|
|
Uses ~4 characters per token heuristic.
|
|
Less accurate but faster for rough estimates.
|
|
|
|
Args:
|
|
text: Text to estimate
|
|
|
|
Returns:
|
|
Estimated token count
|
|
"""
|
|
return len(text) // 4
|