NLP_to_SQL / app.py
atmiLLM's picture
Update app.py
4e6e76d verified
import streamlit as st
import psycopg2
import google.generativeai as genai
# Function to connect to PostgreSQL
def connect_to_db():
try:
conn = psycopg2.connect(
dbname="postgres",
user="postgres",
password="database",
host="localhost",
port="5432"
)
return conn
except Exception as e:
st.error(f"Error connecting to the database: {e}")
return None
# Function to generate SQL query from natural language using Google Gemini
def generate_sql_from_text(prompt):
try:
# Initialize the Gemini client
genai.configure(api_key="AIzaSyCzQ1u1RfAOEgYy2RVyFGGGSzEsiKEtWMk")
model = genai.GenerativeModel('gemini-pro')
# Send the prompt to the Gemini model
response = model.generate_text(prompt)
# Extract the SQL query from the response
sql_query = response["generated_text"]
return sql_query
except Exception as e:
st.error(f"Error generating SQL query: {e}")
return None
# Function to execute SQL query and return results
def execute_sql_query(conn, query):
try:
cursor = conn.cursor()
cursor.execute(query)
results = cursor.fetchall()
cursor.close()
return results
except Exception as e:
st.error(f"Error executing SQL query: {e}")
return None
# Streamlit UI
st.title("Natural Language to SQL with Google Gemini")
# Input prompt from the user
user_prompt = st.text_input("Enter your query in natural language:")
# Connect to the PostgreSQL database
conn = connect_to_db()
if conn and user_prompt:
# Generate SQL query from natural language prompt
st.write("Generating SQL query using Google Gemini...")
sql_query = generate_sql_from_text(user_prompt)
if sql_query:
st.write(f"Generated SQL Query: {sql_query}")
# Execute the generated SQL query
results = execute_sql_query(conn, sql_query)
# Display the results
if results:
st.write("Query Results:")
st.dataframe(results)
# Close the database connection
conn.close()