Update app.py

app.py

@@ -6,10 +6,9 @@ from functools import lru_cache
 import pandas as pd
 import plotly.express as px
 import os
-from openai import OpenAI
 
 # Set OpenAI API key
-client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # =========================
 # Configuration and Setup
@@ -84,13 +83,13 @@ async def parse_query(nl_query):
     ]
 
     try:
-        response =  openai.chat.completions.create(
-            model="gpt-4o-mini",
+        response = await openai.ChatCompletion.acreate(
+            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}"
@@ -169,9 +168,12 @@ with gr.Blocks() as demo:
     ## Instructions
 
     1. **Describe the data you want**: e.g., `Show awards over 1M in CA`
-    2. **Generate SQL**: Click "Generate SQL" to see the SQL query.
-    3. **Execute Query**: Run the query to view results and plots.
-    4. **Dataset Schema**: See available columns and types in the "Schema" tab.
+    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():
@@ -179,6 +181,15 @@ with gr.Blocks() as demo:
             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")
@@ -219,6 +230,50 @@ with gr.Blocks() as demo:
         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])