Update app.py

app.py

@@ -4,20 +4,45 @@ import sqlite3
 import pandas as pd
 import os
 
-# 1. API CONFIG
+# --- 1. SETTINGS & API CONFIG ---
+st.set_page_config(page_title="Agentic SQL", page_icon="🕵️‍♂️")
+
 api_key = os.getenv("GEMINI_API_KEY") 
 if api_key:
     genai.configure(api_key=api_key)
 else:
     st.error("⚠️ API Key missing! Go to Settings -> Secrets and add GEMINI_API_KEY.")
 
-# 2. DB INITIALIZATION
-if not os.path.exists("sales.db"):
-    import generate_db
-    generate_db.init_db()
+# --- 2. DATABASE INITIALIZATION (Self-Contained) ---
+DB_NAME = "sales.db"
+
+def init_db():
+    """Creates a dummy database if it doesn't exist."""
+    conn = sqlite3.connect(DB_NAME)
+    c = conn.cursor()
+    # Create Tables
+    c.execute('''CREATE TABLE IF NOT EXISTS customers 
+                 (id INTEGER PRIMARY KEY, name TEXT, region TEXT)''')
+    c.execute('''CREATE TABLE IF NOT EXISTS products 
+                 (id INTEGER PRIMARY KEY, name TEXT, category TEXT, price REAL)''')
+    c.execute('''CREATE TABLE IF NOT EXISTS orders 
+                 (id INTEGER PRIMARY KEY, customer_id INTEGER, product_id INTEGER, 
+                  quantity INTEGER, order_date DATE)''')
+    # Insert Sample Data
+    c.execute("INSERT OR IGNORE INTO customers VALUES (1, 'Acme Corp', 'North')")
+    c.execute("INSERT OR IGNORE INTO customers VALUES (2, 'Globex', 'West')")
+    c.execute("INSERT OR IGNORE INTO products VALUES (1, 'AI Widget', 'Software', 1000.0)")
+    c.execute("INSERT OR IGNORE INTO products VALUES (2, 'Cloud Server', 'Hardware', 5000.0)")
+    c.execute("INSERT OR IGNORE INTO orders VALUES (101, 1, 1, 5, '2023-10-01')") 
+    conn.commit()
+    conn.close()
+
+if not os.path.exists(DB_NAME):
+    init_db()
 
+# --- 3. HELPER FUNCTIONS ---
 def get_schema():
-    conn = sqlite3.connect('sales.db')
+    conn = sqlite3.connect(DB_NAME)
     cursor = conn.cursor()
     cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
     tables = cursor.fetchall()
@@ -30,7 +55,7 @@ def get_schema():
     return schema_str
 
 def run_sql(sql):
-    conn = sqlite3.connect('sales.db')
+    conn = sqlite3.connect(DB_NAME)
     try:
         df = pd.read_sql_query(sql, conn)
         return df, None
@@ -39,42 +64,63 @@ def run_sql(sql):
     finally:
         conn.close()
 
-# 3. AGENTIC LOGIC
+# --- 4. AGENTIC LOGIC ---
 def sql_agent(user_prompt):
-    # Force the model to use the most stable string
-# Try 'gemini-1.5-flash' first, but if it fails, 'gemini-pro' is the universal fallback
-try:
-    model = genai.GenerativeModel('gemini-1.5-flash')
-    # Test a small call to see if it's alive
-    model.generate_content("test") 
-except:
-    # Fallback to the older, most stable naming convention
-    model = genai.GenerativeModel('gemini-pro')
+    # Attempt to use Flash, fallback to Pro if the environment is restricted
+    try:
+        model = genai.GenerativeModel('gemini-1.5-flash')
+    except:
+        model = genai.GenerativeModel('gemini-pro')
+    
     schema = get_schema()
     
-    # ATTEMPT 1
-    full_prompt = f"Schema:\n{schema}\n\nUser: {user_prompt}\n\nOutput ONLY valid SQL."
-    res = model.generate_content(full_prompt).text.strip().replace('```sql', '').replace('```', '')
+    # SYSTEM PROMPT
+    prompt = f"""
+    You are an expert SQL assistant. Based on the schema below, write a SQLite query to answer the user's request.
+    Schema:
+    {schema}
+    
+    User Request: {user_prompt}
     
-    df, error = run_sql(res)
+    Return ONLY the raw SQL code. No markdown, no backticks, no explanations.
+    """
     
-    # REFLECTION LOOP (Self-Correction)
-    if error:
-        st.warning("🔄 Query failed. Agent is self-correcting...")
-        fix_prompt = f"The query {res} failed with error: {error}. Fix it based on this schema: {schema}. Output ONLY SQL."
-        res = model.generate_content(fix_prompt).text.strip().replace('```sql', '').replace('```', '')
-        df, error = run_sql(res)
+    try:
+        response = model.generate_content(prompt)
+        raw_sql = response.text.strip().replace('```sql', '').replace('```', '')
+        
+        # EXECUTION
+        data, error = run_sql(raw_sql)
         
-    return df, error, res
+        # SELF-REFLECTION LOOP (If error, try once more)
+        if error:
+            st.warning("🔄 Query failed. Agent is self-correcting...")
+            fix_prompt = f"The SQL query '{raw_sql}' failed with error: {error}. Rewrite the SQL correctly based on this schema: {schema}. Return ONLY SQL."
+            response = model.generate_content(fix_prompt)
+            raw_sql = response.text.strip().replace('```sql', '').replace('```', '')
+            data, error = run_sql(raw_sql)
+            
+        return data, error, raw_sql
+    except Exception as e:
+        return None, f"AI Error: {str(e)}", "None"
 
-# 4. UI
+# --- 5. UI ---
 st.title("🕵️‍♂️ Agentic Text-to-SQL")
-query = st.text_input("Ask your Database:", "Show me all software products")
+st.markdown("This agent writes SQL, executes it, and **fixes its own mistakes**.")
+
+query = st.text_input("Ask your database:", value="Show all software products")
 
 if st.button("Run Agent"):
-    data, err, final_sql = sql_agent(query)
-    if err:
-        st.error(f"Final Attempt Failed: {err}")
-    else:
-        st.code(final_sql, language="sql")
-        st.dataframe(data)
+    with st.spinner("Agent is thinking..."):
+        results, err, final_sql = sql_agent(query)
+        
+        if err:
+            st.error(f"Execution Error: {err}")
+        else:
+            st.success("Query Successful!")
+            st.code(final_sql, language="sql")
+            st.dataframe(results, use_container_width=True)
+
+st.divider()
+with st.expander("See Database Schema"):
+    st.text(get_schema())