Update app.py

app.py

@@ -1,770 +1,247 @@
+import asyncio
 import os
 import json
-import asyncio
+import uuid
 import hashlib
 import shutil
-import zipfile
-import uuid
-import requests
-from datetime import datetime, timedelta
+import logging
+from datetime import datetime
 from pathlib import Path
-from typing import List, Dict, Optional
+from typing import List, Optional, Dict, Any
 
-# FastAPI and dependencies
-from fastapi import FastAPI, HTTPException, Depends, File, UploadFile
+from fastapi import FastAPI, HTTPException, Depends, UploadFile, File, Form
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel, EmailStr
-import aiohttp
-import jwt
 
-# Database imports
-import asyncpg
-import redis.asyncio as redis
-from contextlib import asynccontextmanager
+# Import LightRAG manager
+from lightrag_manager import (
+    CloudflareWorker,
+    LightRAGManager,
+    initialize_lightrag_manager,
+    get_lightrag_manager
+)
 
-# Pydantic models
-class UserRegister(BaseModel):
-    email: EmailStr
-    name: str
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-class UserLogin(BaseModel):
-    email: EmailStr
+# Get environment variables from HF Spaces secrets
+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", "")
+JWT_SECRET_KEY = os.getenv("JWT_SECRET_KEY", "your-secret-key")
+
+# Models
+EMBEDDING_MODEL = "@cf/baai/bge-m3"
+LLM_MODEL = "@cf/meta/llama-3.2-3b-instruct"
+
+# API Base URL
+API_BASE_URL = f"https://api.cloudflare.com/client/v4/accounts/{CLOUDFLARE_ACCOUNT_ID}/ai/run/"
 
+# File upload settings
+USER_DATA_DIR = os.getenv("USER_DATA_DIR", "./user_data")
+MAX_UPLOAD_SIZE = int(os.getenv("MAX_UPLOAD_SIZE", "10485760"))  # 10MB
+
+# Create FastAPI app
+app = FastAPI(title="YourAI - LightRAG Powered", version="2.0.0")
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Pydantic models
 class QuestionRequest(BaseModel):
     question: str
-    mode: str = "hybrid"
+    mode: Optional[str] = "hybrid"
     conversation_id: Optional[str] = None
 
-class CustomAIRequest(BaseModel):
-    name: str
-    description: str
-
 class QuestionResponse(BaseModel):
     answer: str
     mode: str
     status: str
     conversation_id: Optional[str] = None
 
+class UserRegisterRequest(BaseModel):
+    email: EmailStr
+    name: str
+
+class UserLoginRequest(BaseModel):
+    email: EmailStr
+
+class CustomAIRequest(BaseModel):
+    name: str
+    description: str
+
 class FileUploadResponse(BaseModel):
     filename: str
     size: int
     message: str
 
-# Database connection management
+# Global variables
+cloudflare_worker: Optional[CloudflareWorker] = None
+lightrag_manager: Optional[LightRAGManager] = None
+
+# Database Manager (simplified version - you'll need to implement based on your DB)
 class DatabaseManager:
-    def __init__(self):
-        self.pool = None
-        self.redis = None
-    
-    async def connect(self):
-        """Initialize database connections"""
-        # PostgreSQL connection
-        database_url = os.getenv('DATABASE_URL')
-        if database_url:
-            try:
-                self.pool = await asyncpg.create_pool(database_url, max_size=20)
-                print("✅ PostgreSQL connected successfully")
-            except Exception as e:
-                print(f"❌ PostgreSQL connection failed: {e}")
-                self.pool = None
+    def __init__(self, database_url: str):
+        self.database_url = database_url
+        # Initialize your database connection here
         
-        # Redis connection
-        redis_url = os.getenv('REDIS_URL')
-        if redis_url:
-            try:
-                self.redis = redis.from_url(redis_url, decode_responses=True)
-                await self.redis.ping()
-                print("✅ Redis connected successfully")
-            except Exception as e:
-                print(f"❌ Redis connection failed: {e}")
-                self.redis = None
-    
-    async def close(self):
-        """Close database connections"""
-        if self.pool:
-            await self.pool.close()
-        if self.redis:
-            await self.redis.close()
-    
     async def execute_query(self, query: str, *args):
-        """Execute a PostgreSQL query"""
-        if not self.pool:
-            raise HTTPException(status_code=503, detail="Database not available")
-        
-        async with self.pool.acquire() as connection:
-            try:
-                if query.strip().upper().startswith('SELECT'):
-                    return await connection.fetch(query, *args)
-                else:
-                    return await connection.execute(query, *args)
-            except Exception as e:
-                print(f"Database query error: {e}")
-                raise HTTPException(status_code=500, detail="Database operation failed")
+        # Implement your database query execution
+        # For now, return empty result
+        return []
     
-    async def cache_set(self, key: str, value: any, ttl: int = 3600):
-        """Set cache value with TTL"""
-        if self.redis:
-            try:
-                await self.redis.setex(key, ttl, json.dumps(value))
-                return True
-            except Exception as e:
-                print(f"Cache set error: {e}")
-        return False
+    async def cache_set(self, key: str, value: Any, ttl: int = 3600):
+        # Implement Redis cache set
+        pass
     
     async def cache_get(self, key: str):
-        """Get cache value"""
-        if self.redis:
-            try:
-                value = await self.redis.get(key)
-                return json.loads(value) if value else None
-            except Exception as e:
-                print(f"Cache get error: {e}")
+        # Implement Redis cache get
         return None
     
     async def cache_delete(self, key: str):
-        """Delete cache value"""
-        if self.redis:
-            try:
-                await self.redis.delete(key)
-                return True
-            except Exception as e:
-                print(f"Cache delete error: {e}")
-        return False
+        # Implement Redis cache delete
+        pass
 
 # Initialize database manager
-db_manager = DatabaseManager()
+db_manager = DatabaseManager(DATABASE_URL)
 
-# Cloudflare Worker class (keep your existing implementation)
-class CloudflareWorker:
-    def __init__(self, cloudflare_api_key: str, api_base_url: str, llm_model_name: str, max_tokens: int = 1024):
-        self.cloudflare_api_key = cloudflare_api_key
-        self.api_base_url = api_base_url
-        self.llm_model_name = llm_model_name
-        self.max_tokens = max_tokens
-
-    async def _send_request(self, model_name: str, input_data: dict) -> str:
-        headers = {"Authorization": f"Bearer {self.cloudflare_api_key}"}
-        url = f"{self.api_base_url}{model_name}"
-
-        async with aiohttp.ClientSession() as session:
-            async with session.post(url, headers=headers, json=input_data) as response:
-                if response.status == 200:
-                    result = await response.json()
-                    return result["result"]["response"]
-                else:
-                    error_text = await response.text()
-                    raise HTTPException(status_code=response.status, detail=f"Cloudflare API error: {error_text}")
-
-    async def query(self, prompt: str, system_prompt: str = '') -> str:
-        message = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": prompt}
-        ]
-        
-        input_ = {
-            "messages": message,
-            "max_tokens": self.max_tokens,
-        }
-        
-        result = await self._send_request(self.llm_model_name, input_)
-        return result
-
-# User Management with Database
-class UserManager:
-    @staticmethod
-    def hash_email(email: str) -> str:
-        return hashlib.md5(email.encode()).hexdigest()[:12]
-    
-    @staticmethod
-    async def create_user(email: str, name: str) -> dict:
-        hashed_email = UserManager.hash_email(email)
-        user_id = str(uuid.uuid4())
-        
-        # Check if user exists
-        existing_query = "SELECT id FROM users WHERE email = $1"
-        existing = await db_manager.execute_query(existing_query, email)
-        if existing:
-            raise HTTPException(status_code=400, detail="User already exists")
-        
-        # Create user
-        query = """
-        INSERT INTO users (id, email, name, hashed_email, created_at, updated_at)
-        VALUES ($1, $2, $3, $4, $5, $6)
-        RETURNING id, email, name, created_at
-        """
-        
-        now = datetime.now()
-        result = await db_manager.execute_query(
-            query, user_id, email, name, hashed_email, now, now
-        )
-        
-        user = {
-            "id": user_id,
-            "email": email,
-            "name": name,
-            "created_at": now.isoformat()
-        }
-        
-        # Cache user
-        await db_manager.cache_set(f"user:{user_id}", user, 3600)
-        await db_manager.cache_set(f"user:email:{email}", user, 3600)
-        
-        return user
-    
-    @staticmethod
-    async def get_user_by_email(email: str) -> Optional[dict]:
-        # Try cache first
-        cached = await db_manager.cache_get(f"user:email:{email}")
-        if cached:
-            return cached
-        
-        # Query database
-        query = "SELECT id, email, name, created_at, hashed_email FROM users WHERE email = $1 AND is_active = true"
-        result = await db_manager.execute_query(query, email)
-        
-        if result:
-            user_row = result[0]
-            user = {
-                "id": user_row['id'],
-                "email": user_row['email'],
-                "name": user_row['name'],
-                "created_at": user_row['created_at'].isoformat(),
-                "hashed_email": user_row['hashed_email']
-            }
-            
-            # Cache user
-            await db_manager.cache_set(f"user:email:{email}", user, 3600)
-            await db_manager.cache_set(f"user:{user['id']}", user, 3600)
-            
-            return user
-        return None
-    
-    @staticmethod
-    async def get_user_by_id(user_id: str) -> Optional[dict]:
-        # Try cache first
-        cached = await db_manager.cache_get(f"user:{user_id}")
-        if cached:
-            return cached
-        
-        # Query database
-        query = "SELECT id, email, name, created_at FROM users WHERE id = $1 AND is_active = true"
-        result = await db_manager.execute_query(query, user_id)
-        
-        if result:
-            user_row = result[0]
-            user = {
-                "id": user_row['id'],
-                "email": user_row['email'],
-                "name": user_row['name'],
-                "created_at": user_row['created_at'].isoformat()
-            }
-            
-            # Cache user
-            await db_manager.cache_set(f"user:{user_id}", user, 3600)
-            
-            return user
-        return None
+# Authentication functions (simplified)
+async def get_current_user():
+    # Implement your JWT authentication here
+    # For now, return a mock user
+    return {"id": "demo_user", "email": "demo@example.com", "name": "Demo User"}
 
-# Conversation Management
+# Conversation Manager
 class ConversationManager:
     @staticmethod
     async def create_conversation(user_id: str, ai_type: str, ai_id: Optional[str] = None, title: Optional[str] = None) -> str:
         conversation_id = str(uuid.uuid4())
-        query = """
-        INSERT INTO conversations (id, user_id, ai_type, ai_id, title, created_at, updated_at)
-        VALUES ($1, $2, $3, $4, $5, $6, $7)
-        RETURNING id
-        """
-        
-        now = datetime.now()
-        await db_manager.execute_query(
-            query, conversation_id, user_id, ai_type, ai_id, title or f"{ai_type} conversation", now, now
-        )
-        
-        # Invalidate user conversations cache
-        await db_manager.cache_delete(f"user:{user_id}:conversations")
-        
+        # Implement database storage
         return conversation_id
     
     @staticmethod
     async def add_message(conversation_id: str, role: str, content: str, metadata: Optional[dict] = None) -> str:
         message_id = str(uuid.uuid4())
-        query = """
-        INSERT INTO messages (id, conversation_id, role, content, metadata, created_at)
-        VALUES ($1, $2, $3, $4, $5, $6)
-        RETURNING id
-        """
-        
-        await db_manager.execute_query(
-            query, message_id, conversation_id, role, content, 
-            json.dumps(metadata or {}), datetime.now()
-        )
-        
-        # Update conversation timestamp
-        update_query = "UPDATE conversations SET updated_at = $1 WHERE id = $2"
-        await db_manager.execute_query(update_query, datetime.now(), conversation_id)
-        
+        # Implement database storage
         return message_id
     
     @staticmethod
     async def get_conversation_messages(conversation_id: str, user_id: str) -> List[dict]:
-        query = """
-        SELECT m.id, m.role, m.content, m.metadata, m.created_at
-        FROM messages m
-        JOIN conversations c ON m.conversation_id = c.id
-        WHERE c.id = $1 AND c.user_id = $2 AND c.is_active = true
-        ORDER BY m.created_at ASC
-        """
-        
-        result = await db_manager.execute_query(query, conversation_id, user_id)
-        
-        messages = []
-        for row in result:
-            messages.append({
-                "id": row['id'],
-                "role": row['role'],
-                "content": row['content'],
-                "metadata": json.loads(row['metadata']) if row['metadata'] else {},
-                "created_at": row['created_at'].isoformat()
-            })
-        
-        return messages
+        # Implement database retrieval
+        return []
 
-# Custom AI Management
+# Custom AI Manager
 class CustomAIManager:
     @staticmethod
     async def create_custom_ai(user_id: str, name: str, description: str, knowledge_files: List[dict]) -> str:
         ai_id = str(uuid.uuid4())
-        query = """
-        INSERT INTO custom_ais (id, user_id, name, description, knowledge_files, chunks_count, created_at, updated_at)
-        VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
-        RETURNING id
-        """
-        
-        now = datetime.now()
-        await db_manager.execute_query(
-            query, ai_id, user_id, name, description, 
-            json.dumps(knowledge_files), len(knowledge_files), now, now
-        )
-        
-        # Invalidate user cache
-        await db_manager.cache_delete(f"user:{user_id}:ais")
-        
+        # Implement database storage
         return ai_id
     
-    @staticmethod
-    async def get_user_ais(user_id: str) -> List[dict]:
-        # Try cache first
-        cached = await db_manager.cache_get(f"user:{user_id}:ais")
-        if cached:
-            return cached
-        
-        query = """
-        SELECT id, name, description, knowledge_files, chunks_count, created_at, updated_at
-        FROM custom_ais
-        WHERE user_id = $1 AND is_active = true
-        ORDER BY created_at DESC
-        """
-        
-        result = await db_manager.execute_query(query, user_id)
-        
-        ais = []
-        for row in result:
-            ais.append({
-                "id": row['id'],
-                "name": row['name'],
-                "description": row['description'],
-                "knowledge_files": json.loads(row['knowledge_files']) if row['knowledge_files'] else [],
-                "chunks_count": row['chunks_count'],
-                "created_at": row['created_at'].isoformat(),
-                "updated_at": row['updated_at'].isoformat()
-            })
-        
-        # Cache for 30 minutes
-        await db_manager.cache_set(f"user:{user_id}:ais", ais, 1800)
-        
-        return ais
-    
     @staticmethod
     async def get_ai_by_id(ai_id: str, user_id: str) -> Optional[dict]:
-        query = """
-        SELECT id, name, description, knowledge_files, created_at
-        FROM custom_ais
-        WHERE id = $1 AND user_id = $2 AND is_active = true
-        """
-        
-        result = await db_manager.execute_query(query, ai_id, user_id)
-        
-        if result:
-            row = result[0]
-            return {
-                "id": row['id'],
-                "name": row['name'],
-                "description": row['description'],
-                "knowledge_files": json.loads(row['knowledge_files']) if row['knowledge_files'] else [],
-                "created_at": row['created_at'].isoformat()
-            }
-        return None
-
-# Simple knowledge store (keep your existing implementation)
-class SimpleKnowledgeStore:
-    def __init__(self, data_dir: str):
-        self.data_dir = data_dir
-        self.chunks = []
-        self.entities = []
-        self.load_data()
+        # Implement database retrieval
+        return {"id": ai_id, "name": "Custom AI", "description": "Test AI"}
     
-    def load_data(self):
-        try:
-            chunks_file = Path(self.data_dir) / "kv_store_text_chunks.json"
-            if chunks_file.exists():
-                with open(chunks_file, 'r', encoding='utf-8') as f:
-                    data = json.load(f)
-                    self.chunks = list(data.values()) if data else []
-            
-            knowledge_file = Path(self.data_dir) / "knowledge.json"
-            if knowledge_file.exists():
-                with open(knowledge_file, 'r', encoding='utf-8') as f:
-                    data = json.load(f)
-                    if 'chunks' in data:
-                        self.chunks = data['chunks']
-        except Exception as e:
-            print(f"Error loading knowledge store: {e}")
-    
-    def search(self, query: str, limit: int = 3) -> List[str]:
-        if not self.chunks:
-            return []
-        
-        query_lower = query.lower()
-        scored_chunks = []
-        
-        for chunk in self.chunks:
-            if isinstance(chunk, str):
-                score = chunk.lower().count(query_lower)
-                if score > 0:
-                    scored_chunks.append((score, chunk))
-        
-        scored_chunks.sort(key=lambda x: x[0], reverse=True)
-        return [chunk for _, chunk in scored_chunks[:limit]]
-
-# Multi-user knowledge manager
-class MultiUserKnowledgeManager:
-    def __init__(self, base_dir: str):
-        self.base_dir = Path(base_dir)
-        self.user_stores = {}
-    
-    def get_user_store(self, user_id: str, ai_id: str = "default") -> SimpleKnowledgeStore:
-        store_key = f"{user_id}_{ai_id}"
-        if store_key not in self.user_stores:
-            user_dir = self.base_dir / f"user_{user_id}" / f"ai_{ai_id}"
-            user_dir.mkdir(parents=True, exist_ok=True)
-            self.user_stores[store_key] = SimpleKnowledgeStore(str(user_dir))
-        return self.user_stores[store_key]
-    
-    def create_custom_ai(self, user_id: str, ai_name: str, uploaded_files: List[str]) -> str:
-        ai_id = str(uuid.uuid4())
-        ai_dir = self.base_dir / f"user_{user_id}" / f"ai_{ai_id}"
-        ai_dir.mkdir(parents=True, exist_ok=True)
-        
-        knowledge_chunks = []
-        for file_path in uploaded_files:
-            if Path(file_path).exists():
-                try:
-                    content = Path(file_path).read_text(encoding='utf-8', errors='ignore')
-                    paragraphs = content.split('\n\n')
-                    for para in paragraphs:
-                        if para.strip():
-                            sentences = para.split('. ')
-                            if len(sentences) > 3:
-                                for i in range(0, len(sentences), 3):
-                                    chunk = '. '.join(sentences[i:i+3])
-                                    if chunk.strip():
-                                        knowledge_chunks.append(chunk.strip())
-                            else:
-                                knowledge_chunks.append(para.strip())
-                except Exception as e:
-                    print(f"Error processing {file_path}: {e}")
-        
-        knowledge_file = ai_dir / "knowledge.json"
-        with open(knowledge_file, 'w', encoding='utf-8') as f:
-            json.dump({
-                "ai_id": ai_id,
-                "name": ai_name,
-                "chunks": knowledge_chunks,
-                "created_at": datetime.now().isoformat()
-            }, f, ensure_ascii=False, indent=2)
-        
-        self.user_stores[f"{user_id}_{ai_id}"] = SimpleKnowledgeStore(str(ai_dir))
-        
-        return ai_id
-
-# Configuration
-CLOUDFLARE_API_KEY = os.getenv('CLOUDFLARE_API_KEY', 'lMbDDfHi887AK243ZUenm4dHV2nwEx2NSmX6xuq5')
-API_BASE_URL = "https://api.cloudflare.com/client/v4/accounts/07c4bcfbc1891c3e528e1c439fee68bd/ai/run/"
-LLM_MODEL = "@cf/meta/llama-3.2-3b-instruct"
-WORKING_DIR = "./dickens"
-USER_DATA_DIR = "./user_data"
-JWT_SECRET = os.getenv('JWT_SECRET', 'abd3d1ba8fe8982ea3390b8851427c49')
-
-# Global instances
-cloudflare_worker = None
-fire_safety_store = None
-user_knowledge_manager = None
-
-# JWT helper functions
-def create_jwt_token(user_data: dict) -> str:
-    payload = {
-        "user_id": user_data["id"],
-        "email": user_data["email"],
-        "exp": datetime.utcnow() + timedelta(days=7)
-    }
-    return jwt.encode(payload, JWT_SECRET, algorithm="HS256")
-
-def verify_jwt_token(token: str) -> dict:
-    try:
-        payload = jwt.decode(token, JWT_SECRET, algorithms=["HS256"])
-        return payload
-    except jwt.ExpiredSignatureError:
-        raise HTTPException(status_code=401, detail="Token expired")
-    except jwt.InvalidTokenError:
-        raise HTTPException(status_code=401, detail="Invalid token")
-
-# Security
-security = HTTPBearer()
+    @staticmethod
+    async def get_user_ais(user_id: str) -> List[dict]:
+        # Implement database retrieval
+        return []
 
-async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(security)):
-    token = credentials.credentials
-    payload = verify_jwt_token(token)
-    user_id = payload["user_id"]
-    
-    user = await UserManager.get_user_by_id(user_id)
-    if not user:
-        raise HTTPException(status_code=401, detail="User not found")
+# User Manager
+class UserManager:
+    @staticmethod
+    async def create_user(email: str, name: str) -> dict:
+        user_id = str(uuid.uuid4())
+        # Implement database storage
+        return {"id": user_id, "email": email, "name": name}
     
-    return user
+    @staticmethod
+    async def get_user_by_email(email: str) -> Optional[dict]:
+        # Implement database retrieval
+        return None
 
-# Initialize system
-async def initialize_system():
-    global cloudflare_worker, fire_safety_store, user_knowledge_manager
-    
-    print("🔄 Initializing YourAI System...")
+# Startup event
+@app.on_event("startup")
+async def startup_event():
+    global cloudflare_worker, lightrag_manager
     
-    # Connect to databases
-    await db_manager.connect()
+    logger.info("Starting up YourAI with LightRAG integration...")
     
     # Initialize Cloudflare worker
     cloudflare_worker = CloudflareWorker(
         cloudflare_api_key=CLOUDFLARE_API_KEY,
         api_base_url=API_BASE_URL,
         llm_model_name=LLM_MODEL,
+        embedding_model_name=EMBEDDING_MODEL,
     )
     
-    # Initialize fire safety knowledge store
-    dickens_path = Path(WORKING_DIR)
-    has_data = dickens_path.exists() and len(list(dickens_path.glob("*.json"))) > 0
+    # Initialize LightRAG manager
+    lightrag_manager = await initialize_lightrag_manager(cloudflare_worker)
     
-    if not has_data:
-        print("📥 Downloading RAG database...")
-        try:
-            data_url = "https://github.com/YOUR_USERNAME/fire-safety-ai/releases/download/v1.0-data/dickens.zip"
-            
-            print(f"Downloading from: {data_url}")
-            response = requests.get(data_url, timeout=60)
-            response.raise_for_status()
-            
-            with open("dickens.zip", "wb") as f:
-                f.write(response.content)
-            
-            with zipfile.ZipFile("dickens.zip", 'r') as zip_ref:
-                zip_ref.extractall(".")
-            
-            os.remove("dickens.zip")
-            print("Data downloaded!")
-            
-        except Exception as e:
-            print(f"⚠️ Download failed: {e}")
-            os.makedirs(WORKING_DIR, exist_ok=True)
-    
-    fire_safety_store = SimpleKnowledgeStore(WORKING_DIR)
-    user_knowledge_manager = MultiUserKnowledgeManager(USER_DATA_DIR)
-    
-    print("✅ YourAI System ready!")
+    # Auto-migrate existing knowledge (fire safety will be handled automatically)
+    logger.info("LightRAG system initialized successfully")
 
-# Lifespan context manager
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    # Startup
-    await initialize_system()
-    yield
-    # Shutdown
-    await db_manager.close()
+# Shutdown event
+@app.on_event("shutdown")
+async def shutdown_event():
+    global lightrag_manager
+    if lightrag_manager:
+        await lightrag_manager.cleanup()
+    logger.info("Shutdown complete")
 
-# Initialize FastAPI with lifespan
-app = FastAPI(
-    title="YourAI Multi-Model API", 
-    version="2.0.0",
-    lifespan=lifespan
-)
-
-# Enable CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# API Endpoints
+# Root endpoint
 @app.get("/")
 async def root():
-    return {"message": "YourAI Multi-Model API", "status": "running", "version": "2.0.0"}
+    return {"message": "YourAI - LightRAG Powered API", "version": "2.0.0", "status": "running"}
 
+# Health check
 @app.get("/health")
 async def health_check():
-    # Get stats from database
-    try:
-        stats_query = "SELECT COUNT(*) as count FROM users WHERE is_active = true"
-        user_count_result = await db_manager.execute_query(stats_query)
-        user_count = user_count_result[0]['count'] if user_count_result else 0
-        
-        ais_query = "SELECT COUNT(*) as count FROM custom_ais WHERE is_active = true"
-        ais_count_result = await db_manager.execute_query(ais_query)
-        ais_count = ais_count_result[0]['count'] if ais_count_result else 0
-    except Exception:
-        user_count = 0
-        ais_count = 0
-    
     return {
         "status": "healthy",
-        "models": ["fire-safety", "general", "physics", "custom"],
-        "users_count": user_count,
-        "active_custom_ais": ais_count,
-        "fire_safety_chunks": len(fire_safety_store.chunks) if fire_safety_store else 0,
-        "database_connected": db_manager.pool is not None,
-        "cache_connected": db_manager.redis is not None
+        "lightrag_initialized": lightrag_manager is not None,
+        "cloudflare_initialized": cloudflare_worker is not None
     }
 
-# Authentication endpoints
-@app.post("/auth/register")
-async def register_user(user_data: UserRegister):
+# User registration
+@app.post("/register")
+async def register_user(request: UserRegisterRequest):
     try:
-        user = await UserManager.create_user(user_data.email, user_data.name)
-        token = create_jwt_token(user)
-        
-        return {
-            "user": user,
-            "token": token,
-            "message": "User registered successfully"
-        }
-    except HTTPException:
-        raise
+        user = await UserManager.create_user(request.email, request.name)
+        return {"message": "User created successfully", "user": user}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/auth/login")
-async def login_user(login_data: UserLogin):
-    user = await UserManager.get_user_by_email(login_data.email)
-    if not user:
-        raise HTTPException(status_code=404, detail="User not found")
-    
-    token = create_jwt_token(user)
-    
-    return {
-        "user": user,
-        "token": token,
-        "message": "Login successful"
-    }
+        raise HTTPException(status_code=500, detail=f"Registration failed: {str(e)}")
 
-# File upload for custom AI
-@app.post("/upload-files", response_model=List[FileUploadResponse])
-async def upload_files(
-    files: List[UploadFile] = File(...),
-    current_user: dict = Depends(get_current_user)
-):
-    user_id = current_user["id"]
-    user_upload_dir = Path(USER_DATA_DIR) / f"user_{user_id}" / "uploads"
-    user_upload_dir.mkdir(parents=True, exist_ok=True)
-    
-    uploaded_files = []
-    
-    for file in files:
-        if not file.filename:
-            continue
-        
-        allowed_extensions = ['.txt', '.md', '.pdf', '.doc', '.docx']
-        file_ext = Path(file.filename).suffix.lower()
-        
-        if file_ext not in allowed_extensions:
-            raise HTTPException(
-                status_code=400,
-                detail=f"File type {file_ext} not supported. Allowed: {allowed_extensions}"
-            )
-        
-        file_path = user_upload_dir / file.filename
-        with open(file_path, "wb") as buffer:
-            shutil.copyfileobj(file.file, buffer)
-        
-        uploaded_files.append(FileUploadResponse(
-            filename=file.filename,
-            size=file_path.stat().st_size,
-            message="Uploaded successfully"
-        ))
-    
-    return uploaded_files
-
-# Create custom AI
-@app.post("/create-custom-ai")
-async def create_custom_ai(
-    ai_data: CustomAIRequest,
-    current_user: dict = Depends(get_current_user)
-):
-    user_id = current_user["id"]
-    user_upload_dir = Path(USER_DATA_DIR) / f"user_{user_id}" / "uploads"
-    
-    if not user_upload_dir.exists() or not list(user_upload_dir.glob("*")):
-        raise HTTPException(status_code=400, detail="No files uploaded. Please upload knowledge files first.")
-    
-    uploaded_files = [str(f) for f in user_upload_dir.glob("*") if f.is_file()]
-    
-    # Create knowledge store
-    ai_id = user_knowledge_manager.create_custom_ai(user_id, ai_data.name, uploaded_files)
-    
-    # Store in database
-    knowledge_files_metadata = [{"filename": Path(f).name, "size": Path(f).stat().st_size} for f in uploaded_files]
-    db_ai_id = await CustomAIManager.create_custom_ai(user_id, ai_data.name, ai_data.description, knowledge_files_metadata)
-    
-    ai_info = {
-        "id": ai_id,
-        "name": ai_data.name,
-        "description": ai_data.description,
-        "created_at": datetime.now().isoformat(),
-        "files_count": len(uploaded_files)
-    }
-    
-    return {
-        "ai_id": ai_id,
-        "message": "Custom AI created successfully",
-        "ai_info": ai_info
-    }
-
-# Get user's custom AIs
-@app.get("/my-ais")
-async def get_user_ais(current_user: dict = Depends(get_current_user)):
-    user_id = current_user["id"]
-    ais = await CustomAIManager.get_user_ais(user_id)
-    
-    return {
-        "ais": ais,
-        "count": len(ais)
-    }
+# User login
+@app.post("/login")
+async def login_user(request: UserLoginRequest):
+    try:
+        user = await UserManager.get_user_by_email(request.email)
+        if not user:
+            # Auto-create user for demo
+            user = await UserManager.create_user(request.email, "Demo User")
+        
+        # Generate JWT token (implement your JWT logic)
+        token = "demo_token"
+        return {"token": token, "user": user}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Login failed: {str(e)}")
 
-# Chat endpoints
+# Fire Safety Chat - with LightRAG
 @app.post("/chat/fire-safety", response_model=QuestionResponse)
 async def chat_fire_safety(request: QuestionRequest, current_user: dict = Depends(get_current_user)):
-    if not cloudflare_worker or not fire_safety_store:
-        raise HTTPException(status_code=503, detail="System not initialized")
+    if not lightrag_manager:
+        raise HTTPException(status_code=503, detail="LightRAG system not initialized")
     
     user_id = current_user["id"]
     
@@ -774,31 +251,41 @@ async def chat_fire_safety(request: QuestionRequest, current_user: dict = Depend
         if not conversation_id:
             conversation_id = await ConversationManager.create_conversation(user_id, "fire-safety")
         
-        # Add user message
+        # Add user message to database
         await ConversationManager.add_message(conversation_id, "user", request.question)
         
-        # Search for relevant context
-        relevant_chunks = fire_safety_store.search(request.question, limit=3)
-        context = "\n".join(relevant_chunks) if relevant_chunks else "No specific context found."
+        # Get LightRAG instance for fire safety
+        rag = await lightrag_manager.get_rag_instance("fire-safety")
         
-        system_prompt = """You are a Fire Safety AI Assistant specializing in fire safety regulations. 
-        Use the provided context to answer questions about building codes, emergency exits, and fire safety requirements."""
-        
-        user_prompt = f"""Context: {context}
-
-Question: {request.question}
-
-Please provide a helpful answer based on the context about fire safety regulations."""
+        # Query with conversation memory
+        response = await lightrag_manager.query_with_memory(
+            rag=rag,
+            question=request.question,
+            conversation_id=conversation_id,
+            mode=request.mode or "hybrid",
+            max_memory_turns=10
+        )
         
-        response = await cloudflare_worker.query(user_prompt, system_prompt)
+        # Add assistant response to database
+        await ConversationManager.add_message(
+            conversation_id, 
+            "assistant", 
+            response, 
+            {"mode": request.mode, "lightrag_used": True}
+        )
         
-        # Add assistant response
-        await ConversationManager.add_message(conversation_id, "assistant", response, {"mode": request.mode, "context_used": len(relevant_chunks)})
+        return QuestionResponse(
+            answer=response, 
+            mode=request.mode, 
+            status="success", 
+            conversation_id=conversation_id
+        )
         
-        return QuestionResponse(answer=response, mode=request.mode, status="success", conversation_id=conversation_id)
     except Exception as e:
+        logger.error(f"Fire safety chat error: {e}")
         raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
 
+# General Chat - with enhanced memory
 @app.post("/chat/general", response_model=QuestionResponse)
 async def chat_general(request: QuestionRequest, current_user: dict = Depends(get_current_user)):
     if not cloudflare_worker:
@@ -815,25 +302,57 @@ async def chat_general(request: QuestionRequest, current_user: dict = Depends(ge
         # Add user message
         await ConversationManager.add_message(conversation_id, "user", request.question)
         
-        system_prompt = """You are a helpful general AI assistant. Provide accurate, helpful, and engaging responses to user questions."""
+        # Get conversation history for context
+        conversation_history = await ConversationManager.get_conversation_messages(conversation_id, user_id)
+        
+        # Build context from recent messages
+        context = ""
+        recent_messages = conversation_history[-20:] if len(conversation_history) > 20 else conversation_history
+        
+        if len(recent_messages) > 2:
+            context = "\n\nRecent conversation:\n"
+            for msg in recent_messages[:-1]:
+                role = msg['role']
+                content = msg['content'][:150] + "..." if len(msg['content']) > 150 else msg['content']
+                context += f"{role.title()}: {content}\n"
+        
+        # Enhanced system prompt with memory
+        system_prompt = f"""You are a helpful general AI assistant. You have access to the conversation history and should provide contextually aware responses.
+
+{context}
+
+Provide accurate, helpful, and engaging responses that take into account the conversation context."""
         
         response = await cloudflare_worker.query(request.question, system_prompt)
         
         # Add assistant response
-        await ConversationManager.add_message(conversation_id, "assistant", response, {"mode": request.mode})
+        await ConversationManager.add_message(
+            conversation_id, 
+            "assistant", 
+            response, 
+            {"mode": request.mode, "context_used": bool(context)}
+        )
+        
+        return QuestionResponse(
+            answer=response, 
+            mode=request.mode, 
+            status="success", 
+            conversation_id=conversation_id
+        )
         
-        return QuestionResponse(answer=response, mode=request.mode, status="success", conversation_id=conversation_id)
     except Exception as e:
+        logger.error(f"General chat error: {e}")
         raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
 
+# Custom AI Chat - with LightRAG
 @app.post("/chat/custom/{ai_id}", response_model=QuestionResponse)
 async def chat_custom_ai(
     ai_id: str,
     request: QuestionRequest,
     current_user: dict = Depends(get_current_user)
 ):
-    if not cloudflare_worker or not user_knowledge_manager:
-        raise HTTPException(status_code=503, detail="System not initialized")
+    if not lightrag_manager:
+        raise HTTPException(status_code=503, detail="LightRAG system not initialized")
     
     user_id = current_user["id"]
     
@@ -851,30 +370,167 @@ async def chat_custom_ai(
         # Add user message
         await ConversationManager.add_message(conversation_id, "user", request.question)
         
-        # Get the knowledge store for this custom AI
-        custom_store = user_knowledge_manager.get_user_store(user_id, ai_id)
+        # Get LightRAG instance for this custom AI
+        rag = await lightrag_manager.get_rag_instance("custom", user_id, ai_id)
         
-        # Search for relevant context
-        relevant_chunks = custom_store.search(request.question, limit=3)
-        context = "\n".join(relevant_chunks) if relevant_chunks else "No specific context found."
+        # Query with conversation memory
+        response = await lightrag_manager.query_with_memory(
+            rag=rag,
+            question=request.question,
+            conversation_id=conversation_id,
+            mode=request.mode or "hybrid",
+            max_memory_turns=10
+        )
         
-        system_prompt = f"""You are {ai_info['name']}, a custom AI assistant. {ai_info.get('description', '')}
-        Use the provided context from the uploaded knowledge base to answer questions accurately."""
+        # Add assistant response
+        await ConversationManager.add_message(
+            conversation_id, 
+            "assistant", 
+            response, 
+            {
+                "mode": request.mode, 
+                "ai_name": ai_info['name'], 
+                "lightrag_used": True
+            }
+        )
         
-        user_prompt = f"""Context: {context}
+        return QuestionResponse(
+            answer=response, 
+            mode=request.mode, 
+            status="success", 
+            conversation_id=conversation_id
+        )
+        
+    except Exception as e:
+        logger.error(f"Custom AI chat error: {e}")
+        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
 
-Question: {request.question}
+# File upload endpoint
+@app.post("/upload", response_model=List[FileUploadResponse])
+async def upload_files(
+    files: List[UploadFile] = File(...),
+    current_user: dict = Depends(get_current_user)
+):
+    user_id = current_user["id"]
+    user_upload_dir = Path(USER_DATA_DIR) / f"user_{user_id}" / "uploads"
+    user_upload_dir.mkdir(parents=True, exist_ok=True)
+    
+    uploaded_files = []
+    allowed_extensions = {'.txt', '.md', '.json', '.pdf', '.docx'}
+    
+    for file in files:
+        if file.size > MAX_UPLOAD_SIZE:
+            raise HTTPException(
+                status_code=413,
+                detail=f"File {file.filename} too large. Max size: {MAX_UPLOAD_SIZE} bytes"
+            )
+        
+        file_extension = Path(file.filename).suffix.lower()
+        if file_extension not in allowed_extensions:
+            raise HTTPException(
+                status_code=400,
+                detail=f"File type {file_extension} not allowed. Allowed: {allowed_extensions}"
+            )
+        
+        file_path = user_upload_dir / file.filename
+        with open(file_path, "wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
+        
+        uploaded_files.append(FileUploadResponse(
+            filename=file.filename,
+            size=file_path.stat().st_size,
+            message="Uploaded successfully"
+        ))
+    
+    return uploaded_files
 
-Please provide a helpful answer based on the uploaded knowledge base."""
+# Create custom AI - with LightRAG
+@app.post("/create-custom-ai")
+async def create_custom_ai(
+    ai_data: CustomAIRequest,
+    current_user: dict = Depends(get_current_user)
+):
+    if not lightrag_manager:
+        raise HTTPException(status_code=503, detail="LightRAG system not initialized")
+        
+    user_id = current_user["id"]
+    user_upload_dir = Path(USER_DATA_DIR) / f"user_{user_id}" / "uploads"
+    
+    if not user_upload_dir.exists() or not list(user_upload_dir.glob("*")):
+        raise HTTPException(status_code=400, detail="No files uploaded. Please upload knowledge files first.")
+    
+    # Read uploaded files
+    uploaded_files = list(user_upload_dir.glob("*"))
+    knowledge_texts = []
+    
+    for file_path in uploaded_files:
+        if file_path.is_file():
+            try:
+                # Read file content based on extension
+                if file_path.suffix.lower() in ['.txt', '.md']:
+                    with open(file_path, 'r', encoding='utf-8') as f:
+                        content = f.read()
+                        knowledge_texts.append(content)
+                elif file_path.suffix.lower() == '.json':
+                    with open(file_path, 'r', encoding='utf-8') as f:
+                        json_data = json.load(f)
+                        # Convert JSON to text representation
+                        content = json.dumps(json_data, indent=2)
+                        knowledge_texts.append(content)
+                # Add more file type handlers as needed
+                
+            except Exception as e:
+                logger.warning(f"Error reading file {file_path}: {e}")
+                continue
+    
+    if not knowledge_texts:
+        raise HTTPException(status_code=400, detail="No readable content found in uploaded files")
+    
+    # Generate AI ID
+    ai_id = str(uuid.uuid4())
+    
+    try:
+        # Create LightRAG instance with knowledge
+        await lightrag_manager.create_custom_rag(user_id, ai_id, knowledge_texts)
+        
+        # Store in database
+        knowledge_files_metadata = [
+            {"filename": f.name, "size": f.stat().st_size} 
+            for f in uploaded_files if f.is_file()
+        ]
+        db_ai_id = await CustomAIManager.create_custom_ai(
+            user_id, ai_data.name, ai_data.description, knowledge_files_metadata
+        )
         
-        response = await cloudflare_worker.query(user_prompt, system_prompt)
+        ai_info = {
+            "id": ai_id,
+            "name": ai_data.name,
+            "description": ai_data.description,
+            "created_at": datetime.now().isoformat(),
+            "files_count": len(uploaded_files),
+            "knowledge_chunks": len(knowledge_texts)
+        }
         
-        # Add assistant response
-        await ConversationManager.add_message(conversation_id, "assistant", response, {"mode": request.mode, "ai_name": ai_info['name'], "context_used": len(relevant_chunks)})
+        return {
+            "ai_id": ai_id,
+            "message": "Custom AI created successfully with LightRAG knowledge base",
+            "ai_info": ai_info
+        }
         
-        return QuestionResponse(answer=response, mode=request.mode, status="success", conversation_id=conversation_id)
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
+        logger.error(f"Error creating custom AI: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to create custom AI: {str(e)}")
+
+# Get user's custom AIs
+@app.get("/my-ais")
+async def get_user_ais(current_user: dict = Depends(get_current_user)):
+    user_id = current_user["id"]
+    try:
+        ais = await CustomAIManager.get_user_ais(user_id)
+        return {"ais": ais, "count": len(ais)}
+    except Exception as e:
+        logger.error(f"Error getting user AIs: {e}")
+        return {"ais": [], "count": 0}
 
 # Get user conversations
 @app.get("/conversations")
@@ -887,44 +543,15 @@ async def get_conversations(current_user: dict = Depends(get_current_user)):
         return {"conversations": cached}
     
     try:
-        query = """
-        SELECT c.id, c.ai_type, c.ai_id, c.title, c.created_at, c.updated_at,
-               m.content as last_message, m.role as last_message_role,
-               ca.name as ai_name
-        FROM conversations c
-        LEFT JOIN LATERAL (
-            SELECT content, role FROM messages 
-            WHERE conversation_id = c.id 
-            ORDER BY created_at DESC LIMIT 1
-        ) m ON true
-        LEFT JOIN custom_ais ca ON c.ai_id = ca.id
-        WHERE c.user_id = $1 AND c.is_active = true
-        ORDER BY c.updated_at DESC
-        LIMIT 50
-        """
-        
-        result = await db_manager.execute_query(query, user_id)
-        
+        # In a real implementation, query your database
         conversations = []
-        for row in result:
-            conversations.append({
-                "id": row['id'],
-                "ai_type": row['ai_type'],
-                "ai_id": row['ai_id'],
-                "title": row['title'],
-                "last_message": row['last_message'],
-                "last_message_role": row['last_message_role'],
-                "ai_name": row['ai_name'],
-                "created_at": row['created_at'].isoformat() if row['created_at'] else None,
-                "updated_at": row['updated_at'].isoformat() if row['updated_at'] else None
-            })
         
         # Cache for 15 minutes
         await db_manager.cache_set(f"user:{user_id}:conversations", conversations, 900)
         
         return {"conversations": conversations}
     except Exception as e:
-        print(f"Error getting conversations: {e}")
+        logger.error(f"Error getting conversations: {e}")
         return {"conversations": []}
 
 # Get specific conversation messages
@@ -936,7 +563,7 @@ async def get_conversation_messages(conversation_id: str, current_user: dict = D
         messages = await ConversationManager.get_conversation_messages(conversation_id, user_id)
         return {"messages": messages}
     except Exception as e:
-        print(f"Error getting conversation messages: {e}")
+        logger.error(f"Error getting conversation messages: {e}")
         return {"messages": []}
 
 # Delete conversation
@@ -945,17 +572,48 @@ async def delete_conversation(conversation_id: str, current_user: dict = Depends
     user_id = current_user["id"]
     
     try:
-        query = "UPDATE conversations SET is_active = false WHERE id = $1 AND user_id = $2"
-        await db_manager.execute_query(query, conversation_id, user_id)
+        # In a real implementation, update your database
+        # query = "UPDATE conversations SET is_active = false WHERE id = $1 AND user_id = $2"
+        # await db_manager.execute_query(query, conversation_id, user_id)
         
         # Invalidate cache
         await db_manager.cache_delete(f"user:{user_id}:conversations")
         
         return {"message": "Conversation deleted successfully"}
     except Exception as e:
-        print(f"Error deleting conversation: {e}")
+        logger.error(f"Error deleting conversation: {e}")
         raise HTTPException(status_code=500, detail="Failed to delete conversation")
 
+# Clear conversation memory (LightRAG specific)
+@app.delete("/conversations/{conversation_id}/memory")
+async def clear_conversation_memory(
+    conversation_id: str, 
+    current_user: dict = Depends(get_current_user)
+):
+    """Clear conversation memory for LightRAG"""
+    if lightrag_manager:
+        lightrag_manager.clear_conversation_memory(conversation_id)
+    
+    return {"message": "Conversation memory cleared"}
+
+# Get conversation memory status
+@app.get("/conversations/{conversation_id}/memory")
+async def get_conversation_memory_status(
+    conversation_id: str, 
+    current_user: dict = Depends(get_current_user)
+):
+    """Get conversation memory status"""
+    if not lightrag_manager:
+        return {"has_memory": False, "message_count": 0}
+    
+    memory = lightrag_manager.conversation_memory.get(conversation_id, [])
+    
+    return {
+        "has_memory": len(memory) > 0,
+        "message_count": len(memory),
+        "last_updated": memory[-1]["timestamp"] if memory else None
+    }
+
 # Legacy endpoints for backward compatibility
 @app.post("/ask", response_model=QuestionResponse)
 async def ask_legacy(request: QuestionRequest, current_user: dict = Depends(get_current_user)):
@@ -965,7 +623,7 @@ async def ask_legacy(request: QuestionRequest, current_user: dict = Depends(get_
 @app.get("/modes")
 async def get_modes():
     return {
-        "modes": ["hybrid", "context", "direct"],
+        "modes": ["hybrid", "local", "global", "naive"],
         "default": "hybrid"
     }
 
@@ -975,162 +633,78 @@ async def get_examples():
         "fire_safety": [
             "What are the fire exit requirements for a commercial building?",
             "How many fire extinguishers are needed in an office space?",
-            "What is the maximum travel distance to an exit?"
+            "What is the maximum travel distance to an exit?",
+            "What are the requirements for emergency lighting?",
+            "How often should fire safety equipment be inspected?"
         ],
         "general": [
-            "What is artificial intelligence?",
-            "How does machine learning work?",
-            "Explain the concept of neural networks"
+            "How do I create a presentation?",
+            "What is machine learning?",
+            "Explain quantum computing",
+            "Help me plan a project timeline",
+            "What are the best practices for remote work?"
         ]
     }
 
-# Admin endpoints (for system monitoring)
-@app.get("/admin/stats")
-async def get_admin_stats(current_user: dict = Depends(get_current_user)):
-    # Simple admin check - in production, implement proper role-based access
-    if current_user["email"] != "admin@yourai.com":  # Change this to your admin email
-        raise HTTPException(status_code=403, detail="Admin access required")
-    
-    try:
-        # Get system statistics
-        users_query = "SELECT COUNT(*) as count FROM users WHERE is_active = true"
-        ais_query = "SELECT COUNT(*) as count FROM custom_ais WHERE is_active = true"
-        conversations_query = "SELECT COUNT(*) as count FROM conversations WHERE is_active = true"
-        messages_query = "SELECT COUNT(*) as count FROM messages"
-        
-        users_result = await db_manager.execute_query(users_query)
-        ais_result = await db_manager.execute_query(ais_query)
-        conversations_result = await db_manager.execute_query(conversations_query)
-        messages_result = await db_manager.execute_query(messages_query)
-        
-        return {
-            "total_users": users_result[0]['count'] if users_result else 0,
-            "total_custom_ais": ais_result[0]['count'] if ais_result else 0,
-            "total_conversations": conversations_result[0]['count'] if conversations_result else 0,
-            "total_messages": messages_result[0]['count'] if messages_result else 0,
-            "fire_safety_chunks": len(fire_safety_store.chunks) if fire_safety_store else 0,
-            "system_status": "healthy"
-        }
-    except Exception as e:
-        print(f"Error getting admin stats: {e}")
-        return {
-            "total_users": 0,
-            "total_custom_ais": 0,
-            "total_conversations": 0,
-            "total_messages": 0,
-            "fire_safety_chunks": 0,
-            "system_status": "error",
-            "error": str(e)
+# System information endpoint
+@app.get("/system/info")
+async def get_system_info():
+    return {
+        "service": "YourAI",
+        "version": "2.0.0",
+        "features": {
+            "lightrag_integration": True,
+            "conversation_memory": True,
+            "custom_ai_support": True,
+            "file_upload": True,
+            "multi_model_support": True
+        },
+        "models": {
+            "llm": LLM_MODEL,
+            "embedding": EMBEDDING_MODEL
+        },
+        "storage": {
+            "graph": "NetworkXStorage",
+            "vector": "NanoVectorDBStorage",
+            "conversation_memory": "In-Memory"
         }
+    }
 
-# Cleanup endpoint for expired sessions
-@app.post("/admin/cleanup")
-async def cleanup_expired_data():
-    try:
-        # Clean up expired sessions
-        cleanup_query = "UPDATE user_sessions SET is_active = false WHERE expires_at < NOW() AND is_active = true"
-        await db_manager.execute_query(cleanup_query)
-        
-        return {"message": "Cleanup completed successfully"}
-    except Exception as e:
-        print(f"Error during cleanup: {e}")
-        return {"message": "Cleanup completed with some errors", "error": str(e)}
-
-# Rate limiting endpoint
-@app.get("/rate-limit/{user_id}")
-async def check_rate_limit(user_id: str, current_user: dict = Depends(get_current_user)):
-    # Only allow users to check their own rate limit or admin
-    if current_user["id"] != user_id and current_user["email"] != "admin@yourai.com":
-        raise HTTPException(status_code=403, detail="Access denied")
-    
-    try:
-        # Simple rate limiting: 100 requests per hour
-        key = f"rate_limit:{user_id}:{datetime.now().strftime('%Y-%m-%d-%H')}"
-        
-        if db_manager.redis:
-            current_count = await db_manager.redis.get(key) or 0
-            current_count = int(current_count)
-            
-            return {
-                "user_id": user_id,
-                "current_requests": current_count,
-                "limit": 100,
-                "remaining": max(0, 100 - current_count),
-                "reset_time": f"{datetime.now().strftime('%Y-%m-%d %H')}:59:59"
-            }
-        else:
-            return {
-                "user_id": user_id,
-                "current_requests": 0,
-                "limit": 100,
-                "remaining": 100,
-                "reset_time": f"{datetime.now().strftime('%Y-%m-%d %H')}:59:59",
-                "note": "Rate limiting not available (Redis not connected)"
-            }
-    except Exception as e:
-        print(f"Error checking rate limit: {e}")
-        raise HTTPException(status_code=500, detail="Rate limit check failed")
-
-# Health check for specific components
-@app.get("/health/detailed")
-async def detailed_health_check():
-    health_status = {
-        "timestamp": datetime.now().isoformat(),
+# System status endpoint
+@app.get("/system/status")
+async def get_system_status():
+    status = {
         "status": "healthy",
-        "components": {}
+        "components": {
+            "lightrag": lightrag_manager is not None,
+            "cloudflare_worker": cloudflare_worker is not None,
+            "database": True,  # Implement actual DB health check
+            "cache": True      # Implement actual cache health check
+        },
+        "memory": {
+            "active_conversations": len(lightrag_manager.conversation_memory) if lightrag_manager else 0,
+            "rag_instances": len(lightrag_manager.rag_instances) if lightrag_manager else 0
+        }
     }
     
-    # Check database
-    try:
-        await db_manager.execute_query("SELECT 1")
-        health_status["components"]["database"] = {"status": "healthy", "type": "PostgreSQL"}
-    except Exception as e:
-        health_status["components"]["database"] = {"status": "unhealthy", "error": str(e), "type": "PostgreSQL"}
-        health_status["status"] = "degraded"
-    
-    # Check Redis cache
-    try:
-        if db_manager.redis:
-            await db_manager.redis.ping()
-            health_status["components"]["cache"] = {"status": "healthy", "type": "Redis"}
-        else:
-            health_status["components"]["cache"] = {"status": "unavailable", "type": "Redis"}
-    except Exception as e:
-        health_status["components"]["cache"] = {"status": "unhealthy", "error": str(e), "type": "Redis"}
-        health_status["status"] = "degraded"
+    # Overall health check
+    all_healthy = all(status["components"].values())
+    status["status"] = "healthy" if all_healthy else "unhealthy"
     
-    # Check Cloudflare Worker
-    try:
-        if cloudflare_worker:
-            # Simple test query
-            test_response = await cloudflare_worker.query("Test", "Respond with 'OK'")
-            if "OK" in test_response or len(test_response) > 0:
-                health_status["components"]["ai_service"] = {"status": "healthy", "type": "Cloudflare Worker"}
-            else:
-                health_status["components"]["ai_service"] = {"status": "degraded", "type": "Cloudflare Worker"}
-        else:
-            health_status["components"]["ai_service"] = {"status": "unavailable", "type": "Cloudflare Worker"}
-            health_status["status"] = "degraded"
-    except Exception as e:
-        health_status["components"]["ai_service"] = {"status": "unhealthy", "error": str(e), "type": "Cloudflare Worker"}
-        health_status["status"] = "degraded"
-    
-    # Check fire safety knowledge store
-    try:
-        if fire_safety_store and len(fire_safety_store.chunks) > 0:
-            health_status["components"]["knowledge_store"] = {
-                "status": "healthy", 
-                "type": "Fire Safety KB",
-                "chunks_loaded": len(fire_safety_store.chunks)
-            }
-        else:
-            health_status["components"]["knowledge_store"] = {"status": "degraded", "type": "Fire Safety KB"}
-            health_status["status"] = "degraded"
-    except Exception as e:
-        health_status["components"]["knowledge_store"] = {"status": "unhealthy", "error": str(e), "type": "Fire Safety KB"}
-        health_status["status"] = "degraded"
-    
-    return health_status
+    return status
+
+# Test endpoint for development
+@app.get("/test")
+async def test_endpoint():
+    return {
+        "message": "Test endpoint working",
+        "timestamp": datetime.now().isoformat(),
+        "environment": {
+            "cloudflare_configured": bool(CLOUDFLARE_API_KEY),
+            "database_configured": bool(DATABASE_URL),
+            "redis_configured": bool(REDIS_URL)
+        }
+    }
 
 if __name__ == "__main__":
     import uvicorn