Update streamlit_app.py

streamlit_app.py

@@ -28,18 +28,21 @@ if st.button("Generate & Run SQL"):
             st.success("✅ File loaded successfully")
             st.dataframe(df.head(3))
 
-            # 📋 Preview for prompt
+            # 🧠 Prompt construction
             preview = df.head(5).to_string(index=False)
             prompt = f"""
-You are an AI assistant. Generate an SQL query for a pandas DataFrame named 'df' based on the user's question.
+You are a SQL expert assistant. Generate an SQL query for a pandas DataFrame named 'df' based on the user's question.
 
 Data preview:
 {preview}
 
-Question:
+User question:
 {user_query}
 
-Return ONLY the SQL query. Do not include any explanations or Python code.
+IMPORTANT:
+- Only return a valid SQL query (no Python, no comments, no markdown, no explanations).
+- Use SQL compatible with SQLite.
+- Table name is 'df'.
 """
 
             # 🔗 Together API request
@@ -53,7 +56,7 @@ Return ONLY the SQL query. Do not include any explanations or Python code.
                     "max_tokens": 256,
                     "temperature": 0.3,
                     "messages": [
-                        {"role": "system", "content": "You generate SQL queries from natural language questions."},
+                        {"role": "system", "content": "You are a SQL assistant. You return only clean SQL code."},
                         {"role": "user", "content": prompt}
                     ]
                 }
@@ -67,7 +70,9 @@ Return ONLY the SQL query. Do not include any explanations or Python code.
                 if response.status_code != 200:
                     raise ValueError(f"API Error: {response.status_code} - {response.text}")
 
+                # 🧼 Clean the response
                 sql_query = response.json()['choices'][0]['message']['content'].strip()
+                sql_query = sql_query.replace("```sql", "").replace("```", "").strip()
                 st.code(sql_query, language='sql')
 
                 # 🧪 Execute SQL