lightrag_manager.py

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 ['<entity_name>', '', '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 ['<entity_name>', '', '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<entity_name><|><entity_type><|><entity_description>
    relationship<source_entity><|><target_entity><|><relationship_description><|><keywords><|><strength>
    content_keywords<high_level_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