Update app.py

app.py

@@ -9,41 +9,23 @@ import urllib.error
 
 import pandas as pd
 
-from fastapi import (
-    FastAPI,
-    File,
-    UploadFile,
-    HTTPException
-)
-
+from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
 
-# ── Configuration ──────────────────────────────────────────────
+# ─────────────────────────────────────────────
+# CONFIG
+# ─────────────────────────────────────────────
 
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
 
-if GEMINI_API_KEY:
-    print("✅ GEMINI_API_KEY Loaded")
-else:
-    print("❌ GEMINI_API_KEY Missing")
-
-
-# ── In-Memory Stores ───────────────────────────────────────────
-
 _db_store = {}
 _schema_store = {}
 
-
-# ── FastAPI Setup ──────────────────────────────────────────────
-
-app = FastAPI(
-    title="QueryMind Gemini",
-    version="4.0.0"
-)
+app = FastAPI(title="AI Data Analyst Agent", version="5.0.0")
 
 app.add_middleware(
     CORSMiddleware,
@@ -52,82 +34,61 @@ app.add_middleware(
     allow_headers=["*"]
 )
 
-
-# ── Request Model ──────────────────────────────────────────────
-
 class QueryRequest(BaseModel):
     session_id: str
     question: str
 
 
-# ── Heuristic SQL Engine ───────────────────────────────────────
+# ─────────────────────────────────────────────
+# AGENT INTENT ENGINE
+# ─────────────────────────────────────────────
 
-def _heuristic_sql(question: str, table: str, columns: list):
-
-    q = question.lower().strip()
-
-    t = f'"{table}"'
+def agent_think(question: str):
+    q = question.lower()
 
-    # ── Count Queries ─────────────────────────────
+    if any(x in q for x in ["chart", "graph", "plot"]):
+        return "VISUALIZE"
 
-    if any(x in q for x in [
-        "count total",
-        "count records",
-        "total records",
-        "how many records",
-        "total rows",
-        "count rows"
-    ]):
+    if any(x in q for x in ["average", "avg", "sum", "max", "min"]):
+        return "ANALYZE"
 
-        return f'SELECT COUNT(*) AS total_rows FROM {t}'
+    if any(x in q for x in ["why", "explain", "reason"]):
+        return "EXPLAIN"
 
-    # ── Preview Queries ───────────────────────────
+    if any(x in q for x in ["count", "how many"]):
+        return "COUNT"
 
-    if any(x in q for x in [
-        "preview",
-        "show head",
-        "data preview",
-        "show first",
-        "first 10",
-        "show rows"
-    ]):
+    return "SQL"
 
-        return f'SELECT * FROM {t} LIMIT 10'
 
-    # ── Show Columns ──────────────────────────────
+# ─────────────────────────────────────────────
+# HEURISTIC SQL ENGINE
+# ─────────────────────────────────────────────
 
-    if "columns" in q or "column names" in q:
-
-        cols = ", ".join(columns)
-
-        return f"SELECT '{cols}' AS columns_list"
-
-    # ── Unique Values ─────────────────────────────
+def _heuristic_sql(question: str, table: str, columns: list):
 
-    if "unique values in" in q:
+    q = question.lower()
+    t = f'"{table}"'
 
-        col = q.replace("unique values in", "").strip()
+    if "count" in q or "how many" in q:
+        return f"SELECT COUNT(*) AS total_rows FROM {t}"
 
-        if col in columns:
+    if any(x in q for x in ["first row", "show first", "preview"]):
+        return f"SELECT * FROM {t} LIMIT 10"
 
-            return f'''
-            SELECT DISTINCT "{col}"
-            FROM {t}
-            LIMIT 100
-            '''
+    if "last row" in q:
+        return f"SELECT * FROM {t} ORDER BY rowid DESC LIMIT 1"
 
-    # ── Group By ──────────────────────────────────
+    if "unique values" in q:
+        for col in columns:
+            if col.lower() in q:
+                return f'SELECT DISTINCT "{col}" FROM {t} LIMIT 100'
 
     if "group by" in q:
-
-        match = re.search(r'group by\s+(\w+)', q)
-
+        match = re.search(r'group by (\w+)', q)
         if match:
-
             col = match.group(1)
-
             if col in columns:
-
                 return f'''
                 SELECT "{col}", COUNT(*) AS count
                 FROM {t}
@@ -138,246 +99,205 @@ def _heuristic_sql(question: str, table: str, columns: list):
     return None
 
 
-# ── Gemini SQL Generator ──────────────────────────────────────
+# ─────────────────────────────────────────────
+# GEMINI SQL GENERATOR
+# ─────────────────────────────────────────────
 
-def _call_gemini(
-    question: str,
-    schema: str,
-    columns: list,
-    table: str
-):
+def _call_gemini(question, schema, columns, table):
 
     if not GEMINI_API_KEY:
-
-        print("❌ GEMINI_API_KEY Missing")
-
         return ""
 
-    col_list = ", ".join(columns[:30])
-
     prompt = f"""
-You are a SQLite expert.
-
-Convert the natural language question into a valid SQLite query.
+You are a strict SQLite expert.
 
 Rules:
 - Output ONLY SQL
-- Use ONLY the provided table
-- Do NOT explain anything
-- Do NOT use markdown
-
-Table Name:
-{table}
+- Use only given table and columns
+- No explanation
 
-Columns:
-{col_list}
+Table: {table}
+Columns: {columns}
+Schema: {schema}
 
-Schema:
-{schema}
-
-Question:
-{question}
+Question: {question}
 """
 
     payload = json.dumps({
-        "contents": [
-            {
-                "parts": [
-                    {
-                        "text": prompt
-                    }
-                ]
-            }
-        ]
+        "contents": [{"parts": [{"text": prompt}]}]
     }).encode("utf-8")
 
-    # Correct Gemini endpoint
     url = (
         "https://generativelanguage.googleapis.com/"
         f"v1beta/models/gemini-1.5-flash:generateContent?key={GEMINI_API_KEY}"
     )
 
     try:
-
         req = urllib.request.Request(
             url,
             data=payload,
-            headers={
-                "Content-Type": "application/json"
-            }
-        )
-
-        with urllib.request.urlopen(req, timeout=20) as resp:
-
-            data = json.loads(resp.read())
-
-        sql = (
-            data["candidates"][0]
-            ["content"]["parts"][0]
-            ["text"]
-            .strip()
+            headers={"Content-Type": "application/json"}
         )
 
-        # Cleanup
-        sql = (
-            sql
-            .replace("```sql", "")
-            .replace("```", "")
-            .strip()
-            .split(";")[0]
-        )
+        res = urllib.request.urlopen(req, timeout=20)
+        data = json.loads(res.read())
 
-        # Force correct table
-        sql = re.sub(
-            r'\bFROM\s+["\'\w\.]+',
-            f'FROM "{table}"',
-            sql,
-            flags=re.IGNORECASE
-        )
+        sql = data["candidates"][0]["content"]["parts"][0]["text"]
 
-        return sql
+        sql = sql.replace("```sql", "").replace("```", "").strip()
 
-    except urllib.error.HTTPError as e:
-
-        error_body = e.read().decode()
-
-        print(f"❌ GEMINI HTTP ERROR: {e.code}")
-        print(error_body)
-
-        return ""
+        return sql.split(";")[0]
 
     except Exception as e:
-
-        print(f"❌ GEMINI ERROR: {str(e)}")
-
+        print("Gemini Error:", e)
         return ""
 
 
-# ── SQLite Execution ──────────────────────────────────────────
+# ─────────────────────────────────────────────
+# SQL EXECUTION
+# ─────────────────────────────────────────────
 
 def execute_sql(sql, db_bytes):
 
     conn = sqlite3.connect(":memory:")
 
     with tempfile.NamedTemporaryFile(delete=False) as f:
-
         f.write(db_bytes)
-
         f.flush()
+        temp = f.name
 
-        temp_name = f.name
+    try:
+        disk = sqlite3.connect(temp)
+        disk.backup(conn)
+        disk.close()
+    finally:
+        if os.path.exists(temp):
+            os.remove(temp)
+
+    conn.row_factory = sqlite3.Row
 
     try:
+        cur = conn.execute(sql)
+        return [dict(r) for r in cur.fetchall()]
+    except Exception as e:
+        return [{"error": str(e)}]
+    finally:
+        conn.close()
 
-        disk_conn = sqlite3.connect(temp_name)
 
-        disk_conn.backup(conn)
+# ─────────────────────────────────────────────
+# ANALYSIS ENGINE
+# ──────────────────���──────────────────────────
 
-        disk_conn.close()
+def analyze_results(results):
 
-    finally:
+    if not results:
+        return {"message": "No data found"}
 
-        if os.path.exists(temp_name):
-            os.remove(temp_name)
+    return {
+        "rows_returned": len(results),
+        "sample": results[:3]
+    }
 
-    conn.row_factory = sqlite3.Row
 
-    try:
+# ─────────────────────────────────────────────
+# EXPLANATION ENGINE (Gemini)
+# ─────────────────────────────────────────────
 
-        cur = conn.execute(sql)
+def explain_results(question, sql, results):
 
-        rows = cur.fetchall()
+    if not GEMINI_API_KEY:
+        return None
 
-        return [dict(r) for r in rows]
+    prompt = f"""
+You are a data analyst.
 
-    except Exception as e:
+Question: {question}
+SQL: {sql}
+Results: {results[:5]}
 
-        return [{"error": str(e)}]
+Explain this in simple words.
+"""
 
-    finally:
+    payload = json.dumps({
+        "contents": [{"parts": [{"text": prompt}]}]
+    }).encode("utf-8")
 
-        conn.close()
+    url = (
+        "https://generativelanguage.googleapis.com/"
+        f"v1beta/models/gemini-1.5-flash:generateContent?key={GEMINI_API_KEY}"
+    )
 
+    try:
+        req = urllib.request.Request(
+            url,
+            data=payload,
+            headers={"Content-Type": "application/json"}
+        )
 
-# ── Upload Endpoint ───────────────────────────────────────────
+        res = urllib.request.urlopen(req)
+        data = json.loads(res.read())
 
-@app.post("/upload")
-async def upload_csv(file: UploadFile = File(...)):
+        return data["candidates"][0]["content"]["parts"][0]["text"]
 
-    try:
+    except:
+        return None
 
-        contents = await file.read()
 
-        df = pd.read_csv(
-            io.BytesIO(contents)
-        ).dropna(how='all')
+# ─────────────────────────────────────────────
+# UPLOAD CSV
+# ─────────────────────────────────────────────
 
-        session_id = os.urandom(8).hex()
+@app.post("/upload")
+async def upload_csv(file: UploadFile = File(...)):
 
-        # Clean table name
-        clean_name = re.sub(
-            r'[^a-zA-Z0-9_]',
-            '_',
-            os.path.splitext(file.filename)[0]
-        )
+    contents = await file.read()
 
-        if clean_name[0].isdigit():
-            clean_name = "t_" + clean_name
+    df = pd.read_csv(io.BytesIO(contents)).dropna(how="all")
 
-        table_name = clean_name[:32]
+    session_id = os.urandom(8).hex()
 
-        with tempfile.NamedTemporaryFile(delete=False) as tf:
+    table_name = re.sub(r"[^a-zA-Z0-9_]", "_", file.filename)
 
-            conn = sqlite3.connect(tf.name)
+    if table_name[0].isdigit():
+        table_name = "t_" + table_name
 
-            df.to_sql(
-                table_name,
-                conn,
-                index=False,
-                if_exists="replace"
-            )
+    table_name = table_name[:32]
 
-            schema = conn.execute(
-                "SELECT sql FROM sqlite_master WHERE type='table'"
-            ).fetchone()[0]
+    with tempfile.NamedTemporaryFile(delete=False) as tf:
 
-            conn.close()
+        conn = sqlite3.connect(tf.name)
 
-            with open(tf.name, "rb") as f:
-                db_data = f.read()
+        df.to_sql(table_name, conn, index=False, if_exists="replace")
 
-        if os.path.exists(tf.name):
-            os.remove(tf.name)
+        schema = conn.execute(
+            "SELECT sql FROM sqlite_master WHERE type='table'"
+        ).fetchone()[0]
 
-        _db_store[session_id] = {
-            "bytes": db_data,
-            "table": table_name,
-            "cols": list(df.columns)
-        }
+        conn.close()
 
-        _schema_store[session_id] = schema
+        db_bytes = open(tf.name, "rb").read()
 
-        return {
-            "session_id": session_id,
-            "columns": list(df.columns),
-            "row_count": len(df),
-            "table_name": table_name,
-            "preview": df.head(5).to_dict(
-                orient="records"
-            )
-        }
+    os.remove(tf.name)
 
-    except Exception as e:
+    _db_store[session_id] = {
+        "bytes": db_bytes,
+        "table": table_name,
+        "cols": list(df.columns)
+    }
 
-        print(f"❌ UPLOAD ERROR: {e}")
+    _schema_store[session_id] = schema
 
-        raise HTTPException(
-            status_code=500,
-            detail=str(e)
-        )
+    return {
+        "session_id": session_id,
+        "rows": len(df),
+        "columns": list(df.columns)
+    }
 
 
-# ── Query Endpoint ────────────────────────────────────────────
+# ─────────────────────────────────────────────
+# QUERY ENGINE (AGENT CORE)
+# ─────────────────────────────────────────────
 
 @app.post("/query")
 async def query(req: QueryRequest):
@@ -385,24 +305,17 @@ async def query(req: QueryRequest):
     data = _db_store.get(req.session_id)
 
     if not data:
-
-        raise HTTPException(
-            status_code=404,
-            detail="Invalid session_id"
-        )
+        raise HTTPException(404, "Invalid session")
 
     schema = _schema_store.get(req.session_id)
 
-    # Step 1 → Heuristic SQL
-    sql = _heuristic_sql(
-        req.question,
-        data["table"],
-        data["cols"]
-    )
+    # 🧠 Agent thinking
+    intent = agent_think(req.question)
 
-    # Step 2 → Gemini fallback
-    if not sql:
+    # 🗄️ SQL generation
+    sql = _heuristic_sql(req.question, data["table"], data["cols"])
 
+    if not sql:
         sql = _call_gemini(
             req.question,
             schema,
@@ -410,46 +323,53 @@ async def query(req: QueryRequest):
             data["table"]
         )
 
-    # Step 3 → Failure
     if not sql:
+        raise HTTPException(400, "SQL generation failed")
 
-        raise HTTPException(
-            status_code=400,
-            detail="Failed to generate SQL query"
-        )
+    # ⚡ Execute SQL
+    results = execute_sql(sql, data["bytes"])
 
-    results = execute_sql(
-        sql,
-        data["bytes"]
-    )
+    # 📊 Analysis
+    analysis = analyze_results(results)
+
+    # 💬 Explanation (only for analytical intent)
+    explanation = None
+
+    if intent in ["ANALYZE", "EXPLAIN"]:
+
+        explanation = explain_results(
+            req.question,
+            sql,
+            results
+        )
 
     return {
+        "intent": intent,
         "sql": sql,
-        "results": results
+        "results": results[:20],
+        "analysis": analysis,
+        "explanation": explanation
     }
 
 
-# ── Health Endpoint ───────────────────────────────────────────
+# ─────────────────────────────────────────────
+# HEALTH CHECK
+# ─────────────────────────────────────────────
 
 @app.get("/health")
 def health():
-
     return {
         "status": "ok",
-        "model": "gemini-1.5-flash"
+        "model": "AI Data Analyst Agent"
     }
 
 
-# ── Static Frontend ───────────────────────────────────────────
-
-app.mount(
-    "/static",
-    StaticFiles(directory="static"),
-    name="static"
-)
+# ─────────────────────────────────────────────
+# FRONTEND
+# ─────────────────────────────────────────────
 
+app.mount("/static", StaticFiles(directory="static"), name="static")
 
 @app.get("/")
 def root():
-
     return FileResponse("static/webapp.html")