Update app.py

app.py

@@ -1,14 +1,12 @@
 import streamlit as st
-from openai import OpenAI
 import sqlite3
-import pandas as pd  # Add to top of file
-
+import pandas as pd
+import google.generativeai as genai
 
 # App UI
 st.set_page_config(page_title="Zero SQL", layout="wide")
 st.title("Zero SQL - Natural Language to SQL Query")
 
-
 # Database initialization with caching
 @st.cache_resource
 def initialize_database():
@@ -61,11 +59,10 @@ def initialize_database():
 # Initialize the database
 initialize_database()
 
-
 # Sidebar for API key
 with st.sidebar:
     st.header("Configuration")
-    api_key = st.text_input("OpenAI API Key", type="password")
+    api_key_iput = st.text_input("Gemini API Key", type="password")
     st.markdown("---")
     st.markdown("**Sample Questions:**")
     st.markdown("- Show total sales per product")
@@ -82,46 +79,48 @@ with st.form("query_form"):
     submitted = st.form_submit_button("🚀 Generate Query")
 
 if submitted:
-    if not api_key:
-        st.error("🔑 API key is required!")
-    elif not user_input:
+    api_key = api_key_iput or os.getenv("GEMINI_API_KEY")
+
+    if not user_input:
         st.error("📝 Please enter your data request!")
     else:
         try:
-            client = OpenAI(api_key=api_key)
+            # Configure Gemini
+            genai.configure(api_key=api_key)
+            model = genai.GenerativeModel('gemini-2.0-flash')
+
+            
+            # Schema context
+            system_context = """Given these SQL tables:
+                CREATE TABLE Produkte (
+                    ProduktID INTEGER PRIMARY KEY,
+                    Produktname TEXT NOT NULL,
+                    Preis REAL NOT NULL
+                );
+                CREATE TABLE Bestellungen (
+                    BestellungID INTEGER PRIMARY KEY,
+                    ProduktID INTEGER,
+                    Menge INTEGER,
+                    Bestelldatum TEXT,
+                    Person TEXT,
+                    FOREIGN KEY (ProduktID) REFERENCES Produkte(ProduktID)
+                );
+                Generate ONLY the raw SQL query for the following request. 
+                Output ONLY the pure SQL statement without any formatting, 
+                explanations, or markdown blocks."""
+            
+            # Combine context with user input
+            full_prompt = f"{system_context}\n\nUser Request: {user_input}"
             
-            # System prompt with schema
-            system_context = """You are a SQL expert. Given these SQL tables:
-            CREATE TABLE Produkte (
-                ProduktID INTEGER PRIMARY KEY,
-                Produktname TEXT NOT NULL,
-                Preis REAL NOT NULL
-            );
-            CREATE TABLE Bestellungen (
-                BestellungID INTEGER PRIMARY KEY,
-                ProduktID INTEGER,
-                Menge INTEGER,
-                Bestelldatum TEXT,
-                Person TEXT,
-                FOREIGN KEY (ProduktID) REFERENCES Produkte(ProduktID)
-            );
-            Generate ONLY the raw SQL query for the user's request. 
-            NEVER use markdown code blocks or any formatting. 
-            ONLY output the pure SQL statement."""
-        
-           # Generate SQL query
-            response = client.chat.completions.create(
-                model="gpt-4o",
-                messages=[
-                    {"role": "system", "content": system_context},
-                    {"role": "user", "content": user_input}
-                ],
-                temperature=0.3
+            # Generate SQL query
+            response = model.generate_content(
+                full_prompt,
+                generation_config={"temperature": 0.3}
             )
             
             # Clean up the response
-            sql_query = response.choices[0].message.content.strip()
-            sql_query = sql_query.replace("```sql", "").replace("```", "").strip()  # Add this line
+            sql_query = response.text.strip()
+            sql_query = sql_query.replace("```sql", "").replace("```", "").strip()
             
             # Execute and display results
             with sqlite3.connect('database.db') as conn: