Update lightrag_manager.py

lightrag_manager.py

@@ -927,7 +927,7 @@ class PersistentLightRAGManager:
         except Exception as e:
             self.logger.error(f"Failed to save RAG to database: {e}")
             raise
-
+    
     async def _load_from_database(self, config: RAGConfig) -> Optional[LightRAG]:
         """Load RAG from database using stored JSON data FIRST, then fallback to Blob URLs"""
         
@@ -938,421 +938,159 @@ class PersistentLightRAGManager:
                 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']})")
+            self.logger.info(f"🔍 Found RAG instance: {instance_data['name']} (ID: {instance_data['id']})")
             
-            # PRIORITY 1: Check if we have storage data in knowledge_files table (DATABASE FIRST)
+            # Check if we have storage data in knowledge_files table
             async with self.db.pool.acquire() as conn:
-                stored_data = await conn.fetchval("""
-                    SELECT content_text FROM knowledge_files 
+                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 stored_data:
-                self.logger.info("🎯 Loading RAG from DATABASE storage (priority method)")
+            if storage_record and storage_record['content_text']:
+                self.logger.info(f"🎯 Found database storage: {storage_record['file_size']} bytes, {storage_record['token_count']} tokens")
                 
-                # Create new RAG instance
-                rag = await self._create_new_rag_instance(config)
-                
-                # Restore storage data from database
                 try:
-                    storage_data = json.loads(stored_data)
+                    # 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())}")
                     
-                    # Write storage files back to temporary directory
+                    # Check if we have actual data
+                    total_items = 0
+                    for key, data in storage_data.items():
+                        if isinstance(data, (list, dict)):
+                            item_count = len(data)
+                            total_items += item_count
+                            self.logger.info(f"   - {key}: {item_count} items")
+                    
+                    if total_items == 0:
+                        self.logger.warning("❌ Storage data exists but contains no items")
+                        return None
+                    
+                    # Create new RAG instance structure
+                    working_dir = f"/tmp/rag_restored_{uuid.uuid4()}"
+                    os.makedirs(working_dir, exist_ok=True)
+                    
+                    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 empty storages
+                    await rag.initialize_storages()
+                    
+                    # Restore storage files from database JSON
                     storage_restored = False
                     for filename, data in storage_data.items():
-                        if data and (isinstance(data, (list, dict))):
-                            file_path = f"{rag.working_dir}/{filename}.json"
+                        if data and isinstance(data, (list, dict)) and len(data) > 0:
+                            file_path = f"{working_dir}/{filename}.json"
                             with open(file_path, 'w') as f:
                                 json.dump(data, f)
                             
                             file_size = os.path.getsize(file_path)
-                            item_count = len(data) if isinstance(data, (list, dict)) else 'unknown'
-                            self.logger.info(f"✅ Restored {filename}.json: {file_size} bytes, {item_count} items")
+                            self.logger.info(f"✅ Restored {filename}.json: {file_size} bytes")
                             storage_restored = True
                     
                     if storage_restored:
-                        # Force reload storage from files
+                        # Re-initialize storages to load the restored files
                         await rag.initialize_storages()
+                        self.logger.info("🔄 Re-initialized storages with restored data")
                         
-                        # Verify the restored RAG has working data
-                        verification_passed = await self._verify_rag_storage(rag)
+                        # Verify the RAG has working data
+                        verification_passed = await self._verify_restored_rag(rag)
                         if verification_passed:
-                            self.logger.info("🎉 Successfully loaded and verified RAG from DATABASE storage")
+                            self.logger.info("🎉 SUCCESS: Loaded and verified RAG from DATABASE storage")
                             return rag
                         else:
-                            self.logger.warning("⚠️ RAG loaded from database but verification failed")
+                            self.logger.warning("⚠️ RAG restored but verification failed")
                             return None
                     else:
-                        self.logger.warning("⚠️ No valid storage data found in database JSON")
+                        self.logger.warning("⚠️ No valid storage files restored")
                         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 storage data from database: {e}")
-                    # Don't return None here, try Blob storage as fallback
-            
-            # PRIORITY 2: Fallback to Vercel Blob storage (only if database fails)
-            if (instance_data.get('graph_blob_url') and 
-                instance_data.get('vector_blob_url') and 
-                instance_data.get('config_blob_url') and
-                not instance_data.get('graph_blob_url').startswith('database://')):
-                
-                self.logger.info("🔄 Falling back to Vercel Blob storage")
-                return await self._load_from_blob_storage(instance_data)
-            
-            self.logger.warning(f"❌ No usable storage data found for RAG instance {instance_data['id']}")
-            return None
+                    self.logger.error(f"❌ Failed to restore from database storage: {e}")
+                    return None
+            else:
+                self.logger.warning(f"⚠️ No storage data found in database for RAG {instance_data['id']}")
+                return None
             
         except Exception as e:
-            self.logger.error(f"Failed to load RAG from database: {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"""
+    
+    async def _verify_restored_rag(self, rag: LightRAG) -> bool:
+        """Verify that restored RAG has actual working data"""
         try:
+            # Check different storage components
+            checks = {
+                'vector_storage': 0,
+                'chunks': 0,
+                'entities': 0,
+                'relationships': 0
+            }
+            
             # Check vector storage
-            vector_count = 0
             if hasattr(rag.vector_storage, '_data') and rag.vector_storage._data:
-                vector_count = len(rag.vector_storage._data)
+                checks['vector_storage'] = len(rag.vector_storage._data)
             
-            # Check chunks storage
-            chunks_count = 0
+            # Check chunks
             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
+                    checks['chunks'] = len(chunks_data) if chunks_data else 0
                 except:
                     pass
             
-            # Check entities storage
-            entities_count = 0
+            # Check entities  
             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
+                    checks['entities'] = len(entities_data) if entities_data else 0
                 except:
                     pass
             
-            # Check relationships storage
-            relationships_count = 0
+            # Check relationships
             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
+                    checks['relationships'] = 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}")
+            # Log verification results
+            self.logger.info(f"📋 RAG verification: {checks}")
             
-            # 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
+            # Consider successful if ANY component has data
+            has_data = any(count > 0 for count in checks.values())
             
             if has_data:
-                self.logger.info("✅ RAG verification PASSED - has working data")
+                self.logger.info("✅ RAG verification PASSED")
             else:
-                self.logger.warning("❌ RAG verification FAILED - no data found")
+                self.logger.warning("❌ RAG verification FAILED - no data in any component")
             
             return has_data
             
         except Exception as e:
-            self.logger.error(f"Failed to verify RAG storage: {e}")
+            self.logger.error(f"❌ RAG verification error: {e}")
             return False
-    
-        async def _load_from_blob_storage(self, instance_data: Dict[str, Any]) -> Optional[LightRAG]:
-            """Load RAG from Vercel Blob storage"""
-            try:
-                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'])
-                
-                rag_state = {
-                    'graph': pickle.loads(gzip.decompress(graph_data)),
-                    'vectors': pickle.loads(gzip.decompress(vector_data)),
-                    'config': pickle.loads(gzip.decompress(config_data))
-                }
-                
-                self.logger.info("Successfully downloaded and deserialized RAG state")
-                
-                return await self._deserialize_rag_state(rag_state)
-                
-            except Exception as e:
-                self.logger.error(f"Failed to load RAG from Vercel Blob: {e}")
-                return None
-        
-        async def _serialize_rag_state(self, rag: LightRAG) -> Dict[str, Any]:
-            """Serialize RAG state for LightRAG 1.3.7"""
-            try:
-                graph_data = {"nodes": [], "edges": [], "graph_attrs": {}}
-                if hasattr(rag.graph_storage, 'graph') and rag.graph_storage.graph:
-                    graph_data = {
-                        "nodes": list(rag.graph_storage.graph.nodes(data=True)),
-                        "edges": list(rag.graph_storage.graph.edges(data=True)),
-                        "graph_attrs": dict(rag.graph_storage.graph.graph)
-                    }
-                
-                vector_data = {"items": [], "metadata": [], "dimension": 1024}
-                
-                if hasattr(rag.vector_storage, '_data') and rag.vector_storage._data:
-                    vector_data["items"] = len(rag.vector_storage._data)
-                    self.logger.info(f"Serializing {len(rag.vector_storage._data)} vector storage items")
-                elif hasattr(rag.vector_storage, 'embeddings') and rag.vector_storage.embeddings is not None:
-                    vector_data["embeddings"] = rag.vector_storage.embeddings.tolist()
-                    self.logger.info(f"Serializing {len(vector_data['embeddings'])} embeddings")
-                
-                chunks_data = {"count": 0, "data": []}
-                if hasattr(rag, 'chunks') and rag.chunks:
-                    try:
-                        all_chunks = await rag.chunks.get_all()
-                        if all_chunks:
-                            chunks_data["count"] = len(all_chunks)
-                            chunks_data["data"] = all_chunks[:100]
-                            self.logger.info(f"Serializing {chunks_data['count']} chunks")
-                    except Exception as e:
-                        self.logger.warning(f"Failed to get chunks data: {e}")
-                
-                entities_data = {"count": 0}
-                if hasattr(rag, 'entities') and rag.entities:
-                    try:
-                        all_entities = await rag.entities.get_all()
-                        if all_entities:
-                            entities_data["count"] = len(all_entities)
-                            self.logger.info(f"Serializing {entities_data['count']} entities")
-                    except Exception as e:
-                        self.logger.warning(f"Failed to get entities data: {e}")
-                
-                config_data = {
-                    "max_parallel_insert": rag.max_parallel_insert,
-                    "llm_model_name": getattr(rag, 'llm_model_name', 'unknown'),
-                    "embedding_dimension": getattr(rag.embedding_func, 'embedding_dim', 1024),
-                    "created_at": datetime.now().isoformat(),
-                    "working_dir": rag.working_dir,
-                    "lightrag_version": "1.3.7",
-                    "storage_stats": {
-                        "vector_items": vector_data.get("items", 0),
-                        "chunks_count": chunks_data["count"],
-                        "entities_count": entities_data["count"]
-                    }
-                }
-                
-                return {
-                    "graph": graph_data,
-                    "vectors": vector_data,
-                    "chunks": chunks_data,
-                    "entities": entities_data,
-                    "config": config_data,
-                    "version": "1.3.7"
-                }
-                
-            except Exception as e:
-                self.logger.error(f"Failed to serialize RAG state: {e}")
-                raise
-    
-        async def _deserialize_rag_state(self, rag_state: Dict[str, Any]) -> LightRAG:
-            """Deserialize RAG state and reconstruct LightRAG"""
-            try:
-                config = rag_state["config"]
-                
-                working_dir = f"/tmp/rag_restored_{uuid.uuid4()}"
-                os.makedirs(working_dir, exist_ok=True)
-                
-                rag = LightRAG(
-                    working_dir=working_dir,
-                    max_parallel_insert=config.get("max_parallel_insert", 2),
-                    llm_model_func=self.cloudflare_worker.query,
-                    llm_model_name=config.get("llm_model_name", self.cloudflare_worker.llm_model_name),
-                    llm_model_max_token_size=4080,
-                    embedding_func=EmbeddingFunc(
-                        embedding_dim=config.get("embedding_dimension", 1024),
-                        max_token_size=2048,
-                        func=self.cloudflare_worker.embedding_chunk,
-                    ),
-                    graph_storage="NetworkXStorage",
-                    vector_storage="NanoVectorDBStorage",
-                )
-                
-                await rag.initialize_storages()
-                
-                graph_data = rag_state["graph"]
-                if graph_data["nodes"] and hasattr(rag.graph_storage, 'graph'):
-                    rag.graph_storage.graph.add_nodes_from(graph_data["nodes"])
-                if graph_data["edges"] and hasattr(rag.graph_storage, 'graph'):
-                    rag.graph_storage.graph.add_edges_from(graph_data["edges"])
-                
-                vector_data = rag_state["vectors"]
-                if vector_data["embeddings"] and hasattr(rag.vector_storage, 'embeddings'):
-                    rag.vector_storage.embeddings = np.array(vector_data["embeddings"])
-                if hasattr(rag.vector_storage, 'metadata'):
-                    rag.vector_storage.metadata = vector_data["metadata"]
-                
-                return rag
-                
-            except Exception as e:
-                self.logger.error(f"Failed to deserialize RAG state: {e}")
-                raise
-    
-        async def query_with_memory(
-            self,
-            ai_type: str,
-            question: str,
-            conversation_id: str,
-            user_id: Optional[str] = None,
-            ai_id: Optional[str] = None,
-            mode: str = "hybrid",
-            max_memory_turns: int = 10
-        ) -> str:
-            """Query RAG with conversation memory"""
-            
-            try:
-                rag = await self.get_or_create_rag_instance(ai_type, user_id, ai_id)
-                
-                messages = await self.db.get_conversation_messages(conversation_id)
-                
-                context_prompt = self._build_context_prompt(question, messages[-max_memory_turns*2:])
-                
-                response = await rag.aquery(context_prompt, QueryParam(mode=mode))
-                
-                await self.db.save_conversation_message(
-                    conversation_id, "user", question, {
-                        "user_id": user_id,
-                        "ai_type": ai_type,
-                        "ai_id": ai_id
-                    }
-                )
-                await self.db.save_conversation_message(
-                    conversation_id, "assistant", response, {
-                        "mode": mode,
-                        "ai_type": ai_type,
-                        "user_id": user_id,
-                        "ai_id": ai_id
-                    }
-                )
-                
-                return response
-                
-            except Exception as e:
-                self.logger.error(f"Query with memory failed: {e}")
-                return "I apologize, but I'm experiencing technical difficulties. Please try again later."
-    
-        def _build_context_prompt(self, question: str, messages: List[Dict[str, Any]]) -> str:
-            """Build context prompt with conversation memory"""
-            if not messages:
-                return question
-            
-            context = "Previous conversation:\n"
-            for msg in messages:
-                role = msg['role']
-                content = msg['content'][:200] + "..." if len(msg['content']) > 200 else msg['content']
-                context += f"{role.title()}: {content}\n"
-            
-            context += f"\nCurrent question: {question}"
-            return context
-    
-        async def create_custom_ai(
-            self,
-            user_id: str,
-            ai_name: str,
-            description: str,
-            uploaded_files: List[Dict[str, Any]]
-        ) -> str:
-            """Create custom AI with uploaded files"""
-            
-            ai_id = str(uuid.uuid4())
-            
-            try:
-                rag = await self.get_or_create_rag_instance(
-                    "custom", user_id, ai_id, ai_name, description
-                )
-                
-                combined_content = ""
-                for file_data in uploaded_files:
-                    content = file_data.get('content', '')
-                    combined_content += f"\n\n=== {file_data['filename']} ===\n\n{content}\n\n"
-                
-                if combined_content.strip():
-                    await rag.ainsert(combined_content)
-                
-                config = RAGConfig(
-                    ai_type="custom",
-                    user_id=user_id,
-                    ai_id=ai_id,
-                    name=ai_name,
-                    description=description
-                )
-                
-                await self._save_to_database(config, rag)
-                
-                return ai_id
-                
-            except Exception as e:
-                self.logger.error(f"Failed to create custom AI: {e}")
-                raise
-    
-        def _estimate_tokens(self, rag_state: Dict[str, Any]) -> int:
-            """Estimate token count from RAG state"""
-            try:
-                content_size = len(json.dumps(rag_state))
-                return content_size // 4
-            except:
-                return 0
-    
-        def get_conversation_memory_status(self, conversation_id: str) -> Dict[str, Any]:
-            """Get conversation memory status"""
-            if conversation_id in self.conversation_memory:
-                return {
-                    "has_memory": True,
-                    "message_count": len(self.conversation_memory[conversation_id]),
-                    "last_updated": datetime.now().isoformat()
-                }
-            return {"has_memory": False, "message_count": 0}
-    
-        def clear_conversation_memory(self, conversation_id: str):
-            """Clear conversation memory"""
-            if conversation_id in self.conversation_memory:
-                del self.conversation_memory[conversation_id]
-    
-        async def cleanup(self):
-            """Clean up resources"""
-            self.rag_instances.clear()
-            self.conversation_memory.clear()
-            self.processing_locks.clear()
-            await self.db.close()
-            self.logger.info("LightRAG manager cleaned up")
-    
-        async def _verify_rag_storage(self, rag: LightRAG) -> bool:
-            """Verify that RAG storage has been properly loaded"""
-            try:
-                vector_count = 0
-                if hasattr(rag.vector_storage, '_data') and rag.vector_storage._data:
-                    vector_count = len(rag.vector_storage._data)
-                
-                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
-                
-                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
-                
-                self.logger.info(f"RAG storage verification: vectors={vector_count}, chunks={chunks_count}, entities={entities_count}")
-                
-                return vector_count > 0 or chunks_count > 0 or entities_count > 0
-                
-            except Exception as e:
-                self.logger.error(f"Failed to verify RAG storage: {e}")
-                return False
 
 # Global instance
 lightrag_manager: Optional[PersistentLightRAGManager] = None