memory_manager.py

"""
Memory Manager - Lưu trữ và quản lý lịch sử hội thoại
"""
import json
import os
from datetime import datetime
from typing import List, Dict, Optional
from pathlib import Path


class MemoryManager:
    def __init__(self, storage_dir: str = "conversation_history"):
        """
        Initialize Memory Manager
        
        Args:
            storage_dir: Thư mục lưu trữ lịch sử hội thoại
        """
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(exist_ok=True)
        self.current_session_file = None
    
    def create_new_session(self) -> str:
        """
        Tạo session mới
        
        Returns:
            Session ID
        """
        session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.current_session_file = self.storage_dir / f"session_{session_id}.json"
        
        # Tạo file session mới
        session_data = {
            "session_id": session_id,
            "created_at": datetime.now().isoformat(),
            "messages": []
        }
        
        self._save_session(session_data)
        return session_id
    
    def save_message(self, role: str, content: str, metadata: Dict = None):
        """
        Lưu tin nhắn vào session hiện tại
        
        Args:
            role: 'user' hoặc 'assistant'
            content: Nội dung tin nhắn
            metadata: Thông tin bổ sung (action, q_values, etc.)
        """
        if not self.current_session_file:
            self.create_new_session()
        
        session_data = self._load_session()
        
        message = {
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        }
        
        if metadata:
            message.update(metadata)
        
        session_data["messages"].append(message)
        self._save_session(session_data)
    
    def load_session(self, session_id: str) -> Optional[Dict]:
        """
        Load session theo ID
        
        Args:
            session_id: ID của session cần load
            
        Returns:
            Session data hoặc None nếu không tìm thấy
        """
        session_file = self.storage_dir / f"session_{session_id}.json"
        
        if not session_file.exists():
            return None
        
        with open(session_file, 'r', encoding='utf-8') as f:
            return json.load(f)
    
    def get_all_sessions(self) -> List[Dict]:
        """
        Lấy danh sách tất cả sessions
        
        Returns:
            List các session info
        """
        sessions = []
        
        for file in self.storage_dir.glob("session_*.json"):
            try:
                with open(file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    sessions.append({
                        "session_id": data["session_id"],
                        "created_at": data["created_at"],
                        "message_count": len(data["messages"]),
                        "file": str(file)
                    })
            except Exception as e:
                print(f"Error loading session {file}: {e}")
        
        # Sắp xếp theo thời gian tạo (mới nhất trước)
        sessions.sort(key=lambda x: x["created_at"], reverse=True)
        return sessions
    
    def delete_session(self, session_id: str) -> bool:
        """
        Xóa session
        
        Args:
            session_id: ID của session cần xóa
            
        Returns:
            True nếu xóa thành công
        """
        session_file = self.storage_dir / f"session_{session_id}.json"
        
        if session_file.exists():
            session_file.unlink()
            return True
        
        return False
    
    def _load_session(self) -> Dict:
        """Load session hiện tại"""
        if not self.current_session_file or not self.current_session_file.exists():
            return {
                "session_id": "default",
                "created_at": datetime.now().isoformat(),
                "messages": []
            }
        
        with open(self.current_session_file, 'r', encoding='utf-8') as f:
            return json.load(f)
    
    def _save_session(self, session_data: Dict):
        """Lưu session data"""
        if not self.current_session_file:
            return
        
        with open(self.current_session_file, 'w', encoding='utf-8') as f:
            json.dump(session_data, f, ensure_ascii=False, indent=2)
    
    def get_current_messages(self) -> List[Dict]:
        """Lấy tất cả messages của session hiện tại"""
        session_data = self._load_session()
        return session_data.get("messages", [])
    
    def export_session(self, session_id: str, output_file: str):
        """
        Export session ra file
        
        Args:
            session_id: ID của session
            output_file: Đường dẫn file output
        """
        session_data = self.load_session(session_id)
        
        if not session_data:
            return False
        
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(session_data, f, ensure_ascii=False, indent=2)
        
        return True