File size: 6,285 Bytes
fb44830 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 |
import io
import tempfile
import numpy as np
from PIL import Image
from PIL.ImageFile import ImageFile
from matplotlib import pyplot as plt
from smolagents.tools import Tool
def _plot_line_diagram(
x_values: list,
y_values_list: list[list[int | float]],
labels: list[str] | None = None,
title: str = 'Bar Diagram',
xlabel: str = 'X-axis',
ylabel: str = 'Y-axis'
) -> str:
"""
Plot a line diagram with one or more y-values and save the image to a temporary file.
Return the path to the saved image file.
:param x_values: List of x-values.
:param y_values_list: List of lists containing y-values.
Each inner list represents a separate line.
:param labels: List of labels for each line (optional).
:param title: Title of the plot (default: 'Line Diagram').
:param xlabel: Label for the X-axis (default: 'X-axis').
:param ylabel: Label for the Y-axis (default: 'Y-axis').
:return: Path to the saved image file.
"""
plt.figure(figsize=(10, 6))
for i, y_values in enumerate(y_values_list):
label = labels[i] if labels and i < len(labels) else f'Line {i+1}'
plt.plot(x_values, y_values, label=label)
plt.title(title)
plt.xlabel(xlabel)
plt.ylabel(ylabel)
plt.legend()
plt.grid(True)
# Save the plot as an image file in a temporary directory
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
plt.savefig(temp_file.name)
plt.close()
return temp_file.name
def _plot_bar_diagram(
x_values: list,
y_values_list: list[list[int | float]],
labels: list[str] | None = None,
title: str = 'Bar Diagram',
xlabel: str = 'X-axis',
ylabel: str = 'Y-axis'
) -> str:
"""
Plot a bar diagram with one or more y-values and save the image to a temporary file.
Return the path to the saved image file.
:param x_values: List of x-values.
:param y_values_list: List of lists containing y-values.
Each inner list represents a separate set of bars.
:param labels: List of labels for each set of bars (optional).
:param title: Title of the plot (default: 'Bar Diagram').
:param xlabel: Label for the X-axis (default: 'X-axis').
:param ylabel: Label for the Y-axis (default: 'Y-axis').
:return: Path to the saved image file.
"""
bar_width = 0.2
n = len(y_values_list)
# Set positions of bars on X axis
r = [np.arange(len(x_values))]
for i in range(1, n):
r.append([x + bar_width for x in r[i - 1]])
plt.figure(figsize=(10, 6))
for i, y_values in enumerate(y_values_list):
label = labels[i] if labels and i < len(labels) else f'Set {i + 1}'
plt.bar(r[i], y_values, width=bar_width, label=label)
# Adding xticks
plt.xlabel(xlabel)
plt.ylabel(ylabel)
plt.title(title)
plt.xticks(
[r + bar_width * (n - 1) / 2 for r in range(len(x_values))],
x_values,
rotation=45,
ha='right'
)
plt.legend()
plt.grid(True, axis='y')
plt.tight_layout()
# Save the plot as an image file in a temporary directory
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
plt.savefig(temp_file.name)
plt.close()
return temp_file.name
class PlotTool(Tool):
"""
A tool to plot bar and line diagrams and return the images.
"""
name = 'plot_bar_line_diagrams'
description = (
'Plot a bar or line diagram with one or more y-values and save the image.'
' Return the saved image file as `ImageFile`.'
' An agent must take this `ImageFile` and display the image.'
)
inputs = {
'plot_type': {
'type': 'string',
'description': 'The type of plot. Only two values are valid: `bar` and `line`'
},
'x_values': {'type': 'array', 'description': 'List of x-values.'},
'y_values_list': {
'type': 'array',
'description': (
'A list of lists containing y-values (numbers).'
' Each inner list represents a separate set of bars.'
' The input type is list[list[int | float]]'
)
},
'labels': {
'type': 'array',
'nullable': True,
'description': (
'A list of labels for each set of bars (optional). Defaults to `None`.'
' If provided, the length of `labels` must be equal to the length of `x_values`.'
)
},
'title': {
'type': 'string',
'nullable': True,
'description': 'Title of the plot (default: "Bar Diagram")'
},
'xlabel': {
'type': 'string',
'nullable': True,
'description': 'Label for the X-axis (default: "X-axis")'
},
'ylabel': {
'type': 'string',
'nullable': True,
'description': 'Label for the Y-axis (default: "Y-axis").'
},
}
output_type = 'image'
def __init__(self, **kwargs):
super().__init__()
def forward(
self,
plot_type: str,
x_values: list,
y_values_list: list[list[int | float]],
labels: list[str] | None = None,
title: str = 'Bar Diagram',
xlabel: str = 'X-axis',
ylabel: str = 'Y-axis'
) -> ImageFile:
if plot_type == 'bar':
img_file_name = _plot_bar_diagram(
x_values=x_values,
y_values_list=y_values_list,
labels=labels,
xlabel=xlabel,
ylabel=ylabel,
title=title
)
elif plot_type == 'line':
img_file_name = _plot_line_diagram(
x_values=x_values,
y_values_list=y_values_list,
labels=labels,
xlabel=xlabel,
ylabel=ylabel,
title=title
)
else:
img_file_name = None
if img_file_name:
with open(img_file_name, 'rb') as in_file:
return Image.open(io.BytesIO(in_file.read()))
else:
return None
|