Update app.py

app.py

@@ -46,6 +46,10 @@ def get_schema():
 
 COLUMN_TYPES = {col['column_name']: col['column_type'] for col in get_schema()}
 
+# =========================
+# Database Interaction
+# =========================
+
 def load_dataset_schema():
     con = duckdb.connect()
     try:
@@ -58,8 +62,6 @@ def load_dataset_schema():
     finally:
         con.close()
 
-load_dataset_schema()
-
 # =========================
 # OpenAI API Integration
 # =========================
@@ -71,68 +73,40 @@ def parse_query(nl_query):
     ]
 
     try:
-        response = openai.chat.completions.create(
+        response = openai.ChatCompletion.create(
             model="gpt-4",
             messages=messages,
             temperature=0,
             max_tokens=150,
         )
-        sql_query = response.choices[0].message.content.strip()
+        sql_query = response.choices[0].message['content'].strip()
         return sql_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', 'trend', 'histogram', 'bar', 'line', 'scatter', 'pie']
+    plot_keywords = ['plot', 'graph', 'chart', 'distribution', 'visualize']
     return any(keyword in nl_query.lower() for keyword in plot_keywords)
 
-
-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}"
-
-def generate_plot(nl_query, result_df):
-    if not detect_plot_intent(nl_query):
-        return None, ""
-
-    columns = result_df.columns.tolist()
-    if len(columns) < 2:
-        return None, "Not enough data to generate a plot."
-
-    if 'bar' in nl_query.lower():
-        fig = px.bar(result_df, x=columns[0], y=columns[1], title='Bar Chart')
-    elif 'line' in nl_query.lower():
-        fig = px.line(result_df, x=columns[0], y=columns[1], title='Line Chart')
-    elif 'scatter' in nl_query.lower():
-        fig = px.scatter(result_df, x=columns[0], y=columns[1], title='Scatter Plot')
-    elif 'pie' in nl_query.lower():
-        fig = px.pie(result_df, names=columns[0], values=columns[1], title='Pie Chart')
-    else:
-        fig = px.bar(result_df, x=columns[0], y=columns[1], title='Bar Chart')
-
-    fig.update_layout(title_x=0.5)
-    return fig, ""
-
 # =========================
 # Gradio Application UI
 # =========================
 
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
+with gr.Blocks() as demo:
     gr.Markdown("""
-    <h1 style="text-align: center; font-size: 2.5em; color: #333333;">Parquet Data Explorer</h1>
-    <p style="text-align: center; color: #666666;">Query and visualize your data effortlessly.</p>
+    # 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.Row():
@@ -142,74 +116,39 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
                 placeholder='e.g., "What are the total awards over 1M in California?"',
                 lines=1
             )
-            gr.Markdown("### Example Queries")
-            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;")
-
-            error_out = gr.Markdown(
-                value="",
-                visible=False
-            )
-        with gr.Column(scale=2):
-            results_out = gr.DataFrame(
-                label="Results"
-            )
-
-    # Instructions
-    gr.Markdown("""
-    ## Instructions
-    1. **Enter a query**: Type in a natural language query in the textbox.
-    2. **Use Example Queries**: Click on any example query button above.
-    3. **Generate SQL and Plot**: Click "Execute" to see results.
-    """)
+            # Button to generate the SQL query from NL
+            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
+            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")
 
     # =========================
-    # Click Event Handlers
+    # Event Functions
     # =========================
 
-    def on_query_submit(nl_query):
-        if not nl_query.strip():
-            return gr.update(visible=True, value="Please enter a query."), None, None
-
+    def generate_sql(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 gr.update(visible=True, value=error_msg), None, None
-
-        fig, plot_error = generate_plot(nl_query, result_df)
-        if plot_error:
-            return gr.update(visible=True, value=plot_error), None, None
-
-        return gr.update(visible=False, value=""), result_df, fig
-
-    def on_example_click(query_text):
-        sql_query = parse_query(query_text)
-        result_df, error_msg = execute_query(sql_query)
-        if error_msg:
-            return sql_query, None, None, error_msg
-        fig, plot_error = generate_plot(query_text, result_df)
-        return sql_query, result_df, fig, plot_error if plot_error else ""
-
-    btn_example1.click(lambda: on_example_click("Show awards over 1M in CA"), outputs=[results_out, error_out])
-    btn_example2.click(lambda: on_example_click("List all contracts in New York"), outputs=[results_out, error_out])
-    btn_example3.click(lambda: on_example_click("Show top 5 departments by award amount"), outputs=[results_out, error_out])
-    btn_example4.click(lambda: on_example_click("SELECT * from contract_data LIMIT 10;"), outputs=[results_out, error_out])
-
-    query.submit(
-        fn=on_query_submit,
-        inputs=query,
-        outputs=[error_out, results_out]
-    )
+        return sql_query
 
-# =========================
-# Launch the Gradio App
-# =========================
+    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])
 
-if __name__ == "__main__":
-    demo.launch()
+# Launch the Gradio App
+demo.launch()