File size: 4,213 Bytes
bae13dc | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | from dotenv import load_dotenv
import streamlit as st
import os
import sqlite3
import google.generativeai as genai
load_dotenv()
env_var = os.getenv('GOOGLE_API_KEY')
if env_var is None:
raise EnvironmentError("GOOGLE_API_KEY not found in environment variables.")
genai.configure(api_key=env_var)
def get_gemini_response(question, prompt):
try:
# Initialize the generative model
model = genai.GenerativeModel("gemini-pro")
# Generate content using the model
response = model.generate_content(prompt + question)
# Return the response text
return response.text
except genai.api_core.exceptions.GoogleAPICallError as e:
print(f"A Google API call error occurred: {e}")
return None
except genai.api_core.exceptions.RetryError as e:
print(f"A retry error occurred: {e}")
return None
except Exception as e:
print(f"An unexpected error occurred: {e}")
return None
# Function to retrieve query from database
def read_sql_query(sql, db):
try:
con = sqlite3.connect(db)
cur = con.cursor()
cur.execute(sql)
rows = cur.fetchall()
except sqlite3.Error as e:
print(f"An error occurred: {e}")
rows = []
finally:
if con:
con.close()
return rows
prompt = """
You are an expert in converting complex English questions to SQL queries!
The SQL database has the name STUDENTS and has the following columns - NAME, CLASS, LEVEL, MARKS.
For example,
Example 1 - How many entries of records are present?
The SQL command will be something like this: SELECT COUNT(*) FROM STUDENTS;
Example 2 - List all student names.
The SQL command will be something like this: SELECT NAME FROM STUDENTS;
Example 3 - Find the number of students in each class.
The SQL command will be something like this: SELECT CLASS, COUNT(*) FROM STUDENTS GROUP BY CLASS;
Example 4 - Retrieve all information for students in level 'LEVEL 100'.
The SQL command will be something like this: SELECT * FROM STUDENTS WHERE LEVEL = 'LEVEL 100';
Example 5 - Get the names of students in class 'AI'.
The SQL command will be something like this: SELECT NAME FROM STUDENTS WHERE CLASS = 'AI';
Example 6 - Count the number of students in level 'LEVEL 300'.
The SQL command will be something like this: SELECT COUNT(*) FROM STUDENTS WHERE LEVEL = 'LEVEL 300';
Example 7 - Find the average marks of students in each class.
The SQL command will be something like this: SELECT CLASS, AVG(MARKS) FROM STUDENTS GROUP BY CLASS;
Example 8 - Retrieve the names and marks of students who have marks greater than 90.
The SQL command will be something like this: SELECT NAME, MARKS FROM STUDENTS WHERE MARKS > 90;
Example 9 - List the names of students in 'Data Science' class and 'Level 100'.
The SQL command will be something like this: SELECT NAME FROM STUDENTS WHERE CLASS = 'Data Science' AND LEVEL = 'Level 100';
Example 10 - Find the student with the highest marks.
The SQL command will be something like this: SELECT NAME, MARKS FROM STUDENTS WHERE MARKS = (SELECT MAX(MARKS) FROM STUDENTS);
also the sql code should not have ``` in beginning or end and sql word in output
"""
st.set_page_config(page_title="I Can Retrieve Any SQL Query")
st.header("Gemini App To Retrieve SQL Data")
question = st.text_input("Input: ", key="input")
submit = st.button("Ask the Question")
# If submit is clicked
if submit:
response = get_gemini_response(question, prompt)
if response:
print(response)
try:
response = read_sql_query(response, "students.db")
st.subheader("The Response is")
for row in response:
print(row)
st.write(row)
except (sqlite3.OperationalError, sqlite3.IntegrityError) as e:
st.error(f"A database error occurred: {e}")
except Exception as e:
st.error(f"An unexpected error occurred while executing the SQL query: {e}")
else:
st.error("Failed to generate a valid SQL query.")
|