Update Prompt

README.md

@@ -4,7 +4,7 @@ emoji: 📊
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 5.0.1
+sdk_version: 4.42.0
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -58,16 +58,16 @@ def get_table_schema(table):
     return ddl_create, full_path
 
 
-def get_visualization(question, tool):
+def get_visualization(question, tool, schema, table_name):
     agent = ReactCodeAgent(tools=[tool], llm_engine=llm_engine, add_base_tools=True,
                            additional_authorized_imports=['matplotlib.pyplot',
                                                  'pandas', 'plotly.express',
                                                  'seaborn'], max_iterations=10)
     fig = agent.run(
         
-        task=f'''
-    
-    You are an expert in writing Visualization code. Here are the steps you should follow while writing code for Visualization:
+      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).
@@ -77,13 +77,12 @@ def get_visualization(question, tool):
     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. Do not write functions.
-    11. Use seaborn.
-    12. 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.
-    DO NOT PRINT THE FINAL ANSWER OR FINAL OUTPUT
-    Here is the task:
-    task: {question}
-    ''',
+    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}',
+    schema= f'{schema}',
+    table_name= f'{table_name}',
     )
 
     return fig
@@ -110,13 +109,12 @@ def main(table, text_query):
     fig, ax = plt.subplots()
     ax.set_axis_off()  
     
-    schema, _ = get_table_schema(table)
-    tool.description = f"""Allows you to perform SQL queries on the table. Returns a pandas dataframe representation of the result.
-    The table schema is as follows: \n{schema}"""
+    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)
+        output = get_visualization(question=text_query, tool=tool, schema=schema, table_name=table_name)
         fig = output.get('fig', None)
         generated_sql = output.get('sql', None)
         data = output.get('data', None)