Spaces:
Runtime error
Runtime error
russkiyximik
Overhauled Code Interpreter tool by wrapping it in a new class to make it robust to agent output
138bf3e | import io | |
| from contextlib import redirect_stderr, redirect_stdout | |
| from typing import Any, Dict, Type | |
| import numpy as np | |
| import pandas as pd | |
| from crewai.tools import BaseTool | |
| from crewai_tools import CodeInterpreterTool | |
| from crewai_tools.tools.code_interpreter_tool.code_interpreter_tool import SandboxPython | |
| from pydantic import BaseModel, Field | |
| class ExtractSchemaInput(BaseModel): | |
| """Input schema for ExtractSchema.""" | |
| file_path: str = Field( | |
| ..., description="Path to the local file from which to extract the schema, " \ | |
| "relative to the agent's working directory." | |
| ) | |
| # This file path is relative to ..? (find working dir of agent) | |
| class ExtractSchema(BaseTool): | |
| name: str = "ExtractSchema" | |
| description: str = ( | |
| "This tool is useful for extracting schema of tabular data, which is important to understand" | |
| "the structure of a dataset for analysis." | |
| ) | |
| args_schema: Type[BaseModel] = ExtractSchemaInput | |
| def _run(self, file_path: str) -> Dict[str, Any]: | |
| # Implementation goes here | |
| df = pd.read_csv(file_path) | |
| return { | |
| "shape": df.shape, | |
| "dtypes": df.dtypes.to_dict(), | |
| "columns": df.columns.tolist(), | |
| "numeric_cols": df.select_dtypes(include=[np.number]).columns.tolist(), | |
| "categorical_cols": df.select_dtypes(include=['object']).columns.tolist() | |
| } | |
| class ResultAwareCodeInterpreterTool(CodeInterpreterTool): | |
| """Returns `result` when present and falls back to captured stdout/stderr.""" | |
| def _resolve_execution_output( | |
| exec_locals: Dict[str, Any], | |
| stdout_buffer: io.StringIO, | |
| stderr_buffer: io.StringIO, | |
| ) -> str: | |
| result = exec_locals.get("result") | |
| if result is not None: | |
| return str(result) | |
| stdout_text = stdout_buffer.getvalue().strip() | |
| stderr_text = stderr_buffer.getvalue().strip() | |
| if stdout_text and stderr_text: | |
| return f"{stdout_text}\n\n[stderr]\n{stderr_text}" | |
| if stdout_text: | |
| return stdout_text | |
| if stderr_text: | |
| return stderr_text | |
| return "Code executed successfully, but it did not set `result` or print output." | |
| def run_code_in_restricted_sandbox(code: str) -> str: | |
| stdout_buffer = io.StringIO() | |
| stderr_buffer = io.StringIO() | |
| exec_locals: Dict[str, Any] = {} | |
| try: | |
| with redirect_stdout(stdout_buffer), redirect_stderr(stderr_buffer): | |
| SandboxPython.exec(code=code, locals_=exec_locals) | |
| except Exception as e: | |
| return f"An error occurred: {e!s}" | |
| return ResultAwareCodeInterpreterTool._resolve_execution_output( | |
| exec_locals, | |
| stdout_buffer, | |
| stderr_buffer, | |
| ) | |
| def run_code_unsafe(code: str, libraries_used: list[str]) -> str: | |
| stdout_buffer = io.StringIO() | |
| stderr_buffer = io.StringIO() | |
| for library in libraries_used: | |
| import os | |
| os.system(f"pip install {library}") # noqa: S605 | |
| try: | |
| exec_locals: Dict[str, Any] = {} | |
| with redirect_stdout(stdout_buffer), redirect_stderr(stderr_buffer): | |
| exec(code, {}, exec_locals) # noqa: S102 | |
| except Exception as e: | |
| return f"An error occurred: {e!s}" | |
| return ResultAwareCodeInterpreterTool._resolve_execution_output( | |
| exec_locals, | |
| stdout_buffer, | |
| stderr_buffer, | |
| ) | |