Spaces:
Sleeping
Sleeping
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() | |