Update lightrag_manager.py

lightrag_manager.py

@@ -9,15 +9,57 @@ 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
 
-# Import enhanced database manager
-from enhanced_database_manager import EnhancedDatabaseManager
+# 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:
+    """Cloudflare Workers AI integration with proper LightRAG compatibility"""
+    
     def __init__(
         self,
         cloudflare_api_key: str,
@@ -33,100 +75,348 @@ class CloudflareWorker:
         self.embedding_model_name = embedding_model_name
         self.max_tokens = max_tokens
         self.max_response_tokens = max_response_tokens
+        self.logger = logging.getLogger(__name__)
 
-    async def _send_request(self, model_name: str, input_: dict):
+    async def _send_request(self, model_name: str, input_: dict) -> Any:
+        """Send request to Cloudflare Workers AI"""
         headers = {"Authorization": f"Bearer {self.cloudflare_api_key}"}
         
         try:
-            async with httpx.AsyncClient() as client:
+            async with httpx.AsyncClient(timeout=30.0) as client:
                 response = await client.post(
                     f"{self.api_base_url}{model_name}", 
                     headers=headers, 
-                    json=input_,
-                    timeout=30.0
+                    json=input_
                 )
                 response.raise_for_status()
                 response_data = response.json()
             
             result = response_data.get("result", {})
             
+            # Handle embedding response
             if "data" in result:
                 return np.array(result["data"])
             
+            # Handle LLM response
             if "response" in result:
                 return result["response"]
                 
             raise ValueError("Unexpected Cloudflare response format")
             
         except Exception as e:
-            logging.error(f"Cloudflare API error: {e}")
+            self.logger.error(f"Cloudflare API error: {e}")
             raise
 
     async def query(self, prompt: str, system_prompt: str = "", **kwargs) -> str:
-        # Clean kwargs to avoid LightRAG-specific parameters
-        clean_kwargs = {k: v for k, v in kwargs.items() if k not in ['hashing_kv', 'history_messages']}
+        """
+        LightRAG-compatible query method
+        Fixed to handle LightRAG's parameter expectations
+        """
+        # Filter out LightRAG-specific parameters that shouldn't go to Cloudflare
+        filtered_kwargs = {
+            k: v for k, v in kwargs.items() 
+            if k not in ['hashing_kv', 'history_messages', 'global_kv', 'text_chunks']
+        }
         
-        message = [
-            {"role": "system", "content": system_prompt},
+        messages = [
+            {"role": "system", "content": system_prompt or "You are a helpful AI assistant."},
             {"role": "user", "content": prompt},
         ]
         
-        input_ = {
-            "messages": message,
+        input_data = {
+            "messages": messages,
             "max_tokens": self.max_tokens,
-            "response_token_limit": self.max_response_tokens,
+            **filtered_kwargs
         }
         
-        return await self._send_request(self.llm_model_name, input_)
+        return await self._send_request(self.llm_model_name, input_data)
 
     async def embedding_chunk(self, texts: List[str]) -> np.ndarray:
-        input_ = {
+        """Generate embeddings for text chunks"""
+        input_data = {
             "text": texts,
             "max_tokens": self.max_tokens,
-            "response_token_limit": self.max_response_tokens,
         }
         
-        return await self._send_request(self.embedding_model_name, input_)
+        return await self._send_request(self.embedding_model_name, input_data)
 
 class VercelBlobClient:
-    """Vercel Blob storage client for backup storage"""
+    """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 as backup"""
+        """Upload data to Vercel Blob"""
         try:
-            async with httpx.AsyncClient() as client:
+            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,
-                    timeout=120.0
+                    content=data
                 )
                 response.raise_for_status()
-                return f"https://blob.vercel-storage.com/{filename}"
+                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
 
-class ProductionLightRAGManager:
-    """Final Production LightRAG Manager with Complete Database Storage"""
+    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:
+            # PostgreSQL connection pool
+            self.pool = await asyncpg.create_pool(
+                self.database_url,
+                min_size=2,
+                max_size=20,
+                command_timeout=60
+            )
+            
+            # Redis connection
+            self.redis = Redis.from_url(self.redis_url, decode_responses=True)
+            
+            self.logger.info("Database connections established successfully")
+            
+            # Create tables if they don't exist
+            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,
+                    
+                    -- Blob storage URLs
+                    graph_blob_url TEXT,
+                    vector_blob_url TEXT,
+                    config_blob_url TEXT,
+                    
+                    -- Metadata
+                    total_chunks INTEGER DEFAULT 0,
+                    total_tokens INTEGER DEFAULT 0,
+                    file_count INTEGER DEFAULT 0,
+                    
+                    -- Timestamps
+                    created_at TIMESTAMP DEFAULT NOW(),
+                    updated_at TIMESTAMP DEFAULT NOW(),
+                    last_accessed_at TIMESTAMP DEFAULT NOW(),
+                    
+                    -- Status
+                    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()
+                );
+                
+                -- Indexes for performance
+                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 save_rag_instance(
+        self,
+        config: RAGConfig,
+        graph_blob_url: str,
+        vector_blob_url: str,
+        config_blob_url: str,
+        metadata: Dict[str, Any]
+    ) -> str:
+        """Save RAG instance metadata to database"""
+        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
+                ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
+                ON CONFLICT (ai_type, user_id, ai_id) DO UPDATE SET
+                    name = EXCLUDED.name,
+                    description = EXCLUDED.description,
+                    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 get_rag_instance(self, config: RAGConfig) -> Optional[Dict[str, Any]]:
+        """Get RAG instance from database"""
+        async with self.pool.acquire() as conn:
+            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'
+            """, 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'])
+                
+                return dict(result)
+            
+            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:
+            # Create conversation if it doesn't exist
+            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')
+            )
+            
+            # Save message
+            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: EnhancedDatabaseManager,
-        vercel_blob_client: Optional[VercelBlobClient] = None
+        database_manager: DatabaseManager,
+        blob_client: VercelBlobClient
     ):
         self.cloudflare_worker = cloudflare_worker
         self.db = database_manager
-        self.blob_client = vercel_blob_client
+        self.blob_client = blob_client
         self.rag_instances: Dict[str, LightRAG] = {}
-        self.processing_lock: Dict[str, asyncio.Lock] = {}
+        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,
@@ -135,20 +425,28 @@ class ProductionLightRAGManager:
         name: Optional[str] = None,
         description: Optional[str] = None
     ) -> LightRAG:
-        """Get or create a LightRAG instance with complete database storage"""
+        """Get or create RAG instance with complete Vercel persistence"""
         
-        cache_key = self._get_cache_key(ai_type, user_id, ai_id)
+        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()
         
-        # Check memory cache first
+        # Check memory cache
         if cache_key in self.rag_instances:
             self.logger.info(f"Returning cached RAG instance: {cache_key}")
             return self.rag_instances[cache_key]
         
-        # Ensure only one instance is created/loaded at a time
-        if cache_key not in self.processing_lock:
-            self.processing_lock[cache_key] = asyncio.Lock()
+        # Ensure thread safety
+        if cache_key not in self.processing_locks:
+            self.processing_locks[cache_key] = asyncio.Lock()
         
-        async with self.processing_lock[cache_key]:
+        async with self.processing_locks[cache_key]:
             # Double-check after acquiring lock
             if cache_key in self.rag_instances:
                 return self.rag_instances[cache_key]
@@ -157,7 +455,7 @@ class ProductionLightRAGManager:
             
             # Try to load from database
             try:
-                rag_instance = await self._load_from_database(ai_type, user_id, ai_id)
+                rag_instance = await self._load_from_database(config)
                 if rag_instance:
                     self.rag_instances[cache_key] = rag_instance
                     self.logger.info(f"Loaded RAG instance from database: {cache_key}")
@@ -166,16 +464,10 @@ class ProductionLightRAGManager:
                 self.logger.warning(f"Failed to load RAG from database: {e}")
             
             # Create new instance
-            rag_instance = await self._create_new_rag_instance(
-                ai_type, user_id, ai_id, name or f"{ai_type} AI", description
-            )
+            rag_instance = await self._create_new_rag_instance(config)
             
             # Save to database
-            await self._save_to_database(
-                ai_type, user_id, ai_id, 
-                name or f"{ai_type} AI", description, 
-                rag_instance
-            )
+            await self._save_to_database(config, rag_instance)
             
             # Cache in memory
             self.rag_instances[cache_key] = rag_instance
@@ -183,88 +475,14 @@ class ProductionLightRAGManager:
             
             return rag_instance
 
-    async def _load_from_database(
-        self,
-        ai_type: str,
-        user_id: Optional[str],
-        ai_id: Optional[str]
-    ) -> Optional[LightRAG]:
-        """Load RAG instance from database"""
-        
-        # Get complete RAG instance from database
-        rag_data = await self.db.load_complete_rag_instance(ai_type, user_id, ai_id)
-        if not rag_data:
-            return None
+    async def _create_new_rag_instance(self, config: RAGConfig) -> LightRAG:
+        """Create new RAG instance with in-memory storage"""
         
-        try:
-            # Reconstruct RAG instance from database data
-            rag_instance = await self._deserialize_rag_state(rag_data['rag_state'])
-            
-            self.logger.info(f"Successfully loaded RAG from database: {ai_type}")
-            return rag_instance
-            
-        except Exception as e:
-            self.logger.error(f"Failed to reconstruct RAG from database: {e}")
-            return None
-
-    async def _save_to_database(
-        self,
-        ai_type: str,
-        user_id: Optional[str],
-        ai_id: Optional[str],
-        name: str,
-        description: Optional[str],
-        rag_instance: LightRAG
-    ):
-        """Save RAG instance completely to database"""
-        
-        try:
-            # Serialize RAG state
-            rag_state = await self._serialize_rag_state(rag_instance)
-            
-            # Optional: Create backup in Vercel Blob
-            blob_url = None
-            if self.blob_client:
-                try:
-                    compressed_data = gzip.compress(pickle.dumps(rag_state))
-                    blob_filename = f"lightrag_backup/{ai_type}_{user_id or 'system'}_{ai_id or 'default'}_{uuid.uuid4()}.pkl.gz"
-                    blob_url = await self.blob_client.put(blob_filename, compressed_data)
-                    self.logger.info(f"Created backup in Vercel Blob: {blob_url}")
-                except Exception as e:
-                    self.logger.warning(f"Failed to create Vercel Blob backup: {e}")
-            
-            # Save everything to database
-            await self.db.save_complete_rag_instance(
-                ai_type=ai_type,
-                user_id=user_id,
-                ai_id=ai_id,
-                name=name,
-                description=description,
-                rag_state=rag_state,
-                blob_url=blob_url
-            )
-            
-            self.logger.info(f"Successfully saved RAG to database: {ai_type}")
-            
-        except Exception as e:
-            self.logger.error(f"Failed to save RAG to database: {e}")
-            raise
-
-    async def _create_new_rag_instance(
-        self,
-        ai_type: str,
-        user_id: Optional[str],
-        ai_id: Optional[str],
-        name: str,
-        description: Optional[str]
-    ) -> LightRAG:
-        """Create a new LightRAG instance"""
-        
-        # Create temporary working directory (will be serialized to database)
-        working_dir = f"/tmp/rag_temp_{ai_type}_{user_id or 'system'}_{ai_id or 'default'}_{uuid.uuid4()}"
+        # Create in-memory working directory structure
+        working_dir = f"/tmp/rag_memory_{config.get_cache_key()}_{uuid.uuid4()}"
         os.makedirs(working_dir, exist_ok=True)
         
-        # Initialize LightRAG
+        # Initialize LightRAG with memory-based storage
         rag = LightRAG(
             working_dir=working_dir,
             max_parallel_insert=2,
@@ -280,91 +498,157 @@ class ProductionLightRAGManager:
             vector_storage="NanoVectorDBStorage",
         )
         
-        # Wait for initialization to complete
+        # Initialize storage components
         await rag.initialize_storages()
         
         # Load knowledge based on AI type
-        if ai_type == "fire-safety":
+        if config.ai_type == "fire-safety":
             await self._load_fire_safety_knowledge(rag)
         
         return rag
 
     async def _load_fire_safety_knowledge(self, rag: LightRAG):
-        """Load fire safety knowledge from files"""
-        knowledge_sources = [
-            "/app/book.txt",
-            "/app/book.pdf",
-            "/app/fire_safety.txt",
-            "./book.txt",
-            "./book.pdf"
-        ]
+        """Load fire safety knowledge from available sources"""
         
-        combined_content = ""
-        processed_files = []
+        # Fire safety knowledge content
+        fire_safety_content = """
+        Fire Safety Regulations and Building Codes:
         
-        for source_file in knowledge_sources:
-            try:
-                if os.path.exists(source_file):
-                    if source_file.endswith('.txt'):
-                        with open(source_file, 'r', encoding='utf-8') as f:
-                            content = f.read()
-                    elif source_file.endswith('.pdf'):
-                        content = await self._extract_pdf_content(source_file)
-                    else:
-                        continue
-                    
-                    if content.strip():
-                        combined_content += f"\n\n=== Content from {source_file} ===\n\n{content}\n\n"
-                        processed_files.append(source_file)
-                        
-            except Exception as e:
-                self.logger.warning(f"Failed to process {source_file}: {e}")
+        1. Emergency Exits:
+        - 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 travel
+        - Exits must be clearly marked and illuminated
+        
+        2. Fire Extinguishers:
+        - Type A: Ordinary combustibles (wood, paper, cloth)
+        - Type B: Flammable liquids (gasoline, oil, paint)
+        - Type C: Electrical equipment
+        - Type D: Combustible metals
+        - Type K: Cooking oils and fats
+        
+        3. Fire Detection Systems:
+        - Smoke detectors required in all sleeping areas
+        - Heat detectors in areas where smoke detectors are not suitable
+        - Manual fire alarm pull stations near exits
+        - Central monitoring systems in commercial buildings
+        
+        4. Sprinkler Systems:
+        - Required in buildings over certain heights
+        - Wet pipe systems most common
+        - Dry pipe systems in areas subject to freezing
+        - Deluge systems for high-hazard areas
         
-        if combined_content.strip():
-            self.logger.info(f"Inserting fire safety knowledge from {len(processed_files)} files")
-            await rag.ainsert(combined_content)
-        else:
-            self.logger.warning("No fire safety knowledge found")
+        5. Emergency Lighting:
+        - Required in all exit routes
+        - Must provide minimum 1 foot-candle illumination
+        - Battery backup required for minimum 90 minutes
+        - Monthly testing required
+        
+        6. Fire Doors:
+        - Must be self-closing and self-latching
+        - Fire rating must match wall rating
+        - Annual inspection required
+        - No propping open unless connected to fire alarm system
+        
+        7. Occupancy Limits:
+        - Based on building type and exit capacity
+        - Assembly: 7 sq ft per person (concentrated use)
+        - Business: 100 sq ft per person
+        - Educational: 20 sq ft per person
+        - Industrial: 100 sq ft per person
+        """
+        
+        self.logger.info("Loading fire safety knowledge base")
+        await rag.ainsert(fire_safety_content)
+        self.logger.info("Fire safety knowledge loaded successfully")
 
-    async def _extract_pdf_content(self, pdf_path: str) -> str:
-        """Extract text from PDF file"""
+    async def _save_to_database(self, config: RAGConfig, rag: LightRAG):
+        """Save RAG instance to Vercel Blob + Database"""
+        
         try:
-            import PyPDF2
-            content = ""
-            with open(pdf_path, 'rb') as file:
-                pdf_reader = PyPDF2.PdfReader(file)
-                for page in pdf_reader.pages:
-                    text = page.extract_text()
-                    if text:
-                        content += text + "\n"
-            return content
+            # Serialize RAG state
+            rag_state = await self._serialize_rag_state(rag)
+            
+            # Create blob files
+            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}"
+            
+            # Save to Vercel Blob
+            graph_data = gzip.compress(pickle.dumps(rag_state['graph']))
+            vector_data = gzip.compress(pickle.dumps(rag_state['vectors']))
+            config_data = gzip.compress(pickle.dumps(rag_state['config']))
+            
+            graph_blob_url = await self.blob_client.put(f"{base_filename}_graph.pkl.gz", graph_data)
+            vector_blob_url = await self.blob_client.put(f"{base_filename}_vectors.pkl.gz", vector_data)
+            config_blob_url = await self.blob_client.put(f"{base_filename}_config.pkl.gz", config_data)
+            
+            # Save metadata to database
+            metadata = {
+                'total_chunks': len(rag_state['vectors'].get('embeddings', [])),
+                'total_tokens': self._estimate_tokens(rag_state),
+                'file_count': 1 if config.ai_type == 'fire-safety' else 0
+            }
+            
+            await self.db.save_rag_instance(
+                config, graph_blob_url, vector_blob_url, config_blob_url, metadata
+            )
+            
+            self.logger.info(f"Successfully saved RAG to Vercel storage: {config.ai_type}")
+            
         except Exception as e:
-            self.logger.warning(f"Failed to extract PDF content: {e}")
-            return ""
+            self.logger.error(f"Failed to save RAG to database: {e}")
+            raise
 
-    async def _serialize_rag_state(self, rag: LightRAG) -> Dict[str, Any]:
-        """Serialize LightRAG state for database storage"""
+    async def _load_from_database(self, config: RAGConfig) -> Optional[LightRAG]:
+        """Load RAG instance from database + Vercel Blob"""
+        
         try:
-            graph_storage = rag.graph_storage
-            vector_storage = rag.vector_storage
+            # Get metadata from database
+            instance_data = await self.db.get_rag_instance(config)
+            if not instance_data:
+                return None
+            
+            # Download 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'])
+            
+            # Deserialize
+            rag_state = {
+                'graph': pickle.loads(gzip.decompress(graph_data)),
+                'vectors': pickle.loads(gzip.decompress(vector_data)),
+                'config': pickle.loads(gzip.decompress(config_data))
+            }
+            
+            # Reconstruct RAG instance
+            rag = await self._deserialize_rag_state(rag_state)
+            
+            self.logger.info(f"Successfully loaded RAG from Vercel storage: {config.ai_type}")
+            return rag
             
+        except Exception as e:
+            self.logger.error(f"Failed to load RAG from database: {e}")
+            return None
+
+    async def _serialize_rag_state(self, rag: LightRAG) -> Dict[str, Any]:
+        """Serialize RAG state for storage"""
+        try:
             # Extract graph data
             graph_data = {"nodes": [], "edges": [], "graph_attrs": {}}
-            if hasattr(graph_storage, 'graph') and graph_storage.graph:
+            if hasattr(rag.graph_storage, 'graph') and rag.graph_storage.graph:
                 graph_data = {
-                    "nodes": list(graph_storage.graph.nodes(data=True)),
-                    "edges": list(graph_storage.graph.edges(data=True)),
-                    "graph_attrs": dict(graph_storage.graph.graph)
+                    "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)
                 }
             
             # Extract vector data
             vector_data = {"embeddings": [], "metadata": [], "dimension": 1024}
-            if hasattr(vector_storage, 'embeddings') and vector_storage.embeddings is not None:
-                vector_data["embeddings"] = vector_storage.embeddings.tolist()
-            if hasattr(vector_storage, 'metadata'):
-                vector_data["metadata"] = getattr(vector_storage, 'metadata', [])
-            if hasattr(vector_storage, 'dimension'):
-                vector_data["dimension"] = getattr(vector_storage, 'dimension', 1024)
+            if hasattr(rag.vector_storage, 'embeddings') and rag.vector_storage.embeddings is not None:
+                vector_data["embeddings"] = rag.vector_storage.embeddings.tolist()
+            if hasattr(rag.vector_storage, 'metadata'):
+                vector_data["metadata"] = getattr(rag.vector_storage, 'metadata', [])
             
             # Configuration
             config_data = {
@@ -387,7 +671,7 @@ class ProductionLightRAGManager:
             raise
 
     async def _deserialize_rag_state(self, rag_state: Dict[str, Any]) -> LightRAG:
-        """Deserialize RAG state and reconstruct LightRAG instance"""
+        """Deserialize RAG state and reconstruct LightRAG"""
         try:
             config = rag_state["config"]
             
@@ -395,7 +679,7 @@ class ProductionLightRAGManager:
             working_dir = f"/tmp/rag_restored_{uuid.uuid4()}"
             os.makedirs(working_dir, exist_ok=True)
             
-            # Create new RAG instance
+            # Create RAG instance
             rag = LightRAG(
                 working_dir=working_dir,
                 max_parallel_insert=config.get("max_parallel_insert", 2),
@@ -419,8 +703,6 @@ class ProductionLightRAGManager:
                 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"])
-            if graph_data["graph_attrs"] and hasattr(rag.graph_storage, 'graph'):
-                rag.graph_storage.graph.graph.update(graph_data["graph_attrs"])
             
             # Restore vectors
             vector_data = rag_state["vectors"]
@@ -428,8 +710,6 @@ class ProductionLightRAGManager:
                 rag.vector_storage.embeddings = np.array(vector_data["embeddings"])
             if hasattr(rag.vector_storage, 'metadata'):
                 rag.vector_storage.metadata = vector_data["metadata"]
-            if hasattr(rag.vector_storage, 'dimension'):
-                rag.vector_storage.dimension = vector_data["dimension"]
             
             return rag
             
@@ -447,43 +727,43 @@ class ProductionLightRAGManager:
         mode: str = "hybrid",
         max_memory_turns: int = 10
     ) -> str:
-        """Query RAG with database-persisted conversation memory"""
+        """Query RAG with conversation memory"""
         
         try:
             # Get RAG instance
             rag = await self.get_or_create_rag_instance(ai_type, user_id, ai_id)
             
-            # Get conversation messages from database
+            # Get conversation memory
             messages = await self.db.get_conversation_messages(conversation_id)
             
-            # Build context with conversation memory
+            # Build context with memory
             context_prompt = self._build_context_prompt(question, messages[-max_memory_turns*2:])
             
             # Query LightRAG
             response = await rag.aquery(context_prompt, QueryParam(mode=mode))
             
-            # Save conversation messages to database
-            await self.db.save_conversation_message(conversation_id, "user", question)
-            await self.db.save_conversation_message(conversation_id, "assistant", response, {
-                "mode": mode,
-                "ai_type": ai_type,
-                "user_id": user_id,
-                "ai_id": ai_id
-            })
+            # Save conversation
+            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}")
-            
-            # Fallback to direct query
-            try:
-                rag = await self.get_or_create_rag_instance(ai_type, user_id, ai_id)
-                response = await rag.aquery(question, QueryParam(mode=mode))
-                return response
-            except Exception as fallback_error:
-                self.logger.error(f"Fallback query also failed: {fallback_error}")
-                return "I apologize, but I'm experiencing technical difficulties. Please try again later."
+            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"""
@@ -506,7 +786,7 @@ class ProductionLightRAGManager:
         description: str,
         uploaded_files: List[Dict[str, Any]]
     ) -> str:
-        """Create custom AI with database storage"""
+        """Create custom AI with uploaded files"""
         
         ai_id = str(uuid.uuid4())
         
@@ -526,30 +806,16 @@ class ProductionLightRAGManager:
             if combined_content.strip():
                 await rag.ainsert(combined_content)
             
-            # Save to database with updated knowledge
-            await self._save_to_database(
+            # Save to database
+            config = RAGConfig(
                 ai_type="custom",
                 user_id=user_id,
                 ai_id=ai_id,
                 name=ai_name,
-                description=description,
-                rag_instance=rag
+                description=description
             )
             
-            # Save file metadata to database
-            rag_data = await self.db.load_complete_rag_instance("custom", user_id, ai_id)
-            if rag_data:
-                rag_instance_id = rag_data['metadata']['id']
-                
-                for file_data in uploaded_files:
-                    await self.db.save_knowledge_file(
-                        rag_instance_id=rag_instance_id,
-                        filename=file_data['filename'],
-                        original_filename=file_data['filename'],
-                        file_type=file_data.get('type', 'unknown'),
-                        file_size=file_data.get('size', 0),
-                        content_text=file_data.get('content', '')
-                    )
+            await self._save_to_database(config, rag)
             
             return ai_id
             
@@ -557,39 +823,78 @@ class ProductionLightRAGManager:
             self.logger.error(f"Failed to create custom AI: {e}")
             raise
 
-    async def get_user_ais(self, user_id: str) -> List[Dict[str, Any]]:
-        return await self.db.list_user_rag_instances(user_id)
+    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  # Rough estimate: 4 chars per token
+        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_lock.clear()
+        self.processing_locks.clear()
         await self.db.close()
         self.logger.info("LightRAG manager cleaned up")
 
 # Global instance
-lightrag_manager: Optional[ProductionLightRAGManager] = None
-
-async def initialize_lightrag_manager(
-    cloudflare_worker: CloudflareWorker,
-    database_url: str,
-    vercel_blob_token: str
-) -> ProductionLightRAGManager:
-    """Initialize the production LightRAG manager"""
+lightrag_manager: Optional[PersistentLightRAGManager] = None
+
+async def initialize_lightrag_manager() -> PersistentLightRAGManager:
+    """Initialize the persistent LightRAG manager"""
     global lightrag_manager
     
     if lightrag_manager is None:
-        # Initialize database
-        db_manager = DatabaseManager(database_url)
+        # 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
+        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.2-3b-instruct",
+            embedding_model_name="@cf/baai/bge-m3"
+        )
+        
+        # Initialize database manager
+        db_manager = DatabaseManager(database_url, redis_url)
         await db_manager.connect()
         
         # Initialize blob client
-        blob_client = VercelBlobClient(vercel_blob_token)
+        blob_client = VercelBlobClient(blob_token)
         
         # Create manager
-        lightrag_manager = ProductionLightRAGManager(
+        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