Update tools/normal_distribution.py

tools/normal_distribution.py

@@ -1,6 +1,7 @@
 import random
 import math
 import matplotlib.pyplot as plt
+import matplotlib
 from smolagents import tool
 
 @tool
@@ -35,7 +36,7 @@ def generate_normal_distribution(mean: float, std_dev: float, count: int = 10000
     return samples
 
 @tool
-def create_histogram_and_theorical_pdf(mean: float, std_dev:float, random_numbers:list)->None:
+def create_histogram_and_theorical_pdf(mean: float, std_dev:float, random_numbers:list)->matplotlib.figure.Figure:
     """Generate a histogram of random numbers and overlay the theoretical 
     probability density function (PDF) of a normal distribution. 
     Return the histogram as a base64-encoded string.
@@ -71,13 +72,21 @@ def create_histogram_and_theorical_pdf(mean: float, std_dev:float, random_number
             for x in x_values
     ]
 
-    # Plotting
-    plt.figure(figsize=(10, 6))
-    plt.bar(x_values, hist_data, width=bin_width, alpha=0.6, color='g', label='Histogram')
-    plt.plot(x_values, pdf_values, 'k', linewidth=2, label='Theoretical PDF')
+    # Create a figure
+    fig = plt.figure(figsize=(10, 6))
+    ax = fig.add_subplot(111)  # Create an Axes object on the figure
 
-    plt.title(f'Histogram of {len(random_numbers)} Random Numbers\nMean: {mean}, Std Dev: {std_dev}')
-    plt.xlabel('Value')
-    plt.ylabel('Density')
-    plt.legend()
-    plt.grid()
+    # Add histogram as a bar plot to the Axes
+    ax.bar(x_values, hist_data, width=bin_width, alpha=0.6, color='g', label='Histogram')
+
+    # Add theoretical PDF as a line plot to the same Axes
+    ax.plot(x_values, pdf_values, 'k', linewidth=2, label='Theoretical PDF')
+
+    # Set title and labels
+    ax.set_title(f'Histogram of {len(random_numbers)} Random Numbers\nMean: {mean:.2f}, Std Dev: {std_dev:.2f}')
+    ax.set_xlabel('Value')
+    ax.set_ylabel('Density')
+    ax.legend()
+    ax.grid()
+
+    return fig