""" Optimized conversation memory manager that only loads conversations on demand """ import os import json import time import uuid import traceback import logging from typing import Dict, List, Any, Optional from utils.conversation_memory import Message, ConversationMemoryManager from utils.conversation_memory import Conversation as BaseConversation # Set up logger logger = logging.getLogger("mango_ai_agent") class OptimizedConversation(BaseConversation): """Extended Conversation class with additional functionality""" @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'OptimizedConversation': """Create a conversation from dictionary""" conversation = cls( conversation_id=data["conversation_id"], max_history=len(data.get("messages", [])) ) conversation.created_at = data.get("created_at", time.time()) conversation.last_updated = data.get("last_updated", time.time()) conversation.messages = [Message.from_dict(msg) for msg in data.get("messages", [])] return conversation class OptimizedConversationMemoryManager(ConversationMemoryManager): """Manages conversations with persistence to disk and on-demand loading""" def __init__(self, storage_dir: str = "conversation_data", max_conversations: int = 1000, max_history: int = 50): """ Initialize the optimized conversation memory manager Args: storage_dir: Directory to store conversation data max_conversations: Maximum number of conversations to store max_history: Maximum number of messages per conversation """ super().__init__(max_conversations, max_history) self.storage_dir = storage_dir self.conversations = {} # Start with empty dictionary, load on demand # Create storage directory if it doesn't exist os.makedirs(self.storage_dir, exist_ok=True) def _get_conversation_path(self, conversation_id: str) -> str: """Get the file path for a conversation""" return os.path.join(self.storage_dir, f"{conversation_id}.json") def _load_conversation(self, conversation_id: str) -> Optional[OptimizedConversation]: """Load a single conversation from disk""" try: file_path = self._get_conversation_path(conversation_id) if os.path.exists(file_path): with open(file_path, 'r', encoding='utf-8') as f: data = json.load(f) conversation = OptimizedConversation.from_dict(data) self.conversations[conversation_id] = conversation return conversation return None except Exception as e: logger.error(f"Error loading conversation {conversation_id}: {str(e)}") logger.debug(traceback.format_exc()) return None def _save_conversation(self, conversation_id: str): """Save a conversation to disk""" try: conversation = self.get_conversation(conversation_id) if not conversation: return # Convert to dictionary data = conversation.to_dict() # Save to file with open(self._get_conversation_path(conversation_id), 'w', encoding='utf-8') as f: json.dump(data, f, ensure_ascii=False, indent=2) except Exception as e: logger.error(f"Error saving conversation {conversation_id}: {str(e)}") logger.debug(traceback.format_exc()) def get_conversation(self, conversation_id: str) -> Optional[OptimizedConversation]: """Get a conversation by ID, loading from disk if necessary""" if conversation_id in self.conversations: return self.conversations[conversation_id] # Try to load from disk return self._load_conversation(conversation_id) def get_context_string(self, conversation_id: str, max_messages: int = None) -> str: """Get conversation history as a string, loading from disk if necessary""" conversation = self.get_conversation(conversation_id) if not conversation: return "" # Use the specified max_messages or the default max_msg = max_messages if max_messages is not None else self.max_history # Format messages messages = conversation.messages[-max_msg:] if max_msg > 0 else conversation.messages formatted_messages = [] # Log the number of messages for debugging logger.info(f"Formatting {len(messages)} messages for conversation {conversation_id}") for message in messages: # Use the agent name from config if available try: from utils.agent_config import agent_config role_prefix = "User" if message.role == "user" else agent_config.AGENT_NAME except ImportError: role_prefix = "User" if message.role == "user" else "Assistant" # Log each message for debugging logger.debug(f"Message role: {message.role}, content: {message.content[:50]}...") formatted_messages.append(f"{role_prefix}: {message.content}") return "\n\n".join(formatted_messages) def create_conversation(self, conversation_id: str = None) -> OptimizedConversation: """Create a new conversation with optional ID and save it""" # Use provided ID or generate a new one conv_id = conversation_id if conversation_id else str(uuid.uuid4()) # Create conversation conversation = OptimizedConversation(conversation_id=conv_id) conversation.created_at = time.time() conversation.last_updated = time.time() # Add to memory and save self.conversations[conv_id] = conversation self._save_conversation(conv_id) return conversation def add_message(self, conversation_id: str, content: str, role: str) -> Optional[Message]: """Add a message to a conversation and save it""" # Get or load conversation conversation = self.get_conversation(conversation_id) # Create conversation if it doesn't exist if not conversation: conversation = self.create_conversation(conversation_id) # Create message message = Message( content=content, role=role, timestamp=time.time() ) # Add message to conversation conversation.messages.append(message) conversation.last_updated = time.time() # Trim if needed if len(conversation.messages) > self.max_history: conversation.messages = conversation.messages[-self.max_history:] # Save conversation self._save_conversation(conversation_id) return message def delete_conversation(self, conversation_id: str) -> bool: """Delete a conversation""" # Remove from memory if present if conversation_id in self.conversations: del self.conversations[conversation_id] # Delete the file if it exists file_path = self._get_conversation_path(conversation_id) if os.path.exists(file_path): os.remove(file_path) return True return False def list_conversations(self) -> List[str]: """List all available conversation IDs without loading them""" try: # Get all JSON files in the storage directory without loading content files = [f.replace('.json', '') for f in os.listdir(self.storage_dir) if f.endswith('.json')] logger.info(f"Found {len(files)} conversation files on disk") return files except Exception as e: logger.error(f"Error listing conversations: {str(e)}") logger.debug(traceback.format_exc()) return [] def cleanup_old_conversations(self, max_age_days: int = 30) -> int: """Delete conversations older than max_age_days without loading them all""" try: current_time = time.time() max_age_seconds = max_age_days * 24 * 60 * 60 deleted_count = 0 # Get list of conversation files files = [f for f in os.listdir(self.storage_dir) if f.endswith('.json')] for file in files: try: # Get file path file_path = os.path.join(self.storage_dir, file) conv_id = file.replace('.json', '') # Check if already loaded in memory if conv_id in self.conversations: conversation = self.conversations[conv_id] if (current_time - conversation.last_updated) > max_age_seconds: self.delete_conversation(conv_id) deleted_count += 1 else: # Check file modification time without loading content file_mod_time = os.path.getmtime(file_path) if (current_time - file_mod_time) > max_age_seconds: # Delete without loading os.remove(file_path) deleted_count += 1 except Exception as file_error: logger.error(f"Error processing file {file}: {str(file_error)}") return deleted_count except Exception as e: logger.error(f"Error cleaning up old conversations: {str(e)}") logger.debug(traceback.format_exc()) return 0 # Create a singleton instance optimized_conversation_memory = OptimizedConversationMemoryManager()