memory/memory_store.py

import os
import json
from datetime import datetime
from typing import List, Dict, Optional
from utils.logger import setup_logger
from utils.encryption import encryption_manager

logger = setup_logger("memory_store")

class MemoryStore:
    def __init__(self, sessions_dir: str = "memory/sessions"):
        self.sessions_dir = sessions_dir
        if not os.path.exists(self.sessions_dir):
            os.makedirs(self.sessions_dir)
            logger.info(f"Created sessions directory: {self.sessions_dir}")

    def _get_session_path(self, session_id: str) -> str:
        return os.path.join(self.sessions_dir, f"{session_id}.json")

    def _get_log_path(self, session_id: str) -> str:
        return os.path.join(self.sessions_dir, f"{session_id}.log.json")

    def save_interaction(self, session_id: str, user_text: str, assistant_text: str):
        """Saves a turn to a delta log to avoid full file rewrite."""
        log_path = self._get_log_path(session_id)
        
        # 1. Load delta log
        log_data = encryption_manager.safe_read(log_path)
        if not log_data:
            log_data = []

        # 2. Add entry
        log_data.append({
            "user": user_text,
            "assistant": assistant_text,
            "timestamp": datetime.now().isoformat()
        })
        
        # 3. Save delta log
        try:
            encryption_manager.safe_write(log_path, log_data)
            
            # 4. Selective Consolidation (Every 50 interactions)
            if len(log_data) >= 50:
                self.consolidate(session_id)
                
            logger.info(f"Interaction logged (delta) for session {session_id}")
        except Exception as e:
            logger.error(f"Error saving delta log for {session_id}: {e}")

    def consolidate(self, session_id: str):
        """Merges delta log into main session file."""
        main_path = self._get_session_path(session_id)
        log_path = self._get_log_path(session_id)
        
        main_data = encryption_manager.safe_read(main_path) or {"session_id": session_id, "history": []}
        log_data = encryption_manager.safe_read(log_path) or []
        
        if not log_data:
            return

        logger.info(f"Consolidating memory for session {session_id}...")
        main_data["history"].extend(log_data)
        
        # Add summary markers or logic could go here
        
        try:
            encryption_manager.safe_write(main_path, main_data)
            # Remove the log file (safe_write works on .json but handles .enc)
            # To delete, we need the actual path with .enc
            full_log_path = log_path + ".enc"
            if os.path.exists(full_log_path):
                os.remove(full_log_path)
            logger.info(f"Consolidation complete for {session_id}")
        except Exception as e:
            logger.error(f"Consolidation failed for {session_id}: {e}")

    def get_recent_history(self, session_id: str, limit: int = 3) -> List[Dict[str, str]]:
        """Retrieves history by merging main file and delta log."""
        main_path = self._get_session_path(session_id)
        log_path = self._get_log_path(session_id)
        
        history = []
        try:
            main_data = encryption_manager.safe_read(main_path)
            if main_data:
                history.extend(main_data.get("history", []))
            
            log_data = encryption_manager.safe_read(log_path)
            if log_data:
                history.extend(log_data)
                
            return history[-limit:]
        except Exception as e:
            logger.error(f"Error reading history for {session_id}: {e}")
            return []