Added Prompt 2

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import matplotlib.pyplot as plt
 from transformers import HfEngine, ReactCodeAgent
 from transformers.agents import Tool
+from prompt import PROMPT
 
 # Height of the Tabs Text Area
 TAB_LINES = 8
@@ -64,23 +65,9 @@ def get_visualization(question, tool, schema, table_name):
                                                  'pandas', 'plotly.express',
                                                  'seaborn'], max_iterations=10)
     fig = agent.run(
-        
-      instruction='''
-    THINK STEP BY STEP
-    Here are the steps you should follow while writing code for Visualization:
-    1. Select the most appropriate chart type for data. Use bar charts for categorical comparisons, line charts for trends over time, scatter plots for relationships between variables, pie charts for proportions, histograms for distribution analysis, and box plots for visualizing data spread and outliers.
-    2. Ensure clear and appropriate labels, colors, and design elements, keeping visual elements legible and uncluttered.
-    3. Follow best practices, avoiding unnecessary visual distractions (chartjunk).
-    4. Ensure the code is error-free, with correct fields, transformations, and aesthetics.
-    5. Use descriptive and accurate x and y axis labels that reflect the data.
-    6. Ensure units of measurement are clearly indicated on axes (e.g., %, $, cm).
-    7. Ensure that categorical data is plotted on one axis and numerical data on the other, with appropriate labels that clearly represent the data being visualized.
-    8. When plotting categorical data, arrange categories in a meaningful order (e.g., by size, time, or frequency) rather than randomly.
-    9. Ensure that the categorical data are plotted on the x-axis, and the frequencies (numerical data) are plotted on the y-axis.
-    10. Use seaborn
-    11. In the end you have to return a dict which contain final fig as fig key, Generated SQL as sql key, Data as a dataframe with data key using the `final_answer` tool e.g. final_answer(answer={"fig": fig, "sql": sql, "data": data})''',
-    
-    task= f'{question}',
+
+    task=PROMPT,
+    query_description= f'{question}',
     schema= f'{schema}',
     table_name= f'{table_name}',
     )
@@ -96,6 +83,7 @@ class SQLExecutorTool(Tool):
             "description": f"The query to perform. This should be correct DuckDB SQL.",
         }
     }
+    description = """Allows you to perform SQL queries on the table. Returns a pandas dataframe representation of the result."""
     output_type = "pandas.core.frame.DataFrame"
 
     def forward(self, query: str) -> str:
@@ -110,8 +98,6 @@ def main(table, text_query):
     ax.set_axis_off()  
     
     schema, table_name = get_table_schema(table)
-    tool.description = f"""Allows you to perform SQL queries on the table. Returns a pandas dataframe representation of the result."""
-    
     
     try:
         output = get_visualization(question=text_query, tool=tool, schema=schema, table_name=table_name)
@@ -121,7 +107,6 @@ def main(table, text_query):
     except Exception as e:
         gr.Warning(f"❌ Unable to generate the visualization. {e}")
         
-        
     return fig, generated_sql, data
     
     

prompt.py

@@ -0,0 +1,62 @@
+PROMPT = '''
+    Here are the steps you should follow while writing code for Visualization:
+    1. You have access to the database with the `sql_engine` tool, which allows you to run DuckDB SQL queries and return results as a df.
+    2. Query the database using `sql_engine`, print the first 5 rows to inspect the data.
+    3. Select the most appropriate chart type for the data:
+       - Use bar charts for categorical comparisons, line charts for trends over time, scatter plots for relationships between variables, pie charts for proportions, histograms for distribution, and box plots for data spread and outliers.
+    4. Analyze the data and choose the best visualization type to answer the query.
+    5. Always include a plot in your answer.
+    6. Use Seaborn for the plots.
+    7. In the end, return a dictionary containing the final figure (`fig` key), the generated SQL (`sql` key), and the data as a dataframe (`data` key) using the `final_answer` tool, e.g. `final_answer(answer={{"fig": 'fig.png', "sql": sql, "data": data}})`.
+
+    Example:
+
+    ```python
+    # Input query
+    query_description = 'Average tip amount based on the ride time length in minutes.'
+
+    # SQL Query to get ride time length and average tip amount
+    query = """
+        SELECT 
+            EXTRACT(EPOCH FROM (tpep_dropoff_datetime - tpep_pickup_datetime)) / 60 AS ride_time_length,
+            AVG(tip_amount) AS avg_tip_amount
+        FROM 
+            sample_data.nyc.taxi
+        GROUP BY 
+            EXTRACT(EPOCH FROM (tpep_dropoff_datetime - tpep_pickup_datetime)) / 60
+    """
+
+    # Execute the query using the sql_engine tool
+    df = sql_engine(query=query)
+
+    # Print the result to observe the data
+    print(df)
+
+    # Create a line plot using seaborn
+    import seaborn as sns
+    import matplotlib.pyplot as plt
+
+    plt.figure(figsize=(10,6))
+    sns.lineplot(x="ride_time_length", y="avg_tip_amount", data=df)
+
+    # Set the title and labels
+    plt.title("Average Tip Amount vs Ride Time Length")
+    plt.xlabel("Ride Time Length (minutes)")
+    plt.ylabel("Average Tip Amount")
+
+    # Print the plot to observe the results
+    print("Plot created")
+
+    # Since we are required to return a fig, sql, and data, let's store the plot in a variable
+    fig = plt.gcf()
+
+    # Store the query in a variable
+    sql = query
+
+    # Store the dataframe in a variable
+    data = df
+
+    # Return the final answer
+    final_answer(answer={"fig": fig, "sql": sql, "data": data})
+    ```
+    '''