Spaces:
Sleeping
Sleeping
File size: 2,442 Bytes
e9bcde7 |
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 |
import re
import sqlite3
from dotenv import load_dotenv
import os
import streamlit as st
import google.generativeai as genai
load_dotenv()
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
def get_gemini_response(question, prompt):
model = genai.GenerativeModel("gemini-1.5-pro-latest")
response = model.generate_content([prompt[0], question])
return response.text
def read_sql_query(response, db):
conn = sqlite3.connect(db)
cur = conn.cursor()
# Extract SQL query from the response
sql_match = re.search(r'```sql\n(.*?)\n```', response, re.DOTALL)
if sql_match:
sql = sql_match.group(1).strip()
else:
# If no SQL code block is found, try to find a SQL-like statement
sql_match = re.search(r'SELECT.*?FROM.*?;', response, re.DOTALL | re.IGNORECASE)
if sql_match:
sql = sql_match.group(0).strip()
else:
st.header("Error: No valid SQL query found in the response.\nPlease write more clear retrieval query.")
exit()
try:
cur.execute(sql)
rows = cur.fetchall()
conn.close()
return rows
except sqlite3.Error as e:
conn.close()
return f"SQLite error: {e}"
prompt = [
"""
You are an expert in converting English questions to SQL code.
The SQL database has the name STUDENT and has the following columns - NAME, CLASS,
SECTION \n\nFor example, \nExample 1: How many entries of records are present?,
the SQL command will be something like this SELECT COUNT(*) FROM STUDENT ;
\nExample 2: What are the names of students who study in AI?, the SQL command will be
something like this SELECT NAME FROM STUDENT WHERE CLASS = "AI";
also the sql code should not have ''' in beginning or end or a 'sql' word in output
"""
]
st.set_page_config(page_title="SQL Retrieval App")
st.info("This is a demo the default db is a student database containing a table of student name, class and section")
st.header("Gemini App - Retrieve SQL Data")
question = st.text_input("Input a retrieval question about the database (e.g. get all student names and class) ", key="input")
submit = st.button("Ask SQL")
if submit:
response = get_gemini_response(question, prompt)
print(response)
response = read_sql_query(response, "student.db")
st.subheader("Response: ")
for row in response:
print(row)
st.header(row)
|