rachitdani/Text-to-SQL-GenAI

1from dotenv import load_dotenv2 3load_dotenv() # load all envirnment variables4 5import streamlit as st6import os7import sqlite38 9import google.generativeai as genai10 11#Configure Genai Key12 13genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))14 15# Function to Load Google Gemini Model16 17def get_gemini_response(question,prompt):18    model=genai.GenerativeModel('gemini-pro')19    response=model.generate_content([prompt[0],question])20    return response.text21 22#Function to retrieve query from database23 24def read_sql_query(sql,db):25    conn = sqlite3.connect(db)26    cur = conn.cursor()27    cur.execute(sql)28    rows=cur.fetchall()29    conn.commit()30    conn.close()31    for row in rows:32        print(row)33    return rows34 35# Define your prompt 36prompt=[37    """38    You are an expert in converting English questions to SQL query!39    The SQL database has the name STUDENT and has the following columns - NAME, CLASS, 40    SECTION \n\nFor example,\nExample 1 - How many entries of records are present?, 41    the SQL command will be something like this SELECT COUNT(*) FROM STUDENT ;42    \nExample 2 - Tell me all the students studying in Data Science class?, 43    the SQL command will be something like this SELECT * FROM STUDENT 44    where CLASS="Data Science"; 45    also the sql code should not have ``` in beginning or end and sql word in output46    47    """48]49 50# Streamlit app51 52st.set_page_config(page_title="I Can Retrieve any SQL query")53st.header("Gemini App to Retrieve SQL Data")54 55question=st.text_input("Input: ",key="input")56 57submit = st.button("Ask the question")58 59# if submit is clicked60if submit:61    response=get_gemini_response(question, prompt)62    print(response)63    response=read_sql_query(response, "student.db")64    st.subheader("The Response is")65    for row in response:66        print(row)67        st.text(row)