""" Database Manager for Epic 2 Demo Persistent Storage ================================================== Handles database connections, operations, and high-level persistence management for the Epic 2 demo to achieve <5 second initialization times. """ import logging import hashlib import time import uuid from datetime import datetime, timedelta from pathlib import Path from typing import Dict, Any, List, Optional, Tuple, Union from contextlib import contextmanager import numpy as np from sqlalchemy import create_engine, text from sqlalchemy.orm import sessionmaker, Session from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.pool import StaticPool from .database_schema import Base, Document, DocumentChunk, SystemCache, ProcessingSession, DatabaseSchema logger = logging.getLogger(__name__) class DatabaseManager: """Manages database operations for Epic 2 demo persistence""" def __init__(self, database_url: str = "sqlite:///demo/epic2_demo.db", echo: bool = False): """ Initialize database manager Args: database_url: Database connection URL echo: Whether to echo SQL statements (for debugging) """ self.database_url = database_url self.echo = echo # Create database directory if using SQLite if database_url.startswith("sqlite:///"): db_path = Path(database_url.replace("sqlite:///", "")) db_path.parent.mkdir(parents=True, exist_ok=True) # Create engine with optimized settings self.engine = create_engine( database_url, echo=echo, poolclass=StaticPool if "sqlite" in database_url else None, connect_args={"check_same_thread": False} if "sqlite" in database_url else {}, pool_pre_ping=True, pool_recycle=3600 # 1 hour ) # Create session factory self.SessionLocal = sessionmaker( autocommit=False, autoflush=False, bind=self.engine ) # Initialize database self._initialize_database() def _initialize_database(self) -> None: """Initialize database tables and indexes""" try: logger.info("Initializing database schema...") DatabaseSchema.create_all_tables(self.engine) # Optimize SQLite if using it if "sqlite" in self.database_url: self._optimize_sqlite() logger.info("Database initialization complete") except Exception as e: logger.error(f"Database initialization failed: {e}") raise def _optimize_sqlite(self) -> None: """Apply SQLite-specific optimizations""" try: with self.engine.connect() as conn: # Performance optimizations conn.execute(text("PRAGMA journal_mode = WAL")) conn.execute(text("PRAGMA synchronous = NORMAL")) conn.execute(text("PRAGMA cache_size = 10000")) conn.execute(text("PRAGMA temp_store = MEMORY")) conn.execute(text("PRAGMA mmap_size = 268435456")) # 256MB conn.commit() logger.info("SQLite optimizations applied") except Exception as e: logger.warning(f"SQLite optimization failed: {e}") @contextmanager def get_session(self): """Context manager for database sessions""" session = self.SessionLocal() try: yield session session.commit() except Exception as e: session.rollback() logger.error(f"Database session error: {e}") raise finally: session.close() def get_database_stats(self) -> Dict[str, Any]: """Get comprehensive database statistics""" with self.get_session() as session: stats = DatabaseSchema.get_database_stats(session) # Add database file size if SQLite if "sqlite" in self.database_url: try: db_path = Path(self.database_url.replace("sqlite:///", "")) if db_path.exists(): stats['database_size_mb'] = db_path.stat().st_size / (1024 * 1024) except: pass return stats def is_database_populated(self) -> bool: """Check if database has any processed documents""" try: with self.get_session() as session: count = session.query(Document).filter( Document.processing_status == 'completed' ).count() return count > 0 except: return False def is_cache_valid(self, pdf_files: List[Path], processor_config: Dict[str, Any], embedder_config: Dict[str, Any]) -> bool: """ Check if database cache is valid for given files and configurations Args: pdf_files: List of PDF files to check processor_config: Document processor configuration embedder_config: Embedder configuration Returns: True if cache is valid and up-to-date """ try: with self.get_session() as session: # Simple check: do we have any completed documents in database? total_docs = session.query(Document).filter( Document.processing_status == 'completed' ).count() total_chunks = session.query(DocumentChunk).filter( DocumentChunk.embedding_vector != None ).count() logger.info(f"Database validation: {total_docs} documents, {total_chunks} chunks with embeddings") if total_docs == 0 or total_chunks == 0: logger.info("No valid documents/chunks in database") return False # Check if we have any matching files available_files = session.query(Document.filename).filter( Document.processing_status == 'completed' ).all() available_filenames = [doc.filename for doc in available_files] requested_filenames = [pdf_file.name for pdf_file in pdf_files] # Option A: If no corpus files requested but database has documents, load ALL if len(requested_filenames) == 0 and len(available_filenames) > 0: logger.info(f"No corpus files specified, but database has {len(available_filenames)} documents - will load ALL") return True matching_files = [f for f in requested_filenames if f in available_filenames] logger.info(f"File matching: {len(matching_files)}/{len(requested_filenames)} files available in database") # Accept if we have at least some matching files if len(matching_files) > 0: logger.info("Database cache validation successful (partial match)") return True else: logger.info("No matching files in database") return False except Exception as e: logger.error(f"Cache validation error: {e}") return False def load_documents_and_embeddings(self, pdf_files: List[Path]) -> Tuple[List[Any], Optional[np.ndarray]]: """ Load documents and embeddings from database Args: pdf_files: List of PDF files to load Returns: Tuple of (documents, embeddings) or (None, None) if failed """ try: with self.get_session() as session: # Load all chunks for the specified files file_names = [f.name for f in pdf_files] # First check if we have any documents at all total_docs = session.query(Document).count() logger.info(f"Total documents in database: {total_docs}") if total_docs == 0: logger.warning("No documents found in database") return None, None # Check which files we have available_docs = session.query(Document.filename).filter( Document.processing_status == 'completed' ).all() available_files = [doc.filename for doc in available_docs] logger.info(f"Available files in database: {available_files[:5]}...") # Show first 5 # Option A: If no specific files requested, load ALL documents if len(file_names) == 0: logger.info(f"No specific files requested - loading ALL {len(available_files)} documents from database") chunks = session.query(DocumentChunk).join(Document).filter( Document.processing_status == 'completed', DocumentChunk.embedding_vector != None ).order_by(Document.id, DocumentChunk.chunk_index).all() else: # Find intersection of requested and available files matching_files = [f for f in file_names if f in available_files] logger.info(f"Matching files: {len(matching_files)}/{len(file_names)}") if not matching_files: logger.warning("No matching files found in database") return None, None chunks = session.query(DocumentChunk).join(Document).filter( Document.filename.in_(matching_files), Document.processing_status == 'completed', DocumentChunk.embedding_vector != None ).order_by(Document.id, DocumentChunk.chunk_index).all() if not chunks: logger.warning("No chunks found in database") return None, None # Convert chunks to document objects and collect embeddings documents = [] embeddings = [] for chunk in chunks: # Create document-like object doc = { 'id': chunk.id, 'content': chunk.content, 'metadata': chunk.chunk_metadata or {}, 'confidence': chunk.confidence_score or 0.8, 'embedding': chunk.get_embedding() } # Add document metadata if doc['metadata'] is None: doc['metadata'] = {} doc['metadata'].update({ 'source': chunk.document.filename, 'page': chunk.chunk_metadata.get('page', 1) if chunk.chunk_metadata else 1, 'chunk_index': chunk.chunk_index }) documents.append(doc) # Collect embedding embedding = chunk.get_embedding() if embedding is not None: embeddings.append(embedding) else: logger.warning(f"Missing embedding for chunk {chunk.id}") if not embeddings: logger.warning("No embeddings found in database") return documents, None embeddings_array = np.array(embeddings) logger.info(f"Loaded {len(documents)} documents and {embeddings_array.shape} embeddings from database") return documents, embeddings_array except Exception as e: logger.error(f"Failed to load from database: {e}") return None, None def save_documents_and_embeddings(self, documents: List[Any], pdf_files: List[Path], processor_config: Dict[str, Any], embedder_config: Dict[str, Any]) -> bool: """ Save documents and embeddings to database Args: documents: List of processed document objects pdf_files: List of source PDF files processor_config: Document processor configuration embedder_config: Embedder configuration Returns: True if save successful """ try: processor_hash = self._hash_config(processor_config) embedder_hash = self._hash_config(embedder_config) # Create processing session session_id = str(uuid.uuid4()) processing_start = time.time() with self.get_session() as session: # Create processing session record proc_session = ProcessingSession( session_id=session_id, processor_config_hash=processor_hash, embedder_config_hash=embedder_hash, documents_processed=len(pdf_files), chunks_created=len(documents) ) session.add(proc_session) session.flush() # Group documents by source file docs_by_file = {} for doc in documents: # Get source and extract filename metadata = doc.get('metadata', {}) if isinstance(doc, dict) else getattr(doc, 'metadata', {}) source = metadata.get('source', 'unknown') # Extract filename from full path import os if source != 'unknown': source_filename = os.path.basename(source) else: source_filename = metadata.get('source_name', 'unknown') if source_filename not in docs_by_file: docs_by_file[source_filename] = [] docs_by_file[source_filename].append(doc) logger.info(f"Grouped documents by file: {list(docs_by_file.keys())[:5]}...") # Show first 5 # Process each file for pdf_file in pdf_files: file_docs = docs_by_file.get(pdf_file.name, []) if not file_docs: logger.warning(f"No documents found for file: {pdf_file.name}") continue # Create or update document record file_hash = self._hash_file(pdf_file) file_mtime = pdf_file.stat().st_mtime doc_record = session.query(Document).filter( Document.filename == pdf_file.name ).first() if not doc_record: doc_record = Document( filename=pdf_file.name, file_path=str(pdf_file), file_hash=file_hash, file_size=pdf_file.stat().st_size, file_mtime=file_mtime, processor_config_hash=processor_hash, chunk_count=len(file_docs), processing_status='completed', doc_metadata={} # Initialize with empty metadata ) session.add(doc_record) session.flush() else: # Update existing record doc_record.file_hash = file_hash doc_record.file_mtime = file_mtime doc_record.processor_config_hash = processor_hash doc_record.chunk_count = len(file_docs) doc_record.processing_status = 'completed' doc_record.processed_at = datetime.utcnow() # Delete old chunks session.query(DocumentChunk).filter( DocumentChunk.document_id == doc_record.id ).delete() # Save chunks for idx, doc in enumerate(file_docs): # Get content and metadata properly if isinstance(doc, dict): content = doc.get('content', '') metadata = doc.get('metadata', {}) confidence = doc.get('confidence', 0.8) else: content = getattr(doc, 'content', '') metadata = getattr(doc, 'metadata', {}) confidence = getattr(doc, 'confidence', 0.8) chunk = DocumentChunk( document_id=doc_record.id, chunk_index=idx, content=content, content_hash=self._hash_text(content), chunk_metadata=metadata, embedding_model=embedder_config.get('model', {}).get('model_name', 'unknown'), embedder_config_hash=embedder_hash, confidence_score=confidence ) # Set embedding if available embedding = None if hasattr(doc, 'embedding') and doc.embedding is not None: embedding = doc.embedding elif isinstance(doc, dict) and 'embedding' in doc and doc['embedding'] is not None: embedding = doc['embedding'] if embedding is not None: # Convert to numpy array if it's a list if isinstance(embedding, list): embedding = np.array(embedding, dtype=np.float32) elif not isinstance(embedding, np.ndarray): embedding = np.array(embedding, dtype=np.float32) chunk.set_embedding(embedding) session.add(chunk) # Update processing session processing_time = (time.time() - processing_start) * 1000 proc_session.completed_at = datetime.utcnow() proc_session.status = 'completed' proc_session.total_processing_time_ms = processing_time proc_session.chunks_created = len(documents) session.commit() logger.info(f"Successfully saved {len(documents)} documents to database in {processing_time:.0f}ms") return True except Exception as e: logger.error(f"Failed to save to database: {e}") return False def cleanup_old_data(self, retention_days: int = 30) -> None: """Clean up old processing sessions and orphaned data""" try: cutoff_date = datetime.utcnow() - timedelta(days=retention_days) with self.get_session() as session: # Clean up old processing sessions old_sessions = session.query(ProcessingSession).filter( ProcessingSession.started_at < cutoff_date ).delete() # Clean up invalid cache entries invalid_cache = session.query(SystemCache).filter( SystemCache.is_valid == False ).delete() session.commit() logger.info(f"Cleaned up {old_sessions} old sessions and {invalid_cache} invalid cache entries") except Exception as e: logger.error(f"Cleanup failed: {e}") def get_processing_history(self, limit: int = 10) -> List[Dict[str, Any]]: """Get recent processing session history""" try: with self.get_session() as session: sessions = session.query(ProcessingSession).order_by( ProcessingSession.started_at.desc() ).limit(limit).all() return [ { 'session_id': s.session_id, 'started_at': s.started_at.isoformat(), 'completed_at': s.completed_at.isoformat() if s.completed_at else None, 'status': s.status, 'documents_processed': s.documents_processed, 'chunks_created': s.chunks_created, 'processing_time_ms': s.total_processing_time_ms, 'documents_per_second': s.documents_per_second } for s in sessions ] except Exception as e: logger.error(f"Failed to get processing history: {e}") return [] def clear_database(self) -> bool: """Clear all data from database (for testing/reset)""" try: with self.get_session() as session: session.query(DocumentChunk).delete() session.query(Document).delete() session.query(ProcessingSession).delete() session.query(SystemCache).delete() session.commit() logger.info("Database cleared successfully") return True except Exception as e: logger.error(f"Failed to clear database: {e}") return False def _hash_file(self, file_path: Path) -> str: """Generate hash of file content""" try: with open(file_path, 'rb') as f: return hashlib.md5(f.read()).hexdigest() except Exception as e: logger.warning(f"Failed to hash file {file_path}: {e}") return "" def _hash_text(self, text: str) -> str: """Generate hash of text content""" return hashlib.md5(text.encode('utf-8')).hexdigest() def _hash_config(self, config: Dict[str, Any]) -> str: """Generate hash of configuration dictionary""" try: import json # Convert config to string, handling any non-serializable objects config_str = json.dumps(config, sort_keys=True, default=str) return hashlib.md5(config_str.encode('utf-8')).hexdigest() except Exception as e: logger.warning(f"Config hash generation failed: {e}") # Fallback to string representation config_str = str(sorted(config.items())) return hashlib.md5(config_str.encode('utf-8')).hexdigest() # Global database manager instance _db_manager = None def get_database_manager(database_url: str = "sqlite:///demo/epic2_demo.db") -> DatabaseManager: """Get global database manager instance""" global _db_manager if _db_manager is None: _db_manager = DatabaseManager(database_url) return _db_manager def reset_database_manager(): """Reset global database manager (for testing)""" global _db_manager _db_manager = None