""" Session Manager Module - Enhanced Version Handles per-user state isolation, persistence, and recovery for multi-user deployments. Features: - Save/load sessions to/from disk - Session restoration on reconnect - Graph state persistence - RLHF data collection with detailed interaction tracking - Session archiving before cleanup - Thread-safe operations """ import json import logging import threading import glob from datetime import datetime, timedelta from pathlib import Path from dataclasses import dataclass, field, asdict from typing import Dict, List, Optional, Any logger = logging.getLogger(__name__) @dataclass class ChatMessage: """A single chat message.""" role: str # "user", "assistant", "error" content: str timestamp: str = field(default_factory=lambda: datetime.now().isoformat()) language: str = "en" def to_dict(self) -> Dict: return asdict(self) @classmethod def from_dict(cls, data: Dict) -> "ChatMessage": return cls(**data) @dataclass class UserInteraction: """Records a single user interaction for RLHF data collection.""" timestamp: str action_type: str # 'query', 'prune', 'resurrect', 'inject', 'feedback' node_id: Optional[str] = None content: Optional[str] = None feedback: Optional[str] = None # 'correct', 'incorrect', 'partial' metadata: Dict[str, Any] = field(default_factory=dict) def to_dict(self) -> Dict: return { 'timestamp': self.timestamp, 'action_type': self.action_type, 'node_id': self.node_id, 'content': self.content, 'feedback': self.feedback, 'metadata': self.metadata } @classmethod def from_dict(cls, data: Dict) -> 'UserInteraction': return cls(**data) @dataclass class Session: """User session state with RLHF tracking.""" session_id: str created_at: str = field(default_factory=lambda: datetime.now().isoformat()) updated_at: str = field(default_factory=lambda: datetime.now().isoformat()) language: str = "en" chat_history: List[ChatMessage] = field(default_factory=list) interactions: List[UserInteraction] = field(default_factory=list) graph_state: Optional[Dict] = None metadata: Dict[str, Any] = field(default_factory=dict) feedback_count: Dict[str, int] = field( default_factory=lambda: {'correct': 0, 'incorrect': 0, 'partial': 0} ) def add_message(self, role: str, content: str, language: str = "en"): """Add a message to chat history.""" self.chat_history.append(ChatMessage( role=role, content=content, language=language )) self.updated_at = datetime.now().isoformat() def add_interaction( self, action_type: str, node_id: Optional[str] = None, content: Optional[str] = None, feedback: Optional[str] = None, metadata: Optional[Dict] = None ): """Record an interaction for RLHF.""" interaction = UserInteraction( timestamp=datetime.now().isoformat(), action_type=action_type, node_id=node_id, content=content, feedback=feedback, metadata=metadata or {} ) self.interactions.append(interaction) self.updated_at = datetime.now().isoformat() if feedback and feedback in self.feedback_count: self.feedback_count[feedback] += 1 def get_title(self, max_length: int = 35) -> str: """Get session title from first user message.""" for msg in self.chat_history: if msg.role == "user": title = msg.content[:max_length] return title + "..." if len(msg.content) > max_length else title return "New Chat" def to_dict(self) -> Dict: return { "session_id": self.session_id, "created_at": self.created_at, "updated_at": self.updated_at, "language": self.language, "chat_history": [m.to_dict() for m in self.chat_history], "interactions": [i.to_dict() for i in self.interactions], "graph_state": self.graph_state, "metadata": self.metadata, "feedback_summary": self.feedback_count, } def export_rlhf_data(self) -> Dict: """Export session data formatted for RLHF training.""" return { "session_id": self.session_id, "created_at": self.created_at, "language": self.language, "interactions": [i.to_dict() for i in self.interactions], "feedback_summary": self.feedback_count, "chat_history": [m.to_dict() for m in self.chat_history], "graph_state": self.graph_state, } @classmethod def from_dict(cls, data: Dict) -> "Session": session = cls( session_id=data["session_id"], created_at=data.get("created_at", datetime.now().isoformat()), updated_at=data.get("updated_at", datetime.now().isoformat()), language=data.get("language", "en"), chat_history=[ChatMessage.from_dict(m) for m in data.get("chat_history", [])], graph_state=data.get("graph_state"), metadata=data.get("metadata", {}), feedback_count=data.get("feedback_summary", {'correct': 0, 'incorrect': 0, 'partial': 0}), ) session.interactions = [ UserInteraction.from_dict(i) for i in data.get("interactions", []) ] return session class SessionManager: """ Manages user sessions with persistence, RLHF tracking, and cleanup. Thread-safe for multi-user deployments. """ def __init__( self, sessions_dir: str = "./data/sessions", max_sessions: int = 1000, session_max_age_hours: int = 24 ): self.sessions_dir = Path(sessions_dir) self.sessions_dir.mkdir(parents=True, exist_ok=True) # Archive directory for RLHF data self.archive_dir = self.sessions_dir / "archive" self.archive_dir.mkdir(parents=True, exist_ok=True) self.max_sessions = max_sessions self.session_max_age = timedelta(hours=session_max_age_hours) # In-memory session cache self._sessions: Dict[str, Session] = {} self._lock = threading.RLock() def get_session(self, session_id: str) -> Optional[Session]: """Get session by ID, loading from disk if needed.""" with self._lock: # Check memory cache if session_id in self._sessions: session = self._sessions[session_id] # Check if expired try: updated = datetime.fromisoformat(session.updated_at) if datetime.now() - updated > self.session_max_age: self._archive_session(session) del self._sessions[session_id] return None except Exception: pass return session # Try loading from disk session = self._load_from_disk(session_id) if session: self._sessions[session_id] = session return session def create_session(self, session_id: str) -> Session: """Create a new session.""" with self._lock: session = Session(session_id=session_id) self._sessions[session_id] = session self._save_to_disk(session) # Cleanup if too many sessions if len(self._sessions) > self.max_sessions: self._cleanup_old_sessions() logger.info(f"Created new session: {session_id}") return session def get_or_create(self, session_id: str) -> Session: """Get existing session or create new one.""" session = self.get_session(session_id) if session is None: session = self.create_session(session_id) return session def save_session(self, session_id: str): """Save session to disk.""" with self._lock: session = self._sessions.get(session_id) if session: session.updated_at = datetime.now().isoformat() self._save_to_disk(session) def update_graph_state(self, session_id: str, state: Dict): """Update session's graph state.""" with self._lock: session = self.get_or_create(session_id) session.graph_state = state session.updated_at = datetime.now().isoformat() self._save_to_disk(session) def update_language(self, session_id: str, language: str): """Update the detected language for a session.""" session = self.get_session(session_id) if session: session.language = language def record_interaction( self, session_id: str, action_type: str, node_id: Optional[str] = None, content: Optional[str] = None, feedback: Optional[str] = None, metadata: Optional[Dict] = None ): """Record a user interaction for RLHF.""" session = self.get_session(session_id) if session: session.add_interaction(action_type, node_id, content, feedback, metadata) # Auto-save on significant interactions if action_type in ['feedback', 'prune', 'inject', 'resurrect']: self.save_session(session_id) def add_feedback( self, session_id: str, node_id: str, feedback_type: str, context: str = "" ): """Record user feedback on reasoning.""" self.record_interaction( session_id=session_id, action_type='feedback', node_id=node_id, feedback=feedback_type, metadata={'context': context} ) def delete_session(self, session_id: str, archive: bool = True): """Delete a session, optionally archiving for RLHF.""" with self._lock: session = self._sessions.pop(session_id, None) if session and archive: self._archive_session(session) session_file = self.sessions_dir / f"{session_id}.json" if session_file.exists(): session_file.unlink() def list_sessions(self, limit: int = 20) -> List[Dict]: """List recent sessions with titles.""" sessions = [] # Get from disk for session_file in self.sessions_dir.glob("*.json"): try: with open(session_file) as f: data = json.load(f) sessions.append({ "id": data["session_id"], "title": self._get_title(data), "timestamp": data.get("updated_at", ""), "language": data.get("language", "en"), }) except Exception: continue # Sort by timestamp sessions.sort(key=lambda x: x["timestamp"], reverse=True) return sessions[:limit] def cleanup_stale_sessions(self): """Remove expired sessions, archiving for RLHF.""" cutoff = datetime.now() - self.session_max_age with self._lock: # Clean memory cache to_remove = [] for sid, session in self._sessions.items(): try: updated = datetime.fromisoformat(session.updated_at) if updated < cutoff: self._archive_session(session) to_remove.append(sid) except Exception: continue for sid in to_remove: del self._sessions[sid] # Clean disk for session_file in self.sessions_dir.glob("*.json"): if session_file.name == "archive": continue try: mtime = datetime.fromtimestamp(session_file.stat().st_mtime) if mtime < cutoff: # Load and archive before deleting try: with open(session_file) as f: data = json.load(f) session = Session.from_dict(data) self._archive_session(session) except Exception: pass session_file.unlink() logger.debug(f"Cleaned up stale session: {session_file.stem}") except Exception: continue if to_remove: logger.info(f"Cleaned up {len(to_remove)} stale sessions") def export_all_rlhf_data(self) -> List[Dict]: """Export all active sessions' RLHF data.""" with self._lock: return [session.export_rlhf_data() for session in self._sessions.values()] def get_stats(self) -> Dict: """Get session manager statistics.""" with self._lock: total_interactions = sum(len(s.interactions) for s in self._sessions.values()) total_feedback = sum(sum(s.feedback_count.values()) for s in self._sessions.values()) return { 'active_sessions': len(self._sessions), 'total_interactions': total_interactions, 'total_feedback': total_feedback, } def _archive_session(self, session: Session): """Archive session data for RLHF before deletion.""" try: filename = self.archive_dir / f"session_{session.session_id}_{session.created_at[:10]}.json" with open(filename, 'w', encoding='utf-8') as f: json.dump(session.export_rlhf_data(), f, indent=2, ensure_ascii=False) logger.info(f"Archived session {session.session_id}") except Exception as e: logger.error(f"Failed to archive session {session.session_id}: {e}") def _load_from_disk(self, session_id: str) -> Optional[Session]: """Load session from disk.""" session_file = self.sessions_dir / f"{session_id}.json" if not session_file.exists(): return None try: with open(session_file) as f: data = json.load(f) logger.info(f"Loaded session {session_id} from disk") return Session.from_dict(data) except Exception as e: logger.warning(f"Failed to load session {session_id}: {e}") return None def _save_to_disk(self, session: Session): """Save session to disk.""" session_file = self.sessions_dir / f"{session.session_id}.json" try: with open(session_file, 'w', encoding='utf-8') as f: json.dump(session.to_dict(), f, indent=2, ensure_ascii=False) logger.debug(f"Saved session {session.session_id} to disk") except Exception as e: logger.warning(f"Failed to save session {session.session_id}: {e}") def _cleanup_old_sessions(self): """Remove oldest sessions when over limit.""" sessions = list(self._sessions.values()) sessions.sort(key=lambda s: s.updated_at) # Remove oldest 10% to_remove = sessions[:len(sessions) // 10] for session in to_remove: self.delete_session(session.session_id, archive=True) def _get_title(self, data: Dict) -> str: """Extract title from session data.""" for msg in data.get("chat_history", []): if msg.get("role") == "user": content = msg.get("content", "") return content[:35] + "..." if len(content) > 35 else content return "New Chat" # Global instance _session_manager: Optional[SessionManager] = None def get_session_manager() -> SessionManager: """Get global session manager instance.""" global _session_manager if _session_manager is None: from .config import get_config config = get_config() _session_manager = SessionManager( sessions_dir=str(config.sessions_dir), max_sessions=config.max_sessions, session_max_age_hours=config.session_max_age_hours ) return _session_manager def generate_session_id() -> str: """Generate a unique session ID for a new user.""" import uuid return str(uuid.uuid4())[:12]