Update app.py

app.py

@@ -4,18 +4,18 @@ import sqlite3
 import pandas as pd
 import os
 
-# --- 1. CONFIG & API ---
-st.set_page_config(page_title="Agentic SQL", page_icon="🕵️‍♂️")
+# --- 1. SETTINGS ---
+st.set_page_config(page_title="SQL AI Agent", layout="wide")
 
-# CRITICAL: This is where you use the key you just copied
-api_key = os.getenv("GEMINI_API_KEY") 
+# Get Key from Secrets
+api_key = os.getenv("GEMINI_API_KEY")
 
 if api_key:
     genai.configure(api_key=api_key)
 else:
-    st.error("⚠️ API Key missing! Add it to Hugging Face Settings > Secrets.")
+    st.error("⚠️ API Key missing! Go to Settings > Secrets and add 'GEMINI_API_KEY'")
 
-# --- 2. DB SETUP ---
+# --- 2. DATABASE SETUP ---
 DB_NAME = "sales.db"
 def init_db():
     conn = sqlite3.connect(DB_NAME)
@@ -30,37 +30,47 @@ def init_db():
 if not os.path.exists(DB_NAME):
     init_db()
 
-# --- 3. AGENT LOGIC ---
+# --- 3. AGENTIC LOGIC ---
 def sql_agent(user_prompt):
-    # FALLBACK LOGIC: If 'flash' fails due to the 404 error, we try 'gemini-pro'
-    try:
-        model = genai.GenerativeModel('gemini-1.5-flash')
-    except:
-        model = genai.GenerativeModel('gemini-pro')
-    
-    conn = sqlite3.connect(DB_NAME)
-    schema = "Table: customers (id, name, region); Table: products (id, name, category, price)"
+    # Forced stable model names
+    model_names = ['gemini-1.5-flash', 'gemini-1.5-pro', 'gemini-pro']
+    model = None
     
-    prompt = f"System: Output ONLY SQLite SQL. Schema: {schema}\nUser: {user_prompt}"
+    for name in model_names:
+        try:
+            model = genai.GenerativeModel(name)
+            # Test if this model works
+            model.generate_content("ping")
+            break
+        except:
+            continue
+            
+    if not model:
+        return None, "Could not connect to any Gemini models.", ""
+
+    schema = "Tables: customers (id, name, region), products (id, name, category, price)"
+    prompt = f"System: Output ONLY clean SQLite SQL. No backticks. Schema: {schema}\nUser: {user_prompt}"
     
     try:
         response = model.generate_content(prompt)
         sql = response.text.strip().replace('```sql', '').replace('```', '')
+        
+        conn = sqlite3.connect(DB_NAME)
         df = pd.read_sql_query(sql, conn)
+        conn.close()
         return df, None, sql
     except Exception as e:
         return None, str(e), "Error"
-    finally:
-        conn.close()
 
 # --- 4. UI ---
-st.title("🕵️‍♂️ Agentic SQL Fix")
-query = st.text_input("Ask your DB:", "List all North customers")
+st.title("🕵️‍♂️ Agentic Text-to-SQL")
+query = st.text_input("Ask a question:", "Show me all products")
 
-if st.button("Run"):
-    data, err, sql = sql_agent(query)
-    if err:
-        st.error(f"Error: {err}")
-    else:
-        st.code(sql)
-        st.dataframe(data)
+if st.button("Run Analysis"):
+    with st.spinner("Agent is querying..."):
+        data, err, sql = sql_agent(query)
+        if err:
+            st.error(f"Error: {err}")
+        else:
+            st.code(sql, language="sql")
+            st.dataframe(data)