""" Conversation Management Handles message history and context window management. """ from typing import List, Dict, Optional from datetime import datetime import uuid from loguru import logger class ConversationManager: """ Manages conversation history and context. Features: - Track multiple conversations - Automatic context window management - Message summarization when context grows too large """ def __init__(self, max_messages: int = 50, max_tokens: int = 8000): self.conversations: Dict[str, List[Dict]] = {} self.max_messages = max_messages self.max_tokens = max_tokens def create_conversation(self) -> str: """Create a new conversation and return its ID.""" conv_id = str(uuid.uuid4()) self.conversations[conv_id] = [] logger.debug(f"Created conversation: {conv_id}") return conv_id def get_conversation(self, conv_id: str) -> List[Dict]: """Get messages for a conversation.""" return self.conversations.get(conv_id, []) def add_message( self, conv_id: str, role: str, content: str, metadata: Optional[Dict] = None ) -> None: """Add a message to a conversation.""" if conv_id not in self.conversations: self.conversations[conv_id] = [] message = { "role": role, "content": content, "timestamp": datetime.now().isoformat(), } if metadata: message["metadata"] = metadata self.conversations[conv_id].append(message) # Trim if needed self._trim_conversation(conv_id) def _trim_conversation(self, conv_id: str) -> None: """Trim conversation if it exceeds limits.""" messages = self.conversations.get(conv_id, []) if len(messages) > self.max_messages: # Keep system messages and last N messages system_messages = [m for m in messages if m["role"] == "system"] other_messages = [m for m in messages if m["role"] != "system"] # Keep last N-1 messages (plus system) keep_count = self.max_messages - len(system_messages) - 1 trimmed = system_messages + other_messages[-keep_count:] self.conversations[conv_id] = trimmed logger.debug(f"Trimmed conversation {conv_id} to {len(trimmed)} messages") def delete_conversation(self, conv_id: str) -> None: """Delete a conversation.""" if conv_id in self.conversations: del self.conversations[conv_id] logger.debug(f"Deleted conversation: {conv_id}") def list_conversations(self) -> List[str]: """List all conversation IDs.""" return list(self.conversations.keys()) def estimate_tokens(self, messages: List[Dict]) -> int: """Estimate token count for messages.""" # Rough estimate: ~4 characters per token total_chars = sum( len(m.get("content", "")) + len(m.get("role", "")) for m in messages ) return total_chars // 4