import asyncio import os import json import logging import numpy as np import pickle import gzip from typing import Dict, List, Optional, Any, Tuple from datetime import datetime import uuid import httpx import base64 from dataclasses import dataclass # LightRAG imports from lightrag import LightRAG, QueryParam from lightrag.utils import EmbeddingFunc from lightrag.kg.shared_storage import initialize_pipeline_status # Database imports import asyncpg from redis import Redis # Environment validation REQUIRED_ENV_VARS = [ 'CLOUDFLARE_API_KEY', 'CLOUDFLARE_ACCOUNT_ID', 'DATABASE_URL', 'BLOB_READ_WRITE_TOKEN', 'REDIS_URL', 'JWT_SECRET' ] class EnvironmentError(Exception): """Raised when required environment variables are missing""" pass def validate_environment(): """Validate all required environment variables are present""" missing_vars = [] for var in REQUIRED_ENV_VARS: if not os.getenv(var): missing_vars.append(var) if missing_vars: raise EnvironmentError(f"Missing required environment variables: {', '.join(missing_vars)}") @dataclass class RAGConfig: """Configuration for RAG instances""" ai_type: str user_id: Optional[str] = None ai_id: Optional[str] = None name: Optional[str] = None description: Optional[str] = None def get_cache_key(self) -> str: """Generate cache key for this RAG configuration""" return f"rag_{self.ai_type}_{self.user_id or 'system'}_{self.ai_id or 'default'}" class CloudflareWorker: def __init__(self, cloudflare_api_key: str, api_base_url: str, llm_model_name: str, embedding_model_name: str, max_tokens: int = 4080): self.cloudflare_api_key = cloudflare_api_key self.api_base_url = api_base_url self.max_tokens = max_tokens self.logger = logging.getLogger(__name__) self.llm_model_name = llm_model_name self.embedding_model_name = embedding_model_name self.llm_models = [ "@cf/meta/llama-3.1-8b-instruct", "@cf/deepseek-ai/deepseek-r1-distill-qwen-32b", "@cf/mistralai/mistral-small-3.1-24b-instruct", "@cf/meta/llama-4-scout-17b-16e-instruct", "@cf/meta/llama-3.2-11b-vision-instruct", "@cf/meta/llama-3-8b-instruct", # ✅ VERY GOOD - Llama 3, 8B params "@cf/mistral/mistral-7b-instruct-v0.1", # ✅ GOOD - Mistral, excellent reasoning "@cf/meta/llama-2-7b-chat-int8", # ✅ RELIABLE - Stable Llama 2 "@cf/microsoft/phi-2", # ✅ FAST - Microsoft's small but powerful "@cf/meta/llama-3.2-3b-instruct", # ✅ CURRENT - Your working model "@cf/google/gemma-3-12b-it", "@cf/google/gemma-7b-it", # ✅ GOOD - Google's model "@cf/qwen/qwen1.5-7b-chat-awq", # ✅ ALTERNATIVE - Chinese but works "@cf/tiiuae/falcon-7b-instruct", "@cf/microsoft/dialoGPT-medium", ] # VERIFIED WORKING embedding models self.embedding_models = [ "@cf/baai/bge-large-en-v1.5", # 🏆 BEST - Largest, most accurate "@cf/baai/bge-base-en-v1.5", # ✅ GOOD - Standard choice "@cf/baai/bge-small-en-v1.5", # ✅ FAST - Smaller but decent "@cf/baai/bge-m3", # ✅ CURRENT - Multilingual ] self.current_llm_index = 0 self.current_embedding_index = 0 async def query(self, prompt: str, system_prompt: str = "", **kwargs) -> str: """Enhanced query with better entity extraction prompting""" # ENHANCED: Better system prompt for entity extraction if not system_prompt: system_prompt = """You are an expert technical document analyzer. Your main goal is to identify and extract important technical entities, concepts, and objects from specialized documents. Focus on: - Technical terms and concepts - Equipment and devices - Procedures and processes - Standards and requirements - Physical objects and systems Be precise and technical in your analysis.""" filtered_kwargs = {k: v for k, v in kwargs.items() if k not in ['hashing_kv', 'history_messages', 'global_kv', 'text_chunks']} messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt[:self.max_tokens]}, ] # ENHANCED: Better parameters for technical content input_data = { "messages": messages, "max_tokens": min(self.max_tokens, 4096), "temperature": 0.2, # Lower temperature for more focused, technical responses "top_p": 0.85, # Slightly focused sampling **filtered_kwargs } response, new_index = await self._send_request_with_fallback(self.llm_models, self.current_llm_index, input_data) self.current_llm_index = new_index # Log which model was used model_used = self.llm_models[new_index] self.logger.info(f"🤖 Used model: {model_used}") return response async def _send_request_with_fallback(self, model_list: List[str], current_index: int, input_: dict) -> Tuple[ Any, int]: """Send request with model fallback""" for i in range(len(model_list)): model_index = (current_index + i) % len(model_list) model_name = model_list[model_index] try: headers = {"Authorization": f"Bearer {self.cloudflare_api_key}"} async with httpx.AsyncClient(timeout=30.0) as client: response = await client.post( f"{self.api_base_url}{model_name}", headers=headers, json=input_ ) response.raise_for_status() result = response.json().get("result", {}) if "data" in result: return np.array(result["data"]), model_index elif "response" in result: return result["response"], model_index else: continue except Exception as e: self.logger.warning(f"Model {model_name} failed: {e}") continue raise Exception("All models failed") async def query(self, prompt: str, system_prompt: str = "", **kwargs) -> str: filtered_kwargs = {k: v for k, v in kwargs.items() if k not in ['hashing_kv', 'history_messages', 'global_kv', 'text_chunks']} messages = [ {"role": "system", "content": system_prompt or "You are a helpful AI assistant. Your main goal is to help with the knowledge you have from LightRAG files"}, {"role": "user", "content": prompt[:self.max_tokens]}, ] input_data = {"messages": messages, "max_tokens": min(self.max_tokens, 4096), **filtered_kwargs} response, new_index = await self._send_request_with_fallback(self.llm_models, self.current_llm_index, input_data) self.current_llm_index = new_index return response async def embedding_chunk(self, texts: List[str]) -> np.ndarray: truncated_texts = [text[:2000] for text in texts] input_data = {"text": truncated_texts} response, new_index = await self._send_request_with_fallback(self.embedding_models, self.current_embedding_index, input_data) self.current_embedding_index = new_index return response class VercelBlobClient: """Vercel Blob storage client for RAG state persistence""" def __init__(self, token: str): self.token = token self.logger = logging.getLogger(__name__) async def put(self, filename: str, data: bytes) -> str: """Upload data to Vercel Blob""" try: async with httpx.AsyncClient(timeout=120.0) as client: response = await client.put( f"https://blob.vercel-storage.com/{filename}", headers={"Authorization": f"Bearer {self.token}"}, content=data ) response.raise_for_status() result = response.json() return result.get('url', f"https://blob.vercel-storage.com/{filename}") except Exception as e: self.logger.error(f"Failed to upload to Vercel Blob: {e}") raise async def get(self, url: str) -> bytes: """Download data from Vercel Blob""" try: async with httpx.AsyncClient(timeout=120.0) as client: response = await client.get(url) response.raise_for_status() return response.content except Exception as e: self.logger.error(f"Failed to download from Vercel Blob: {e}") raise class DatabaseManager: """Database manager with complete RAG persistence""" def __init__(self, database_url: str, redis_url: str): self.database_url = database_url self.redis_url = redis_url self.pool = None self.redis = None self.logger = logging.getLogger(__name__) async def connect(self): """Initialize database connections""" try: self.pool = await asyncpg.create_pool( self.database_url, min_size=2, max_size=20, command_timeout=60 ) self.redis = Redis.from_url(self.redis_url, decode_responses=True) self.logger.info("Database connections established successfully") await self._create_tables() except Exception as e: self.logger.error(f"Database connection failed: {e}") raise async def _create_tables(self): """Create necessary tables for RAG persistence""" async with self.pool.acquire() as conn: await conn.execute(""" CREATE TABLE IF NOT EXISTS rag_instances ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), ai_type VARCHAR(50) NOT NULL, user_id VARCHAR(100), ai_id VARCHAR(100), name VARCHAR(255) NOT NULL, description TEXT, graph_blob_url TEXT, vector_blob_url TEXT, config_blob_url TEXT, total_chunks INTEGER DEFAULT 0, total_tokens INTEGER DEFAULT 0, file_count INTEGER DEFAULT 0, created_at TIMESTAMP DEFAULT NOW(), updated_at TIMESTAMP DEFAULT NOW(), last_accessed_at TIMESTAMP DEFAULT NOW(), status VARCHAR(20) DEFAULT 'active', UNIQUE(ai_type, user_id, ai_id) ); CREATE TABLE IF NOT EXISTS knowledge_files ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), rag_instance_id UUID REFERENCES rag_instances(id) ON DELETE CASCADE, filename VARCHAR(255) NOT NULL, original_filename VARCHAR(255), file_type VARCHAR(50), file_size INTEGER, blob_url TEXT, content_text TEXT, processed_at TIMESTAMP DEFAULT NOW(), processing_status VARCHAR(20) DEFAULT 'processed', token_count INTEGER DEFAULT 0, created_at TIMESTAMP DEFAULT NOW() ); CREATE TABLE IF NOT EXISTS conversations ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), user_id VARCHAR(100) NOT NULL, ai_type VARCHAR(50) NOT NULL, ai_id VARCHAR(100), title VARCHAR(255), created_at TIMESTAMP DEFAULT NOW(), updated_at TIMESTAMP DEFAULT NOW(), is_active BOOLEAN DEFAULT TRUE ); CREATE TABLE IF NOT EXISTS conversation_messages ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), conversation_id UUID REFERENCES conversations(id) ON DELETE CASCADE, role VARCHAR(20) NOT NULL, content TEXT NOT NULL, metadata JSONB DEFAULT '{}', created_at TIMESTAMP DEFAULT NOW() ); CREATE TABLE IF NOT EXISTS system_stats ( id VARCHAR(50) PRIMARY KEY DEFAULT gen_random_uuid()::text, total_users INTEGER NOT NULL DEFAULT 0, total_ais INTEGER NOT NULL DEFAULT 0, total_messages INTEGER NOT NULL DEFAULT 0, date TIMESTAMP WITHOUT TIME ZONE NOT NULL DEFAULT NOW() ); CREATE INDEX IF NOT EXISTS idx_system_stats_date ON system_stats(date DESC); CREATE UNIQUE INDEX IF NOT EXISTS idx_system_stats_date_unique ON system_stats(DATE(date)); -- Insert initial stats if table is empty INSERT INTO system_stats (id, total_users, total_ais, total_messages, date) SELECT 'initial', 0, 0, 0, NOW() WHERE NOT EXISTS (SELECT 1 FROM system_stats); CREATE INDEX IF NOT EXISTS idx_rag_instances_lookup ON rag_instances(ai_type, user_id, ai_id); CREATE INDEX IF NOT EXISTS idx_conversations_user ON conversations(user_id); CREATE INDEX IF NOT EXISTS idx_conversation_messages_conv ON conversation_messages(conversation_id); """) self.logger.info("Database tables created/verified successfully") async def initialize_system_stats(self): """Initialize system stats with current counts from database""" try: async with self.pool.acquire() as conn: # Count actual users user_count = await conn.fetchval(""" SELECT COUNT(*) FROM users WHERE is_active = true """) or 0 # Count actual custom AIs ai_count = await conn.fetchval(""" SELECT COUNT(*) FROM custom_ais WHERE is_active = true """) or 0 # Count actual messages from both tables message_count = await conn.fetchval(""" SELECT (SELECT COUNT(*) FROM messages) + (SELECT COUNT(*) FROM conversation_messages) """) or 0 # Check if today's stats already exist today = datetime.now().date() existing_stats = await conn.fetchrow(""" SELECT id FROM system_stats WHERE DATE(date) = $1 """, today) if existing_stats: # Update existing record await conn.execute(""" UPDATE system_stats SET total_users = $1, total_ais = $2, total_messages = $3, date = NOW() WHERE DATE(date) = $4 """, user_count, ai_count, message_count, today) else: # Insert new record for today await conn.execute(""" INSERT INTO system_stats (id, total_users, total_ais, total_messages, date) VALUES ($1, $2, $3, $4, NOW()) """, f"stats_{today}", user_count, ai_count, message_count) self.logger.info( f"📊 Initialized system stats: {user_count} users, {ai_count} AIs, {message_count} messages") except Exception as e: self.logger.error(f"Failed to initialize system stats: {e}") async def update_system_stat(self, stat_type: str, increment: int = 1): """Update a specific system statistic""" try: async with self.pool.acquire() as conn: today = datetime.now().date() # Map stat types to column names column_map = { 'users': 'total_users', 'ais': 'total_ais', 'messages': 'total_messages' } if stat_type not in column_map: self.logger.warning(f"Unknown stat type: {stat_type}") return column_name = column_map[stat_type] # Upsert today's record await conn.execute(f""" INSERT INTO system_stats (id, total_users, total_ais, total_messages, date) VALUES ($1, CASE WHEN '{column_name}' = 'total_users' THEN $2 ELSE 0 END, CASE WHEN '{column_name}' = 'total_ais' THEN $2 ELSE 0 END, CASE WHEN '{column_name}' = 'total_messages' THEN $2 ELSE 0 END, NOW()) ON CONFLICT (DATE(date)) DO UPDATE SET {column_name} = system_stats.{column_name} + $2, date = NOW() """, f"stats_{today}", increment) self.logger.debug(f"📈 Updated {stat_type} by {increment}") except Exception as e: self.logger.error(f"Failed to update {stat_type} stat: {e}") async def get_current_stats(self): """Get current system statistics""" try: async with self.pool.acquire() as conn: # Get latest stats stats_row = await conn.fetchrow(""" SELECT total_users, total_ais, total_messages, date FROM system_stats ORDER BY date DESC LIMIT 1 """) if not stats_row: # Initialize if no stats exist await self.initialize_system_stats() return await self.get_current_stats() # Calculate total characters (lines of code) total_characters = await conn.fetchval(""" SELECT COALESCE( (SELECT SUM(LENGTH(content)) FROM messages) + (SELECT SUM(LENGTH(content)) FROM conversation_messages), 0 ) """) return { 'total_users': stats_row['total_users'], 'total_ais': stats_row['total_ais'], 'total_messages': stats_row['total_messages'], 'lines_of_code_generated': total_characters or 0, 'last_updated': stats_row['date'].isoformat() } except Exception as e: self.logger.error(f"Failed to get current stats: {e}") # Return default stats on error return { 'total_users': 0, 'total_ais': 0, 'total_messages': 0, 'lines_of_code_generated': 0, 'last_updated': datetime.now().isoformat() } async def save_rag_instance(self, config: RAGConfig, graph_blob_url: str, vector_blob_url: str, config_blob_url: str, metadata: Dict[str, Any]) -> str: async with self.pool.acquire() as conn: rag_instance_id = await conn.fetchval(""" INSERT INTO rag_instances ( ai_type, user_id, ai_id, name, description, graph_blob_url, vector_blob_url, config_blob_url, total_chunks, total_tokens, file_count, created_at, updated_at, last_accessed_at ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, NOW(), NOW(), NOW()) ON CONFLICT (ai_type, user_id, ai_id) DO UPDATE SET graph_blob_url = EXCLUDED.graph_blob_url, vector_blob_url = EXCLUDED.vector_blob_url, config_blob_url = EXCLUDED.config_blob_url, total_chunks = EXCLUDED.total_chunks, total_tokens = EXCLUDED.total_tokens, file_count = EXCLUDED.file_count, updated_at = NOW() RETURNING id; """, config.ai_type, config.user_id, config.ai_id, config.name, config.description, graph_blob_url, vector_blob_url, config_blob_url, metadata.get('total_chunks', 0), metadata.get('total_tokens', 0), metadata.get('file_count', 0) ) return str(rag_instance_id) async def cleanup_duplicate_rag_instances(self, ai_type: str, keep_latest: bool = True): """Clean up duplicate RAG instances, keeping only the latest one""" async with self.pool.acquire() as conn: if keep_latest: # Keep the latest instance, deactivate others await conn.execute(""" UPDATE rag_instances SET status = 'duplicate_cleanup' WHERE ai_type = $1 AND user_id IS NULL AND ai_id IS NULL AND status = 'active' AND id NOT IN ( SELECT id FROM rag_instances WHERE ai_type = $1 AND user_id IS NULL AND ai_id IS NULL AND status = 'active' ORDER BY created_at DESC LIMIT 1 ) """, ai_type) count = await conn.fetchval(""" SELECT COUNT(*) FROM rag_instances WHERE ai_type = $1 AND status = 'duplicate_cleanup' """, ai_type) self.logger.info(f"🧹 Cleaned up {count} duplicate {ai_type} RAG instances") # Return the active instance info active_instance = await conn.fetchrow(""" SELECT id, name, created_at FROM rag_instances WHERE ai_type = $1 AND user_id IS NULL AND ai_id IS NULL AND status = 'active' ORDER BY created_at DESC LIMIT 1 """, ai_type) if active_instance: self.logger.info(f"✅ Active {ai_type} RAG: {active_instance['name']} (ID: {active_instance['id']})") return active_instance async def get_rag_instance(self, config: RAGConfig) -> Optional[Dict[str, Any]]: """Get RAG instance from database with FIXED cache key matching""" async with self.pool.acquire() as conn: # Handle NULL/None matching properly for PostgreSQL if config.user_id is None and config.ai_id is None: # System-level RAG (fire-safety, general, etc.) result = await conn.fetchrow(""" SELECT id, ai_type, user_id, ai_id, name, description, graph_blob_url, vector_blob_url, config_blob_url, total_chunks, total_tokens, file_count, created_at, updated_at, last_accessed_at, status FROM rag_instances WHERE ai_type = $1 AND user_id IS NULL AND ai_id IS NULL AND status = 'active' ORDER BY created_at DESC LIMIT 1 """, config.ai_type) elif config.user_id is not None and config.ai_id is None: # User-specific RAG (user's general AI) result = await conn.fetchrow(""" SELECT id, ai_type, user_id, ai_id, name, description, graph_blob_url, vector_blob_url, config_blob_url, total_chunks, total_tokens, file_count, created_at, updated_at, last_accessed_at, status FROM rag_instances WHERE ai_type = $1 AND user_id = $2 AND ai_id IS NULL AND status = 'active' ORDER BY created_at DESC LIMIT 1 """, config.ai_type, config.user_id) else: # Custom AI RAG result = await conn.fetchrow(""" SELECT id, ai_type, user_id, ai_id, name, description, graph_blob_url, vector_blob_url, config_blob_url, total_chunks, total_tokens, file_count, created_at, updated_at, last_accessed_at, status FROM rag_instances WHERE ai_type = $1 AND user_id = $2 AND ai_id = $3 AND status = 'active' ORDER BY created_at DESC LIMIT 1 """, config.ai_type, config.user_id, config.ai_id) if result: # Update last accessed time await conn.execute(""" UPDATE rag_instances SET last_accessed_at = NOW() WHERE id = $1 """, result['id']) self.logger.info(f"🎯 Database lookup SUCCESS: Found {result['name']} (ID: {result['id']})") return dict(result) self.logger.info( f"🔍 Database lookup: No RAG found for ai_type='{config.ai_type}', user_id={config.user_id}, ai_id={config.ai_id}") return None async def save_conversation_message( self, conversation_id: str, role: str, content: str, metadata: Optional[Dict[str, Any]] = None ) -> str: """Save conversation message to database""" async with self.pool.acquire() as conn: await conn.execute(""" INSERT INTO conversations (id, user_id, ai_type, ai_id, title) VALUES ($1, $2, $3, $4, $5) ON CONFLICT (id) DO NOTHING """, conversation_id, metadata.get('user_id', 'anonymous'), metadata.get('ai_type', 'unknown'), metadata.get('ai_id'), metadata.get('title', 'New Conversation') ) message_id = await conn.fetchval(""" INSERT INTO conversation_messages (conversation_id, role, content, metadata) VALUES ($1, $2, $3, $4) RETURNING id """, conversation_id, role, content, json.dumps(metadata or {})) return str(message_id) async def get_conversation_messages( self, conversation_id: str, limit: int = 50 ) -> List[Dict[str, Any]]: """Get conversation messages from database""" async with self.pool.acquire() as conn: messages = await conn.fetch(""" SELECT id, role, content, metadata, created_at FROM conversation_messages WHERE conversation_id = $1 ORDER BY created_at DESC LIMIT $2 """, conversation_id, limit) return [dict(msg) for msg in reversed(messages)] async def close(self): """Close database connections""" if self.pool: await self.pool.close() if self.redis: self.redis.close() class PersistentLightRAGManager: """ Complete LightRAG manager with Vercel-only persistence Zero dependency on HuggingFace ephemeral storage """ def __init__( self, cloudflare_worker: CloudflareWorker, database_manager: DatabaseManager, blob_client: VercelBlobClient ): self.cloudflare_worker = cloudflare_worker self.db = database_manager self.blob_client = blob_client self.rag_instances: Dict[str, LightRAG] = {} self.processing_locks: Dict[str, asyncio.Lock] = {} self.conversation_memory: Dict[str, List[Dict[str, Any]]] = {} self.logger = logging.getLogger(__name__) async def get_or_create_rag_instance(self, ai_type: str, user_id: Optional[str] = None, ai_id: Optional[str] = None, name: Optional[str] = None, description: Optional[str] = None) -> LightRAG: config = RAGConfig(ai_type=ai_type, user_id=user_id, ai_id=ai_id, name=name or f"{ai_type} AI", description=description) cache_key = config.get_cache_key() if cache_key in self.rag_instances: self.logger.info(f"Returning cached RAG instance: {cache_key}") return self.rag_instances[cache_key] if cache_key not in self.processing_locks: self.processing_locks[cache_key] = asyncio.Lock() async with self.processing_locks[cache_key]: if cache_key in self.rag_instances: return self.rag_instances[cache_key] try: self.logger.info(f"Checking for existing RAG instance: {cache_key}") instance_data = await self.db.get_rag_instance(config) if instance_data: self.logger.info( f"Found existing RAG instance: {instance_data['name']} (ID: {instance_data['id']})") async with self.db.pool.acquire() as conn: storage_check = await conn.fetchrow(""" SELECT filename, file_size, processing_status, token_count FROM knowledge_files WHERE rag_instance_id = $1 AND filename = 'lightrag_storage.json' LIMIT 1 """, instance_data['id']) if storage_check: self.logger.info( f"Found storage data: {storage_check['file_size']} bytes, {storage_check['token_count']} tokens, status: {storage_check['processing_status']}") rag_instance = await self._load_from_database(config) if rag_instance: self.rag_instances[cache_key] = rag_instance self.logger.info(f"Successfully loaded existing RAG from database: {cache_key}") return rag_instance else: self.logger.error(f"Failed to load RAG from database despite having storage data") else: self.logger.warning(f"RAG instance exists but no storage data found") else: self.logger.info(f"No existing RAG instance found in database for: {cache_key}") except Exception as e: self.logger.error(f"Error checking/loading existing RAG instance: {e}") self.logger.info(f"Creating new RAG instance: {cache_key}") rag_instance = await self._create_new_rag_instance(config) await self._save_to_database(config, rag_instance) self.rag_instances[cache_key] = rag_instance return rag_instance async def _create_new_rag_instance(self, config: RAGConfig) -> LightRAG: """Create new RAG instance with CORRECT LightRAG 1.3.7 configuration""" working_dir = f"/tmp/rag_memory_{config.get_cache_key()}_{uuid.uuid4()}" os.makedirs(working_dir, exist_ok=True) # FIXED: Use only valid LightRAG 1.3.7 parameters rag = LightRAG( working_dir=working_dir, max_parallel_insert=1, # Reduce for stability llm_model_func=self.cloudflare_worker.query, llm_model_name=self.cloudflare_worker.llm_models[0], llm_model_max_token_size=4080, embedding_func=EmbeddingFunc( embedding_dim=1024, max_token_size=2048, func=self.cloudflare_worker.embedding_chunk, ), graph_storage="NetworkXStorage", vector_storage="NanoVectorDBStorage", # REMOVED invalid parameters: # enable_entity_extraction=True, # NOT VALID IN 1.3.7 # chunk_token_size=1200, # NOT VALID IN 1.3.7 # entity_extract_max_gleaning=1, # NOT VALID IN 1.3.7 # entity_summarization_enabled=True, # NOT VALID IN 1.3.7 ) # Initialize storages await rag.initialize_storages() await initialize_pipeline_status() self.logger.info(f"✅ Initialized LightRAG 1.3.7 with working directory: {working_dir}") # Verify configuration self.logger.info(f"🔧 LightRAG Configuration:") self.logger.info(f" - Working dir: {rag.working_dir}") self.logger.info(f" - LLM model: {rag.llm_model_name}") self.logger.info(f" - Graph storage: {type(rag.graph_storage).__name__}") self.logger.info(f" - Vector storage: {type(rag.vector_storage).__name__}") # Initialize pipeline status properly if not hasattr(rag, 'pipeline_status') or rag.pipeline_status is None: rag.pipeline_status = {"history_messages": []} elif "history_messages" not in rag.pipeline_status: rag.pipeline_status["history_messages"] = [] self.logger.info(f"✅ Pipeline status initialized for {config.get_cache_key()}") # Load knowledge based on AI type if config.ai_type == "fire-safety": self.logger.info(f"🔥 Loading fire safety knowledge for {config.get_cache_key()}") success = await self._load_fire_safety_knowledge(rag) if success: # CRITICAL: Wait for entity extraction to complete self.logger.info("⏳ Waiting for entity extraction to complete...") await asyncio.sleep(10) # Give LightRAG time to process entities # Check what was actually created await self._check_storage_contents(rag) else: self.logger.warning("⚠️ Fire safety knowledge loading reported failure") return rag async def _check_storage_contents(self, rag: LightRAG): """Check what was actually stored after document insertion""" try: self.logger.info("🔍 Checking storage contents after insertion...") # Check all storage files storage_files = { 'vdb_entities.json': 'entities', 'vdb_chunks.json': 'chunks', 'vdb_relationships.json': 'relationships' } total_items = 0 for filename, storage_type in storage_files.items(): file_path = f"{rag.working_dir}/{filename}" if os.path.exists(file_path): try: file_size = os.path.getsize(file_path) with open(file_path, 'r') as f: data = json.load(f) item_count = len(data.get('data', [])) has_matrix = bool(data.get('matrix', '')) total_items += item_count if item_count > 0: self.logger.info( f"✅ {storage_type}: {item_count} items, {file_size} bytes, matrix: {has_matrix}") # Show sample for debugging if item_count > 0 and len(data['data']) > 0: sample_item = data['data'][0] if isinstance(sample_item, dict): sample_keys = list(sample_item.keys())[:5] # Show first 5 keys self.logger.info(f" Sample item keys: {sample_keys}") else: self.logger.info(f" Sample item type: {type(sample_item)}") else: self.logger.warning(f"⚠️ {storage_type}: EMPTY ({file_size} bytes)") except Exception as e: self.logger.error(f"❌ Failed to read {filename}: {e}") else: self.logger.warning(f"⚠️ {filename} doesn't exist") self.logger.info(f"📊 Total items across all storage: {total_items}") # Test if entity extraction is working by checking entities specifically if total_items > 0: await self._test_entity_extraction_quality(rag) except Exception as e: self.logger.error(f"❌ Storage content check failed: {e}") async def _test_entity_extraction_quality(self, rag: LightRAG): """Test the quality of entity extraction""" try: self.logger.info("🧪 Testing entity extraction quality...") # Check entities file specifically entities_file = f"{rag.working_dir}/vdb_entities.json" if os.path.exists(entities_file): with open(entities_file, 'r') as f: entities_data = json.load(f) entities_count = len(entities_data.get('data', [])) if entities_count > 0: self.logger.info(f"✅ Found {entities_count} entities") # Show some sample entities for i, entity in enumerate(entities_data['data'][:3]): # Show first 3 if isinstance(entity, dict): entity_name = entity.get('content', entity.get('name', str(entity))) self.logger.info(f" Entity {i + 1}: {entity_name}") return True else: self.logger.warning("⚠️ No entities found - this will break HYBRID mode") return False else: self.logger.warning("⚠️ Entities file doesn't exist") return False except Exception as e: self.logger.error(f"❌ Entity extraction test failed: {e}") return False async def debug_entity_extraction(self, rag: LightRAG): """Debug why entities aren't being extracted""" try: self.logger.info("🔍 Debugging entity extraction process...") # Check if entity extraction is working at all test_content = """ Fire safety regulations require that all commercial buildings have fire extinguishers. Emergency exits must be clearly marked with illuminated signs. Sprinkler systems are mandatory in buildings over 15,000 square feet. """ # Try manual entity extraction try: # This should trigger entity extraction await rag.ainsert(test_content) # Wait for processing await asyncio.sleep(3) # Check what was created entities_file = f"{rag.working_dir}/vdb_entities.json" relationships_file = f"{rag.working_dir}/vdb_relationships.json" for file_path in [entities_file, relationships_file]: if os.path.exists(file_path): with open(file_path, 'r') as f: data = json.load(f) filename = os.path.basename(file_path) item_count = len(data.get('data', [])) self.logger.info(f"📊 {filename}: {item_count} items") if item_count > 0: # Show sample data sample = data['data'][0] self.logger.info(f"📝 Sample {filename} item: {sample}") else: self.logger.warning(f"⚠️ {filename} is still empty after insertion") else: self.logger.warning(f"⚠️ {file_path} doesn't exist") except Exception as e: self.logger.error(f"❌ Entity extraction test failed: {e}") # Check LightRAG configuration self.logger.info(f"🔧 LightRAG config:") self.logger.info(f" - Working dir: {rag.working_dir}") self.logger.info(f" - LLM model: {getattr(rag, 'llm_model_name', 'unknown')}") self.logger.info(f" - Graph storage: {type(rag.graph_storage).__name__}") self.logger.info(f" - Vector storage: {type(rag.vector_storage).__name__}") # Check if extraction is enabled if hasattr(rag, 'enable_entity_extraction'): self.logger.info(f" - Entity extraction enabled: {rag.enable_entity_extraction}") return True except Exception as e: self.logger.error(f"❌ Debug entity extraction failed: {e}") return False async def validate_extracted_entities(self, rag: LightRAG, original_content: str) -> bool: """Validate that extracted entities actually exist in the source content""" try: entities_file = f"{rag.working_dir}/vdb_entities.json" if not os.path.exists(entities_file): return True # No entities to validate with open(entities_file, 'r') as f: entities_data = json.load(f) entities = entities_data.get('data', []) invalid_entities = [] valid_entities = [] self.logger.info(f"🔍 Validating {len(entities)} extracted entities against source content...") for entity in entities: if isinstance(entity, dict): entity_name = entity.get('entity_name', '').strip() # Skip empty or placeholder entities if not entity_name or entity_name in ['', '', 'Unknown']: invalid_entities.append(f"Empty/placeholder: '{entity_name}'") continue # Check if entity name appears in the original content if entity_name.lower() in original_content.lower(): valid_entities.append(entity_name) self.logger.info(f" ✅ Valid entity: '{entity_name}'") else: invalid_entities.append(f"Not found in content: '{entity_name}'") self.logger.warning(f" ❌ INVALID entity: '{entity_name}' - NOT FOUND in source content!") self.logger.info(f"📊 Entity validation results:") self.logger.info(f" ✅ Valid entities: {len(valid_entities)}") self.logger.info(f" ❌ Invalid entities: {len(invalid_entities)}") if invalid_entities: self.logger.error(f"🚨 ENTITY HALLUCINATION DETECTED!") for invalid in invalid_entities[:5]: # Show first 5 self.logger.error(f" {invalid}") if len(invalid_entities) > 5: self.logger.error(f" ... and {len(invalid_entities) - 5} more invalid entities") return False return True except Exception as e: self.logger.error(f"❌ Entity validation failed: {e}") return False async def clean_hallucinated_entities(self, rag: LightRAG, original_content: str): """Remove entities that don't exist in the source content""" try: entities_file = f"{rag.working_dir}/vdb_entities.json" if not os.path.exists(entities_file): return with open(entities_file, 'r') as f: entities_data = json.load(f) original_entities = entities_data.get('data', []) cleaned_entities = [] removed_count = 0 self.logger.info(f"🧹 Cleaning hallucinated entities from {len(original_entities)} total entities...") for entity in original_entities: if isinstance(entity, dict): entity_name = entity.get('entity_name', '').strip() # Remove empty/placeholder entities if not entity_name or entity_name in ['', '', 'Unknown']: removed_count += 1 continue # Remove entities not found in content if entity_name.lower() not in original_content.lower(): self.logger.warning(f" 🗑️ Removing hallucinated entity: '{entity_name}'") removed_count += 1 continue # Keep valid entities cleaned_entities.append(entity) # Update the entities file with cleaned data entities_data['data'] = cleaned_entities with open(entities_file, 'w') as f: json.dump(entities_data, f) self.logger.info(f"✅ Entity cleaning complete:") self.logger.info(f" 📊 Original entities: {len(original_entities)}") self.logger.info(f" 🗑️ Removed: {removed_count}") self.logger.info(f" ✅ Remaining: {len(cleaned_entities)}") except Exception as e: self.logger.error(f"❌ Entity cleaning failed: {e}") async def _load_fire_safety_knowledge(self, rag: LightRAG): """Load fire safety knowledge with FIXED insertion process""" self.logger.info(f"🔥 Loading fire safety knowledge for {rag.working_dir}") # Prepare knowledge content base_knowledge = """ FIRE SAFETY REGULATIONS AND BUILDING CODES 1. Emergency Exit Requirements: - All buildings must have at least two exits on each floor - Maximum travel distance to exit: 75 feet in unsprinklered buildings, 100 feet in sprinklered buildings - Exit doors must swing in direction of egress travel - All exits must be clearly marked with illuminated exit signs - Exit routes must be free of obstructions at all times - Minimum width for exits: 32 inches for single doors, 64 inches for double doors 2. Fire Extinguisher Requirements: - Type A: For ordinary combustible materials (wood, paper, cloth, rubber, plastic) - Type B: For flammable and combustible liquids (gasoline, oil, paint, grease) - Type C: For energized electrical equipment (motors, generators, switches) - Type D: For combustible metals (magnesium, titanium, zirconium, lithium) - Type K: For cooking oils and fats in commercial kitchen equipment - Distribution: Maximum travel distance of 75 feet to nearest extinguisher - Inspection: Monthly visual inspections and annual professional service 3. Fire Detection and Alarm Systems: - Smoke detectors required in all sleeping areas and hallways - Heat detectors required in areas where smoke detectors unsuitable - Manual fire alarm pull stations required near all exits - Central monitoring systems required in commercial buildings over 10,000 sq ft - Backup power systems required for all alarm components - Testing schedule: Monthly for batteries, annually for full system 4. Sprinkler System Requirements: - Required in all buildings over 3 stories or 15,000 sq ft - Wet pipe systems: Most common, water-filled pipes - Dry pipe systems: For areas subject to freezing temperatures - Deluge systems: For high-hazard areas with rapid fire spread potential - Inspection: Quarterly for valves, annually for full system testing """ all_content = [base_knowledge] # Load additional files if they exist book_files = ['/app/book.pdf', '/app/book.txt'] for file_path in book_files: if os.path.exists(file_path): try: if file_path.endswith('.pdf'): try: import PyPDF2 with open(file_path, 'rb') as file: pdf_reader = PyPDF2.PdfReader(file) for page_num in range(min(20, len(pdf_reader.pages))): # Reduced from 50 to 20 page_text = pdf_reader.pages[page_num].extract_text() if page_text and len(page_text.strip()) > 100: all_content.append( f"PDF Page {page_num + 1}: {page_text[:3000]}") # Reduced chunk size except Exception as e: self.logger.warning(f"PDF processing failed: {e}") continue else: with open(file_path, 'r', encoding='utf-8', errors='ignore') as file: txt_content = file.read() # Split into smaller chunks for i in range(0, min(len(txt_content), 60000), 3000): # Reduced chunk size and total chunk = txt_content[i:i + 3000] if chunk.strip(): all_content.append(f"TXT Section {i // 3000 + 1}: {chunk}") self.logger.info(f"✅ Loaded {file_path}") except Exception as e: self.logger.error(f"❌ Failed to load {file_path}: {e}") self.logger.info(f"📚 Starting insertion of {len(all_content)} documents") # CRITICAL: Insert documents with better error handling and timeout successful_insertions = 0 for i, content in enumerate(all_content): try: self.logger.info(f"📝 Inserting document {i + 1}/{len(all_content)} ({len(content)} chars)") entities_before = await self._count_entities(rag) # Insert with timeout insertion_task = asyncio.create_task(rag.ainsert(content)) try: await asyncio.wait_for(insertion_task, timeout=45.0) # 30 second timeout per document successful_insertions += 1 self.logger.info(f"✅ Document {i + 1} inserted successfully") # Brief pause between insertions await asyncio.sleep(2) entities_after = await self._count_entities(rag) entities_added = entities_after - entities_before self.logger.info( f"✅ Document {i + 1} inserted - Entities added: {entities_added} (total: {entities_after})") await asyncio.sleep(1) except asyncio.TimeoutError: self.logger.error(f"⏰ Document {i + 1} insertion timed out after 30 seconds") insertion_task.cancel() continue except Exception as e: self.logger.error(f"❌ Failed to insert document {i + 1}: {e}") continue self.logger.info(f"📊 Insertion complete: {successful_insertions}/{len(all_content)} documents successful") # Force storage verification with timeout if successful_insertions > 0: self.logger.info("🔍 Final validation and cleaning...") await asyncio.sleep(5) # Wait for processing is_valid = await self.validate_extracted_entities(rag, all_content) if not is_valid: self.logger.warning("🧹 Cleaning hallucinated entities...") await self.clean_hallucinated_entities(rag, all_content) # Final verification final_entities = await self._count_entities(rag) final_relationships = await self._count_relationships(rag) final_chunks = await self._count_chunks(rag) self.logger.info( f"📊 Final counts after cleaning: {final_chunks} chunks, {final_entities} entities, {final_relationships} relationships") if final_entities > 0: self.logger.info("🎉 Entity extraction SUCCESS - HYBRID mode should work!") else: self.logger.warning("⚠️ No entities extracted - HYBRID mode will fail") # Show cleaned entities try: # Check storage files storage_verified = False for storage_file in ['vdb_chunks.json', 'vdb_entities.json', 'vdb_relationships.json']: file_path = f"{rag.working_dir}/{storage_file}" if os.path.exists(file_path) and os.path.getsize(file_path) > 100: with open(file_path, 'r') as f: data = json.load(f) if data.get('data') and len(data['data']) > 0: storage_verified = True self.logger.info(f"✅ {storage_file}: {len(data['data'])} items") if storage_verified: self.logger.info("🎉 Storage verification PASSED") else: self.logger.error("❌ Storage verification FAILED") except Exception as e: self.logger.error(f"❌ Storage verification error: {e}") self.logger.info("🔍 Starting entity extraction debugging...") await self.debug_entity_extraction(rag) return successful_insertions > 0 async def _count_entities(self, rag: LightRAG) -> int: """Count entities in storage""" try: entities_file = f"{rag.working_dir}/vdb_entities.json" if os.path.exists(entities_file): with open(entities_file, 'r') as f: data = json.load(f) return len(data.get('data', [])) return 0 except: return 0 async def _count_relationships(self, rag: LightRAG) -> int: """Count relationships in storage""" try: relationships_file = f"{rag.working_dir}/vdb_relationships.json" if os.path.exists(relationships_file): with open(relationships_file, 'r') as f: data = json.load(f) return len(data.get('data', [])) return 0 except: return 0 async def _count_chunks(self, rag: LightRAG) -> int: """Count chunks in storage""" try: chunks_file = f"{rag.working_dir}/vdb_chunks.json" if os.path.exists(chunks_file): with open(chunks_file, 'r') as f: data = json.load(f) return len(data.get('data', [])) return 0 except: return 0 # Add this method to your PersistentLightRAGManager class async def fix_entity_extraction_for_custom_ai(self, rag: LightRAG, content_list: List[str]): """Fix entity extraction issues for custom AI""" try: self.logger.info("🔧 Starting entity extraction fix...") # Clear existing corrupted data for storage_file in ['vdb_entities.json', 'vdb_chunks.json', 'vdb_relationships.json']: file_path = f"{rag.working_dir}/{storage_file}" if os.path.exists(file_path): # Backup the file backup_path = f"{file_path}.backup" shutil.copy2(file_path, backup_path) self.logger.info(f"📋 Backed up {storage_file}") # Try different LLM models for entity extraction successful_extractions = 0 for i, content in enumerate(content_list): if len(content.strip()) < 50: # Skip very short content continue try: self.logger.info(f"🔍 Processing content chunk {i + 1}/{len(content_list)}") # Use a more explicit prompt for entity extraction enhanced_prompt = f""" You are an expert at extracting entities and relationships from text. Extract entities and relationships from this text and return ONLY the extracted information in the exact format requested. Text to analyze: {content[:2000]} Requirements: 1. Extract ALL important entities (people, organizations, locations, concepts, objects) 2. For each entity provide: name, type (person/organization/geo/event/category), description 3. Extract relationships between entities with descriptions and strength (1-10) 4. Identify high-level keywords that summarize main concepts 5. Ignore all entities and do not add if they are: Alex, Taylor, Jordan, Cruz, The Device. 6. Make sure to double check if this entity is actually in the text or are you just hallucinating it. Return the results in this exact format: entity<|><|> relationship<|><|><|><|> content_keywords Use ## to separate each item. """ # Use timeout for entity extraction extraction_task = asyncio.create_task( rag.ainsert(content[:3000]) # Limit content size ) try: await asyncio.wait_for(extraction_task, timeout=45.0) successful_extractions += 1 self.logger.info(f"✅ Successfully processed chunk {i + 1}") # Add delay between insertions await asyncio.sleep(2) except asyncio.TimeoutError: self.logger.warning(f"⏰ Chunk {i + 1} timed out") extraction_task.cancel() continue except Exception as e: self.logger.error(f"❌ Failed to process chunk {i + 1}: {e}") continue # Verify the results await asyncio.sleep(5) # Wait for processing entities_count = await self._count_entities(rag) chunks_count = await self._count_chunks(rag) relationships_count = await self._count_relationships(rag) self.logger.info( f"📊 Extraction results: {chunks_count} chunks, {entities_count} entities, {relationships_count} relationships") if entities_count > 0: self.logger.info("🎉 Entity extraction fix SUCCESS!") return True else: self.logger.error("❌ Entity extraction fix FAILED - no entities found") return False except Exception as e: self.logger.error(f"❌ Entity extraction fix failed: {e}") return False # Also add this helper method async def force_rebuild_custom_ai(self, ai_id: str, user_id: str): """Force rebuild a custom AI from scratch""" try: self.logger.info(f"🔧 Force rebuilding custom AI: {ai_id}") # Get the AI details and uploaded files async with self.db.pool.acquire() as conn: # Get custom AI info ai_info = await conn.fetchrow(""" SELECT * FROM rag_instances WHERE ai_id = $1 AND user_id = $2 AND ai_type = 'custom' """, ai_id, user_id) if not ai_info: self.logger.error(f"❌ Custom AI not found: {ai_id}") return False # Get uploaded files files = await conn.fetch(""" SELECT content_text, original_name FROM knowledge_files WHERE rag_instance_id = $1 AND filename != 'lightrag_storage.json' AND processing_status = 'processed' """, ai_info['id']) if not files: self.logger.error(f"❌ No files found for custom AI: {ai_id}") return False # Create new RAG config config = RAGConfig( ai_type="custom", user_id=user_id, ai_id=ai_id, name=ai_info['name'], description=ai_info['description'] ) # Create fresh RAG instance rag = await self._create_new_rag_instance(config) # Re-process all files with fixed entity extraction content_list = [] for file_record in files: if file_record['content_text']: content_list.append(file_record['content_text']) # Use the fixed entity extraction success = await self.fix_entity_extraction_for_custom_ai(rag, content_list) if success: # Save the rebuilt RAG await self._save_to_database(config, rag) # Clear cache cache_key = config.get_cache_key() if cache_key in self.rag_instances: del self.rag_instances[cache_key] self.logger.info(f"✅ Successfully rebuilt custom AI: {ai_id}") return True else: self.logger.error(f"❌ Failed to rebuild custom AI: {ai_id}") return False except Exception as e: self.logger.error(f"❌ Force rebuild failed: {e}") return False async def _force_storage_to_database(self, rag: LightRAG, rag_instance_id: str): try: entities_file = f"{rag.working_dir}/vdb_entities.json" chunks_file = f"{rag.working_dir}/vdb_chunks.json" relationships_file = f"{rag.working_dir}/vdb_relationships.json" storage_data = {} total_items = 0 storage_files = { 'vdb_entities': entities_file, 'vdb_chunks': chunks_file, 'vdb_relationships': relationships_file } for storage_key, file_path in storage_files.items(): if os.path.exists(file_path): try: with open(file_path, 'r') as f: file_data = json.load(f) if isinstance(file_data, dict) and 'data' in file_data: item_count = len(file_data.get('data', [])) total_items += item_count storage_data[storage_key] = file_data self.logger.info(f"✅ Read {storage_key}: {item_count} items") else: self.logger.warning(f"⚠️ Invalid format in {file_path}") storage_data[storage_key] = {"data": [], "matrix": ""} except Exception as e: self.logger.error(f"❌ Failed to read {file_path}: {e}") storage_data[storage_key] = {"data": [], "matrix": ""} else: self.logger.warning(f"⚠️ Storage file not found: {file_path}") storage_data[storage_key] = {"data": [], "matrix": ""} # Only proceed if we have some data if total_items > 0 and storage_data: try: async with self.db.pool.acquire() as conn: # Check if rag_instance exists instance_exists = await conn.fetchval(""" SELECT COUNT(*) FROM rag_instances WHERE id = $1::uuid """, rag_instance_id) if not instance_exists: self.logger.error(f"❌ RAG instance {rag_instance_id} does not exist") return False # Insert the knowledge file record with complete storage data await conn.execute(""" INSERT INTO knowledge_files ( id, user_id, rag_instance_id, filename, original_name, file_type, file_size, blob_url, content_text, processing_status, token_count, created_at, updated_at ) VALUES ( gen_random_uuid(), 'system', $1::uuid, 'lightrag_storage.json', 'LightRAG Storage Data', 'json', $2, 'database://storage', $3, 'processed', $4, NOW(), NOW() ) ON CONFLICT (rag_instance_id, filename) DO UPDATE SET content_text = EXCLUDED.content_text, file_size = EXCLUDED.file_size, token_count = EXCLUDED.token_count, updated_at = NOW() """, rag_instance_id, len(json.dumps(storage_data)), json.dumps(storage_data), total_items) self.logger.info( f"✅ Stored LightRAG data for instance {rag_instance_id}: {total_items} total items") # Log detailed breakdown for key, data in storage_data.items(): item_count = len(data.get('data', [])) self.logger.info(f" - {key}: {item_count} items") return True except Exception as e: self.logger.error(f"❌ Database storage failed: {e}") import traceback self.logger.error(f"Full traceback: {traceback.format_exc()}") return False else: if not storage_data: self.logger.warning("⚠️ No storage data to save") else: self.logger.warning(f"⚠️ No items found in storage data (total: {total_items})") return False except Exception as e: self.logger.error(f"❌ Failed to store to database: {e}") import traceback self.logger.error(f"Full traceback: {traceback.format_exc()}") return False async def _wait_for_pipeline_completion(self, rag: LightRAG, doc_name: str, max_wait_time: int = 30): """Wait for LightRAG 1.3.7 pipeline to complete processing""" for attempt in range(max_wait_time): try: await asyncio.sleep(1) if hasattr(rag, 'doc_status') and rag.doc_status: status_data = await rag.doc_status.get_all() if status_data: completed_docs = [doc for doc in status_data if 'completed' in str(doc).lower()] if completed_docs: self.logger.info(f"Pipeline processing detected for {doc_name}") return True if hasattr(rag.vector_storage, '_data') and rag.vector_storage._data: data_count = len(rag.vector_storage._data) if data_count > 0: self.logger.info(f"Vector storage contains {data_count} items after {doc_name}") return True if hasattr(rag, 'chunks') and rag.chunks: chunks_data = await rag.chunks.get_all() if chunks_data and len(chunks_data) > 0: self.logger.info(f"Chunks storage contains {len(chunks_data)} items after {doc_name}") return True except Exception as e: self.logger.debug(f"Pipeline check attempt {attempt + 1} failed: {e}") continue self.logger.warning(f"Pipeline completion check timed out for {doc_name}") return False async def _verify_knowledge_base_state(self, rag: LightRAG): """Verify the final state of the knowledge base""" try: storage_stats = {} if hasattr(rag.vector_storage, '_data'): storage_stats['vector_items'] = len(rag.vector_storage._data) if rag.vector_storage._data else 0 if hasattr(rag, 'chunks') and rag.chunks: try: chunks_data = await rag.chunks.get_all() storage_stats['chunks'] = len(chunks_data) if chunks_data else 0 except: storage_stats['chunks'] = 0 if hasattr(rag, 'entities') and rag.entities: try: entities_data = await rag.entities.get_all() storage_stats['entities'] = len(entities_data) if entities_data else 0 except: storage_stats['entities'] = 0 if hasattr(rag, 'relationships') and rag.relationships: try: relationships_data = await rag.relationships.get_all() storage_stats['relationships'] = len(relationships_data) if relationships_data else 0 except: storage_stats['relationships'] = 0 self.logger.info(f"Knowledge base state: {storage_stats}") return any(count > 0 for count in storage_stats.values()) except Exception as e: self.logger.error(f"Failed to verify knowledge base state: {e}") return False def _intelligent_chunk_split(self, content: str, max_chunk_size: int = 8000) -> List[str]: """Split content intelligently on sentence and paragraph boundaries""" if len(content) <= max_chunk_size: return [content] chunks = [] current_chunk = "" paragraphs = content.split('\n\n') for paragraph in paragraphs: if len(paragraph) > max_chunk_size: sentences = paragraph.split('. ') for sentence in sentences: if len(current_chunk) + len(sentence) + 2 <= max_chunk_size: current_chunk += sentence + '. ' else: if current_chunk: chunks.append(current_chunk.strip()) current_chunk = sentence + '. ' else: if len(current_chunk) + len(paragraph) + 2 <= max_chunk_size: current_chunk += paragraph + '\n\n' else: if current_chunk: chunks.append(current_chunk.strip()) current_chunk = paragraph + '\n\n' if current_chunk: chunks.append(current_chunk.strip()) return chunks def _split_content_into_chunks(self, content: str, max_length: int) -> List[str]: """Split content into manageable chunks""" chunks = [] words = content.split() current_chunk = [] current_length = 0 for word in words: if current_length + len(word) + 1 <= max_length: current_chunk.append(word) current_length += len(word) + 1 else: if current_chunk: chunks.append(' '.join(current_chunk)) current_chunk = [word] current_length = len(word) if current_chunk: chunks.append(' '.join(current_chunk)) return chunks async def _save_to_database(self, config: RAGConfig, rag: LightRAG): """Save RAG instance to Database with CORRECT order of operations""" try: self.logger.info("💾 Starting database save process...") # STEP 1: Calculate metadata from actual storage files metadata = await self._calculate_storage_metadata(rag) # STEP 2: Create empty blob URLs (we're using database storage) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") base_filename = f"rag_{config.ai_type}_{config.user_id or 'system'}_{config.ai_id or 'default'}_{timestamp}" fake_blob_urls = { 'graph_blob_url': f"database://graph_{base_filename}", 'vector_blob_url': f"database://vector_{base_filename}", 'config_blob_url': f"database://config_{base_filename}" } # STEP 3: CREATE the rag_instance record FIRST rag_instance_id = await self.db.save_rag_instance( config, fake_blob_urls['graph_blob_url'], fake_blob_urls['vector_blob_url'], fake_blob_urls['config_blob_url'], metadata ) self.logger.info(f"✅ Created RAG instance in database: {rag_instance_id}") # STEP 4: NOW store the RAG data (with existing rag_instance_id) storage_success = await self._force_storage_to_database(rag, str(rag_instance_id)) if storage_success: self.logger.info(f"✅ Successfully saved complete RAG to database: {rag_instance_id}") else: self.logger.warning(f"⚠️ RAG instance created but storage data save failed: {rag_instance_id}") except Exception as e: self.logger.error(f"❌ Failed to save RAG to database: {e}") import traceback self.logger.error(f"Full traceback: {traceback.format_exc()}") raise async def _calculate_storage_metadata(self, rag: LightRAG) -> Dict[str, Any]: """Calculate metadata from RAG storage""" try: total_chunks = 0 total_tokens = 0 file_count = 0 # Check storage files for storage_file in ['vdb_chunks.json', 'vdb_entities.json', 'vdb_relationships.json']: file_path = f"{rag.working_dir}/{storage_file}" if os.path.exists(file_path): try: with open(file_path, 'r') as f: data = json.load(f) if data.get('data'): chunk_count = len(data['data']) total_chunks += chunk_count # Estimate tokens (rough calculation) for item in data['data']: if isinstance(item, dict) and 'content' in item: # Rough token estimation: ~4 chars per token total_tokens += len(str(item['content'])) // 4 file_count += 1 except Exception as e: self.logger.warning(f"Failed to read {storage_file}: {e}") return { 'total_chunks': total_chunks, 'total_tokens': max(total_tokens, 100), # Minimum 100 tokens 'file_count': file_count } except Exception as e: self.logger.error(f"Failed to calculate metadata: {e}") return { 'total_chunks': 0, 'total_tokens': 100, 'file_count': 0 } async def _load_from_database(self, config: RAGConfig) -> Optional[LightRAG]: """Load RAG from database with PROPER NanoVectorDB restoration""" try: # Get RAG instance metadata instance_data = await self.db.get_rag_instance(config) if not instance_data: self.logger.info(f"No RAG instance found in database for {config.get_cache_key()}") return None self.logger.info(f"🔍 Found RAG instance: {instance_data['name']} (ID: {instance_data['id']})") # Check if we have storage data in knowledge_files table async with self.db.pool.acquire() as conn: storage_record = await conn.fetchrow(""" SELECT content_text, file_size, token_count FROM knowledge_files WHERE rag_instance_id = $1 AND filename = 'lightrag_storage.json' AND processing_status = 'processed' ORDER BY created_at DESC LIMIT 1 """, instance_data['id']) if not storage_record or not storage_record['content_text']: self.logger.warning(f"⚠️ No storage data found in database for RAG {instance_data['id']}") return None self.logger.info( f"🎯 Found database storage: {storage_record['file_size']} bytes, {storage_record['token_count']} tokens") try: # Parse the JSON storage data storage_data = json.loads(storage_record['content_text']) self.logger.info(f"📊 Parsed storage data with keys: {list(storage_data.keys())}") # CRITICAL: Check if we have chunks with actual content chunks_data = storage_data.get('vdb_chunks', {}) if not chunks_data.get('data') or len(chunks_data['data']) == 0: self.logger.warning("❌ No chunk data found in storage") return None chunk_count = len(chunks_data['data']) self.logger.info(f"📦 Found {chunk_count} chunks in storage") # Create working directory working_dir = f"/tmp/rag_restored_{uuid.uuid4()}" os.makedirs(working_dir, exist_ok=True) # Write storage files with PROPER NanoVectorDB format for filename, file_data in storage_data.items(): try: file_path = f"{working_dir}/{filename}.json" # CRITICAL: Ensure proper NanoVectorDB format if isinstance(file_data, dict) and 'data' in file_data: # NanoVectorDB expects the EXACT format from your data with open(file_path, 'w') as f: json.dump(file_data, f) file_size = os.path.getsize(file_path) self.logger.info(f"✅ Wrote {filename}.json: {file_size} bytes") else: self.logger.warning(f"⚠️ Skipping {filename}: invalid format") except Exception as e: self.logger.error(f"❌ Failed to write {filename}: {e}") # Create LightRAG instance self.logger.info("🚀 Creating LightRAG instance with restored files") rag = LightRAG( working_dir=working_dir, max_parallel_insert=2, llm_model_func=self.cloudflare_worker.query, llm_model_name=self.cloudflare_worker.llm_models[0], llm_model_max_token_size=4080, embedding_func=EmbeddingFunc( embedding_dim=1024, max_token_size=2048, func=self.cloudflare_worker.embedding_chunk, ), graph_storage="NetworkXStorage", vector_storage="NanoVectorDBStorage", ) # Initialize storages await rag.initialize_storages() self.logger.info("🔄 Initialized storages") # Initialize pipeline status if not hasattr(rag, 'pipeline_status') or rag.pipeline_status is None: rag.pipeline_status = {"history_messages": []} elif "history_messages" not in rag.pipeline_status: rag.pipeline_status["history_messages"] = [] # CRITICAL: Test with multiple query types self.logger.info("🧪 Testing RAG with comprehensive queries...") # Test 1: Simple fire safety query try: from lightrag import QueryParam test_response = await rag.aquery( "What are fire exit requirements?", QueryParam(mode="hybrid") ) if test_response and len(test_response.strip()) > 50 and not test_response.startswith("Sorry"): self.logger.info(f"🎉 SUCCESS: Hybrid query test passed - {len(test_response)} chars") return rag else: self.logger.warning(f"⚠️ Hybrid query failed: '{test_response[:100]}'") except Exception as e: self.logger.error(f"❌ Hybrid query test failed: {e}") # Test 2: Try local mode try: local_response = await rag.aquery("fire safety", QueryParam(mode="local")) if local_response and len(local_response.strip()) > 20 and not local_response.startswith("Sorry"): self.logger.info(f"✅ LOCAL query worked: {local_response[:100]}...") return rag except Exception as e: self.logger.error(f"❌ Local query failed: {e}") # Test 3: Try naive mode try: naive_response = await rag.aquery("fire", QueryParam(mode="naive")) if naive_response and len(naive_response.strip()) > 10 and not naive_response.startswith("Sorry"): self.logger.info(f"✅ NAIVE query worked: {naive_response[:100]}...") return rag except Exception as e: self.logger.error(f"❌ Naive query failed: {e}") # If all queries fail, return None self.logger.error("❌ ALL query tests failed - RAG is not functional") return None except json.JSONDecodeError as e: self.logger.error(f"❌ Failed to parse JSON storage data: {e}") return None except Exception as e: self.logger.error(f"❌ Failed to restore from database storage: {e}") import traceback self.logger.error(f"Full traceback: {traceback.format_exc()}") return None except Exception as e: self.logger.error(f"❌ Database loading failed: {e}") return None async def _verify_rag_storage(self, rag: LightRAG) -> bool: """Verify that RAG storage has been properly loaded with actual data""" try: # Check vector storage vector_count = 0 if hasattr(rag.vector_storage, '_data') and rag.vector_storage._data: vector_count = len(rag.vector_storage._data) # Check chunks storage chunks_count = 0 if hasattr(rag, 'chunks') and rag.chunks: try: chunks_data = await rag.chunks.get_all() chunks_count = len(chunks_data) if chunks_data else 0 except: pass # Check entities storage entities_count = 0 if hasattr(rag, 'entities') and rag.entities: try: entities_data = await rag.entities.get_all() entities_count = len(entities_data) if entities_data else 0 except: pass # Check relationships storage relationships_count = 0 if hasattr(rag, 'relationships') and rag.relationships: try: relationships_data = await rag.relationships.get_all() relationships_count = len(relationships_data) if relationships_data else 0 except: pass self.logger.info( f"📋 RAG storage verification: vectors={vector_count}, chunks={chunks_count}, entities={entities_count}, relationships={relationships_count}") # Consider RAG loaded if ANY storage has data has_data = vector_count > 0 or chunks_count > 0 or entities_count > 0 or relationships_count > 0 if has_data: self.logger.info("✅ RAG verification PASSED - has working data") else: self.logger.warning("❌ RAG verification FAILED - no data found") return has_data except Exception as e: self.logger.error(f"Failed to verify RAG storage: {e}") return False async def _serialize_rag_state(self, rag: LightRAG) -> Dict[str, Any]: """Serialize RAG state for storage in Vercel Blob + Database""" try: rag_state = { 'graph': {}, 'vectors': {}, 'config': {} } # Serialize graph storage (NetworkX) if hasattr(rag, 'graph_storage') and rag.graph_storage: try: # Get the NetworkX graph data if hasattr(rag.graph_storage, '_graph'): import networkx as nx graph_data = nx.node_link_data(rag.graph_storage._graph) rag_state['graph'] = graph_data self.logger.info( f"📊 Serialized graph: {len(graph_data.get('nodes', []))} nodes, {len(graph_data.get('links', []))} edges") else: rag_state['graph'] = {} except Exception as e: self.logger.warning(f"Failed to serialize graph storage: {e}") rag_state['graph'] = {} # Serialize vector storage (NanoVectorDB) if hasattr(rag, 'vector_storage') and rag.vector_storage: try: vectors_data = { 'embeddings': [], 'metadata': [], 'config': { 'embedding_dim': getattr(rag.vector_storage, 'embedding_dim', 1024), 'metric': getattr(rag.vector_storage, 'metric', 'cosine') } } # Get vector data if hasattr(rag.vector_storage, '_data') and rag.vector_storage._data: vectors_data['embeddings'] = rag.vector_storage._data.tolist() if hasattr( rag.vector_storage._data, 'tolist') else list(rag.vector_storage._data) if hasattr(rag.vector_storage, '_metadata') and rag.vector_storage._metadata: vectors_data['metadata'] = rag.vector_storage._metadata rag_state['vectors'] = vectors_data self.logger.info(f"📊 Serialized vectors: {len(vectors_data['embeddings'])} embeddings") except Exception as e: self.logger.warning(f"Failed to serialize vector storage: {e}") rag_state['vectors'] = {'embeddings': [], 'metadata': [], 'config': {}} # Serialize configuration and metadata rag_state['config'] = { 'working_dir': rag.working_dir, 'llm_model_name': getattr(rag, 'llm_model_name', ''), 'llm_model_max_token_size': getattr(rag, 'llm_model_max_token_size', 4080), 'graph_storage_type': 'NetworkXStorage', 'vector_storage_type': 'NanoVectorDBStorage', 'embedding_dim': 1024, 'created_at': datetime.now().isoformat() } # Add pipeline status if available if hasattr(rag, 'pipeline_status') and rag.pipeline_status: rag_state['config']['pipeline_status'] = rag.pipeline_status self.logger.info(f"✅ Successfully serialized RAG state") return rag_state except Exception as e: self.logger.error(f"Failed to serialize RAG state: {e}") # Return minimal state to avoid complete failure return { 'graph': {}, 'vectors': {'embeddings': [], 'metadata': [], 'config': {}}, 'config': { 'working_dir': getattr(rag, 'working_dir', '/tmp/unknown'), 'created_at': datetime.now().isoformat() } } async def _deserialize_rag_state(self, rag_state: Dict[str, Any], working_dir: str) -> LightRAG: """Deserialize RAG state from Vercel Blob storage""" try: # Create new RAG instance rag = LightRAG( working_dir=working_dir, max_parallel_insert=2, llm_model_func=self.cloudflare_worker.query, llm_model_name=self.cloudflare_worker.llm_models[0], llm_model_max_token_size=4080, embedding_func=EmbeddingFunc( embedding_dim=1024, max_token_size=2048, func=self.cloudflare_worker.embedding_chunk, ), graph_storage="NetworkXStorage", vector_storage="NanoVectorDBStorage", ) # Initialize storages await rag.initialize_storages() # Restore graph data if rag_state.get('graph') and hasattr(rag, 'graph_storage'): try: import networkx as nx graph_data = rag_state['graph'] if graph_data and 'nodes' in graph_data: restored_graph = nx.node_link_graph(graph_data) rag.graph_storage._graph = restored_graph self.logger.info(f"🔄 Restored graph: {len(graph_data.get('nodes', []))} nodes") except Exception as e: self.logger.warning(f"Failed to restore graph: {e}") # Restore vector data if rag_state.get('vectors') and hasattr(rag, 'vector_storage'): try: vectors_data = rag_state['vectors'] if vectors_data.get('embeddings'): embeddings = np.array(vectors_data['embeddings']) rag.vector_storage._data = embeddings if vectors_data.get('metadata'): rag.vector_storage._metadata = vectors_data['metadata'] self.logger.info(f"🔄 Restored vectors: {len(vectors_data.get('embeddings', []))} embeddings") except Exception as e: self.logger.warning(f"Failed to restore vectors: {e}") # Restore configuration if rag_state.get('config'): config = rag_state['config'] if config.get('pipeline_status'): rag.pipeline_status = config['pipeline_status'] # Ensure pipeline status is initialized if not hasattr(rag, 'pipeline_status') or rag.pipeline_status is None: rag.pipeline_status = {"history_messages": []} self.logger.info("✅ Successfully deserialized RAG state") return rag except Exception as e: self.logger.error(f"Failed to deserialize RAG state: {e}") raise async def _estimate_tokens(self, rag_state: Dict[str, Any]) -> int: """Estimate token count from RAG state""" try: token_count = 0 # Count tokens from vector embeddings if rag_state.get('vectors', {}).get('embeddings'): embeddings = rag_state['vectors']['embeddings'] token_count += len(embeddings) * 10 # Rough estimate: 10 tokens per embedding # Count tokens from graph nodes if rag_state.get('graph', {}).get('nodes'): nodes = rag_state['graph']['nodes'] token_count += len(nodes) * 5 # Rough estimate: 5 tokens per node # Count tokens from graph edges if rag_state.get('graph', {}).get('links'): links = rag_state['graph']['links'] token_count += len(links) * 3 # Rough estimate: 3 tokens per edge return max(token_count, 100) # Minimum 100 tokens except Exception as e: self.logger.warning(f"Failed to estimate tokens: {e}") return 100 async def query_with_memory( self, ai_type: str, question: str, conversation_id: str, user_id: str, ai_id: Optional[str] = None, mode: str = "hybrid" ) -> str: """Query RAG with conversation memory""" try: # Get or create RAG instance rag_instance = await self.get_or_create_rag_instance( ai_type=ai_type, user_id=user_id if ai_type == "custom" else None, ai_id=ai_id, name=f"{ai_type.title()} AI", description=f"AI assistant for {ai_type}" ) # Save user message to database await self.db.save_conversation_message( conversation_id, "user", question, { "user_id": user_id, "ai_type": ai_type, "ai_id": ai_id } ) # Query RAG with LightRAG QueryParam from lightrag import QueryParam response = await rag_instance.aquery(question, QueryParam(mode=mode)) # Save assistant response to database await self.db.save_conversation_message( conversation_id, "assistant", response, { "mode": mode, "ai_type": ai_type, "ai_id": ai_id, "user_id": user_id } ) return response except Exception as e: self.logger.error(f"Query with memory failed: {e}") # Fallback to direct Cloudflare query fallback_response = await self.cloudflare_worker.query( question, f"You are a helpful {ai_type} AI assistant." ) # Save fallback response await self.db.save_conversation_message( conversation_id, "assistant", fallback_response, { "mode": "fallback", "ai_type": ai_type, "user_id": user_id, "error": str(e) } ) return fallback_response async def _load_from_blob_storage(self, instance_data: Dict[str, Any]) -> Optional[LightRAG]: """Load RAG from Vercel Blob storage (fallback method)""" try: self.logger.info("🔄 Loading RAG from Vercel Blob storage") # Download RAG state from Vercel Blob self.logger.info("📥 Downloading RAG state from Vercel Blob...") graph_data = await self.blob_client.get(instance_data['graph_blob_url']) vector_data = await self.blob_client.get(instance_data['vector_blob_url']) config_data = await self.blob_client.get(instance_data['config_blob_url']) # Decompress and deserialize graph_state = pickle.loads(gzip.decompress(graph_data)) vector_state = pickle.loads(gzip.decompress(vector_data)) config_state = pickle.loads(gzip.decompress(config_data)) rag_state = { 'graph': graph_state, 'vectors': vector_state, 'config': config_state } self.logger.info("✅ Successfully downloaded and deserialized RAG state") # Create working directory working_dir = f"/tmp/rag_restored_{uuid.uuid4()}" os.makedirs(working_dir, exist_ok=True) # Deserialize RAG instance rag = await self._deserialize_rag_state(rag_state, working_dir) return rag except Exception as e: self.logger.error(f"❌ Failed to load RAG from Vercel Blob: {e}") return None async def test_model_entity_extraction(self): """Test different models to see which extracts entities best""" test_content = """ Fire extinguishers are required in commercial buildings. Type A fire extinguishers are used for ordinary combustible materials like wood and paper. Emergency exits must be clearly marked with illuminated exit signs. Sprinkler systems are mandatory in buildings over 15,000 square feet. Building codes require fire-resistant construction materials. """ results = {} for i, model in enumerate(self.llm_models[:5]): # Test top 5 models try: self.logger.info(f"🧪 Testing entity extraction with {model}") # Temporarily switch to this model original_index = self.current_llm_index self.current_llm_index = i # Test entity extraction response = await self.query( f"Extract all important technical entities, concepts, and objects from this text. List each entity with a brief description:\n\n{test_content}", "You are an expert at identifying technical entities and concepts in specialized documents." ) # Count how many entities it found (rough estimate) entity_count = response.count('\n') if response else 0 results[model] = { "response_length": len(response) if response else 0, "estimated_entities": entity_count, "response_preview": response[:200] if response else "No response" } self.logger.info( f" 📊 {model}: {entity_count} estimated entities, {len(response) if response else 0} chars") # Restore original index self.current_llm_index = original_index except Exception as e: results[model] = {"error": str(e)} self.logger.error(f" ❌ {model} failed: {e}") # Find the best model best_model = None best_score = 0 for model, result in results.items(): if "error" not in result: score = result.get("estimated_entities", 0) + (result.get("response_length", 0) // 100) if score > best_score: best_score = score best_model = model if best_model: self.logger.info(f"🏆 Best model for entity extraction: {best_model}") # Switch to the best model self.current_llm_index = self.llm_models.index(best_model) return results # Global instance lightrag_manager: Optional[PersistentLightRAGManager] = None # Replace the initialize_lightrag_manager function with correct logger usage async def initialize_lightrag_manager() -> PersistentLightRAGManager: """Initialize with OPTIMIZED models for entity extraction""" global lightrag_manager if lightrag_manager is None: # Get logger for this function func_logger = logging.getLogger(__name__) # Validate environment validate_environment() # Get environment variables cloudflare_api_key = os.getenv("CLOUDFLARE_API_KEY") cloudflare_account_id = os.getenv("CLOUDFLARE_ACCOUNT_ID") database_url = os.getenv("DATABASE_URL") redis_url = os.getenv("REDIS_URL") blob_token = os.getenv("BLOB_READ_WRITE_TOKEN") # Initialize Cloudflare worker with BEST models api_base_url = f"https://api.cloudflare.com/client/v4/accounts/{cloudflare_account_id}/ai/run/" cloudflare_worker = CloudflareWorker( cloudflare_api_key=cloudflare_api_key, api_base_url=api_base_url, llm_model_name="@cf/meta/llama-3.1-8b-instruct", # Start with BEST model embedding_model_name="@cf/baai/bge-large-en-v1.5" # Start with BEST embedding ) # Test the enhanced model func_logger.info("🧪 Testing enhanced model configuration...") try: test_response = await cloudflare_worker.query( "Extract entities from: Fire extinguishers are required in commercial buildings.", "You are an expert at identifying technical entities and concepts." ) func_logger.info(f"✅ Model test successful: {test_response[:100]}...") except Exception as e: func_logger.warning(f"⚠️ Model test failed: {e}") # Initialize database manager db_manager = DatabaseManager(database_url, redis_url) await db_manager.connect() # Initialize blob client blob_client = VercelBlobClient(blob_token) # Create manager lightrag_manager = PersistentLightRAGManager( cloudflare_worker, db_manager, blob_client ) return lightrag_manager def get_lightrag_manager() -> Optional[PersistentLightRAGManager]: """Get the current LightRAG manager instance""" return lightrag_manager