Equinox-M/Text-To-SQL-Gemini-Model

1import streamlit as st2import sqlite33import os4from dotenv import load_dotenv 5import google.generativeai as genai6#Load all the environment variables7load_dotenv()8# Configuration Genai key9genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))10 11# Function to Load Google Gemini 12def get_gemini_response(question, prompt):13    model = genai.GenerativeModel('gemini-pro')14    response = model.generate_content([prompt[0], question])15    return response.text16 17# Function to retrieve query from the database18def read_sql_query(sql, db):19    connection = sqlite3.connect(db)20    try:21        cursor = connection.cursor()22        cursor.execute(sql)23        data_rows = cursor.fetchall()24        connection.commit()  # Commit changes if any25        return data_rows26    finally:27        connection.close()28 29# Define Your Prompt30prompt = [31    """32You are an expert in converting English questions to SQL code!33The SQL database has the name STUDENT and has the following Columns - NAME, CLASS, SECTION34 35For example:36- How many entries of records are present? SQL command: SELECT COUNT(*) FROM STUDENT;37- Tell me all the students studying in Data Science class? SQL command: SELECT * FROM STUDENT where CLASS="Data Science";38 39Also, the SQL code should not have ''' in the beginning or at the end, and SQL word in output.40    """41]42 43# Setting up the Streamlit App 44st.set_page_config(page_title="Gemini App: Convert Text to SQL", layout="wide")45 46# Define layout47st.markdown("# Gemini App For Converting Text to SQL 🌐🗄️")48st.markdown("### Enter Your Question")49question = st.text_input("Input: ", key="input")50submit = st.button("Submit")51 52# Show generated SQL query and query results53if submit:54    ai_response = get_gemini_response(question, prompt)55    st.markdown("## Generated SQL Query")56    st.code(ai_response)57 58    try:59        sql_response = read_sql_query(ai_response, "student.db")60 61        st.markdown("## Query Results")62        if len(sql_response) > 0:63            # Display results as a table64            st.table(sql_response)65        else:66            st.write("No results found.")67        68        # Check if the SQL query is a data manipulation query69        if ai_response.strip().lower().startswith(("insert", "update", "delete")):70            # Add message indicating successful execution for data manipulation queries71            st.success("Query execution successful.")72    except Exception as e:73        st.error("Error executing SQL query:", e)74