Update app.py

app.py

@@ -8,7 +8,7 @@ import plotly.express as px
 import os
 
 # Set OpenAI API key
-client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # =========================
 # Configuration and Setup
@@ -73,7 +73,7 @@ def load_dataset_schema():
 # OpenAI API Integration
 # =========================
 
-async def parse_query(nl_query):
+def parse_query(nl_query):
     """
     Converts a natural language query into a SQL query using OpenAI's API.
     """
@@ -102,180 +102,94 @@ def detect_plot_intent(nl_query):
     """
     Detects if the user's query involves plotting.
     """
-    plot_keywords = ['plot', 'graph', 'chart', 'distribution', 'visualize']
+    plot_keywords = ['plot', 'graph', 'chart', 'distribution', 'visualize', 'trend', 'histogram', 'bar', 'line']
     return any(keyword in nl_query.lower() for keyword in plot_keywords)
 
-async def generate_sql_and_plot_code(query):
+def generate_plot_code(sql_query, result_df):
     """
-    Generates SQL query and optional plotting code.
+    Generates plotting code based on the SQL query and result DataFrame.
     """
-    is_plot = detect_plot_intent(query)
-    sql_query = await parse_query(query)
-    plot_code = ""
-    if is_plot and not sql_query.startswith("Error"):
-        plot_code = """
-import plotly.express as px
-fig = px.bar(result_df, x='x_column', y='y_column', title='Generated Plot')
-fig.update_layout(title_x=0.5)
-"""
-    return sql_query, plot_code
-
-def execute_query(sql_query):
-    """
-    Executes the SQL query and returns the results.
-    """
-    if sql_query.startswith("Error"):
-        return None, 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}"
+    if not detect_plot_intent(sql_query):
+        return None
 
-def generate_plot(plot_code, result_df):
-    """
-    Executes the plot code to generate a plot from the result DataFrame.
-    """
-    if not plot_code.strip():
-        return None, "No plot code provided."
-    try:
-        columns = result_df.columns.tolist()
-        if len(columns) < 2:
-            return None, "Not enough columns to plot."
-        plot_code = plot_code.replace('x_column', columns[0])
-        plot_code = plot_code.replace('y_column', columns[1])
-        local_vars = {'result_df': result_df, 'px': px}
-        exec(plot_code, {}, local_vars)
-        fig = local_vars.get('fig', None)
-        return fig, "" if fig else "Plot could not be generated."
-    except Exception as e:
-        return None, f"Error generating plot: {e}"
+    columns = result_df.columns.tolist()
+    if len(columns) >= 2:
+        fig = px.bar(result_df, x=columns[0], y=columns[1], title='Generated Plot')
+        fig.update_layout(title_x=0.5)
+        return fig
+    else:
+        return None
 
 # =========================
 # Gradio Application UI
 # =========================
 
-with gr.Blocks() as demo:
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
     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"')
+    ## Parquet Data Explorer
 
-                    # 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;")
+    **Query and visualize data effortlessly.**
 
-                    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")
+    """, elem_id="main-title")
 
-        with gr.TabItem("Dataset Schema"):
-            gr.Markdown("### Dataset Schema")
-            schema_display = gr.JSON(label="Schema", value=json.loads(json.dumps(get_schema(), indent=2)))
+    with gr.Row():
+        with gr.Column(scale=1):
+            query = gr.Textbox(
+                label="Ask a question about the data",
+                placeholder='e.g., "What are the total awards over 1M in California?"',
+                lines=1
+            )
+            # Display schema next to the input
+            schema_display = gr.JSON(value=json.loads(json.dumps(get_schema(), indent=2)), visible=False)
 
-    # =========================
-    # Click Event Handlers
-    # =========================
+            error_out = gr.Alert(variant="error", visible=False)
 
-    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
+        with gr.Column(scale=2):
+            results_out = gr.DataFrame(label="Results")
+            plot_out = gr.Plot()
 
-    def on_execute_click(sql_query, plot_code):
-        """
-        Handles the "Execute Query" button click event.
-        """
+    def on_query_submit(nl_query):
+        sql_query = parse_query(nl_query)
+        if sql_query.startswith("Error"):
+            return gr.update(visible=True, value=sql_query), None, None
         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, ""
+            return gr.update(visible=True, value=error_msg), None, None
+        fig = generate_plot_code(nl_query, result_df)
+        return gr.update(visible=False), result_df, fig
+
+    def on_focus():
+        return gr.update(visible=True)
+
+    query.submit(
+        fn=on_query_submit,
+        inputs=query,
+        outputs=[error_out, results_out, plot_out]
+    )
+    query.focus(
+        fn=on_focus,
+        outputs=schema_display
+    )
 
-    # 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;")
+# =========================
+# Helper Functions
+# =========================
 
-    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])
+def execute_query(sql_query):
+    """
+    Executes the SQL query and returns the results.
+    """
+    if sql_query.startswith("Error"):
+        return None, sql_query
 
-    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])
+    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}"
 
 # =========================
 # Launch the Gradio App