Adith29/Text_To_SQL_Generative_AI_model

1from dotenv import load_dotenv
2
3load_dotenv() ## load all the environment variables
4import streamlit as st
5import os
6import sqlite3
7
8import google.generativeai as genai
9
10
11## CONFIGURE GENAI API KEY
12
13genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
14
15
16## Funciton to load gemini model && provide Queries as response
17
18def get_gemini_response(question,prompt):
19    model=genai.GenerativeModel("gemini-pro")
20    response=model.generate_content([prompt[0],question])
21    return response.text
22
23
24##functio to retrieve query from database
25
26def read_sql_query(sql,db):
27    conn = sqlite3.connect(db)
28    cur = conn.cursor()
29    cur.execute(sql)
30    rows=cur.fetchall()
31    conn.commit()
32    conn.close()
33    for r in rows:
34        print(r)
35    return rows
36
37##prompt designing
38
39prompt=[
40    '''
41    You are an expert in converting English questions to SQL query!
42    The SQL database has the name STUDENT and  has the following columns - NAME,CLASS,
43    SECTION\n\n For example ,\n Example1- How many entries of records are present?,
44    the SQL command will be something like this SELECT COUNT(*) FROM STUDENT ;
45    Example2- Tell me how many students study in ComputerScience class?,
46    the SQL command will be something like this SELECT * FROM STUDENT where CLASS = "ComputerScience";
47    also the sql code should not have ``` in the beginning or end and sql word in output 
48    '''
49]
50
51##Streamlit Setup 
52
53st.set_page_config(page_title="SQL_Helper")
54st.header("GEMINI APP TO RETRIEVE SQL DATA")
55
56question = st.text_input("input",key="input");
57
58submit = st.button("ASK ME")
59
60## Submit clicked --
61
62if submit:
63    response = get_gemini_response(question,prompt)
64    if response:
65        st.write("Generated SQL Query:", response)
66        query_results = read_sql_query(response, "student.db")
67        st.subheader("Response:")
68        if query_results:
69            for r in query_results:
70                st.write(r)
71        else:
72            st.write("No results found.")