milad818/sql-query-retriever

1from dotenv import load_dotenv
2
3import streamlit as st
4import os
5
6import google.generativeai as genai
7
8from helper import *
9
10
11
12load_dotenv()
13
14# configure google api key
15genai.configure(api_key=os.getenv("GOOGLE_API-KEY"))
16
17prompt = [
18    """
19        You are an expert in converting English Questions to SQL query.
20        The SQL database is named as STUDENT and has the following columns NAME,
21        CLASS, SECTION \n\n Below are the examples: \n\n Example 1: How many entries are present in the table?
22        The SQL command will be something like this SELECT COUNT(*) FROM STUDENT; \n
23        Example 2: How many students are taking the Computer Science course? The SQL command will be something like
24        SELECT * FROM STUDENT WHERE CLASS="Computer Science"; 
25        also the SQL code should not have ''' in the beginning or end neither the SQL word in output.
26    """
27]
28
29
30# streamlit app
31st.set_page_config(page_title="Retrieve SQL Query via Gemini Model")
32st.header("Ask me about your database..")
33
34question = st.text_input("Input: ", key="input")
35submit = st.button("Ask Question")
36
37if submit:
38    db = "student.db"
39    response = get_response(prompt, question)
40    response = read_sql_query(response, db)
41    st.subheader("The response is: ")
42
43    st.markdown("""
44                <style>
45                .card {
46                background-color: #f0f2f6;
47                padding: 1rem;
48                margin-bottom: 1rem;
49                border-radius: 10px;
50                box-shadow: 0 2px 5px rgba(0,0,0,0.1);
51                }
52                </style>
53                """, unsafe_allow_html=True)
54
55    for row in response:
56        print(row)
57        st.markdown(f"<div class='card'>📝 {row}</div>", unsafe_allow_html=True)