feat: add conversation persistence and context management layer
- SQLAlchemy async database layer (SQLite dev, PostgreSQL prod) - Conversation and Message models with UUID primary keys - Token counting utilities using litellm - Context summarization at 80% token threshold - REST API endpoints for multi-turn conversations - 19 conversation tests, 6 token tests (176 total passing) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,74 @@
|
||||
"""
|
||||
Token counting utilities for context management.
|
||||
|
||||
Uses litellm for accurate multi-model token counting.
|
||||
"""
|
||||
from src.shared.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Default model for token counting (Mistral Nemo)
|
||||
DEFAULT_MODEL = "mistral/mistral-nemo"
|
||||
|
||||
|
||||
def count_tokens(text: str, model: str = DEFAULT_MODEL) -> 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)
|
||||
except Exception as e:
|
||||
# Fallback to rough estimate if litellm 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:
|
||||
"""
|
||||
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)
|
||||
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
|
||||
Reference in New Issue
Block a user