Update app.py

app.py

@@ -6,57 +6,104 @@ import os
 # Load environment variables
 load_dotenv()
 
-# Configure GenerativeAI with API key
-genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
+os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")
+genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
 
 def main():
     # Function to generate SQL query based on prompt input
     def gemini_ans(prompt_input):
-        model = genai.GenerativeModel('gemini-pro')
+        model = genai.GenerativeModel('gemini-1.5-flash-latest')
         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
+    # Custom CSS for styling
     st.markdown(
         """
         <style>
+            body {
+                background-color: #f5f5f5;
+                font-family: 'Arial', sans-serif;
+            }
             .header {
+                background-color: #ffffff;
                 padding: 20px;
-                background-color: #f0f0f0;
                 border-radius: 10px;
+                box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
                 margin-bottom: 20px;
-                color: black;
             }
-            .header h1, .header h3 {
+            .header h1 {
+                margin: 0;
+                color: #333333;
+                text-align: center;
+            }
+            .header p {
                 margin: 0;
-                color: black;
+                text-align: center;
+                color: #666666;
+                font-size: 14px;
+            }
+            .input-container {
+                background-color: #ffffff;
+                padding: 20px;
+                border-radius: 10px;
+                box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
+                margin-bottom: 20px;
+            }
+            .stTextArea, .stTextInput, .stButton {
+                margin-top: 15px;
+            }
+            .stButton > button {
+                background-color: #4CAF50;
+                color: white;
+                border: none;
+                padding: 10px 20px;
+                font-size: 16px;
+                border-radius: 5px;
+                cursor: pointer;
+            }
+            .stButton > button:hover {
+                background-color: #45a049;
             }
         </style>
-        """
-    , unsafe_allow_html=True)
+        """,
+        unsafe_allow_html=True
+    )
 
+    # Header section
     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>
+            <h1>SQL Query Generator</h1>
+            <p>Effortlessly transform your natural language input into SQL queries!</p>
         </div>
         """
     , unsafe_allow_html=True)
-    
-    # Text area for input
-    input_text = st.text_area('Enter your query...')
-    
+
+    # Input section with a card-like design
+    st.markdown(
+        """
+        <div class="input-container">
+            <h4>Enter your details below:</h4>
+        </div>
+        """,
+        unsafe_allow_html=True
+    )
+
+    # Inputs
+    input_text = st.text_area('Enter your natural language query...')
+    schema_name = st.text_input('Enter the schema name (optional):', placeholder="e.g., public")
+    table_name = st.text_input('Enter the table name:', placeholder="e.g., users")
+
     # 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},
+        You are an English to SQL language translator. The schema name is {schema}, 
+        and the table name is {table}. Using the given text here {en},
         write a SQL query only without making any mistakes.
     """
 
@@ -78,12 +125,20 @@ def main():
 
     # Handle button click event
     if submit:
+        if not table_name:
+            st.error("Please enter the table name!")
+            return
+        
         with st.spinner('Generating SQL Query...'):
             # Generate SQL query
-            sql_query = gemini_ans(prompt.format(en=input_text))
+            sql_query = gemini_ans(prompt.format(
+                schema=schema_name or 'default schema', 
+                table=table_name, 
+                en=input_text
+            ))
             st.header('Model Response')
             st.success("Generated SQL Query Successfully!")
-            st.write(sql_query)
+            st.code(sql_query, language="sql")
 
             # Generate sample expected output
             sql_table = gemini_ans(prompt1.format(query=sql_query))