Update app.py

app.py

@@ -1,63 +1,100 @@
-from dotenv import load_dotenv
 import streamlit as st
-import os
-import sqlite3
 import google.generativeai as genai
+from dotenv import load_dotenv
+import os
+
+# Load environment variables
+load_dotenv()
+
+# Configure GenerativeAI with API key
+genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
+
+def main():
+    # Function to generate SQL query based on prompt input
+    def gemini_ans(prompt_input):
+        model = genai.GenerativeModel('gemini-pro')
+        response = model.generate_content([prompt_input])
+        return response.text
+
+    # Set page configuration
+    st.set_page_config(page_title='SQL Query Generator', page_icon=':robot:')
+    
+    # Header section
+    st.markdown(
+        """
+        <style>
+            .header {
+                padding: 20px;
+                background-color: #f0f0f0;
+                border-radius: 10px;
+                margin-bottom: 20px;
+                color: black;
+            }
+            .header h1, .header h3 {
+                margin: 0;
+                color: black;
+            }
+        </style>
+        """
+    , unsafe_allow_html=True)
 
-# Configure API key
-genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
-
-def get_gemini_response(question, prompt):
-    model = genai.GenerativeModel('gemini-pro')
-    response = model.generate_content([prompt[0], question])
-    return response.text.strip()
-
-def read_sql_query(sql, db):
-    conn = sqlite3.connect(db)
-    cur = conn.cursor()
-    cur.execute(sql)
-    rows = cur.fetchall()
-    conn.commit()
-    conn.close()
-    return rows
-
-# Prompt for the application
-prompt = [
+    st.markdown(
+        """
+        <div class="header">
+            <h1 style="text-align: center;">SQL Query Generator 🤖</h1>
+            <h3 style="text-align: center;">Generate SQL queries with ease!</h3>
+            <p style="text-align: center;">This tool allows you to generate SQL queries based on your prompt.</p>
+        </div>
+        """
+    , unsafe_allow_html=True)
+    
+    # Text area for input
+    input_text = st.text_area('Enter your query...')
+    
+    # Generate SQL Query button
+    submit = st.button('Generate SQL Query', key='generate_button')
+    
+    # Prompts for model responses
+    prompt = """
+        You are an English to SQL language translator. Using the given text here {en},
+        write a SQL query only without making any mistakes.
     """
-    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 \n\nFor example,\nExample 1 - How many entries of records are present?, 
-    the SQL command will be something like this SELECT COUNT(*) FROM STUDENT ;
-    \nExample 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"; 
+
+    prompt1 = """
+        What would be the expected response of this query snippet:
+        ```
+        {query}
+        ```
+        Provide a sample tabular response with no explanation.
+    """
+
+    prompt2 = """
+        Explain the SQL query:
+        ```
+        {query}
+        ```
+        Please provide the simplest explanation.
     """
-]
-
-# Streamlit app
-st.set_page_config(page_title="Text to SQL Query Converter")
-st.title("Text to SQL Query Converter")
-
-# User input
-question = st.text_input("Enter your question:", key="input")
-
-# Submit button
-if st.button("Convert to SQL Query"):
-    if not question:
-        st.error("Please enter a question.")
-    else:
-        # Generate SQL query from the question
-        sql_query = get_gemini_response(question, prompt)
-        st.write("SQL Query:")
-        st.code(sql_query)
-
-        # Execute the SQL query and display results
-        try:
-            results = read_sql_query(sql_query, "student.db")
-            if results:
-                st.success("Query executed successfully. Results:")
-                st.table(results)
-            else:
-                st.warning("No results found.")
-        except Exception as e:
-            st.error(f"An error occurred: {e}")
+
+    # Handle button click event
+    if submit:
+        with st.spinner('Generating SQL Query...'):
+            # Generate SQL query
+            sql_query = gemini_ans(prompt.format(en=input_text))
+            st.header('Model Response')
+            st.success("Generated SQL Query Successfully!")
+            st.write(sql_query)
+
+            # Generate sample expected output
+            sql_table = gemini_ans(prompt1.format(query=sql_query))
+            st.success('Sample Expected Output')
+            st.markdown(sql_table)
+
+            # Generate explanation of the given query
+            sql_explanation = gemini_ans(prompt2.format(query=sql_query))
+            st.success('Explanation of the Given Query')
+            st.markdown(sql_explanation)
+
+# Execute main function
+if __name__ == "__main__":
+    main()