Added Prompt

app.py

@@ -64,15 +64,80 @@ def get_visualization(question, tool, schema, table_name):
                            additional_authorized_imports=['matplotlib.pyplot',
                                                  'pandas', 'plotly.express',
                                                  'seaborn'], max_iterations=10)
-    fig = agent.run(
+    results = agent.run(
 
-    task=PROMPT,
-    query_description= f'{question}',
-    schema= f'{schema}',
-    table_name= f'{table_name}',
+    task= f'''
+
+    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}})
+    ```
+
+
+    Here is the query you should generate a plot for: '{question}'.
+    Here is the schema: '{schema}' and here is the table name: '{table_name}
+
+      
+    '''
     )
 
-    return fig
+    return results
 
 
 class SQLExecutorTool(Tool):