# context_manager.py import sqlite3 import json import logging import uuid import hashlib import threading from contextlib import contextmanager from datetime import datetime, timedelta logger = logging.getLogger(__name__) class TransactionManager: """Manage database transactions with proper locking""" def __init__(self, db_path): self.db_path = db_path self._lock = threading.RLock() self._connections = {} @contextmanager def transaction(self, session_id=None): """Context manager for database transactions with automatic rollback""" conn = None cursor = None try: with self._lock: conn = sqlite3.connect(self.db_path, isolation_level='IMMEDIATE') conn.execute('PRAGMA journal_mode=WAL') # Write-Ahead Logging for better concurrency conn.execute('PRAGMA busy_timeout=5000') # 5 second timeout for locks cursor = conn.cursor() yield cursor conn.commit() logger.debug(f"Transaction committed for session {session_id}") except Exception as e: if conn: conn.rollback() logger.error(f"Transaction rolled back for session {session_id}: {e}") raise finally: if conn: conn.close() class EfficientContextManager: def __init__(self, llm_router=None): self.session_cache = {} # In-memory for active sessions self.cache_config = { "max_session_size": 10, # MB per session "ttl": 3600, # 1 hour "compression": "gzip", "eviction_policy": "LRU" } self.db_path = "sessions.db" self.llm_router = llm_router # For generating context summaries logger.info(f"Initializing ContextManager with DB path: {self.db_path}") self.transaction_manager = TransactionManager(self.db_path) self._init_database() self.optimize_database_indexes() def _init_database(self): """Initialize database and create tables""" try: logger.info("Initializing database...") conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Create sessions table if not exists cursor.execute(""" CREATE TABLE IF NOT EXISTS sessions ( session_id TEXT PRIMARY KEY, user_id TEXT DEFAULT 'Test_Any', created_at TIMESTAMP, last_activity TIMESTAMP, context_data TEXT, user_metadata TEXT ) """) # Add user_id column to existing sessions table if it doesn't exist try: cursor.execute("ALTER TABLE sessions ADD COLUMN user_id TEXT DEFAULT 'Test_Any'") logger.info("✓ Added user_id column to sessions table") except sqlite3.OperationalError: # Column already exists pass logger.info("✓ Sessions table ready") # Create interactions table cursor.execute(""" CREATE TABLE IF NOT EXISTS interactions ( id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT REFERENCES sessions(session_id), user_input TEXT, context_snapshot TEXT, created_at TIMESTAMP, FOREIGN KEY(session_id) REFERENCES sessions(session_id) ) """) logger.info("✓ Interactions table ready") # Create user_contexts table (persistent user persona summaries) cursor.execute(""" CREATE TABLE IF NOT EXISTS user_contexts ( user_id TEXT PRIMARY KEY, persona_summary TEXT, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) logger.info("✓ User contexts table ready") # Create session_contexts table (session summaries) cursor.execute(""" CREATE TABLE IF NOT EXISTS session_contexts ( session_id TEXT PRIMARY KEY, user_id TEXT, session_summary TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY(session_id) REFERENCES sessions(session_id), FOREIGN KEY(user_id) REFERENCES user_contexts(user_id) ) """) logger.info("✓ Session contexts table ready") # Create interaction_contexts table (individual interaction summaries) cursor.execute(""" CREATE TABLE IF NOT EXISTS interaction_contexts ( interaction_id TEXT PRIMARY KEY, session_id TEXT, user_input TEXT, system_response TEXT, interaction_summary TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY(session_id) REFERENCES sessions(session_id) ) """) logger.info("✓ Interaction contexts table ready") conn.commit() conn.close() # Update schema with new columns and tables for user change tracking self._update_database_schema() logger.info("Database initialization complete") except Exception as e: logger.error(f"Database initialization error: {e}", exc_info=True) def _update_database_schema(self): """Add missing columns and tables for user change tracking""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Add needs_refresh column to interaction_contexts try: cursor.execute(""" ALTER TABLE interaction_contexts ADD COLUMN needs_refresh INTEGER DEFAULT 0 """) logger.info("✓ Added needs_refresh column to interaction_contexts") except sqlite3.OperationalError: pass # Column already exists # Create user change log table cursor.execute(""" CREATE TABLE IF NOT EXISTS user_change_log ( id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT, old_user_id TEXT, new_user_id TEXT, timestamp TIMESTAMP, FOREIGN KEY(session_id) REFERENCES sessions(session_id) ) """) conn.commit() conn.close() logger.info("✓ Database schema updated successfully for user change tracking") # Update interactions table for deduplication self._update_interactions_table() except Exception as e: logger.error(f"Schema update error: {e}", exc_info=True) def _update_interactions_table(self): """Add interaction_hash column for deduplication""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Check if column already exists cursor.execute("PRAGMA table_info(interactions)") columns = [row[1] for row in cursor.fetchall()] # Add interaction_hash column if it doesn't exist if 'interaction_hash' not in columns: try: cursor.execute(""" ALTER TABLE interactions ADD COLUMN interaction_hash TEXT """) logger.info("✓ Added interaction_hash column to interactions table") except sqlite3.OperationalError: pass # Column already exists # Create unique index for deduplication (this enforces uniqueness) try: cursor.execute(""" CREATE UNIQUE INDEX IF NOT EXISTS idx_interaction_hash_unique ON interactions(interaction_hash) """) logger.info("✓ Created unique index on interaction_hash") except sqlite3.OperationalError: # Index might already exist, try non-unique index as fallback cursor.execute(""" CREATE INDEX IF NOT EXISTS idx_interaction_hash ON interactions(interaction_hash) """) conn.commit() conn.close() logger.info("✓ Interactions table updated for deduplication") except Exception as e: logger.error(f"Error updating interactions table: {e}", exc_info=True) async def manage_context(self, session_id: str, user_input: str, user_id: str = "Test_Any") -> dict: """ Efficient context management with separated session/user caching STEP 1: Fetch User Context (if available) STEP 2: Get Previous Interaction Contexts STEP 3: Combine for workflow use """ # Use session-only cache key to prevent user_id conflicts session_cache_key = f"session_{session_id}" user_cache_key = f"user_{user_id}" # Get session context from cache session_context = self._get_from_memory_cache(session_cache_key) # Check if cached session context matches current user_id if session_context and session_context.get("user_id") != user_id: # User changed, invalidate session cache logger.info(f"User mismatch in cache for session {session_id}, invalidating cache") session_context = None if session_cache_key in self.session_cache: del self.session_cache[session_cache_key] # Get user context separately user_context = self._get_from_memory_cache(user_cache_key) if not session_context: # Retrieve from database with user context session_context = await self._retrieve_from_db(session_id, user_input, user_id) # Cache session context (cache invalidation for user changes is handled in _retrieve_from_db) self._warm_memory_cache(session_cache_key, session_context) # Handle user context separately to prevent loops if not user_context or not user_context.get("user_context_loaded"): user_context_data = await self.get_user_context(user_id) user_context = { "user_context": user_context_data, "user_context_loaded": True, "user_id": user_id } # Cache user context separately self._warm_memory_cache(user_cache_key, user_context) # Merge contexts without duplication merged_context = { **session_context, "user_context": user_context.get("user_context", ""), "user_context_loaded": True, "user_id": user_id # Ensure current user_id is used } # Update context with new interaction updated_context = self._update_context(merged_context, user_input, user_id=user_id) return self._optimize_context(updated_context) async def get_user_context(self, user_id: str) -> str: """ STEP 1: Fetch or generate User Context (500-token persona summary) Available for all interactions except first time per user """ try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Check if user context exists cursor.execute(""" SELECT persona_summary FROM user_contexts WHERE user_id = ? """, (user_id,)) row = cursor.fetchone() if row and row[0]: # Existing user context found conn.close() logger.info(f"✓ User context loaded for {user_id}") return row[0] # Generate new user context from all historical data logger.info(f"Generating new user context for {user_id}") # Fetch all historical Session and Interaction contexts for this user all_session_summaries = [] all_interaction_summaries = [] # Get all session contexts cursor.execute(""" SELECT session_summary FROM session_contexts WHERE user_id = ? ORDER BY created_at DESC LIMIT 50 """, (user_id,)) for row in cursor.fetchall(): if row[0]: all_session_summaries.append(row[0]) # Get all interaction contexts cursor.execute(""" SELECT ic.interaction_summary FROM interaction_contexts ic JOIN sessions s ON ic.session_id = s.session_id WHERE s.user_id = ? ORDER BY ic.created_at DESC LIMIT 100 """, (user_id,)) for row in cursor.fetchall(): if row[0]: all_interaction_summaries.append(row[0]) conn.close() if not all_session_summaries and not all_interaction_summaries: # First time user - no context to generate logger.info(f"No historical data for {user_id} - first time user") return "" # Generate persona summary using LLM (500 tokens) historical_data = "\n\n".join(all_session_summaries + all_interaction_summaries[:20]) if self.llm_router: prompt = f"""Generate a concise 500-token persona summary for user {user_id} based on their interaction history: Historical Context: {historical_data} Create a persona summary that captures: - Communication style and preferences - Common topics and interests - Interaction patterns - Key information shared across sessions Keep the summary concise and focused (approximately 500 tokens).""" try: persona_summary = await self.llm_router.route_inference( task_type="general_reasoning", prompt=prompt, max_tokens=500, temperature=0.7 ) if persona_summary and isinstance(persona_summary, str) and persona_summary.strip(): # Store in database conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(""" INSERT OR REPLACE INTO user_contexts (user_id, persona_summary, updated_at) VALUES (?, ?, ?) """, (user_id, persona_summary.strip(), datetime.now().isoformat())) conn.commit() conn.close() logger.info(f"✓ Generated and stored user context for {user_id}") return persona_summary.strip() except Exception as e: logger.error(f"Error generating user context: {e}", exc_info=True) # Fallback: Return empty if LLM fails logger.warning(f"Could not generate user context for {user_id} - using empty") return "" except Exception as e: logger.error(f"Error getting user context: {e}", exc_info=True) return "" async def generate_interaction_context(self, interaction_id: str, session_id: str, user_input: str, system_response: str, user_id: str = "Test_Any") -> str: """ STEP 2: Generate Interaction Context (50-token summary) Called after each response """ try: if not self.llm_router: return "" prompt = f"""Summarize this interaction in approximately 50 tokens: User Input: {user_input[:200]} System Response: {system_response[:300]} Provide a brief summary capturing the key exchange.""" try: summary = await self.llm_router.route_inference( task_type="general_reasoning", prompt=prompt, max_tokens=50, temperature=0.7 ) if summary and isinstance(summary, str) and summary.strip(): # Store in database conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(""" INSERT OR REPLACE INTO interaction_contexts (interaction_id, session_id, user_input, system_response, interaction_summary, created_at) VALUES (?, ?, ?, ?, ?, ?) """, ( interaction_id, session_id, user_input[:500], system_response[:1000], summary.strip(), datetime.now().isoformat() )) conn.commit() conn.close() logger.info(f"✓ Generated interaction context for {interaction_id}") return summary.strip() except Exception as e: logger.error(f"Error generating interaction context: {e}", exc_info=True) # Fallback on LLM failure return "" except Exception as e: logger.error(f"Error in generate_interaction_context: {e}", exc_info=True) return "" async def generate_session_context(self, session_id: str, user_id: str = "Test_Any") -> str: """ FINAL STEP: Generate Session Context (100-token summary) Called at session end """ try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Get all interaction contexts for this session cursor.execute(""" SELECT interaction_summary FROM interaction_contexts WHERE session_id = ? ORDER BY created_at ASC """, (session_id,)) interaction_summaries = [row[0] for row in cursor.fetchall() if row[0]] conn.close() if not interaction_summaries: logger.info(f"No interactions to summarize for session {session_id}") return "" # Generate session summary using LLM (100 tokens) if self.llm_router: combined_context = "\n".join(interaction_summaries) prompt = f"""Summarize this session's interactions in approximately 100 tokens: Interaction Summaries: {combined_context} Create a concise session summary capturing: - Main topics discussed - Key outcomes or information shared - User's focus areas Keep the summary concise (approximately 100 tokens).""" try: session_summary = await self.llm_router.route_inference( task_type="general_reasoning", prompt=prompt, max_tokens=100, temperature=0.7 ) if session_summary and isinstance(session_summary, str) and session_summary.strip(): # Store in database conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(""" INSERT OR REPLACE INTO session_contexts (session_id, user_id, session_summary, created_at) VALUES (?, ?, ?, ?) """, (session_id, user_id, session_summary.strip(), datetime.now().isoformat())) conn.commit() conn.close() logger.info(f"✓ Generated session context for {session_id}") return session_summary.strip() except Exception as e: logger.error(f"Error generating session context: {e}", exc_info=True) # Fallback on LLM failure return "" except Exception as e: logger.error(f"Error in generate_session_context: {e}", exc_info=True) return "" async def end_session(self, session_id: str, user_id: str = "Test_Any"): """ FINAL STEP: Generate Session Context and clear cache """ try: # Generate session context await self.generate_session_context(session_id, user_id) # Clear in-memory cache for this session (session-only key) session_cache_key = f"session_{session_id}" if session_cache_key in self.session_cache: del self.session_cache[session_cache_key] logger.info(f"✓ Cleared cache for session {session_id}") except Exception as e: logger.error(f"Error ending session: {e}", exc_info=True) def _clear_user_cache_on_change(self, session_id: str, new_user_id: str, old_user_id: str): """Clear cache entries when user changes""" if new_user_id != old_user_id: # Clear old composite cache keys old_cache_key = f"{session_id}_{old_user_id}" if old_cache_key in self.session_cache: del self.session_cache[old_cache_key] logger.info(f"Cleared old cache for user {old_user_id} on session {session_id}") def _optimize_context(self, context: dict) -> dict: """ Optimize context for LLM consumption Format: [Interaction Context #N, #N-1, ...] + User Context """ user_context = context.get("user_context", "") interaction_contexts = context.get("interaction_contexts", []) # Format interaction contexts as requested formatted_interactions = [] for idx, ic in enumerate(interaction_contexts[:10]): # Last 10 interactions formatted_interactions.append(f"[Interaction Context #{len(interaction_contexts) - idx}]\n{ic.get('summary', '')}") # Combine User Context + Interaction Contexts combined_context = "" if user_context: combined_context += f"[User Context]\n{user_context}\n\n" if formatted_interactions: combined_context += "\n\n".join(formatted_interactions) return { "session_id": context.get("session_id"), "user_id": context.get("user_id", "Test_Any"), "user_context": user_context, "interaction_contexts": interaction_contexts, "combined_context": combined_context, # For direct use in prompts "preferences": context.get("preferences", {}), "active_tasks": context.get("active_tasks", []), "last_activity": context.get("last_activity") } def _get_from_memory_cache(self, cache_key: str) -> dict: """ Retrieve context from in-memory session cache """ return self.session_cache.get(cache_key) async def _retrieve_from_db(self, session_id: str, user_input: str, user_id: str = "Test_Any") -> dict: """ Retrieve session context with proper user_id synchronization Uses transactions to ensure atomic updates of database and cache """ conn = None try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Use transaction to ensure atomic updates cursor.execute("BEGIN TRANSACTION") # Get session data (SQLite doesn't support FOR UPDATE, but transaction ensures consistency) cursor.execute(""" SELECT context_data, user_metadata, last_activity, user_id FROM sessions WHERE session_id = ? """, (session_id,)) row = cursor.fetchone() if row: context_data = json.loads(row[0]) if row[0] else {} user_metadata = json.loads(row[1]) if row[1] else {} last_activity = row[2] session_user_id = row[3] if len(row) > 3 else user_id # Check for user_id change and update atomically user_changed = False if session_user_id != user_id: logger.info(f"User change detected: {session_user_id} -> {user_id} for session {session_id}") user_changed = True # Update session with new user_id cursor.execute(""" UPDATE sessions SET user_id = ?, last_activity = ? WHERE session_id = ? """, (user_id, datetime.now().isoformat(), session_id)) # Clear any cached interaction contexts for old user by marking for refresh try: cursor.execute(""" UPDATE interaction_contexts SET needs_refresh = 1 WHERE session_id = ? """, (session_id,)) except sqlite3.OperationalError: # Column might not exist yet, will be created by schema update pass # Log user change event try: cursor.execute(""" INSERT INTO user_change_log (session_id, old_user_id, new_user_id, timestamp) VALUES (?, ?, ?, ?) """, (session_id, session_user_id, user_id, datetime.now().isoformat())) except sqlite3.OperationalError: # Table might not exist yet, will be created by schema update pass # Clear old cache entries when user changes self._clear_user_cache_on_change(session_id, user_id, session_user_id) cursor.execute("COMMIT") # Get interaction contexts with refresh flag check try: cursor.execute(""" SELECT interaction_summary, created_at, needs_refresh FROM interaction_contexts WHERE session_id = ? AND (needs_refresh IS NULL OR needs_refresh = 0) ORDER BY created_at DESC LIMIT 20 """, (session_id,)) except sqlite3.OperationalError: # Column might not exist yet, fall back to query without needs_refresh cursor.execute(""" SELECT interaction_summary, created_at FROM interaction_contexts WHERE session_id = ? ORDER BY created_at DESC LIMIT 20 """, (session_id,)) interaction_contexts = [] for ic_row in cursor.fetchall(): # Handle both query formats (with and without needs_refresh) if len(ic_row) >= 2: summary = ic_row[0] timestamp = ic_row[1] needs_refresh = ic_row[2] if len(ic_row) > 2 else 0 if summary and not needs_refresh: interaction_contexts.append({ "summary": summary, "timestamp": timestamp }) context = { "session_id": session_id, "user_id": user_id, "interaction_contexts": interaction_contexts, "preferences": user_metadata.get("preferences", {}), "active_tasks": user_metadata.get("active_tasks", []), "last_activity": last_activity, "user_context_loaded": False, "user_changed": user_changed } conn.close() return context else: # Create new session with transaction cursor.execute(""" INSERT INTO sessions (session_id, user_id, created_at, last_activity, context_data, user_metadata) VALUES (?, ?, ?, ?, ?, ?) """, (session_id, user_id, datetime.now().isoformat(), datetime.now().isoformat(), "{}", "{}")) cursor.execute("COMMIT") conn.close() return { "session_id": session_id, "user_id": user_id, "interaction_contexts": [], "preferences": {}, "active_tasks": [], "user_context_loaded": False, "user_changed": False } except sqlite3.Error as e: logger.error(f"Database transaction error: {e}", exc_info=True) if conn: try: conn.rollback() except: pass conn.close() # Return safe fallback return { "session_id": session_id, "user_id": user_id, "interaction_contexts": [], "preferences": {}, "active_tasks": [], "user_context_loaded": False, "error": str(e), "user_changed": False } except Exception as e: logger.error(f"Database retrieval error: {e}", exc_info=True) if conn: try: conn.rollback() except: pass conn.close() # Return safe fallback return { "session_id": session_id, "user_id": user_id, "interaction_contexts": [], "preferences": {}, "active_tasks": [], "user_context_loaded": False, "error": str(e), "user_changed": False } def _warm_memory_cache(self, cache_key: str, context: dict): """ Warm the in-memory cache with retrieved context """ self.session_cache[cache_key] = context def _update_context(self, context: dict, user_input: str, response: str = None, user_id: str = "Test_Any") -> dict: """ Update context with deduplication and idempotency checks Prevents duplicate context updates using interaction hashes """ try: # Generate unique interaction hash to prevent duplicates interaction_hash = self._generate_interaction_hash(user_input, context["session_id"], user_id) # Check if this interaction was already processed if self._is_duplicate_interaction(interaction_hash): logger.info(f"Duplicate interaction detected, skipping update: {interaction_hash[:8]}") return context # Use transaction for atomic updates current_time = datetime.now().isoformat() with self.transaction_manager.transaction(context["session_id"]) as cursor: # Update session activity (only if last_activity is older to prevent unnecessary updates) cursor.execute(""" UPDATE sessions SET last_activity = ?, user_id = ? WHERE session_id = ? AND (last_activity IS NULL OR last_activity < ?) """, (current_time, user_id, context["session_id"], current_time)) # Store interaction with duplicate prevention using INSERT OR IGNORE session_context = { "preferences": context.get("preferences", {}), "active_tasks": context.get("active_tasks", []) } cursor.execute(""" INSERT OR IGNORE INTO interactions ( interaction_hash, session_id, user_input, context_snapshot, created_at ) VALUES (?, ?, ?, ?, ?) """, ( interaction_hash, context["session_id"], user_input, json.dumps(session_context), current_time )) # Mark interaction as processed (outside transaction) self._mark_interaction_processed(interaction_hash) # Update in-memory context context["last_interaction"] = user_input context["last_update"] = current_time logger.info(f"Context updated for session {context['session_id']} with hash {interaction_hash[:8]}") return context except Exception as e: logger.error(f"Error updating context: {e}", exc_info=True) return context def _generate_interaction_hash(self, user_input: str, session_id: str, user_id: str) -> str: """Generate unique hash for interaction to prevent duplicates""" # Use session_id, user_id, and user_input for exact duplicate detection # Normalize user input by stripping whitespace normalized_input = user_input.strip() content = f"{session_id}:{user_id}:{normalized_input}" return hashlib.sha256(content.encode()).hexdigest() def _is_duplicate_interaction(self, interaction_hash: str) -> bool: """Check if interaction was already processed""" # Keep a rolling window of recent interaction hashes in memory if not hasattr(self, '_processed_interactions'): self._processed_interactions = set() # Check in-memory cache first if interaction_hash in self._processed_interactions: return True # Also check database for persistent duplicates try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Check if interaction_hash column exists and query for duplicates cursor.execute("PRAGMA table_info(interactions)") columns = [row[1] for row in cursor.fetchall()] if 'interaction_hash' in columns: cursor.execute(""" SELECT COUNT(*) FROM interactions WHERE interaction_hash IS NOT NULL AND interaction_hash = ? """, (interaction_hash,)) count = cursor.fetchone()[0] conn.close() return count > 0 else: conn.close() return False except sqlite3.OperationalError: # Column might not exist yet, only check in-memory return interaction_hash in self._processed_interactions def _mark_interaction_processed(self, interaction_hash: str): """Mark interaction as processed""" if not hasattr(self, '_processed_interactions'): self._processed_interactions = set() self._processed_interactions.add(interaction_hash) # Limit memory usage by keeping only last 1000 hashes if len(self._processed_interactions) > 1000: # Keep most recent 500 entries (simple truncation) self._processed_interactions = set(list(self._processed_interactions)[-500:]) async def manage_context_optimized(self, session_id: str, user_input: str, user_id: str = "Test_Any") -> dict: """ Efficient context management with transaction optimization """ # Use session-only cache key session_cache_key = f"session_{session_id}" # Try to get from cache first (no DB access) cached_context = self._get_from_memory_cache(session_cache_key) if cached_context and self._is_cache_valid(cached_context): logger.debug(f"Using cached context for session {session_id}") return cached_context # Use transaction for all DB operations with self.transaction_manager.transaction(session_id) as cursor: # Atomic session retrieval and update cursor.execute(""" SELECT s.context_data, s.user_metadata, s.last_activity, s.user_id, COUNT(ic.interaction_id) as interaction_count FROM sessions s LEFT JOIN interaction_contexts ic ON s.session_id = ic.session_id WHERE s.session_id = ? GROUP BY s.session_id """, (session_id,)) row = cursor.fetchone() if row: # Parse existing session data context_data = json.loads(row[0] or '{}') user_metadata = json.loads(row[1] or '{}') last_activity = row[2] stored_user_id = row[3] or user_id interaction_count = row[4] or 0 # Handle user change atomically if stored_user_id != user_id: self._handle_user_change_atomic(cursor, session_id, stored_user_id, user_id) # Get interaction contexts efficiently interaction_contexts = self._get_interaction_contexts_atomic(cursor, session_id) else: # Create new session atomically cursor.execute(""" INSERT INTO sessions (session_id, user_id, created_at, last_activity, context_data, user_metadata) VALUES (?, ?, datetime('now'), datetime('now'), '{}', '{}') """, (session_id, user_id)) context_data = {} user_metadata = {} interaction_contexts = [] interaction_count = 0 # Load user context asynchronously (outside transaction) user_context = await self._load_user_context_async(user_id) # Build final context final_context = { "session_id": session_id, "user_id": user_id, "interaction_contexts": interaction_contexts, "user_context": user_context, "preferences": user_metadata.get("preferences", {}), "active_tasks": user_metadata.get("active_tasks", []), "interaction_count": interaction_count, "cache_timestamp": datetime.now().isoformat() } # Update cache self._warm_memory_cache(session_cache_key, final_context) return self._optimize_context(final_context) def _handle_user_change_atomic(self, cursor, session_id: str, old_user_id: str, new_user_id: str): """Handle user change within transaction""" logger.info(f"Handling user change in transaction: {old_user_id} -> {new_user_id}") # Update session cursor.execute(""" UPDATE sessions SET user_id = ?, last_activity = datetime('now') WHERE session_id = ? """, (new_user_id, session_id)) # Log the change try: cursor.execute(""" INSERT INTO user_change_log (session_id, old_user_id, new_user_id, timestamp) VALUES (?, ?, ?, datetime('now')) """, (session_id, old_user_id, new_user_id)) except sqlite3.OperationalError: # Table might not exist yet pass # Invalidate related caches try: cursor.execute(""" UPDATE interaction_contexts SET needs_refresh = 1 WHERE session_id = ? """, (session_id,)) except sqlite3.OperationalError: # Column might not exist yet pass def _get_interaction_contexts_atomic(self, cursor, session_id: str, limit: int = 20): """Get interaction contexts within transaction""" try: cursor.execute(""" SELECT interaction_summary, created_at, interaction_id FROM interaction_contexts WHERE session_id = ? AND (needs_refresh IS NULL OR needs_refresh = 0) ORDER BY created_at DESC LIMIT ? """, (session_id, limit)) except sqlite3.OperationalError: # Fallback if needs_refresh column doesn't exist cursor.execute(""" SELECT interaction_summary, created_at, interaction_id FROM interaction_contexts WHERE session_id = ? ORDER BY created_at DESC LIMIT ? """, (session_id, limit)) contexts = [] for row in cursor.fetchall(): if row[0]: contexts.append({ "summary": row[0], "timestamp": row[1], "id": row[2] if len(row) > 2 else None }) return contexts async def _load_user_context_async(self, user_id: str): """Load user context asynchronously to avoid blocking""" try: # Check memory cache first user_cache_key = f"user_{user_id}" cached = self._get_from_memory_cache(user_cache_key) if cached: return cached.get("user_context", "") # Load from database return await self.get_user_context(user_id) except Exception as e: logger.error(f"Error loading user context: {e}") return "" def _is_cache_valid(self, cached_context: dict, max_age_seconds: int = 60) -> bool: """Check if cached context is still valid""" if not cached_context: return False cache_timestamp = cached_context.get("cache_timestamp") if not cache_timestamp: return False try: cache_time = datetime.fromisoformat(cache_timestamp) age = (datetime.now() - cache_time).total_seconds() return age < max_age_seconds except: return False def optimize_database_indexes(self): """Create database indexes for better query performance""" try: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() # Create indexes for frequently queried columns indexes = [ "CREATE INDEX IF NOT EXISTS idx_sessions_user_id ON sessions(user_id)", "CREATE INDEX IF NOT EXISTS idx_sessions_last_activity ON sessions(last_activity)", "CREATE INDEX IF NOT EXISTS idx_interactions_session_id ON interactions(session_id)", "CREATE INDEX IF NOT EXISTS idx_interaction_contexts_session_id ON interaction_contexts(session_id)", "CREATE INDEX IF NOT EXISTS idx_interaction_contexts_created_at ON interaction_contexts(created_at)", "CREATE INDEX IF NOT EXISTS idx_user_change_log_session_id ON user_change_log(session_id)", "CREATE INDEX IF NOT EXISTS idx_user_contexts_updated_at ON user_contexts(updated_at)" ] for index in indexes: try: cursor.execute(index) except sqlite3.OperationalError as e: # Table might not exist yet, skip this index logger.debug(f"Skipping index creation (table may not exist): {e}") # Analyze database for query optimization try: cursor.execute("ANALYZE") except sqlite3.OperationalError: # ANALYZE might not be available in all SQLite versions pass conn.commit() conn.close() logger.info("✓ Database indexes optimized successfully") except Exception as e: logger.error(f"Error optimizing database indexes: {e}", exc_info=True) def _extract_entities(self, context: dict) -> list: """ Extract essential entities from context """ # TODO: Implement entity extraction return [] def _generate_summary(self, context: dict) -> str: """ Generate conversation summary """ # TODO: Implement summary generation return ""