Update app.py

app.py

@@ -2,96 +2,92 @@ import torch
 import sqlite3
 import pandas as pd
 import gradio as gr
+import re
 from langchain_community.llms import HuggingFacePipeline
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# ============================================================
-# 🚀 Load SQLCoder model
-# ============================================================
+# Load model
 model_id = "defog/sqlcoder-7b-2"
-
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype="auto",
-    device_map="auto"
-)
-
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=256,
-    do_sample=False
-)
-
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256, do_sample=False)
 sqlcoder_llm = HuggingFacePipeline(pipeline=pipe)
 
-# ============================================================
-# 🧠 Define query function
-# ============================================================
 def ask_question(user_db, question):
-    """Takes an uploaded SQLite database + a question, returns SQL + result"""
     if not user_db:
-        return "❌ Please upload a database file.", None
-
+        return "❌ Upload database", None
+    
     conn = sqlite3.connect(user_db.name)
     cursor = conn.cursor()
+    
+    # Get full schema with columns
+    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+    tables = [row[0] for row in cursor.fetchall()]
+    
+    schema_info = []
+    for table in tables:
+        cursor.execute(f"PRAGMA table_info({table});")
+        columns = [col[1] for col in cursor.fetchall()]
+        schema_info.append(f"{table}({', '.join(columns)})")
+    
+    schema_text = "\n".join(schema_info)
+    
+    # Smart prompt - let model figure out the right table
+    prompt = f"""You are a SQL expert. Generate ONLY the SQL query, nothing else.
 
-    # Create a Text-to-SQL prompt
-    prompt = f"""
-    You are an expert SQL generator.
-    The database follows the Chinook schema with tables:
-    customers, invoices, invoice_items, tracks, albums, artists, employees, genres, media_types, playlists, playlist_track.
-    Translate this question into a valid SQLite query for this schema.
-    Return only SQL (no text).
-    Question: {question}
-    SQL:
-    """
+Database Schema:
+{schema_text}
 
+Instructions:
+- Use the EXACT table and column names from the schema above
+- If user asks about concepts (like "sales", "customers", "products"), find the most relevant table
+- Return ONLY valid SQL with semicolon
+- No explanations, no markdown, just SQL
 
-    # ✅ Use .invoke() instead of calling the object directly
+Question: {question}
+SQL:"""
+    
+    # Generate SQL
     response = sqlcoder_llm.invoke(prompt)
-
-    # Ensure we get plain string
-    if isinstance(response, dict) and "text" in response:
-        response = response["text"]
-    elif isinstance(response, list):
-        response = response[0]["generated_text"]
-
-    # Clean and finalize SQL
-    sql_query = response.strip().split("SQL:")[-1].strip()
-    sql_query = sql_query.split("\n")[0].strip()
-    if not sql_query.endswith(";"):
-        sql_query += ";"
-
+    sql = str(response).strip()
+    
+    # Extract SQL
+    if "SQL:" in sql:
+        sql = sql.split("SQL:")[-1].strip()
+    sql = sql.split("\n")[0].strip()
+    if not sql.endswith(";"): 
+        sql += ";"
+    
+    # Remove common formatting
+    sql = sql.replace("```sql", "").replace("```", "").strip()
+    
+    # Execute
     try:
-        cursor.execute(sql_query)
+        cursor.execute(sql)
         rows = cursor.fetchall()
-        columns = [desc[0] for desc in cursor.description]
-        df = pd.DataFrame(rows, columns=columns)
+        if cursor.description:
+            df = pd.DataFrame(rows, columns=[d[0] for d in cursor.description])
+        else:
+            df = pd.DataFrame()
         conn.close()
-        return sql_query, df
-    except Exception as e:
+        return f"✅ SQL:\n{sql}\n\n📊 {len(df)} rows", df
+    except sqlite3.Error as e:
         conn.close()
-        return f"❌ Error executing query: {e}\n\nGenerated SQL:\n{sql_query}", None
+        return f"❌ Error: {e}\n\nSQL tried:\n{sql}\n\n💡 Available tables:\n{schema_text}", None
 
-# ============================================================
-# 🎨 Gradio UI
-# ============================================================
+# UI
 demo = gr.Interface(
     fn=ask_question,
     inputs=[
-        gr.File(label="Upload SQLite Database (.db)"),
-        gr.Textbox(label="Ask your question")
+        gr.File(label="📁 Upload Database (.db)"),
+        gr.Textbox(label="❓ Ask Question", placeholder="e.g., show all data, highest value, total count")
     ],
     outputs=[
-        gr.Textbox(label="Generated SQL Query"),
-        gr.Dataframe(label="Query Result")
+        gr.Textbox(label="🤖 SQL & Status", lines=6),
+        gr.Dataframe(label="📊 Results")
     ],
-    title="🧠 Text-to-SQL on Your Own Database",
-    description="Upload your SQLite database and ask natural language questions."
+    title="🔮 Universal Text-to-SQL",
+    description="Upload ANY SQLite database and ask questions. The AI will figure out the right tables!"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()