Update app.py

app.py

@@ -6,6 +6,9 @@ import yaml
 from tools.final_answer import FinalAnswerTool
 import matplotlib.pyplot as plt
 import io
+from PIL import Image as PILImage
+from IPython.display import display, Image
+
 
 from Gradio_UI import GradioUI
 
@@ -36,7 +39,7 @@ def get_current_time_in_timezone(timezone: str) -> str:
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
 @tool
-def generate_bar_chart(categories: list[str], values: list[int]) -> bytes:
+def generate_bar_chart(categories: list[str], values: list[int]) -> PILImage:
     """A tool that generates a bar chart from given categories and values, and returns it as an image.
     
     Args:
@@ -44,7 +47,7 @@ def generate_bar_chart(categories: list[str], values: list[int]) -> bytes:
         values: A list of integers or floats representing values (e.g., [10, 20, 15]).
 
     Returns:
-        Image data in bytes format.
+        An image object containing the bar chart.
     """
     try:
         if len(categories) != len(values):
@@ -56,7 +59,7 @@ def generate_bar_chart(categories: list[str], values: list[int]) -> bytes:
         # Add labels and title
         plt.xlabel('Categories')
         plt.ylabel('Values')
-        plt.title('Bar Graph')
+        plt.title('User-Generated Bar Graph')
 
         # Save the plot to a BytesIO buffer
         img_buffer = io.BytesIO()
@@ -65,10 +68,13 @@ def generate_bar_chart(categories: list[str], values: list[int]) -> bytes:
         
         # Get image bytes
         img_buffer.seek(0)
-        return img_buffer.getvalue()
+
+        # Open the image and return as PIL Image object
+        img = PILImage.open(img_buffer)
+        return img
 
     except Exception as e:
-        return f"Error generating bar chart: {str(e)}".encode()
+        print(f"Error generating bar chart: {str(e)}")