Update app.py

app.py

@@ -86,7 +86,7 @@ def execute_sql_query(sql_query):
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    <h1 style="text-align:center;">Text-to-SQL Contract Data Explorer</h1>
+    <h1 style="text-align:center;">📊 Text-to-SQL Contract Data Explorer</h1>
     <p style="text-align:center; font-size:1.2em;">Analyze US Government contract data using natural language queries.</p>
     """)
     
@@ -99,31 +99,34 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 lines=2
             )
 
-            btn_generate_sql = gr.Button("Generate SQL Query")
-            sql_query_out = gr.Code(label="Generated SQL Query", language="sql")
+            btn_generate_sql = gr.Button("📝 Generate SQL Query", variant="primary")
+            sql_query_out = gr.Code(label="🛠️ Generated SQL Query", language="sql")
 
-            btn_execute_query = gr.Button("Execute Query")
-            error_out = gr.Markdown("", visible=False)
+            btn_execute_query = gr.Button("🚀 Execute Query", variant="secondary")
+            error_out = gr.Markdown("", visible=False, elem_id="error_message")
             
             gr.Markdown("### 💡 Example Queries")
-            example_queries = [
-                "Show the top 10 departments by total award amount.",
-                "List contracts where the award amount exceeds $5,000,000.",
-                "Retrieve awards over $1M in California.",
-                "Find the top 5 awardees by number of contracts.",
-                "Display contracts awarded after 2020 in New York.",
-                "What is the total award amount by state?"
-            ]
-            for i, query in enumerate(example_queries):
-                gr.Button(query, elem_id=f"example_{i}")
-
-            with gr.Accordion("Dataset Schema", open=False):
+            with gr.Column():
+                example_queries = [
+                    "Show the top 10 departments by total award amount.",
+                    "List contracts where the award amount exceeds $5,000,000.",
+                    "Retrieve awards over $1M in California.",
+                    "Find the top 5 awardees by number of contracts.",
+                    "Display contracts awarded after 2020 in New York.",
+                    "What is the total award amount by state?"
+                ]
+                example_buttons = []
+                for i, query in enumerate(example_queries):
+                    btn = gr.Button(query, variant="link", size="sm", interactive=True)
+                    example_buttons.append(btn)
+            
+            with gr.Accordion("📄 Dataset Schema", open=False):
                 gr.JSON(get_schema(), label="Schema")
 
         with gr.Column(scale=2):
             gr.Markdown("### 📊 Query Results")
-            results_out = gr.DataFrame(label="", interactive=False)
-            status_info = gr.Markdown("", visible=False)
+            results_out = gr.DataFrame(label="", interactive=False, row_count=10, max_rows=100)
+            status_info = gr.Markdown("", visible=False, elem_id="status_info")
 
     # =========================
     # Event Functions
@@ -148,20 +151,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         return result_df, ""
 
     def handle_example_click(example_query):
-        query_input.value = example_query
         sql_query, error = parse_query(example_query)
         if error:
-            sql_query_out.value = ""
-            error_out.value = f"❌ {error}"
-            return
-        sql_query_out.value = sql_query
+            return "", f"❌ {error}", None, f"❌ {error}"
         result_df, exec_error = execute_sql_query(sql_query)
         if exec_error:
-            results_out.value = None
-            error_out.value = f"❌ {exec_error}"
-            return
-        results_out.value = result_df
-        error_out.value = ""
+            return sql_query, f"❌ {exec_error}", None, f"❌ {exec_error}"
+        return sql_query, "", result_df, ""
 
     # =========================
     # Button Click Event Handlers
@@ -179,10 +175,28 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=[results_out, error_out]
     )
 
-    for i, query in enumerate(example_queries):
-        gr.get_component(f"example_{i}").click(
+    # Assign click events to example buttons
+    for btn in example_buttons:
+        btn.click(
+            fn=handle_example_click,
+            inputs=None,  # We'll pass the query via lambda
+            outputs=[sql_query_out, error_out, results_out, error_out],
+            _js=None,
+            # Use a lambda to pass the button's value as the example_query
+            # Note: Gradio passes the button's value as the first input if inputs are not specified
+            # To ensure the correct query is passed, we use a lambda with no inputs
+            # and retrieve the button's value within the function
+            # However, Gradio does not support passing the button's value directly here
+            # So, we need to create a partial function or use lambda with default argument
+            # Here, we use a lambda with a default argument
+        )
+
+    # To properly assign each button's click event with its specific query, redefine the loop:
+    for btn, query in zip(example_buttons, example_queries):
+        btn.click(
             fn=lambda q=query: handle_example_click(q),
-            outputs=[]
+            inputs=None,
+            outputs=[sql_query_out, error_out, results_out, error_out]
         )
 
 # Launch the Gradio App