File size: 2,144 Bytes
655b738 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
# streamlit_app.py
import streamlit as st
import pandas as pd
from sqlalchemy import create_engine, text
import openai
import os
# ---- CONFIG ----
# Set your API key as an environment variable or in a .env file
openai.api_key = os.getenv("OPENAI_API_KEY")
# Database connection (update these with your credentials)
DB_TYPE = "mysql+pymysql"
DB_USER = "username"
DB_PASS = "password"
DB_HOST = "host"
DB_PORT = "3306"
DB_NAME = "db_name"
DATABASE_URL = f"{DB_TYPE}://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
engine = create_engine(DATABASE_URL)
# ---- FUNCTIONS ----
def generate_sql(user_question, table_names=[]):
"""
Generates SQL query from user question using OpenAI GPT
"""
table_info = ""
if table_names:
table_info = f"These are your tables: {table_names}\n"
prompt = f"""
You are an expert SQL generator.
{table_info}
Write a SQL query that answers the following question:
\"\"\"{user_question}\"\"\"
Only return SQL, do not explain.
"""
response = openai.Completion.create(
engine="text-davinci-003",
prompt=prompt,
temperature=0,
max_tokens=300
)
sql_query = response.choices[0].text.strip()
return sql_query
def run_query(sql_query):
"""
Runs SQL query using SQLAlchemy
"""
try:
with engine.connect() as conn:
result = pd.read_sql(text(sql_query), conn)
return result
except Exception as e:
return f"Error executing query: {e}"
# ---- STREAMLIT UI ----
st.title("🧠 AI SQL Assistant")
st.markdown("Ask a question about your database, and it will generate SQL and show results.")
user_question = st.text_input("Enter your question:")
if st.button("Run Query") and user_question:
with st.spinner("Generating SQL..."):
sql_query = generate_sql(user_question)
st.code(sql_query, language="sql")
with st.spinner("Executing SQL..."):
result = run_query(sql_query)
if isinstance(result, pd.DataFrame):
st.success("Query executed successfully!")
st.dataframe(result)
else:
st.error(result)
|