First_agent_template / tools /visualize_csv.py
Athmiiii's picture
Update tools/visualize_csv.py
c663d92 verified
import pandas as pd
import matplotlib.pyplot as plt
import io
import base64
from smolagents import tool
@tool
def visualize_csv(csv_content: str, x_col: str, y_col: str) -> str:
"""
Generates a plot from CSV content.
Args:
csv_content: The raw CSV file content as a string.
x_col: The name of the column for the x-axis.
y_col: The name of the column for the y-axis.
Returns:
A base64-encoded image tag of the generated plot.
"""
try:
df = pd.read_csv(io.StringIO(csv_content))
plt.figure(figsize=(8, 5))
plt.plot(df[x_col], df[y_col])
plt.xlabel(x_col)
plt.ylabel(y_col)
plt.title(f"{y_col} vs {x_col}")
buf = io.BytesIO()
plt.savefig(buf, format='png')
buf.seek(0)
img_base64 = base64.b64encode(buf.read()).decode('utf-8')
plt.close()
return f'<img src="data:image/png;base64,{img_base64}" />'
except Exception as e:
return f"Error: {str(e)}"