Update app.py

app.py

@@ -1,8 +1,15 @@
+"""
+QueryMind — CSV-to-SQL Engine (v3.0.4)
+Final Production Build: Gemini 1.5 Flash + Hybrid Heuristics
+Author: Nilotpal Dhar
+"""
+
 import os
 import re
 import io
-import json 
+import json
 import sqlite3
+import tempfile
 import pandas as pd
 import urllib.request
 from fastapi import FastAPI, File, UploadFile, HTTPException
@@ -12,51 +19,49 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
 # ── Configuration ──────────────────────────────────────────────────────────────
+# In HF Spaces, set this in Settings -> Variables and Secrets
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
 
 _db_store = {}
 _schema_store = {}
 
-app = FastAPI(title="QueryMind Gemini", version="3.0.3")
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+app = FastAPI(title="QueryMind Gemini", version="3.0.4")
+app.add_middleware(
+    CORSMiddleware, 
+    allow_origins=["*"], 
+    allow_methods=["*"], 
+    allow_headers=["*"]
+)
 
 class QueryRequest(BaseModel):
     session_id: str
     question: str
 
-# ── Logic Helpers ──────────────────────────────────────────────────────────────
-
-def _find_col(question: str, columns: list) -> str or None:
-    q = question.lower()
-    for col in sorted(columns, key=len, reverse=True):
-        if col.lower() in q: return col
-    return None
-
-# ── Improved Heuristics (Less aggressive) ──────────────────────────────────────
+# ── Heuristic Logic (Instant Speed Layer) ─────────────────────────────────────
 
 def _heuristic_sql(question: str, table: str, columns: list) -> str or None:
+    """Handles basic queries locally without calling Gemini to save time/quota."""
     q = question.lower().strip()
     t = f'"{table}"'
     
-    # Only trigger heuristics for VERY specific, simple patterns
-    # "how many records" or "total rows"
+    # Simple counting
     if re.fullmatch(r'(how many records|total rows|count rows|count total)', q):
         return f'SELECT COUNT(*) AS total_rows FROM {t}'
     
-    # "preview" or "show head"
-    if re.fullmatch(r'(preview|show head|data preview)', q):
+    # Simple data preview
+    if re.fullmatch(r'(preview|show head|data preview|show all)', q):
         return f'SELECT * FROM {t} LIMIT 10'
 
-    # If it's anything else, return None so Gemini can handle the logic
     return None
 
-# ── Improved Gemini Prompt (For Hard Queries) ──────────────────────────────────
+# ── Gemini API Call (Neural Logic Layer) ─────────────────────────────────────
 
 def _call_gemini(question: str, schema: str, columns: list, table: str) -> str:
-    if not GEMINI_API_KEY: return ""
+    """Calls Gemini 1.5 Flash to translate Natural Language into SQLite."""
+    if not GEMINI_API_KEY:
+        return ""
     
     col_list = ", ".join(columns[:30])
-    # Improved prompt to handle filtering, grouping, and ordering
     prompt = (
         f"You are a SQLite expert. Convert the question into a single valid SQL query.\n"
         f"Table: {table}\n"
@@ -65,7 +70,7 @@ def _call_gemini(question: str, schema: str, columns: list, table: str) -> str:
         f"Question: {question}\n\n"
         f"Rules:\n"
         f"1. Use double quotes for table and column names.\n"
-        f"2. Output ONLY the SQL code.\n"
+        f"2. Output ONLY the SQL code. No markdown, no explanation.\n"
         f"3. If the question asks for 'the first', use LIMIT 1.\n"
         f"4. If filtering by text, use the LIKE operator for flexibility.\n\n"
         f"SQL:"
@@ -73,7 +78,7 @@ def _call_gemini(question: str, schema: str, columns: list, table: str) -> str:
     
     payload = json.dumps({
         "contents": [{"parts": [{"text": prompt}]}],
-        "generationConfig": {"temperature": 0.1} # Slight temperature for better reasoning
+        "generationConfig": {"temperature": 0.1, "maxOutputTokens": 300}
     }).encode("utf-8")
 
     url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={GEMINI_API_KEY}"
@@ -84,49 +89,57 @@ def _call_gemini(question: str, schema: str, columns: list, table: str) -> str:
             data = json.loads(resp.read())
             sql = data["candidates"][0]["content"]["parts"][0]["text"].strip()
         
-        # Clean response
+        # Clean potential markdown artifacts
         sql = sql.replace("```sql", "").replace("```", "").strip().split(";")[0]
-        # Force correct table name
+        # Safety check: Force the correct table name from our store
         sql = re.sub(r'\bFROM\s+["\'\w\.]+', f'FROM "{table}"', sql, flags=re.IGNORECASE)
         return sql
     except Exception as e:
-        print(f"Gemini Error: {e}")
+        print(f"[GEMINI ERROR] {e}")
         return ""
 
-# 🔥 FIXED: Execution logic to handle bytes correctly
+# ── Database Management ───────────────────────────────────────────────────────
+
 def execute_sql(sql, db_bytes):
-    # Create a fresh in-memory DB and load the bytes into it
+    """Restores the SQLite database from memory and executes the query."""
+    # Create an empty in-memory database
     conn = sqlite3.connect(":memory:")
-    db_file = io.BytesIO(db_bytes)
-    # We use a temporary disk file because SQLite backup doesn't like raw BytesIO directly
+    
+    # We use a temporary file to bridge bytes back into a SQLite connection
     with tempfile.NamedTemporaryFile() as f:
         f.write(db_bytes)
         f.flush()
         disk_conn = sqlite3.connect(f.name)
-        disk_conn.backup(conn)
-        disk_conn.close()
+        try:
+            disk_conn.backup(conn)
+        finally:
+            disk_conn.close()
     
     conn.row_factory = sqlite3.Row
     try:
         cur = conn.execute(sql)
         return [dict(r) for r in cur.fetchall()]
-    finally: conn.close()
+    except Exception as e:
+        return [{"error": str(e)}]
+    finally:
+        conn.close()
 
 # ── API Endpoints ─────────────────────────────────────────────────────────────
-import tempfile # Added this import
 
 @app.post("/upload")
 async def upload_csv(file: UploadFile = File(...)):
+    """Receives CSV, creates session, and prepares in-memory SQL database."""
     try:
         contents = await file.read()
         df = pd.read_csv(io.BytesIO(contents)).dropna(how='all')
         
         session_id = os.urandom(8).hex()
+        # Clean table name for SQLite safety
         clean_name = re.sub(r'[^a-zA-Z0-9_]', '_', os.path.splitext(file.filename)[0])
         if clean_name[0].isdigit(): clean_name = "t_" + clean_name
         table_name = clean_name[:32]
         
-        # 🔥 FIXED: Use a temporary file to bridge Pandas to SQLite bytes
+        # Build the SQL database using a temp file to get raw bytes
         with tempfile.NamedTemporaryFile() as tf:
             conn = sqlite3.connect(tf.name)
             df.to_sql(table_name, conn, index=False, if_exists="replace")
@@ -136,6 +149,7 @@ async def upload_csv(file: UploadFile = File(...)):
             with open(tf.name, "rb") as f:
                 db_data = f.read()
         
+        # Store session data globally (Shared with bot.py)
         _db_store[session_id] = {
             "bytes": db_data, 
             "table": table_name, 
@@ -143,6 +157,7 @@ async def upload_csv(file: UploadFile = File(...)):
         }
         _schema_store[session_id] = schema
 
+        # This response is synchronized with webapp.html logic
         return {
             "session_id": session_id,
             "columns": list(df.columns),
@@ -156,17 +171,35 @@ async def upload_csv(file: UploadFile = File(...)):
 
 @app.post("/query")
 async def query(req: QueryRequest):
+    """Main query handler: Heuristics -> Gemini -> SQL Execution."""
     data = _db_store.get(req.session_id)
-    if not data: raise HTTPException(status_code=404)
-    sql = _heuristic_sql(req.question, data["table"], data["cols"]) or _call_gemini(req.question, _schema_store[req.session_id], data["cols"], data["table"])
-    return {"sql": sql, "results": execute_sql(sql, data["bytes"])}
+    if not data:
+        raise HTTPException(status_code=404, detail="Session expired. Please re-upload your file.")
+    
+    # 1. Try Heuristics
+    sql = _heuristic_sql(req.question, data["table"], data["cols"])
+    
+    # 2. Try Gemini
+    if not sql:
+        sql = _call_gemini(req.question, _schema_store[req.session_id], data["cols"], data["table"])
+    
+    if not sql:
+        raise HTTPException(status_code=400, detail="I couldn't translate that question into a SQL query.")
+
+    results = execute_sql(sql, data["bytes"])
+    return {"sql": sql, "results": results}
+
+# ── Health & Static Assets ──
 
 @app.get("/health")
 def health():
+    """Health check endpoint for Hugging Face and the Web UI."""
     return {"status": "ok", "model": "gemini-1.5-flash"}
 
+# Mount the static directory to serve webapp.html and CSS
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 @app.get("/")
 def root():
+    """Serves the main frontend dashboard."""
     return FileResponse("static/webapp.html")