dukepato2010/Text2SQLGenAIApp

1from dotenv import load_dotenv
2load_dotenv() #loads all the environment variables
3import streamlit as st
4import os
5import sqlite3
6import google.generativeai as genai
7import pandas as pd
8
9genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
10
11#Function to load google gemini model and provide query as response
12
13def get_gemini_response(question, prompt):
14    model = genai.GenerativeModel('gemini-pro')
15    response= model.generate_content([prompt[0], question])
16    return response.text
17
18def get_gemini_formatted(sql, response, prompt):
19    model = genai.GenerativeModel('gemini-pro')
20    response= model.generate_content([prompt[1],sql,response])
21    return response.text
22## Function to retreieve query from database
23def read_sql_query(sql,db):
24    conn = sqlite3.connect(db)
25    cursor=conn.cursor()
26    cursor.execute(sql)
27    rows=cursor.fetchall()
28    conn.commit()
29    conn.close()
30    for row in rows:
31        print(row)
32    return rows
33
34##Define your prompt
35prompt = [
36    """
37    You are an expert in converting English questions to SQL query!
38    The SQL database has the name STUDENT and has the following columns- NAME, CLASS, SECTION \n\n
39    Example 1- How many entries of records are present?, the SQL command will be something lke SELECT COUNT(*) FROM STUDENT;
40    \n\nExample 2- How many students study in Data Science class?
41    SQL query will be SELECT * FROM STUDENT where CLASS='Data Science';
42    Also, the SQL code should not have ''' in the beginning or end, or SQL word in output
43    RETURN ONLY the SQL code as response so it can be used to query a database.
44    """,
45    """
46    You are an expert in converting sql results and the given sql query into readable table format.
47    The SQL database that has been queried the name STUDENT and has the following columns- NAME, CLASS, SECTION \n\n
48    You will receive both the SQL query and the corresponding response from the database. Be smart enough
49    to convert the sql response as a readable format or table depending on the response. When returning
50    your output, if it is a table, returning a pandas dataframe output.\n\n 
51    Example if the response is [('Duke', 'Data Science', 'A'), ('Patrick', 'Data Science', 'B')], then it is a table, return 
52    pd.DataFrame([('Duke', 'Data Science', 'A'), ('Patrick', 'Data Science', 'B')], columns=['NAME', 'CLASS', 'SECTION']) as response
53    """
54
55]
56
57## Setup streamlit app
58st.set_page_config(page_title="I can retrieve any SQL query")
59st.header("Gemini App to Retrieve SQL Query Data")
60
61question = st.text_input("input: ", key="input")
62submit=st.button("Ask the question")
63
64#If submit is clicked
65if submit:
66    response1 = get_gemini_response(question, prompt)
67    response = read_sql_query(response1, "student.db")
68    st.subheader("The response is:")
69    #formatted_answer = get_gemini_formatted(response1, response, prompt)
70    for row in response:
71         print(row)
72         st.write(row)
73    st.write(f"{response1}")
74
75