zero_sql / app.py
Sebbe33's picture
Update app.py
0941bd4 verified
raw
history blame
5 kB
import streamlit as st
from openai import OpenAI
import sqlite3
import pandas as pd # Add to top of file
# 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():
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()
# Sidebar for API key
with st.sidebar:
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 over €500",
height=100
)
submitted = st.form_submit_button("πŸš€ Generate Query")
if submitted:
if not api_key:
st.error("πŸ”‘ API key is required!")
elif not user_input:
st.error("πŸ“ Please enter your data request!")
else:
try:
client = OpenAI(api_key=api_key)
# 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
)
# Clean up the response
sql_query = response.choices[0].message.content.strip()
sql_query = sql_query.replace("```sql", "").replace("```", "").strip() # Add this line
# Execute and display results
with sqlite3.connect('database.db') as conn:
cursor = conn.cursor()
cursor.execute(sql_query)
results = cursor.fetchall()
column_names = [desc[0] for desc in cursor.description]
st.subheader("Generated SQL")
st.code(sql_query, language="sql")
st.subheader("Results")
if results:
df = pd.DataFrame(results, columns=column_names)
st.dataframe(
data=df,
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)}")
except Exception as e:
st.error(f"πŸ’₯ Unexpected Error: {str(e)}")