src/rag/memory.py

"""Chat history and memory management for RAG conversations."""

import json
import os
from typing import List, Dict, Any, Optional, Tuple
from datetime import datetime
from pathlib import Path
from dataclasses import dataclass, asdict
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from src.core.config import config
from src.core.logging_config import get_logger

logger = get_logger(__name__)

@dataclass
class ChatMessage:
    """Represents a single chat message."""
    role: str  # "user" or "assistant"
    content: str
    timestamp: str
    sources: Optional[List[str]] = None  # Source documents used for the response
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'ChatMessage':
        """Create from dictionary."""
        return cls(**data)

@dataclass
class ChatSession:
    """Represents a chat session with history."""
    session_id: str
    created_at: str
    updated_at: str
    messages: List[ChatMessage]
    document_sources: List[str]  # Documents available in this session
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "session_id": self.session_id,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "messages": [msg.to_dict() for msg in self.messages],
            "document_sources": self.document_sources
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'ChatSession':
        """Create from dictionary."""
        messages = [ChatMessage.from_dict(msg) for msg in data.get("messages", [])]
        return cls(
            session_id=data["session_id"],
            created_at=data["created_at"],
            updated_at=data["updated_at"],
            messages=messages,
            document_sources=data.get("document_sources", [])
        )

class ChatMemoryManager:
    """Manages chat history and memory for RAG conversations."""
    
    def __init__(self, persist_directory: Optional[str] = None):
        """
        Initialize the chat memory manager.
        
        Args:
            persist_directory: Directory to persist chat history
        """
        if persist_directory is None:
            persist_directory = config.rag.chat_history_path
        
        self.persist_directory = Path(persist_directory)
        self.persist_directory.mkdir(parents=True, exist_ok=True)
        
        self.current_session: Optional[ChatSession] = None
        
        logger.info(f"ChatMemoryManager initialized with persist_directory={self.persist_directory}")
    
    def create_session(self, document_sources: Optional[List[str]] = None) -> str:
        """
        Create a new chat session.
        
        Args:
            document_sources: List of document sources available for this session
            
        Returns:
            Session ID
        """
        session_id = f"session_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}"
        now = datetime.now().isoformat()
        
        self.current_session = ChatSession(
            session_id=session_id,
            created_at=now,
            updated_at=now,
            messages=[],
            document_sources=document_sources or []
        )
        
        logger.info(f"Created new chat session: {session_id}")
        return session_id
    
    def add_message(self, role: str, content: str, sources: Optional[List[str]] = None) -> None:
        """
        Add a message to the current session.
        
        Args:
            role: "user" or "assistant"
            content: Message content
            sources: Source documents used (for assistant messages)
        """
        if self.current_session is None:
            self.create_session()
        
        message = ChatMessage(
            role=role,
            content=content,
            timestamp=datetime.now().isoformat(),
            sources=sources
        )
        
        self.current_session.messages.append(message)
        self.current_session.updated_at = datetime.now().isoformat()
        
        logger.info(f"Added {role} message to session {self.current_session.session_id}")
    
    def get_conversation_history(self, max_messages: Optional[int] = None) -> List[Tuple[str, str]]:
        """
        Get conversation history in Gradio chat format.
        
        Args:
            max_messages: Maximum number of messages to return
            
        Returns:
            List of (user_message, assistant_message) tuples
        """
        if not self.current_session or not self.current_session.messages:
            return []
        
        messages = self.current_session.messages
        if max_messages:
            messages = messages[-max_messages:]
        
        # Group messages into pairs
        history = []
        user_msg = None
        
        for msg in messages:
            if msg.role == "user":
                user_msg = msg.content
            elif msg.role == "assistant" and user_msg is not None:
                history.append((user_msg, msg.content))
                user_msg = None
        
        return history
    
    def get_context_messages(self, max_context_length: int = 4000) -> List[BaseMessage]:
        """
        Get recent messages formatted for LangChain context.
        
        Args:
            max_context_length: Maximum context length in characters
            
        Returns:
            List of LangChain message objects
        """
        if not self.current_session or not self.current_session.messages:
            return []
        
        context_messages = []
        current_length = 0
        
        # Start from the most recent messages and work backwards
        for msg in reversed(self.current_session.messages):
            msg_length = len(msg.content)
            
            if current_length + msg_length > max_context_length:
                break
            
            if msg.role == "user":
                context_messages.insert(0, HumanMessage(content=msg.content))
            elif msg.role == "assistant":
                context_messages.insert(0, AIMessage(content=msg.content))
            
            current_length += msg_length
        
        logger.info(f"Retrieved {len(context_messages)} context messages ({current_length} chars)")
        return context_messages
    
    def save_session(self) -> bool:
        """
        Save the current session to disk.
        
        Returns:
            True if successful, False otherwise
        """
        if not self.current_session:
            return False
        
        try:
            session_file = self.persist_directory / f"{self.current_session.session_id}.json"
            
            with open(session_file, 'w', encoding='utf-8') as f:
                json.dump(self.current_session.to_dict(), f, indent=2, ensure_ascii=False)
            
            logger.info(f"Saved session {self.current_session.session_id}")
            return True
            
        except Exception as e:
            logger.error(f"Error saving session: {e}")
            return False
    
    def load_session(self, session_id: str) -> bool:
        """
        Load a session from disk.
        
        Args:
            session_id: Session ID to load
            
        Returns:
            True if successful, False otherwise
        """
        try:
            session_file = self.persist_directory / f"{session_id}.json"
            
            if not session_file.exists():
                logger.warning(f"Session file not found: {session_id}")
                return False
            
            with open(session_file, 'r', encoding='utf-8') as f:
                session_data = json.load(f)
            
            self.current_session = ChatSession.from_dict(session_data)
            logger.info(f"Loaded session {session_id}")
            return True
            
        except Exception as e:
            logger.error(f"Error loading session {session_id}: {e}")
            return False
    
    def list_sessions(self) -> List[Dict[str, Any]]:
        """
        List all saved sessions.
        
        Returns:
            List of session metadata
        """
        sessions = []
        
        try:
            for session_file in self.persist_directory.glob("session_*.json"):
                try:
                    with open(session_file, 'r', encoding='utf-8') as f:
                        session_data = json.load(f)
                    
                    sessions.append({
                        "session_id": session_data["session_id"],
                        "created_at": session_data["created_at"],
                        "updated_at": session_data["updated_at"],
                        "message_count": len(session_data.get("messages", [])),
                        "document_sources": session_data.get("document_sources", [])
                    })
                    
                except Exception as e:
                    logger.warning(f"Error reading session file {session_file}: {e}")
                    
        except Exception as e:
            logger.error(f"Error listing sessions: {e}")
        
        # Sort by updated_at (most recent first)
        sessions.sort(key=lambda x: x["updated_at"], reverse=True)
        return sessions
    
    def clear_current_session(self) -> None:
        """Clear the current session."""
        self.current_session = None
        logger.info("Cleared current session")

# Global chat memory manager instance
chat_memory_manager = ChatMemoryManager()