Akashr18/SQL-Query-Retrieval-App

1import streamlit as st2import os3import sqlite34import google.generativeai as genai5 6from dotenv import load_dotenv7load_dotenv()8 9genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))10 11## Function To Load Google Gemini Model and provide queries as response12def gemini_response(system_message, input):13    model = genai.GenerativeModel(model_name= 'gemini-pro')  14    response = model.generate_content([system_message[0], input])15    return response.text16 17## Fucntion To retrieve query from the database18def read_sql_query(sql,db):19    conn=sqlite3.connect(db)20    cur=conn.cursor()21    cur.execute(sql)22    rows=cur.fetchall()23    conn.commit()24    conn.close()25    for row in rows:26        print(row)27    return rows28 29## System message30prompt=[31    """32    You are an expert in converting English questions to SQL query!33    The SQL database has the name STUDENT and has the following columns - NAME, CLASS, 34    SECTION and MARKS \n\nFor example,\nExample 1 - How many entries of records are present?, 35    the SQL command will be something like this SELECT COUNT(*) FROM STUDENT ;36    \nExample 2 - Tell me all the students studying in Data Science class?, 37    the SQL command will be something like this SELECT * FROM STUDENT 38    where CLASS="Data Science"; 39    also the sql code should not have ``` in beginning or end and sql word in output40    """41]42 43st.set_page_config(page_title="I can Retrieve Any SQL query")44st.header("SQL Data Retrieval App")45 46input = st.text_input("Enter your query", key='input')47 48button = st.button("Get Response")49 50if button:51    response = gemini_response(prompt, input)52    print("Response: ", response)53    response_sql = read_sql_query(response,"student.db")54    st.subheader("Response: ")55    for row in response_sql:56        st.write(row)57    st.subheader("SQL Query used: ")58    st.write(response)