memory/memory_manager.py

# Session & Persistent Memory Manager
import json
import os
from datetime import datetime
from typing import Dict, Any, Optional

class MemoryManager:
    """Manages session data and persistent storage for the LinkedIn enhancer"""
    
    def __init__(self, storage_dir: str = "data"):
        self.storage_dir = storage_dir
        self.session_data = {}
        self.persistent_file = os.path.join(storage_dir, "persistent_data.json")
        
        # Create storage directory if it doesn't exist
        os.makedirs(storage_dir, exist_ok=True)
        
        # Load existing persistent data
        self.persistent_data = self._load_persistent_data()
    
    def store_session(self, profile_url: str, data: Dict[str, Any]) -> None:
        """
        Store session data for a specific profile
        
        Args:
            profile_url (str): LinkedIn profile URL as key
            data (Dict[str, Any]): Session data to store
        """
        session_key = self._create_session_key(profile_url)
        
        self.session_data[session_key] = {
            'timestamp': datetime.now().isoformat(),
            'profile_url': profile_url,
            'data': data
        }
    
    def get_session(self, profile_url: str) -> Optional[Dict[str, Any]]:
        """
        Retrieve session data for a specific profile
        
        Args:
            profile_url (str): LinkedIn profile URL
            
        Returns:
            Optional[Dict[str, Any]]: Session data if exists
        """
        session_key = self._create_session_key(profile_url)
        return self.session_data.get(session_key)
    
    def store_persistent(self, key: str, data: Any) -> None:
        """
        Store data persistently to disk
        
        Args:
            key (str): Storage key
            data (Any): Data to store
        """
        self.persistent_data[key] = {
            'timestamp': datetime.now().isoformat(),
            'data': data
        }
        
        self._save_persistent_data()
    
    def get_persistent(self, key: str) -> Optional[Any]:
        """
        Retrieve persistent data
        
        Args:
            key (str): Storage key
            
        Returns:
            Optional[Any]: Stored data if exists
        """
        stored_item = self.persistent_data.get(key)
        return stored_item['data'] if stored_item else None
    
    def store_user_preferences(self, user_id: str, preferences: Dict[str, Any]) -> None:
        """
        Store user preferences
        
        Args:
            user_id (str): User identifier
            preferences (Dict[str, Any]): User preferences
        """
        pref_key = f"user_preferences_{user_id}"
        self.store_persistent(pref_key, preferences)
    
    def get_user_preferences(self, user_id: str) -> Dict[str, Any]:
        """
        Retrieve user preferences
        
        Args:
            user_id (str): User identifier
            
        Returns:
            Dict[str, Any]: User preferences
        """
        pref_key = f"user_preferences_{user_id}"
        preferences = self.get_persistent(pref_key)
        return preferences if preferences else {}
    
    def store_analysis_history(self, profile_url: str, analysis: Dict[str, Any]) -> None:
        """
        Store analysis history for tracking improvements
        
        Args:
            profile_url (str): LinkedIn profile URL
            analysis (Dict[str, Any]): Analysis results
        """
        history_key = f"analysis_history_{self._create_session_key(profile_url)}"
        
        # Get existing history
        history = self.get_persistent(history_key) or []
        
        # Add new analysis with timestamp
        history.append({
            'timestamp': datetime.now().isoformat(),
            'analysis': analysis
        })
        
        # Keep only last 10 analyses
        history = history[-10:]
        
        self.store_persistent(history_key, history)
    
    def get_analysis_history(self, profile_url: str) -> list:
        """
        Retrieve analysis history for a profile
        
        Args:
            profile_url (str): LinkedIn profile URL
            
        Returns:
            list: Analysis history
        """
        history_key = f"analysis_history_{self._create_session_key(profile_url)}"
        return self.get_persistent(history_key) or []
    
    def clear_session(self, profile_url: str = None) -> None:
        """
        Clear session data
        
        Args:
            profile_url (str, optional): Specific profile to clear, or all if None
        """
        if profile_url:
            session_key = self._create_session_key(profile_url)
            self.session_data.pop(session_key, None)
        else:
            self.session_data.clear()
    
    def clear_session_cache(self, profile_url: str = None) -> None:
        """
        Clear session cache for a specific profile or all profiles
        
        Args:
            profile_url (str, optional): URL to clear cache for. If None, clears all.
        """
        if profile_url:
            session_key = self._create_session_key(profile_url)
            if session_key in self.session_data:
                del self.session_data[session_key]
                print(f"🗑️ Cleared session cache for: {profile_url}")
        else:
            self.session_data.clear()
            print("🗑️ Cleared all session cache")
    
    def force_refresh_session(self, profile_url: str) -> None:
        """
        Force refresh by clearing cache for a specific profile
        
        Args:
            profile_url (str): LinkedIn profile URL
        """
        self.clear_session_cache(profile_url)
        print(f"🔄 Forced refresh for: {profile_url}")

    def get_session_summary(self) -> Dict[str, Any]:
        """
        Get summary of current session data
        
        Returns:
            Dict[str, Any]: Session summary
        """
        return {
            'active_sessions': len(self.session_data),
            'sessions': list(self.session_data.keys()),
            'storage_location': self.storage_dir
        }
    
    def _create_session_key(self, profile_url: str) -> str:
        """Create a clean session key from profile URL"""
        # Extract username or create a hash-like key
        import hashlib
        return hashlib.md5(profile_url.encode()).hexdigest()[:16]
    
    def _load_persistent_data(self) -> Dict[str, Any]:
        """Load persistent data from disk"""
        if os.path.exists(self.persistent_file):
            try:
                with open(self.persistent_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except (json.JSONDecodeError, IOError):
                return {}
        return {}
    
    def _save_persistent_data(self) -> None:
        """Save persistent data to disk"""
        try:
            with open(self.persistent_file, 'w', encoding='utf-8') as f:
                json.dump(self.persistent_data, f, indent=2, ensure_ascii=False)
        except IOError as e:
            print(f"Warning: Could not save persistent data: {e}")
    
    def export_data(self, filename: str = None) -> str:
        """
        Export all data to a JSON file
        
        Args:
            filename (str, optional): Custom filename
            
        Returns:
            str: Path to exported file
        """
        if not filename:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"linkedin_enhancer_export_{timestamp}.json"
        
        export_path = os.path.join(self.storage_dir, filename)
        
        export_data = {
            'session_data': self.session_data,
            'persistent_data': self.persistent_data,
            'export_timestamp': datetime.now().isoformat()
        }
        
        with open(export_path, 'w', encoding='utf-8') as f:
            json.dump(export_data, f, indent=2, ensure_ascii=False)
        
        return export_path