""" Simple database client for artifact caching """ import os import json import hashlib import logging from typing import List, Dict, Optional from datetime import datetime from dotenv import load_dotenv from pathlib import Path # Conditional streamlit import to avoid issues when not available try: import streamlit as st STREAMLIT_AVAILABLE = True except ImportError: st = None STREAMLIT_AVAILABLE = False # Load configuration using the configuration manager try: from .config_manager import load_configuration load_configuration() print("✅ Configuration loaded in simple_db.py") except Exception as e: print(f"⚠️ Error loading configuration in simple_db.py: {e}") # Fallback to manual loading project_root = Path(__file__).parent.parent env_path = project_root / ".env" load_dotenv(env_path, override=True) logger = logging.getLogger(__name__) class SimpleArtifactDB: """Simple database client for artifact storage and caching""" def __init__(self): self.supabase_client = None self.enabled = False # Try to initialize Supabase try: url = os.getenv("SUPABASE_URL") key = os.getenv("SUPABASE_ANON_KEY") enable = os.getenv("ENABLE_SUPABASE", "false").lower() logger.info(f"🔍 Environment check - URL: {'SET' if url else 'NOT SET'}, " f"KEY: {'SET' if key else 'NOT SET'}, " f"ENABLE: {enable}") if url and key and enable == "true": from supabase import create_client self.supabase_client = create_client(url, key) self.enabled = True logger.info("✅ Simple database client initialized") else: logger.info(f"📝 Database disabled - URL: {bool(url)}, KEY: {bool(key)}, ENABLE: {enable}") except Exception as e: logger.warning(f"⚠️ Database initialization failed: {e}") self.enabled = False def _hash_file(self, file_path: str) -> str: """Generate SHA-256 hash of file content""" try: hasher = hashlib.sha256() with open(file_path, 'rb') as f: for chunk in iter(lambda: f.read(4096), b""): hasher.update(chunk) return hasher.hexdigest() except Exception as e: logger.error(f"Error hashing file {file_path}: {e}") return "" def _create_content_fingerprint(self, file_path: str) -> str: """Create a fast content fingerprint using file start + end + size This is much faster than full file hashing and sufficient for cache identification. Combines file size, first 1KB, and last 1KB of content. """ try: if not os.path.exists(file_path): logger.warning(f"File does not exist for fingerprinting: {file_path}") return "" file_size = os.path.getsize(file_path) with open(file_path, 'rb') as f: # Read first 1KB start_chunk = f.read(1024) # Read last 1KB if file is large enough if file_size > 2048: f.seek(-1024, 2) # Seek to 1KB from end end_chunk = f.read(1024) else: end_chunk = b"" # Combine: file_size + start_chunk + end_chunk fingerprint_data = f"{file_size}".encode() + start_chunk + end_chunk fingerprint = hashlib.sha256(fingerprint_data).hexdigest() logger.debug(f"Created content fingerprint for {os.path.basename(file_path)}: {fingerprint[:16]}... (size: {file_size})") return fingerprint except Exception as e: logger.error(f"Error creating content fingerprint for {file_path}: {e}") return "" def check_page_level_cache(self, doc_group: dict, start_page: int, end_page, ocr_model: str, extraction_model: str, thresholds: dict) -> tuple: """ Check cache for each page in the requested range Returns: tuple: (cached_artifacts, missing_pages, cache_stats) """ # Handle None end_page by setting to a large number (will be handled by actual processing) if end_page is None: end_page = 9999 # Large number to indicate "till end" if not self.enabled: return [], list(range(start_page, end_page + 1)), {"cached_pages": 0, "missing_pages": end_page - start_page + 1} # Validate inputs if not doc_group or not doc_group.get("EN"): logger.error("Invalid doc_group - missing English document") return [], list(range(start_page, end_page + 1)), {"cached_pages": 0, "missing_pages": end_page - start_page + 1} if start_page > end_page: logger.error(f"Invalid page range: {start_page} > {end_page}") return [], [], {"cached_pages": 0, "missing_pages": 0} try: # FINAL SAFETY CHECK: Handle None end_page even if it somehow got through if end_page is None: logger.warning(f"🚨 EMERGENCY: end_page is None in cache function, converting to 9999") end_page = 9999 requested_pages = list(range(start_page, end_page + 1)) cached_artifacts = [] missing_pages = [] # Create processing params hash processing_params_hash = hashlib.sha256( json.dumps(thresholds, sort_keys=True).encode() ).hexdigest() logger.info(f"🔍 Checking cache for pages {start_page}-{end_page}") logger.info(f"🔍 Cache check parameters: OCR={ocr_model}, Extract={extraction_model}") for page_num in requested_pages: # Create page cache key page_cache_key = self._create_page_cache_key( doc_group, page_num, ocr_model, extraction_model, thresholds ) # Check if this page exists in cache try: result = self.supabase_client.rpc("check_page_cache", { "p_page_cache_key": page_cache_key }).execute() except Exception as rpc_error: logger.warning(f"⚠️ RPC call failed for page {page_num}, falling back to direct query: {rpc_error}") # Fallback to direct table query result = self.supabase_client.table("artifacts").select("*").eq( "page_cache_key", page_cache_key ).execute() if result.data: # Convert database format back to original format page_artifacts = [] for record in result.data: artifact = { "Name_EN": record.get("name_en", ""), "Name_AR": record.get("name_ar", ""), "Name_FR": record.get("name_fr", ""), "Creator": record.get("creator", ""), "Creation Date": record.get("creation_date", ""), "Materials": record.get("materials", ""), "Origin": record.get("origin", ""), "Description": record.get("description", ""), "Category": record.get("category", ""), "source_page": record.get("source_page", page_num), "source_document": record.get("source_document", ""), "Name_validation": record.get("name_validation", "") } page_artifacts.append(artifact) cached_artifacts.extend(page_artifacts) logger.info(f"✅ Cache HIT - Page {page_num}: Found {len(page_artifacts)} cached artifacts") else: missing_pages.append(page_num) logger.info(f"❌ Cache MISS - Page {page_num}: Needs processing") cache_stats = { "cached_pages": len(requested_pages) - len(missing_pages), "missing_pages": len(missing_pages), "total_cached_artifacts": len(cached_artifacts) } logger.info(f"📊 Cache Summary: {cache_stats['cached_pages']} cached, {cache_stats['missing_pages']} need processing") return cached_artifacts, missing_pages, cache_stats except Exception as e: logger.error(f"Error checking page-level cache: {e}") # Fallback: treat all pages as missing # Handle None end_page in fallback as well if end_page is None: end_page = 9999 return [], list(range(start_page, end_page + 1)), {"cached_pages": 0, "missing_pages": end_page - start_page + 1} def _create_page_cache_key(self, doc_group: dict, page_num: int, ocr_model: str, extraction_model: str, thresholds: dict) -> str: """Create unique cache key using content fingerprints instead of file paths""" content_hashes = {} for lang, file_path in doc_group.items(): if file_path and os.path.exists(file_path): # Use content fingerprint for fast, reliable identification content_hashes[lang] = self._create_content_fingerprint(file_path) else: # Fallback to original filename from session state try: if STREAMLIT_AVAILABLE and st and hasattr(st, 'session_state') and hasattr(st.session_state, 'uploaded_file_names'): original_name = st.session_state.uploaded_file_names.get(lang, "") if original_name: content_hashes[lang] = hashlib.sha256(original_name.encode()).hexdigest()[:16] else: content_hashes[lang] = "missing" else: content_hashes[lang] = "missing" except Exception: content_hashes[lang] = "missing" # Create parameter combination params = { 'content_hashes': content_hashes, 'page': page_num, 'ocr_model': ocr_model, 'extraction_model': extraction_model, 'thresholds': json.dumps(thresholds, sort_keys=True) } # Generate SHA-256 hash params_str = json.dumps(params, sort_keys=True) cache_key = hashlib.sha256(params_str.encode()).hexdigest() return cache_key def _create_run_cache_key(self, doc_group: dict, start_page: int, end_page: int, ocr_model: str, extraction_model: str, thresholds: dict) -> str: """Create unique cache key for an entire processing run using content""" en_file = doc_group.get("EN") if en_file and os.path.exists(en_file): # Use content fingerprint for reliable identification content_hash = self._create_content_fingerprint(en_file) logger.debug(f"Using content fingerprint for run cache: {content_hash[:16]}...") else: # Fallback to original filename from session state try: if STREAMLIT_AVAILABLE and st and hasattr(st, 'session_state') and hasattr(st.session_state, 'uploaded_file_names'): original_name = st.session_state.uploaded_file_names.get("EN", "") if original_name: content_hash = hashlib.sha256(original_name.encode()).hexdigest() logger.debug(f"Using filename fallback for run cache: {original_name}") else: content_hash = "unknown_content" logger.warning("No English file or filename available for run cache") else: content_hash = "unknown_content" logger.debug("Session state not available for run cache, using 'unknown_content'") except Exception as e: logger.warning(f"Error accessing session state for run cache: {e}") content_hash = "unknown_content" params = { 'content_hash': content_hash, 'start_page': start_page, 'end_page': end_page, 'ocr_model': ocr_model, 'extraction_model': extraction_model, 'thresholds': json.dumps(thresholds, sort_keys=True) } params_str = json.dumps(params, sort_keys=True) run_cache_key = hashlib.sha256(params_str.encode()).hexdigest() logger.debug(f"🔍 Run cache key generated: {run_cache_key[:16]}...") return run_cache_key def _map_artifact_to_new_schema(self, artifact: dict, page_cache_key: str, page_num: int, ocr_model: str, extraction_model: str, processing_params_hash: str, actual_source_document: str = None) -> dict: """Map artifact fields from old format to new schema with cache keys""" # Use actual source document name if provided, otherwise fall back to artifact data source_doc = actual_source_document if actual_source_document else artifact.get("source_document", "") return { "page_cache_key": page_cache_key, "page_number": page_num, "name_en": artifact.get("Name_EN", artifact.get("Name", "")), "name_ar": artifact.get("Name_AR", ""), "name_fr": artifact.get("Name_FR", ""), "creator": artifact.get("Creator", ""), "creation_date": artifact.get("Creation Date", ""), "materials": artifact.get("Materials", ""), "origin": artifact.get("Origin", ""), "description": artifact.get("Description", ""), "category": artifact.get("Category", ""), "source_page": artifact.get("source_page", page_num), "source_document": source_doc, "name_validation": artifact.get("Name_validation", ""), "ocr_model": ocr_model, "extraction_model": extraction_model, "processing_params_hash": processing_params_hash } def save_page_artifacts(self, doc_group: dict, page_num: int, artifacts: List[Dict], ocr_model: str, extraction_model: str, thresholds: dict, provided_file_hash: str = None) -> bool: """Save artifacts from a specific page with cache key""" if not self.enabled or not artifacts: return False # Debug: Log the models being saved logger.info(f"💾 DB Save - OCR model: '{ocr_model}', Extraction model: '{extraction_model}'") try: # Get main file hash - use provided hash if available, otherwise create content fingerprint if provided_file_hash: file_hash = provided_file_hash logger.info(f"💾 Using provided file hash: {file_hash[:16]}...") else: en_file = doc_group.get("EN", "") if en_file and os.path.exists(en_file): # Use content fingerprint for reliable, fast identification file_hash = self._create_content_fingerprint(en_file) logger.info(f"💾 Created content fingerprint: {file_hash[:16]}...") else: # Content-based fallback using original filename from session state try: if STREAMLIT_AVAILABLE and st and hasattr(st, 'session_state') and hasattr(st.session_state, 'uploaded_file_names'): original_name = st.session_state.uploaded_file_names.get("EN", "") if original_name: file_hash = hashlib.sha256(original_name.encode()).hexdigest() logger.info(f"💾 Using filename-based hash: {file_hash[:16]}... (from: {original_name})") else: logger.error("💾 Cannot create file hash - no file or filename available") return False else: logger.error("💾 Cannot create file hash - session state not available and no file exists") return False except Exception as e: logger.error(f"💾 Error accessing session state for file hash: {e}") return False if not file_hash: logger.error("💾 Cannot create file hash for page artifacts") return False # Create cache keys page_cache_key = self._create_page_cache_key( doc_group, page_num, ocr_model, extraction_model, thresholds ) processing_params_hash = hashlib.sha256( json.dumps(thresholds, sort_keys=True).encode() ).hexdigest() # Check if this page is already cached existing_check = self.supabase_client.table("artifacts").select("id").eq( "page_cache_key", page_cache_key ).limit(1).execute() if existing_check.data: logger.info(f"Page {page_num} already cached, skipping save") return True # Map and save each artifact saved_count = 0 for artifact in artifacts: # Use the source_document from the artifact itself (which is already correct in the UI) artifact_source_document = artifact.get("source_document", "") mapped_artifact = self._map_artifact_to_new_schema( artifact, page_cache_key, page_num, ocr_model, extraction_model, processing_params_hash, artifact_source_document ) # Add file hash mapped_artifact["file_hash"] = file_hash try: self.supabase_client.table("artifacts").insert(mapped_artifact).execute() saved_count += 1 except Exception as e: logger.error(f"Error saving individual artifact: {e}") continue logger.info(f"💾 Saved {saved_count} artifacts for page {page_num}") return saved_count > 0 except Exception as e: logger.error(f"Error saving page artifacts: {e}") return False def save_run_statistics(self, doc_group: dict, start_page: int, end_page, ocr_model: str, extraction_model: str, thresholds: dict, total_artifacts: int, cached_pages: int, processed_pages: int, provided_file_hash: str = None) -> bool: """Save statistics for a complete processing run""" if not self.enabled: return False # Handle None end_page if end_page is None: end_page = 9999 try: # Get main file hash - use provided hash if available, otherwise try to hash the file if provided_file_hash: file_hash = provided_file_hash logger.info(f"📊 Using provided file hash: {file_hash[:16]}...") else: en_file = doc_group.get("EN", "") file_hash = self._hash_file(en_file) if en_file and os.path.exists(en_file) else "" if not file_hash: logger.warning(f"⚠️ Could not generate file hash for {en_file}") # Create run cache key run_cache_key = self._create_run_cache_key( doc_group, start_page, end_page, ocr_model, extraction_model, thresholds ) processing_params_hash = hashlib.sha256( json.dumps(thresholds, sort_keys=True).encode() ).hexdigest() # Save run statistics run_record = { "run_cache_key": run_cache_key, "file_hash": file_hash, "start_page": start_page, "end_page": end_page, "ocr_model": ocr_model, "extraction_model": extraction_model, "processing_params_hash": processing_params_hash, "total_artifacts": total_artifacts, "cached_pages": cached_pages, "processed_pages": processed_pages, "processing_status": "completed" } # Use upsert to avoid duplicates result = self.supabase_client.table("processing_cache").upsert(run_record).execute() logger.info(f"📊 Saved run statistics: {total_artifacts} artifacts, {cached_pages} cached pages, {processed_pages} processed pages") return True except Exception as e: logger.error(f"Error saving run statistics: {e}") return False def save_artifacts_from_data(self, file_hash: str, artifacts_data) -> bool: """Legacy method for backward compatibility - converts to page-based saving""" if not self.enabled: logger.warning("Database not enabled - cannot save artifacts") return False try: # Handle different data structures artifacts = [] if isinstance(artifacts_data, list): artifacts = artifacts_data logger.info(f"Received artifacts as list with {len(artifacts)} items") elif isinstance(artifacts_data, dict): artifacts = artifacts_data.get('artifacts', []) logger.info(f"Received artifacts as dict, extracted {len(artifacts)} artifacts") if not artifacts and artifacts_data: artifacts = [artifacts_data] logger.info("Treating single dict as one artifact") else: logger.error(f"Unexpected data type: {type(artifacts_data)}") return False if not artifacts: logger.warning(f"No artifacts found in data") return False # Group artifacts by page for proper caching artifacts_by_page = {} for artifact in artifacts: page_num = artifact.get("source_page", 1) if page_num not in artifacts_by_page: artifacts_by_page[page_num] = [] artifacts_by_page[page_num].append(artifact) # Save each page separately (using default parameters for legacy data) doc_group = {"EN": f"legacy_file_{file_hash[:8]}.pdf"} default_model = "gpt-4o" default_thresholds = {"EN": 0.05, "AR": 0.10, "FR": 0.07} total_saved = 0 for page_num, page_artifacts in artifacts_by_page.items(): success = self.save_page_artifacts( doc_group, page_num, page_artifacts, default_model, default_model, default_thresholds ) if success: total_saved += len(page_artifacts) logger.info(f"💾 Successfully saved {total_saved} artifacts across {len(artifacts_by_page)} pages") return total_saved > 0 except Exception as e: logger.error(f"Error saving artifacts from data: {e}") return False def search_artifacts(self, query: str = None, limit: int = 100) -> List[Dict]: """Search artifacts using the new schema""" if not self.enabled: return [] try: # Use the database function for search if query: result = self.supabase_client.rpc("search_artifacts", { "search_term": query, "limit_count": limit }).execute() else: result = self.supabase_client.table("artifacts").select("*").order( "created_at", desc=True ).limit(limit).execute() # Convert back to original format for compatibility artifacts = [] for record in result.data: artifact = { "Name_EN": record.get("name_en", ""), "Name_AR": record.get("name_ar", ""), "Name_FR": record.get("name_fr", ""), "Creator": record.get("creator", ""), "Creation Date": record.get("creation_date", ""), "Materials": record.get("materials", ""), "Origin": record.get("origin", ""), "Description": record.get("description", ""), "Category": record.get("category", ""), "source_page": record.get("source_page", record.get("page_number", "")), "source_document": record.get("source_document", ""), "Name_validation": record.get("name_validation", "") } artifacts.append(artifact) logger.info(f"🔍 Found {len(artifacts)} artifacts matching query") return artifacts except Exception as e: logger.error(f"Error searching artifacts: {e}") return [] def get_stats(self) -> Dict: """Get database statistics using the new schema""" if not self.enabled: return {"enabled": False} try: # Use the enhanced statistics function try: stats_result = self.supabase_client.rpc("get_artifact_statistics").execute() except Exception as rpc_error: logger.warning(f"RPC call for statistics failed, falling back to simple counting: {rpc_error}") stats_result = None if stats_result.data: stats = stats_result.data[0] return { "enabled": True, "total_artifacts": stats.get("total_artifacts", 0), "unique_runs": stats.get("unique_runs", 0), "unique_pages": stats.get("unique_pages", 0), "categories_count": stats.get("categories_count", 0), "creators_count": stats.get("creators_count", 0), "origins_count": stats.get("origins_count", 0), "unique_documents_count": stats.get("unique_documents_count", 0), "last_updated": datetime.now().isoformat() } else: # Fallback to simple counting artifacts_count = self.supabase_client.table("artifacts").select( "*", count="exact" ).execute() runs_count = self.supabase_client.table("processing_cache").select( "run_cache_key", count="exact" ).execute() return { "enabled": True, "total_artifacts": artifacts_count.count, "unique_runs": runs_count.count, "last_updated": datetime.now().isoformat() } except Exception as e: logger.error(f"Error getting stats: {e}") return {"enabled": True, "error": str(e)} # Global instance _db_instance = None def get_simple_db() -> SimpleArtifactDB: """Get the global database instance""" global _db_instance if _db_instance is None: _db_instance = SimpleArtifactDB() return _db_instance