app.py

# Import necessary libraries
from dotenv import load_dotenv # For loading environment variables from .env file
import streamlit as st # For building the web application UI
import os # For interacting with the operating system (e.g., getting environment variables)
import sqlite3 # For interacting with SQLite database

import google.generativeai as genai # Google Generative AI SDK

## Load all the environment variables from .env file
# This line will look for a .env file in the same directory as app.py
# and load any key-value pairs found there into the environment.
load_dotenv()


API_KEY = "API Key" # Your provided API key

# Configure the Generative AI library with the retrieved API key.
genai.configure(api_key=API_KEY)


def get_gemini_response(question, prompt):
    model = genai.GenerativeModel('gemini-2.5-pro') # Changed model to gemini-2.5-pro
    
    # Generate content using the model.
    # The input is a list containing the prompt (as a list with one element) and the user's question.
    response = model.generate_content([prompt[0], question])
    
    # Return the text content of the model's response.
    return response.text

## Function to retrieve data from the SQLite database
# This function takes an SQL query and a database file path,
# executes the query, and returns the fetched rows.
def read_sql_query(sql, db):
    conn = None # Initialize connection object to None
    try:
        # Connect to the SQLite database.
        conn = sqlite3.connect(db)
        # Create a cursor object to execute SQL commands.
        cur = conn.cursor()
        
        # Execute the provided SQL query.
        cur.execute(sql)
        
        # Fetch all rows from the executed query result.
        rows = cur.fetchall()
        
        # Commit any changes made to the database (important for INSERT, UPDATE, DELETE).
        conn.commit()
        
        # Return the fetched rows.
        return rows
    except sqlite3.Error as e:
        # Catch any SQLite database errors and display them in Streamlit.
        st.error(f"Database error: {e}")
        return [] # Return an empty list on error to prevent further issues
    finally:
        # Ensure the database connection is closed, even if an error occurred.
        if conn:
            conn.close()

## Define the Prompt for the Gemini Model
# This prompt instructs the Gemini model on how to convert English questions into SQL queries.
# It includes examples and specifies the database schema.
prompt = [
    """
    You are an expert in converting English questions to SQL query!
    The SQL database has the name STUDENT and has the following columns - NAME, CLASS,
    SECTION, MARKS.
    
    For example:
    Example 1 - How many entries of records are present?,
    the SQL command will be something like this: SELECT COUNT(*) FROM STUDENT;
    
    Example 2 - Tell me all the students studying in Data Science class?,
    the SQL command will be something like this: SELECT * FROM STUDENT WHERE CLASS="Data Science";
    
    Example 3 - What are the names of students in section A?,
    the SQL command will be something like this: SELECT NAME FROM STUDENT WHERE SECTION="A";
    
    Example 4 - What are the marks of Krish?,
    the SQL command will be something like this: SELECT MARKS FROM STUDENT WHERE NAME="Krish";

    Please ensure the SQL code does not have ``` in the beginning or end and does not contain the word 'sql' in the output.
    """
]

## Streamlit Application UI
# Set the page configuration for the Streamlit app.
st.set_page_config(page_title="SQL Query Generator with Gemini")
st.header("Prompt to Query - App To Retrieve SQL Data")

# Text input for the user to ask a question.
question = st.text_input("Input your question about the student database:", key="input")

# Button to submit the question.
submit = st.button("Generate SQL and Retrieve Data")

# Check if the submit button is clicked.
if submit:
    try:
        # Get the SQL query generated by the Gemini model.
        sql_query_response = get_gemini_response(question, prompt)
        
        # Display the generated SQL query to the user for transparency.
        st.subheader("Generated SQL Query:")
        st.code(sql_query_response, language='sql') # Use st.code for better SQL formatting

        # Execute the generated SQL query against the database.
        # Ensure 'student.db' exists in the same directory or provide its full path.
        db_rows = read_sql_query(sql_query_response, "student.db")

        st.subheader("The Retrieved Data is:")
        
        # Display the retrieved data.
        if db_rows:
            # Use st.table for a clean, tabular display of results.
            st.table(db_rows)
        else:
            # Inform the user if no data was found or if the query returned no results.
            st.info("No data found for the given query, or the query returned an empty result set.")

    except Exception as e:
        # Catch any other unexpected errors during the process.
        st.error(f"An unexpected error occurred: {e}")