|
|
from smolagents._function_type_hints_utils import FunctionSchema |
|
|
|
|
|
@tool(schema=FunctionSchema( |
|
|
name="execute_code", |
|
|
description="Executes Python code and returns the result as a string.", |
|
|
parameters=[ |
|
|
{ |
|
|
"name": "code_string", |
|
|
"type": "string", |
|
|
"description": "Python code to execute." |
|
|
}, |
|
|
{ |
|
|
"name": "data", |
|
|
"type": "object", |
|
|
"description": "The pandas DataFrame to use.", |
|
|
"properties":{}, |
|
|
} |
|
|
], |
|
|
returns={ |
|
|
"type":"string", |
|
|
"description": "The result of executing the code." |
|
|
} |
|
|
)) |
|
|
def execute_code(code_string: str, data: pd.DataFrame) -> str: |
|
|
"""Executes python code and returns results as a string. |
|
|
|
|
|
Args: |
|
|
code_string (str): Python code to execute. |
|
|
data (pd.DataFrame): The dataframe to use in the code |
|
|
Returns: |
|
|
str: The result of executing the code or an error message |
|
|
""" |
|
|
try: |
|
|
local_vars = {"data": data, "pd": pd, "np": np, "plt": plt, "sns": sns} |
|
|
exec(code_string, local_vars) |
|
|
|
|
|
if "result" in local_vars: |
|
|
if isinstance(local_vars["result"], (pd.DataFrame, pd.Series)): |
|
|
return local_vars["result"].to_string() |
|
|
elif isinstance(local_vars["result"], plt.Figure): |
|
|
buf = io.BytesIO() |
|
|
local_vars["result"].savefig(buf, format="png") |
|
|
plt.close(local_vars["result"]) |
|
|
return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}" |
|
|
else: |
|
|
return str(local_vars["result"]) |
|
|
else: |
|
|
return "Code executed successfully, but no variable called 'result' was assigned." |
|
|
|
|
|
except Exception as e: |
|
|
return f"Error executing code: {str(e)}" |
