Update app.py

app.py

@@ -1,11 +1,51 @@
-@tool
+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
-    """
-    return f"Received code: {code_string}"
+  """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)}"