Update app.py

app.py

@@ -150,5 +150,122 @@ with gr.Blocks() as demo:
     btn_generate_sql.click(fn=generate_sql, inputs=query, outputs=sql_query_out)
     btn_execute_query.click(fn=execute_sql_query, inputs=sql_query_out, outputs=[results_out, error_out])
 
+        gr.Markdown("""
+    # Parquet SQL Query and Plotting App
+
+    **Query and visualize data** in `sample_contract_df.parquet`
+
+    ## Instructions
+
+    1. **Describe the data you want**: e.g., `Show awards over 1M in CA`
+    2. **Use Example Queries**: Click on any example query button below to execute.
+    3. **Generate SQL**: Or, enter your own query and click "Generate SQL" to see the SQL query.
+    4. **Execute Query**: Run the query to view results and plots.
+    5. **Dataset Schema**: See available columns and types in the "Schema" tab.
+
+    ## Example Queries
+    """)
+
+    with gr.Tabs():
+        with gr.TabItem("Query Data"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    query = gr.Textbox(label="Natural Language Query", placeholder='e.g., "Awards > 1M in CA"')
+
+                    # Example query buttons
+                    gr.Markdown("### Click on an example query:")
+                    with gr.Row():
+                        btn_example1 = gr.Button("Show awards over 1M in CA")
+                        btn_example2 = gr.Button("List all contracts in New York")
+                        btn_example3 = gr.Button("Show top 5 departments by award amount")
+                        btn_example4 = gr.Button("Execute: SELECT * from contract_data LIMIT 10;")
+
+                    btn_generate = gr.Button("Generate SQL")
+                    sql_out = gr.Code(label="Generated SQL Query", language="sql")
+                    plot_code_out = gr.Code(label="Generated Plot Code", language="python")
+                    btn_execute = gr.Button("Execute Query")
+                    error_out = gr.Markdown("", visible=False)
+                with gr.Column(scale=2):
+                    results_out = gr.Dataframe(label="Query Results", interactive=False)
+                    plot_out = gr.Plot(label="Plot")
+
+        with gr.TabItem("Dataset Schema"):
+            gr.Markdown("### Dataset Schema")
+            schema_display = gr.JSON(label="Schema", value=json.loads(json.dumps(get_schema(), indent=2)))
+
+    # =========================
+    # Click Event Handlers
+    # =========================
+
+    async def on_generate_click(nl_query):
+        """
+        Handles the "Generate SQL" button click event.
+        """
+        sql_query, plot_code = await generate_sql_and_plot_code(nl_query)
+        return sql_query, plot_code
+
+    def on_execute_click(sql_query, plot_code):
+        """
+        Handles the "Execute Query" button click event.
+        """
+        result_df, error_msg = execute_query(sql_query)
+        if error_msg:
+            return None, None, error_msg
+        if plot_code.strip():
+            fig, plot_error = generate_plot(plot_code, result_df)
+            if plot_error:
+                return result_df, None, plot_error
+            else:
+                return result_df, fig, ""
+        else:
+            return result_df, None, ""
+
+    # Functions for example query buttons
+    async def on_example_nl_click(query_text):
+        sql_query, plot_code = await generate_sql_and_plot_code(query_text)
+        result_df, error_msg = execute_query(sql_query)
+        fig = None
+        if error_msg:
+            return sql_query, plot_code, None, None, error_msg
+        if plot_code.strip():
+            fig, plot_error = generate_plot(plot_code, result_df)
+            if plot_error:
+                error_msg = plot_error
+            else:
+                error_msg = ""
+        else:
+            fig = None
+            error_msg = ""
+        return sql_query, plot_code, result_df, fig, error_msg
+
+    def on_example_sql_click(sql_query):
+        result_df, error_msg = execute_query(sql_query)
+        fig = None
+        plot_code = ""
+        if error_msg:
+            return sql_query, plot_code, None, None, error_msg
+        else:
+            return sql_query, plot_code, result_df, fig, ""
+
+    async def on_example1_click():
+        return await on_example_nl_click("Show awards over 1M in CA")
+
+    async def on_example2_click():
+        return await on_example_nl_click("List all contracts in New York")
+
+    async def on_example3_click():
+        return await on_example_nl_click("Show top 5 departments by award amount")
+
+    def on_example4_click():
+        return on_example_sql_click("SELECT * from contract_data LIMIT 10;")
+
+    btn_example1.click(fn=on_example1_click, inputs=[], outputs=[sql_out, plot_code_out, results_out, plot_out, error_out])
+    btn_example2.click(fn=on_example2_click, inputs=[], outputs=[sql_out, plot_code_out, results_out, plot_out, error_out])
+    btn_example3.click(fn=on_example3_click, inputs=[], outputs=[sql_out, plot_code_out, results_out, plot_out, error_out])
+    btn_example4.click(fn=on_example4_click, inputs=[], outputs=[sql_out, plot_code_out, results_out, plot_out, error_out])
+
+    btn_generate.click(fn=on_generate_click, inputs=query, outputs=[sql_out, plot_code_out])
+    btn_execute.click(fn=on_execute_click, inputs=[sql_out, plot_code_out], outputs=[results_out, plot_out, error_out])
+
 # Launch the Gradio App
 demo.launch()