Update app.py

app.py

@@ -1,28 +1,81 @@
 import streamlit as st
 from openai import OpenAI
 import sqlite3
-from init_db import initialize_database
 
-# Initialize database
+# Database initialization with caching
+@st.cache_resource
+def initialize_database():
+    conn = sqlite3.connect('database.db')
+    cursor = conn.cursor()
+
+    # Create tables
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS Produkte (
+            ProduktID INTEGER PRIMARY KEY AUTOINCREMENT,
+            Produktname TEXT NOT NULL,
+            Preis REAL NOT NULL
+        )
+    ''')
+    
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS Bestellungen (
+            BestellungID INTEGER PRIMARY KEY AUTOINCREMENT,
+            ProduktID INTEGER NOT NULL,
+            Menge INTEGER NOT NULL,
+            Bestelldatum TEXT NOT NULL,
+            Person TEXT NOT NULL,
+            FOREIGN KEY (ProduktID) REFERENCES Produkte(ProduktID)
+        )
+    ''')
+
+    # Insert sample products
+    cursor.execute("SELECT COUNT(*) FROM Produkte")
+    if cursor.fetchone()[0] == 0:
+        products = [
+            ('Laptop', 999.99),
+            ('Smartphone', 699.99),
+            ('Tablet', 399.99)
+        ]
+        cursor.executemany("INSERT INTO Produkte (Produktname, Preis) VALUES (?, ?)", products)
+
+    # Insert sample orders
+    cursor.execute("SELECT COUNT(*) FROM Bestellungen")
+    if cursor.fetchone()[0] == 0:
+        orders = [
+            (1, 2, '2024-10-20', 'Max Mustermann'),
+            (2, 1, '2024-10-21', 'Erika Musterfrau'),
+            (3, 3, '2024-10-22', 'Hans Meier')
+        ]
+        cursor.executemany("INSERT INTO Bestellungen (ProduktID, Menge, Bestelldatum, Person) VALUES (?, ?, ?, ?)", orders)
+
+    conn.commit()
+    conn.close()
+
+# Initialize the database
 initialize_database()
 
-# App configuration
+# App UI
 st.set_page_config(page_title="Zero SQL", layout="wide")
 st.title("Zero SQL - Natural Language to SQL Query")
 
-# Sidebar for API key configuration
+# Sidebar for API key
 with st.sidebar:
-    st.header("API Configuration")
+    st.header("Configuration")
     api_key = st.text_input("OpenAI API Key", type="password")
+    st.markdown("---")
+    st.markdown("**Sample Questions:**")
+    st.markdown("- Show total sales per product")
+    st.markdown("- List orders from Max Mustermann")
+    st.markdown("- Find most popular product by quantity")
 
 # Main form
 with st.form("query_form"):
     user_input = st.text_area(
         "Enter your data request in natural language:",
-        placeholder="e.g. Show all orders from last week",
-        height=150
+        placeholder="e.g. Show all orders over €500",
+        height=100
     )
-    submitted = st.form_submit_button("Generate Query")
+    submitted = st.form_submit_button("🚀 Generate Query")
 
 if submitted:
     if not api_key:
@@ -31,63 +84,59 @@ if submitted:
         st.error("📝 Please enter your data request!")
     else:
         try:
-            # Initialize OpenAI client
             client = OpenAI(api_key=api_key)
             
-            # System context with schema information
-            system_context = """Given the following SQL tables, your job is to write queries given a user's request.
-                                CREATE TABLE Produkte (
-                                    ProduktID INTEGER PRIMARY KEY AUTOINCREMENT,
-                                    Produktname TEXT NOT NULL,
-                                    Preis REAL NOT NULL
-                                );
-
-                                CREATE TABLE Bestellungen (
-                                    BestellungID INTEGER PRIMARY KEY AUTOINCREMENT,
-                                    ProduktID INTEGER NOT NULL,
-                                    Menge INTEGER NOT NULL,
-                                    Bestelldatum TEXT NOT NULL,
-                                    Person TEXT NOT NULL,
-                                    FOREIGN KEY (ProduktID) REFERENCES Produkte(ProduktID)
-                                );"""
+            # 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 SQL query for the user's request. Output raw SQL without any formatting or explanations."""
             
-            # Generate SQL query using OpenAI
+            # Generate SQL query
             response = client.chat.completions.create(
                 model="gpt-4o",
                 messages=[
                     {"role": "system", "content": system_context},
-                    {"role": "user", "content": f"Generate the SQL query for: {user_input}. Only output the raw SQL query without any code block delimiters or markdown."}
+                    {"role": "user", "content": user_input}
                 ],
-                response_format={"type": "text"}
+                temperature=0.3
             )
             
             sql_query = response.choices[0].message.content.strip()
 
-            # Execute query and fetch results
-            conn = sqlite3.connect('database.db')
-            cursor = conn.cursor()
-            cursor.execute(sql_query)
-            
-            results = cursor.fetchall()
-            column_names = [description[0] for description in cursor.description]
-            conn.close()
+            # Execute and display results
+            with sqlite3.connect('database.db') as conn:
+                cursor = conn.cursor()
+                cursor.execute(sql_query)
+                results = cursor.fetchall()
+                columns = [desc[0] for desc in cursor.description]
 
-            # Display results
-            st.subheader("Generated SQL Query")
+            st.subheader("Generated SQL")
             st.code(sql_query, language="sql")
 
-            st.subheader("Query Results")
+            st.subheader("Results")
             if results:
                 st.dataframe(
                     data=results,
-                    columns=column_names,
+                    columns=columns,
                     use_container_width=True,
                     hide_index=True
                 )
             else:
                 st.info("No results found", icon="ℹ️")
-                
+
         except sqlite3.Error as e:
-            st.error(f"SQL Error: {str(e)}")
+            st.error(f"🚨 SQL Error: {str(e)}")
         except Exception as e:
-            st.error(f"An error occurred: {str(e)}")
+            st.error(f"💥 Unexpected Error: {str(e)}")