fix: replace litellm with tiktoken for token counting
- 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>
This commit is contained in:
@@ -1,53 +1,64 @@
|
||||
"""
|
||||
Token counting utilities for context management.
|
||||
|
||||
Uses litellm for accurate multi-model token counting.
|
||||
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__)
|
||||
|
||||
# Default model for token counting (Mistral Nemo)
|
||||
DEFAULT_MODEL = "mistral/mistral-nemo"
|
||||
|
||||
@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, model: str = DEFAULT_MODEL) -> int:
|
||||
def count_tokens(text: str) -> int:
|
||||
"""
|
||||
Count tokens in a text string.
|
||||
|
||||
Args:
|
||||
text: Text to count tokens for
|
||||
model: Model identifier for tokenizer selection
|
||||
|
||||
Returns:
|
||||
Token count
|
||||
"""
|
||||
try:
|
||||
from litellm import token_counter
|
||||
return token_counter(model=model, text=text)
|
||||
encoding = _get_encoding()
|
||||
return len(encoding.encode(text))
|
||||
except Exception as e:
|
||||
# Fallback to rough estimate if litellm fails
|
||||
# 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]],
|
||||
model: str = DEFAULT_MODEL
|
||||
) -> int:
|
||||
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
|
||||
model: Model identifier for tokenizer selection
|
||||
|
||||
Returns:
|
||||
Total token count including message overhead
|
||||
"""
|
||||
try:
|
||||
from litellm import token_counter
|
||||
return token_counter(model=model, messages=messages)
|
||||
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}")
|
||||
|
||||
Reference in New Issue
Block a user