Update app.py

app.py

@@ -3,8 +3,6 @@ import openai
 import gradio as gr
 import duckdb
 from functools import lru_cache
-import pandas as pd
-import plotly.express as px
 import os
 
 # =========================
@@ -44,24 +42,6 @@ schema = [
 def get_schema():
     return schema
 
-COLUMN_TYPES = {col['column_name']: col['column_type'] for col in get_schema()}
-
-# =========================
-# Database Interaction
-# =========================
-
-def load_dataset_schema():
-    con = duckdb.connect()
-    try:
-        con.execute("DROP VIEW IF EXISTS contract_data")
-        con.execute(f"CREATE VIEW contract_data AS SELECT * FROM '{dataset_path}'")
-        return True
-    except Exception as e:
-        print(f"Error loading dataset schema: {e}")
-        return False
-    finally:
-        con.close()
-
 # =========================
 # OpenAI API Integration
 # =========================
@@ -73,7 +53,7 @@ def parse_query(nl_query):
     ]
 
     try:
-        response = openai.chat.completions.create(
+        response = openai.ChatCompletion.create(
             model="gpt-4o-mini",
             messages=messages,
             temperature=0,
@@ -84,9 +64,19 @@ def parse_query(nl_query):
     except Exception as e:
         return f"Error generating SQL query: {e}"
 
-def detect_plot_intent(nl_query):
-    plot_keywords = ['plot', 'graph', 'chart', 'distribution', 'visualize']
-    return any(keyword in nl_query.lower() for keyword in plot_keywords)
+# =========================
+# Database Interaction
+# =========================
+
+def execute_sql_query(sql_query):
+    try:
+        con = duckdb.connect()
+        con.execute(f"CREATE OR REPLACE VIEW contract_data AS SELECT * FROM '{dataset_path}'")
+        result_df = con.execute(sql_query).fetchdf()
+        con.close()
+        return result_df, ""
+    except Exception as e:
+        return None, f"Error executing query: {e}"
 
 # =========================
 # Gradio Application UI
@@ -103,7 +93,7 @@ with gr.Blocks() as demo:
     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.
+    4. **Execute Query**: Run the query to view results.
     5. **Dataset Schema**: See available columns and types in the "Schema" tab.
 
     ## Example Queries
@@ -116,156 +106,61 @@ with gr.Blocks() as demo:
                 placeholder='e.g., "What are the total awards over 1M in California?"',
                 lines=1
             )
-            # Button to generate the SQL query from NL
+
+            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_sql = gr.Button("Generate SQL Query")
-            # Textbox to display generated SQL
-            sql_query_out = gr.Textbox(label="Generated SQL Query", interactive=False)
-            # Execute button
+            sql_query_out = gr.Code(label="Generated SQL Query", language="sql")
             btn_execute_query = gr.Button("Execute Query")
             error_out = gr.Markdown("", visible=False)
-        
-        # Results and Plot output
-        results_out = gr.DataFrame(label="Query Results")
-        plot_out = gr.Plot(label="Plot")
+        with gr.Column(scale=2):
+            results_out = gr.Dataframe(label="Query Results", interactive=False)
+
+    with gr.TabItem("Dataset Schema"):
+        gr.Markdown("### Dataset Schema")
+        schema_display = gr.JSON(label="Schema", value=get_schema())
 
     # =========================
     # Event Functions
     # =========================
 
-    def generate_sql(nl_query):
+    def on_generate_click(nl_query):
         sql_query = parse_query(nl_query)
         return sql_query
 
-    def execute_sql_query(sql_query):
-        try:
-            con = duckdb.connect()
-            con.execute(f"CREATE OR REPLACE VIEW contract_data AS SELECT * FROM '{dataset_path}'")
-            result_df = con.execute(sql_query).fetchdf()
-            con.close()
-            return result_df, ""
-        except Exception as e:
-            return None, f"Error executing query: {e}"
-
-    # Button click event handlers
-    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)
+    def on_execute_click(sql_query):
+        result_df, error_msg = execute_sql_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, ""
+            error_out.update(f"**Error:** {error_msg}")
+            return None
         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
+            error_out.update("")
+            return result_df
+
+    def on_example_click(query_text):
+        query.value = query_text
+        sql_query = parse_query(query_text)
+        sql_query_out.value = sql_query
+        result_df, error_msg = execute_sql_query(sql_query)
         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 = ""
+            error_out.update(f"**Error:** {error_msg}")
+            results_out.value = None
         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;")
+            error_out.update("")
+            results_out.value = result_df
 
-    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_example1.click(fn=lambda: on_example_click("Show awards over 1M in CA"), outputs=[])
+    btn_example2.click(fn=lambda: on_example_click("List all contracts in New York"), outputs=[])
+    btn_example3.click(fn=lambda: on_example_click("Show top 5 departments by award amount"), outputs=[])
+    btn_example4.click(fn=lambda: on_example_click("SELECT * from contract_data LIMIT 10;"), outputs=[])
 
-    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])
+    btn_generate_sql.click(fn=on_generate_click, inputs=query, outputs=sql_query_out)
+    btn_execute_query.click(fn=on_execute_click, inputs=sql_query_out, outputs=results_out)
 
 # Launch the Gradio App
 demo.launch()