Update main.py

main.py

@@ -1,34 +1,16 @@
-import gradio as gr
+# main.py - FastAPI Backend
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 import asyncio
+from lightrag import LightRAG, QueryParam
+from lightrag.utils import EmbeddingFunc
 import os
-import zipfile
 import requests
-from pathlib import Path
 import numpy as np
 from typing import List
 
-# Try different LightRAG imports based on version
-try:
-    from lightrag import LightRAG, QueryParam
-    from lightrag.utils import EmbeddingFunc
-    LIGHTRAG_AVAILABLE = True
-except ImportError:
-    try:
-        from lightrag.lightrag import LightRAG
-        from lightrag.query import QueryParam
-        from lightrag.utils import EmbeddingFunc
-        LIGHTRAG_AVAILABLE = True
-    except ImportError:
-        try:
-            from lightrag.core import LightRAG
-            from lightrag.core import QueryParam
-            from lightrag.utils import EmbeddingFunc
-            LIGHTRAG_AVAILABLE = True
-        except ImportError:
-            print("❌ LightRAG import failed - using fallback mode")
-            LIGHTRAG_AVAILABLE = False
-
-# Fallback CloudflareWorker with simple search
+# Your CloudflareWorker class
 class CloudflareWorker:
     def __init__(self, cloudflare_api_key: str, api_base_url: str, llm_model_name: str, embedding_model_name: str):
         self.cloudflare_api_key = cloudflare_api_key
@@ -38,21 +20,20 @@ class CloudflareWorker:
         self.max_tokens = 4080
         self.max_response_tokens = 4080
 
-    async def _send_request(self, model_name: str, input_: dict, debug_log: str = ""):
+    async def _send_request(self, model_name: str, input_: dict, debug_log: str):
         headers = {"Authorization": f"Bearer {self.cloudflare_api_key}"}
         
         try:
             response_raw = requests.post(
                 f"{self.api_base_url}{model_name}",
                 headers=headers,
-                json=input_,
-                timeout=30
+                json=input_
             ).json()
             
             result = response_raw.get("result", {})
             
             if "data" in result:
-                return np.array(result["data"]) if LIGHTRAG_AVAILABLE else result["data"]
+                return np.array(result["data"])
             if "response" in result:
                 return result["response"]
             
@@ -76,79 +57,23 @@ class CloudflareWorker:
             "response_token_limit": self.max_response_tokens,
         }
         
-        result = await self._send_request(self.llm_model_name, input_)
+        result = await self._send_request(self.llm_model_name, input_, "")
         return result if result is not None else "Error: Failed to get response"
 
-    async def embedding_chunk(self, texts: List[str]):
+    async def embedding_chunk(self, texts: List[str]) -> np.ndarray:
         input_ = {
             "text": texts,
             "max_tokens": self.max_tokens,
             "response_token_limit": self.max_response_tokens,
         }
         
-        result = await self._send_request(self.embedding_model_name, input_)
+        result = await self._send_request(self.embedding_model_name, input_, "")
         
         if result is None:
-            if LIGHTRAG_AVAILABLE:
-                return np.random.rand(len(texts), 1024).astype(np.float32)
-            else:
-                return [[0.0] * 1024 for _ in texts]
+            return np.random.rand(len(texts), 1024).astype(np.float32)
         
         return result
 
-# Simple fallback knowledge store if LightRAG fails
-class SimpleKnowledgeStore:
-    def __init__(self, data_dir: str):
-        self.data_dir = data_dir
-        self.chunks = []
-        self.entities = []
-        self.load_data()
-    
-    def load_data(self):
-        try:
-            import json
-            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 []
-            
-            entities_file = Path(self.data_dir) / "vdb_entities.json"
-            if entities_file.exists():
-                with open(entities_file, 'r', encoding='utf-8') as f:
-                    entities_data = json.load(f)
-                    if isinstance(entities_data, dict) and 'data' in entities_data:
-                        self.entities = entities_data['data']
-                    elif isinstance(entities_data, list):
-                        self.entities = entities_data
-                    else:
-                        self.entities = []
-                    
-            print(f"✅ Loaded {len(self.chunks)} chunks and {len(self.entities)} entities")
-            
-        except Exception as e:
-            print(f"⚠️ Error loading data: {e}")
-            self.chunks = []
-            self.entities = []
-    
-    def search(self, query: str, limit: int = 5) -> List[str]:
-        query_lower = query.lower()
-        results = []
-        
-        for chunk in self.chunks:
-            if isinstance(chunk, dict) and 'content' in chunk:
-                content = chunk['content']
-                if any(word in content.lower() for word in query_lower.split()):
-                    results.append(content)
-        
-        for entity in self.entities:
-            if isinstance(entity, dict):
-                entity_text = str(entity)
-                if any(word in entity_text.lower() for word in query_lower.split()):
-                    results.append(entity_text)
-        
-        return results[:limit]
-
 # Configuration
 CLOUDFLARE_API_KEY = os.getenv('CLOUDFLARE_API_KEY', 'lMbDDfHi887AK243ZUenm4dHV2nwEx2NSmX6xuq5')
 API_BASE_URL = "https://api.cloudflare.com/client/v4/accounts/07c4bcfbc1891c3e528e1c439fee68bd/ai/run/"
@@ -156,181 +81,126 @@ EMBEDDING_MODEL = '@cf/baai/bge-m3'
 LLM_MODEL = "@cf/meta/llama-3.2-3b-instruct"
 WORKING_DIR = "./dickens"
 
-# Global instances
+# Initialize FastAPI
+app = FastAPI(title="Fire Safety AI Assistant API", version="1.0.0")
+
+# Enable CORS for frontend
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # In production, replace with your frontend domain
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Global RAG instance
 rag_instance = None
-knowledge_store = None
-cloudflare_worker = None
 
-async def initialize_system():
-    global rag_instance, knowledge_store, cloudflare_worker
-    
-    print("🔄 Initializing system...")
-    
-    # Download data if needed
-    dickens_path = Path(WORKING_DIR)
-    has_data = dickens_path.exists() and len(list(dickens_path.glob("*.json"))) > 0
+# Pydantic models
+class QuestionRequest(BaseModel):
+    question: str
+    mode: str = "hybrid"  # naive, local, global, hybrid
+
+class QuestionResponse(BaseModel):
+    answer: str
+    mode: str
+    status: str
+
+@app.on_event("startup")
+async def startup_event():
+    """Initialize RAG system on startup"""
+    global rag_instance
     
-    if not has_data:
-        print("📥 Downloading RAG database...")
-        try:
-            # REPLACE YOUR_USERNAME with your actual GitHub username
-            data_url = "https://github.com/YOUR_USERNAME/fire-safety-ai/releases/download/v1.0-data/dickens.zip"
-            
-            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)
+    print("🔄 Initializing RAG system...")
     
-    # Initialize Cloudflare worker
     cloudflare_worker = CloudflareWorker(
         cloudflare_api_key=CLOUDFLARE_API_KEY,
         api_base_url=API_BASE_URL,
         embedding_model_name=EMBEDDING_MODEL,
         llm_model_name=LLM_MODEL,
     )
+
+    rag_instance = LightRAG(
+        working_dir=WORKING_DIR,
+        max_parallel_insert=2,
+        llm_model_func=cloudflare_worker.query,
+        llm_model_name=LLM_MODEL,
+        llm_model_max_token_size=4080,
+        embedding_func=EmbeddingFunc(
+            embedding_dim=1024,
+            max_token_size=2048,
+            func=lambda texts: cloudflare_worker.embedding_chunk(texts),
+        ),
+    )
     
-    # Try to initialize LightRAG, fallback to simple store
-    if LIGHTRAG_AVAILABLE:
-        try:
-            rag_instance = LightRAG(
-                working_dir=WORKING_DIR,
-                max_parallel_insert=2,
-                llm_model_func=cloudflare_worker.query,
-                llm_model_name=LLM_MODEL,
-                llm_model_max_token_size=4080,
-                embedding_func=EmbeddingFunc(
-                    embedding_dim=1024,
-                    max_token_size=2048,
-                    func=lambda texts: cloudflare_worker.embedding_chunk(texts),
-                ),
-            )
-            
-            await rag_instance.initialize_storages()
-            print("✅ LightRAG system initialized!")
-            
-        except Exception as e:
-            print(f"⚠️ LightRAG failed, using fallback: {e}")
-            knowledge_store = SimpleKnowledgeStore(WORKING_DIR)
-    else:
-        print("🔄 Using simple knowledge store...")
-        knowledge_store = SimpleKnowledgeStore(WORKING_DIR)
-    
-    print("✅ System ready!")
+    await rag_instance.initialize_storages()
+    print("✅ RAG system initialized!")
+
+@app.get("/")
+async def root():
+    return {"message": "🔥 Fire Safety AI Assistant API", "status": "running"}
 
-# Initialize on startup
-asyncio.run(initialize_system())
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy", "rag_ready": rag_instance is not None}
 
-async def ask_question(question, mode="hybrid"):
-    if not question.strip():
-        return "❌ Please enter a question."
+@app.post("/ask", response_model=QuestionResponse)
+async def ask_question(request: QuestionRequest):
+    """Ask a question to the Fire Safety AI"""
+    
+    if not rag_instance:
+        raise HTTPException(status_code=503, detail="RAG system not initialized")
+    
+    if not request.question.strip():
+        raise HTTPException(status_code=400, detail="Question cannot be empty")
     
     try:
-        print(f"🔍 Processing question: {question}")
+        # Query the RAG system
+        print(f"🔍 Processing question: {request.question}")
         
-        # Use LightRAG if available, otherwise fallback
-        if rag_instance and LIGHTRAG_AVAILABLE:
-            response = await rag_instance.aquery(
-                question,
-                param=QueryParam(mode=mode)
-            )
-            return response
+        response = await rag_instance.aquery(
+            request.question,
+            param=QueryParam(mode=request.mode)
+        )
         
-        elif knowledge_store and cloudflare_worker:
-            # Fallback: simple search + Cloudflare AI
-            relevant_chunks = knowledge_store.search(question, limit=3)
-            context = "\n".join(relevant_chunks) if relevant_chunks else "No specific context found."
-            
-            system_prompt = """You are a Fire Safety AI Assistant specializing in Vietnamese 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: {question}
-
-Please provide a helpful answer based on the context about Vietnamese fire safety regulations."""
-            
-            response = await cloudflare_worker.query(user_prompt, system_prompt)
-            return response
+        return QuestionResponse(
+            answer=response,
+            mode=request.mode,
+            status="success"
+        )
         
-        else:
-            return "❌ System not initialized yet. Please wait..."
-            
     except Exception as e:
-        return f"❌ Error: {str(e)}"
-
-def sync_ask_question(question, mode):
-    return asyncio.run(ask_question(question, mode))
-
-# Create Gradio interface
-with gr.Blocks(title="🔥 Fire Safety AI Assistant", theme=gr.themes.Soft()) as demo:
-    gr.HTML("<h1 style='text-align: center;'>🔥 Fire Safety AI Assistant</h1>")
-    gr.HTML("<p style='text-align: center;'>Ask questions about Vietnamese fire safety regulations</p>")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            question_input = gr.Textbox(
-                label="Your Question",
-                placeholder="What are the requirements for emergency exits?",
-                lines=3
-            )
-            mode_dropdown = gr.Dropdown(
-                choices=["hybrid", "local", "global", "naive"],
-                value="hybrid",
-                label="Search Mode",
-                info="Hybrid is recommended for best results"
-            )
-            submit_btn = gr.Button("🔍 Ask Question", variant="primary", size="lg")
-        
-        with gr.Column(scale=2):
-            answer_output = gr.Textbox(
-                label="Answer",
-                lines=15,
-                show_copy_button=True
-            )
-    
-    # System status
-    status_text = "✅ LightRAG System" if LIGHTRAG_AVAILABLE else "⚠️ Fallback Mode"
-    gr.HTML(f"<p style='text-align: center; color: gray;'>Status: {status_text}</p>")
-    
-    # Example questions
-    gr.HTML("<h3 style='text-align: center;'>💡 Example Questions:</h3>")
-    
-    with gr.Row():
-        example1 = gr.Button("What are the requirements for emergency exits?", size="sm")
-        example2 = gr.Button("How many exits does a building need?", size="sm")
-    
-    with gr.Row():
-        example3 = gr.Button("What are fire safety rules for stairwells?", size="sm")
-        example4 = gr.Button("What are building safety requirements?", size="sm")
-    
-    # Event handlers
-    submit_btn.click(
-        sync_ask_question,
-        inputs=[question_input, mode_dropdown],
-        outputs=answer_output
-    )
-    
-    question_input.submit(
-        sync_ask_question,
-        inputs=[question_input, mode_dropdown],
-        outputs=answer_output
-    )
-    
-    example1.click(lambda: "What are the requirements for emergency exits?", outputs=question_input)
-    example2.click(lambda: "How many exits does a building need?", outputs=question_input)
-    example3.click(lambda: "What are fire safety rules for stairwells?", outputs=question_input)
-    example4.click(lambda: "What are building safety requirements?", outputs=question_input)
+        print(f"❌ Error processing question: {e}")
+        raise HTTPException(status_code=500, detail=f"Error processing question: {str(e)}")
+
+@app.get("/modes")
+async def get_available_modes():
+    """Get available query modes"""
+    return {
+        "modes": [
+            {"name": "naive", "description": "Simple text search"},
+            {"name": "local", "description": "Search specific document sections"},
+            {"name": "global", "description": "Look at overall document themes"},
+            {"name": "hybrid", "description": "Combined approach (recommended)"}
+        ]
+    }
+
+# Example questions endpoint
+@app.get("/examples")
+async def get_example_questions():
+    """Get example questions users can ask"""
+    return {
+        "examples": [
+            "What are the requirements for emergency exits?",
+            "How many exits does a building need?",
+            "What are fire safety rules for stairwells?",
+            "What are building safety requirements?",
+            "What are the fire safety regulations for high-rise buildings?",
+            "What are the requirements for fire doors?",
+            "How should evacuation routes be designed?"
+        ]
+    }
 
 if __name__ == "__main__":
-    demo.launch()
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)