Spaces:
Sleeping
Sleeping
from dotenv import load_dotenv | |
import streamlit as st | |
import os | |
import sqlite3 | |
import google.generativeai as genai | |
import warnings | |
warnings.filterwarnings("ignore") | |
# Load all the environment variables | |
load_dotenv() | |
# Configure genAI key | |
genai.configure(api_key=os.getenv("GOOGLE_API_KEY")) | |
# Create a function to load the Gemini model and get a response | |
def get_gemini_response(query, prompt): | |
model = genai.GenerativeModel("gemini-pro") | |
response = model.generate_content([prompt[0], query]) | |
print(response.text) | |
return response.text | |
# Function to retrieve query from the database | |
def read_sql_query(sql, db): | |
try: | |
connection = sqlite3.connect(db) | |
cursor = connection.cursor() | |
cursor.execute(sql) | |
rows = cursor.fetchall() | |
connection.commit() | |
for row in rows: | |
print(row) | |
except sqlite3.Error as e: | |
# st.error(f"An error occurred: {e}") | |
rows = [] | |
finally: | |
connection.close() | |
return rows | |
# Define the prompt | |
prompt = [""" | |
You are an expert in converting English questions to SQL code! | |
The SQL database is named STUDENT and has the following columns: Name, Class, Section. | |
For example, | |
Example 1: How many entries of records are present? | |
The SQL command will be something like this: SELECT COUNT(*) FROM STUDENT; | |
Ensure that the SQL code in the response does not have ''' at the beginning or end of the response and does not contain the word 'sql' in the output. | |
"""] | |
# Streamlit app configuration | |
st.set_page_config(page_title="RetrieveGPT") | |
st.header("Retrieve_GPT") | |
question = st.text_input("Input your question:", key='input') | |
submit = st.button("Submit") | |
# If submit is clicked | |
if submit: | |
response = get_gemini_response(question, prompt) | |
if response: | |
results = read_sql_query(response, "Student.db") | |
st.subheader("Output:") | |
if results: | |
for row in results: | |
st.write(row) | |
else: | |
st.write("No results found or an error occurred.") | |