Update app.py

app.py

@@ -2,16 +2,16 @@ import sys
 import os
 import uvicorn
 import re
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import HTMLResponse  # <--- ADDED THIS
+from fastapi.responses import HTMLResponse
 from pydantic import BaseModel
 from typing import List, Optional
 
 # 🚨 FORCE PYTHON TO FIND THE 'src' FOLDER
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
-# Import your custom modules
+# Import custom modules
 from src.db_connector import Database
 from src.rag_manager import RAGSystem 
 from src.sql_generator import SQLGenerator
@@ -32,20 +32,22 @@ class ChatRequest(BaseModel):
 
 # --- HELPER: CLEAN AI OUTPUT ---
 def clean_sql(sql_text: str) -> str:
-    """Removes markdown, 'sql' tags, and extra whitespace."""
-    if not sql_text:
-        return ""
-    # Remove markdown code blocks (```sql ... ```)
+    if not sql_text: return ""
     cleaned = re.sub(r"```sql|```", "", sql_text, flags=re.IGNORECASE).strip()
-    # Remove trailing semicolons for consistency
-    cleaned = cleaned.rstrip(';') 
-    return cleaned
+    return cleaned.rstrip(';')
+
+# --- 🚀 SAFE INITIALIZATION ---
+# Initialize as None so the app doesn't crash if DB fails
+db = None
+rag = None
+generator = None
+startup_error = None
 
-# --- INITIALIZATION SEQUENCE ---
 print("--- 🚀 SYSTEM STARTUP SEQUENCE ---")
 try:
     print("   ...Connecting to Database")
-    db = Database()
+    # ⚠️ IF THIS FAILS, THE APP WILL NOW CATCH IT GRACEFULLY
+    db = Database() 
     print("   ✅ Database Connection: SUCCESS")
     
     print("   ...Initializing RAG System")
@@ -57,85 +59,55 @@ try:
     print("   ✅ AI Model: LOADED")
     
 except Exception as e:
+    startup_error = str(e)
     print(f"   ❌ CRITICAL STARTUP ERROR: {e}")
 
-
-# --- 1. NEW ROUTE TO SERVE THE UI ---
+# --- ROUTES ---
 @app.get("/", response_class=HTMLResponse)
 async def serve_ui():
-    """Serves the index.html file when user visits the root URL"""
     try:
         with open("index.html", "r", encoding="utf-8") as f:
             return f.read()
     except FileNotFoundError:
-        return "Error: index.html not found. Make sure it is in the root folder."
-
+        return "Error: index.html not found."
 
-# --- 2. CHAT API ENDPOINT ---
 @app.post("/chat")
 def chat_endpoint(request: ChatRequest):
-    print(f"\n📨 NEW REQUEST: {request.question}") 
-    
+    # 🚨 CHECK FOR STARTUP ERRORS FIRST
+    if rag is None or db is None:
+        return {
+            "answer": [f"System Error: {startup_error}"],
+            "sql": "-- Database Connection Failed",
+            "message": f"I cannot connect to the database. Reason: {startup_error}. Please check the Logs tab.",
+            "follow_ups": []
+        }
+
     try:
-        # 1. Get Context
         context = rag.get_relevant_schema(request.question)
-        
-        # 2. Generate SQL
         raw_sql, explanation, friendly_msg = generator.generate_sql(request.question, context, request.history)
         
-        # 3. Clean and Validate SQL
-        if "Error:" in raw_sql or "Invalid Query" in raw_sql:
-            return {
-                "answer": [],
-                "sql": raw_sql,
-                "message": "I couldn't generate a safe query for that request. Try asking for specific data like 'Show me users' or 'List orders'.",
-                "follow_ups": ["Show top 10 rows from users", "Count total orders"]
-            }
-
+        # Clean SQL
         cleaned_sql = clean_sql(raw_sql)
-        print(f"   🧹 Cleaned SQL: {cleaned_sql}")
-
-        # Safety check: Ensure it's a SELECT
+        
+        # Security Check
         if not cleaned_sql.upper().startswith("SELECT"):
-             return {
-                "answer": [],
-                "sql": cleaned_sql,
-                "message": "Security Alert: I can only perform READ (SELECT) operations.",
-                "follow_ups": []
-            }
+             return {"answer": [], "sql": cleaned_sql, "message": "Security Alert: Read-Only Mode.", "follow_ups": []}
 
-        # 4. Run Query
+        # Run Query
         try:
             results = db.run_query(cleaned_sql)
         except Exception as db_err:
-            print(f"   ⚠️ DB Error: {db_err}")
-            return {
-                "answer": [f"Error: {str(db_err)}"],
-                "sql": cleaned_sql,
-                "message": "There was a syntax error in the generated SQL. I have displayed the error above.",
-                "follow_ups": []
-            }
-        
-        # 5. Generate Follow-ups
-        follow_ups = generator.generate_followup_questions(request.question, cleaned_sql)
+            return {"answer": [f"SQL Error: {str(db_err)}"], "sql": cleaned_sql, "message": "Syntax Error in SQL.", "follow_ups": []}
         
         return {
             "answer": results,
             "sql": cleaned_sql,
-            "explanation": explanation,
             "message": friendly_msg, 
-            "follow_ups": follow_ups
+            "follow_ups": []
         }
 
     except Exception as e:
-        print(f"❌ General Processing Error: {e}")
-        return {
-            "answer": [],
-            "sql": "-- System Error",
-            "message": f"Critical Error: {str(e)}",
-            "follow_ups": []
-        }
+        return {"answer": [], "sql": "-- Error", "message": f"Processing Error: {str(e)}", "follow_ups": []}
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
-    
+    uvicorn.run(app, host="0.0.0.0", port=7860)