""" Session Management System for Law RAG Chatbot """ import uuid import time import json from typing import Dict, List, Any, Optional from datetime import datetime, timedelta from pathlib import Path import sqlite3 import logging logger = logging.getLogger(__name__) class SessionManager: """Manages user sessions and chat history""" def __init__(self, db_path: str = "chat_sessions.db"): self.db_path = db_path self.sessions: Dict[str, Dict[str, Any]] = {} self._init_database() def _init_database(self): """Initialize SQLite database for session storage""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Create sessions table cursor.execute(''' CREATE TABLE IF NOT EXISTS sessions ( session_id TEXT PRIMARY KEY, created_at TIMESTAMP, last_activity TIMESTAMP, user_info TEXT, metadata TEXT ) ''') # Create chat_history table cursor.execute(''' CREATE TABLE IF NOT EXISTS chat_history ( id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT, question TEXT, answer TEXT, sources TEXT, confidence REAL, processing_time REAL, timestamp TIMESTAMP, FOREIGN KEY (session_id) REFERENCES sessions (session_id) ) ''') # Create search_history table cursor.execute(''' CREATE TABLE IF NOT EXISTS search_history ( id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT, query TEXT, results_count INTEGER, timestamp TIMESTAMP, FOREIGN KEY (session_id) REFERENCES sessions (session_id) ) ''') conn.commit() conn.close() logger.info("Session database initialized successfully") except Exception as e: logger.error(f"Failed to initialize session database: {e}") raise def create_session(self, user_info: Optional[str] = None, metadata: Optional[Dict] = None) -> str: """Create a new session""" try: session_id = str(uuid.uuid4()) current_time = datetime.now() # Store in memory self.sessions[session_id] = { "session_id": session_id, "created_at": current_time, "last_activity": current_time, "user_info": user_info or "anonymous", "metadata": metadata or {}, "chat_count": 0, "search_count": 0 } # Store in database conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' INSERT INTO sessions (session_id, created_at, last_activity, user_info, metadata) VALUES (?, ?, ?, ?, ?) ''', ( session_id, current_time.isoformat(), current_time.isoformat(), user_info or "anonymous", json.dumps(metadata or {}) )) conn.commit() conn.close() logger.info(f"Created new session: {session_id}") return session_id except Exception as e: logger.error(f"Failed to create session: {e}") raise def get_session(self, session_id: str) -> Optional[Dict[str, Any]]: """Get session information""" # Check memory first if session_id in self.sessions: return self.sessions[session_id] # Check database try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' SELECT session_id, created_at, last_activity, user_info, metadata FROM sessions WHERE session_id = ? ''', (session_id,)) row = cursor.fetchone() conn.close() if row: session_data = { "session_id": row[0], "created_at": datetime.fromisoformat(row[1]), "last_activity": datetime.fromisoformat(row[2]), "user_info": row[3], "metadata": json.loads(row[4]) if row[4] else {}, "chat_count": 0, "search_count": 0 } # Load counts session_data["chat_count"] = self._get_chat_count(session_id) session_data["search_count"] = self._get_search_count(session_id) # Store in memory self.sessions[session_id] = session_data return session_data return None except Exception as e: logger.error(f"Failed to get session {session_id}: {e}") return None def update_session_activity(self, session_id: str): """Update session last activity""" current_time = datetime.now() # Update memory if session_id in self.sessions: self.sessions[session_id]["last_activity"] = current_time # Update database try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' UPDATE sessions SET last_activity = ? WHERE session_id = ? ''', (current_time.isoformat(), session_id)) conn.commit() conn.close() except Exception as e: logger.error(f"Failed to update session activity: {e}") def store_chat_response(self, session_id: str, question: str, answer: str, sources: List[Dict], confidence: float, processing_time: float): """Store a chat response in the session""" try: current_time = datetime.now() # Update session activity self.update_session_activity(session_id) # Store in database conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' INSERT INTO chat_history (session_id, question, answer, sources, confidence, processing_time, timestamp) VALUES (?, ?, ?, ?, ?, ?, ?) ''', ( session_id, question, answer, json.dumps(sources), confidence, processing_time, current_time.isoformat() )) conn.commit() conn.close() # Update memory if session_id in self.sessions: self.sessions[session_id]["chat_count"] += 1 logger.info(f"Stored chat response for session {session_id}") except Exception as e: logger.error(f"Failed to store chat response: {e}") def store_search_query(self, session_id: str, query: str, results_count: int): """Store a search query in the session""" try: current_time = datetime.now() # Update session activity self.update_session_activity(session_id) # Store in database conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' INSERT INTO search_history (session_id, query, results_count, timestamp) VALUES (?, ?, ?, ?) ''', ( session_id, query, results_count, current_time.isoformat() )) conn.commit() conn.close() # Update memory if session_id in self.sessions: self.sessions[session_id]["search_count"] += 1 logger.info(f"Stored search query for session {session_id}") except Exception as e: logger.error(f"Failed to store search query: {e}") def get_chat_history(self, session_id: str, limit: int = 10) -> List[Dict[str, Any]]: """Get chat history for a session""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' SELECT question, answer, sources, confidence, processing_time, timestamp FROM chat_history WHERE session_id = ? ORDER BY timestamp DESC LIMIT ? ''', (session_id, limit)) rows = cursor.fetchall() conn.close() history = [] for row in rows: history.append({ "question": row[0], "answer": row[1], "sources": json.loads(row[2]) if row[2] else [], "confidence": row[3], "processing_time": row[4], "timestamp": row[5] }) return history except Exception as e: logger.error(f"Failed to get chat history: {e}") return [] def get_search_history(self, session_id: str, limit: int = 10) -> List[Dict[str, Any]]: """Get search history for a session""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(''' SELECT query, results_count, timestamp FROM search_history WHERE session_id = ? ORDER BY timestamp DESC LIMIT ? ''', (session_id, limit)) rows = cursor.fetchall() conn.close() history = [] for row in rows: history.append({ "query": row[0], "results_count": row[1], "timestamp": row[2] }) return history except Exception as e: logger.error(f"Failed to get search history: {e}") return [] def _get_chat_count(self, session_id: str) -> int: """Get chat count for a session""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute('SELECT COUNT(*) FROM chat_history WHERE session_id = ?', (session_id,)) count = cursor.fetchone()[0] conn.close() return count except Exception as e: logger.error(f"Failed to get chat count: {e}") return 0 def _get_search_count(self, session_id: str) -> int: """Get search count for a session""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute('SELECT COUNT(*) FROM search_history WHERE session_id = ?', (session_id,)) count = cursor.fetchone()[0] conn.close() return count except Exception as e: logger.error(f"Failed to get search count: {e}") return 0 def get_session_stats(self, session_id: str) -> Dict[str, Any]: """Get comprehensive session statistics""" session = self.get_session(session_id) if not session: return {} chat_history = self.get_chat_history(session_id, limit=100) search_history = self.get_search_history(session_id, limit=100) # Calculate average confidence confidences = [chat["confidence"] for chat in chat_history if chat["confidence"] > 0] avg_confidence = sum(confidences) / len(confidences) if confidences else 0 # Calculate average processing time processing_times = [chat["processing_time"] for chat in chat_history if chat["processing_time"] > 0] avg_processing_time = sum(processing_times) / len(processing_times) if processing_times else 0 return { "session_id": session_id, "created_at": session["created_at"].isoformat(), "last_activity": session["last_activity"].isoformat(), "user_info": session["user_info"], "total_chats": len(chat_history), "total_searches": len(search_history), "average_confidence": round(avg_confidence, 3), "average_processing_time": round(avg_processing_time, 3), "recent_questions": [chat["question"] for chat in chat_history[:5]], "recent_searches": [search["query"] for search in search_history[:5]] } def cleanup_old_sessions(self, days: int = 30): """Clean up sessions older than specified days""" try: cutoff_date = datetime.now() - timedelta(days=days) conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Delete old chat history cursor.execute(''' DELETE FROM chat_history WHERE session_id IN ( SELECT session_id FROM sessions WHERE last_activity < ? ) ''', (cutoff_date.isoformat(),)) # Delete old search history cursor.execute(''' DELETE FROM search_history WHERE session_id IN ( SELECT session_id FROM sessions WHERE last_activity < ? ) ''', (cutoff_date.isoformat(),)) # Delete old sessions cursor.execute('DELETE FROM sessions WHERE last_activity < ?', (cutoff_date.isoformat(),)) conn.commit() conn.close() logger.info(f"Cleaned up sessions older than {days} days") except Exception as e: logger.error(f"Failed to cleanup old sessions: {e}") def delete_session(self, session_id: str): """Delete a session and all its data""" try: # Remove from memory if session_id in self.sessions: del self.sessions[session_id] # Remove from database conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute('DELETE FROM chat_history WHERE session_id = ?', (session_id,)) cursor.execute('DELETE FROM search_history WHERE session_id = ?', (session_id,)) cursor.execute('DELETE FROM sessions WHERE session_id = ?', (session_id,)) conn.commit() conn.close() logger.info(f"Deleted session: {session_id}") except Exception as e: logger.error(f"Failed to delete session {session_id}: {e}")