aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/tools/code_sandbox/code_sandbox_tool.py

@@ -0,0 +1,350 @@
+"""Provider-agnostic code-sandbox tool.
+
+Executes Python code in a sandbox selected by ``SANDBOX_PROVIDER`` (or the tool's
+``provider`` field), built entirely on the generic Seam B pieces — ``SandboxRuntime``
+(create/execute/package-install/cleanup) and ``SandboxArtifactWatcher`` (artifact
+detection via ``find -newer``). No backend SDK is referenced directly.
+
+Authors:
+    Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
+    Komang Elang Surya Prawira (komang.e.s.prawira@gdplabs.id)
+    benedikus (benedikus.h.tjuatja@gdplabs.id)
+"""
+
+import asyncio
+import json
+import os
+import time
+from typing import Any
+
+import pandas as pd
+from gllm_tools.code_interpreter.code_template.models import SandboxProvider
+from langchain_core.tools import BaseTool
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
+from aip_agents.a2a.types import get_mime_type_from_filename
+from aip_agents.sandbox.sandbox_runtime import SandboxRuntime
+from aip_agents.utils.artifact_helpers import ArtifactHandler, create_multiple_artifacts_response
+from aip_agents.utils.file_watchers import SandboxArtifactWatcher
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+# The data source is uploaded cwd-relative (the runtime uploads with upload_dir="") so it works
+# on any provider, and the df preamble reads the same relative path.
+DATA_FILE = "data.csv"
+# Directory the tool watches for created artifacts (the model writes its outputs here).
+ARTIFACT_OUTPUT_DIR = "/tmp/output"  # NOSONAR - sandbox-isolated directory
+
+
+class CodeSandboxInput(BaseModel):
+    """Input schema for the code-sandbox tool."""
+
+    code: str = Field(
+        ...,
+        description=(
+            "Python code to execute in the sandbox. "
+            "If `data_source` is provided, the data will be automatically pre-loaded into "
+            "the sandbox as a Pandas DataFrame. "
+            "You can access the data directly via a variable named `df`. "
+            "The final result must be printed to stdout."
+        ),
+        min_length=1,
+    )
+    data_source: str | list[dict[str, Any]] | None = Field(
+        default=None,
+        description=(
+            "The data source used during code execution. "
+            "It can be a tool output reference (using $tool_output) or raw data. "
+            "If the code requires a data source, this field should not be left empty. "
+            "When an output reference is available, prioritize the output reference over raw data."
+        ),
+    )
+    timeout: int = Field(default=30, description="Maximum execution time in seconds", ge=1, le=300)
+    language: str = Field(default="python", description="Programming language for the sandbox (Python only)")
+    additional_packages: list[str] | None = Field(
+        default=None,
+        description="Additional Python packages to install before execution",
+    )
+
+
+class CodeSandboxTool(BaseTool):
+    """Execute Python code in a provider-agnostic cloud sandbox.
+
+    Backend is chosen by ``provider`` (falls back to the ``SANDBOX_PROVIDER`` env, then
+    ``e2b``). Files created during execution under ``/tmp/output`` are returned as artifacts.
+    """
+
+    name: str = "code_sandbox_tool"
+    description: str = (
+        "Useful for executing Python code in a secure cloud sandbox environment. "
+        "Input should include the Python code to execute and optional data source, timeout, and packages. "
+        "Returns execution results including stdout, stderr, and any errors. "
+        "Automatically downloads all files created during execution as artifacts."
+    )
+    save_output_history: bool = Field(default=True)
+    args_schema: type[BaseModel] = CodeSandboxInput
+    provider: SandboxProvider | None = Field(
+        default=None,
+        description="Sandbox backend; defaults to the SANDBOX_PROVIDER env, then 'e2b'.",
+    )
+    api_key: str | None = Field(
+        default=None,
+        description="Backend API key override; the selected provider resolves its own env when None.",
+    )
+    domain: str | None = Field(
+        default=None,
+        description="Backend domain override (e.g. self-hosted endpoint); provider resolves its own env when None.",
+    )
+    template: str | None = Field(
+        default=None,
+        description="Provider-interpreted template alias (E2B) / snapshot name (OpenSandbox); provider default when None.",
+    )
+    default_additional_packages: list[str] = Field(
+        default_factory=lambda: [
+            pkg.strip() for pkg in os.getenv("SANDBOX_TOOL_ADDITIONAL_PACKAGES", "").split(",") if pkg.strip()
+        ],
+        description="Default additional packages from environment variable",
+    )
+    store_final_output: bool = False
+
+    def _run(
+        self,
+        code: str,
+        data_source: str | list[dict[str, Any]] | None = None,
+        timeout: int = 30,
+        language: str = "python",
+        additional_packages: list[str] | None = None,
+    ) -> str | dict[str, Any]:
+        """Synchronous wrapper around :meth:`_arun` (runs in a fresh event loop)."""
+        loop: asyncio.AbstractEventLoop | None = None
+        try:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            try:
+                return loop.run_until_complete(self._arun(code, data_source, timeout, language, additional_packages))
+            finally:
+                try:
+                    pending = [t for t in asyncio.all_tasks(loop) if not t.done()]
+                    for t in pending:
+                        t.cancel()
+                    if pending:
+                        loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+                except Exception:
+                    pass
+                try:
+                    loop.close()
+                except Exception:
+                    pass
+        except Exception as e:
+            logger.warning(f"Error in synchronous execution wrapper: {str(e)}")
+            return self._create_error_result(code, str(e))
+
+    async def _arun(
+        self,
+        code: str,
+        data_source: str | list[dict[str, Any]] | None = None,
+        timeout: int = 30,
+        language: str = "python",  # noqa: ARG002 - accepted for schema compat; backends are Python-only
+        additional_packages: list[str] | None = None,
+    ) -> str | dict[str, Any]:
+        """Execute code in a provider-selected sandbox and return results + artifacts."""
+        runtime: SandboxRuntime | None = None
+        watcher: SandboxArtifactWatcher | None = None
+        try:
+            if isinstance(data_source, str) and data_source.startswith("$tool_output"):
+                logger.warning(f"Reference {data_source!r} not resolved! Skipping the code execution.")
+                return self._create_error_result(code, f"Reference {data_source!r} not resolved!")
+
+            files = self._create_files(data_source)
+            runtime = self._create_runtime(additional_packages)
+            watcher = SandboxArtifactWatcher(runtime)
+
+            await runtime.execute_command(f"mkdir -p {ARTIFACT_OUTPUT_DIR}")
+            await watcher.start(ARTIFACT_OUTPUT_DIR)
+
+            wrapped_code = self._wrap_code_with_dataframe(code, files is not None)
+            start_time = time.time()
+            result = await runtime.execute(wrapped_code, timeout=timeout, files=files)
+            duration_ms = int((time.time() - start_time) * 1000)
+
+            created_files = sorted(await watcher.get_created_files())
+            return await self._process_execution_result(runtime, result, code, created_files, duration_ms)
+
+        except Exception as e:
+            logger.error(f"Error executing code in sandbox: {str(e)}")
+            return self._create_error_result(code, str(e))
+
+        finally:
+            if watcher is not None:
+                try:
+                    await watcher.stop()
+                except Exception as e:
+                    logger.debug(f"Error stopping artifact watcher: {e}")
+            if runtime is not None:
+                await runtime.cleanup()
+
+    def _create_runtime(self, additional_packages: list[str] | None) -> SandboxRuntime:
+        """Build a SandboxRuntime for the configured provider with the requested packages.
+
+        ``api_key``/``domain``/``template`` are forwarded as-is; when None the
+        provider applies its own env fallback (e.g. ``E2B_API_KEY``/``E2B_DOMAIN``).
+        """
+        packages = self._prepare_packages(additional_packages)
+        return SandboxRuntime(
+            provider=self.provider,
+            api_key=self.api_key or None,
+            domain=self.domain,
+            template=self.template,
+            packages=packages or None,
+        )
+
+    def _create_files(self, data_source: Command | list[dict[str, Any]] | None) -> dict[str, str] | None:
+        """Convert a tabular data source into a single cwd-relative CSV upload.
+
+        Returns:
+            ``{DATA_FILE: csv}`` for a non-empty list of dicts, else None.
+        """
+        if isinstance(data_source, Command):
+            data_source = data_source.update["result"]
+
+        if not (isinstance(data_source, list) and data_source and all(isinstance(row, dict) for row in data_source)):
+            if not data_source:
+                logger.info("No data source provided. Ignoring the data source.")
+            else:
+                logger.warning("Unsupported `data_source` type. Expected a non-empty list of dictionaries. Ignoring.")
+            return None
+
+        try:
+            logger.info("Creating a CSV file from the data source to upload to the sandbox...")
+            csv_string = pd.DataFrame(data_source).to_csv(index=False)
+            return {DATA_FILE: csv_string}
+        except Exception as e:
+            logger.warning(f"Error creating files from data source: {str(e)}. Ignoring the data source.")
+            return None
+
+    @staticmethod
+    def _wrap_code_with_dataframe(code: str, has_data: bool) -> str:
+        """Prepend the ``df`` preload when a data file was uploaded (cwd-relative path)."""
+        if not has_data:
+            return code
+        return f"import pandas as pd\ndf = pd.read_csv({DATA_FILE!r})\n{code}"
+
+    def _prepare_packages(self, additional_packages: list[str] | None) -> list[str]:
+        """Combine the default packages with the per-call ones, de-duplicated."""
+        all_packages = list(self.default_additional_packages) + (additional_packages or [])
+        return list(dict.fromkeys(all_packages))
+
+    @staticmethod
+    def _result_to_dict(result: Any, code: str, duration_ms: int) -> dict[str, Any]:
+        """Map a ``SandboxExecutionResult`` to the tool's result dict.
+
+        ``SandboxExecutionResult`` folds the error into ``stderr``; ``status`` is derived
+        from ``exit_code`` and ``duration_ms`` is measured by the tool.
+        """
+        failed = result.exit_code != 0
+        return {
+            "status": "error" if failed else "success",
+            "code": code,
+            "stdout": result.stdout,
+            "stderr": result.stderr,
+            "error": result.stderr if failed else "",
+            "duration_ms": duration_ms,
+        }
+
+    def _create_error_result(self, code: str, error_message: str, duration_ms: int = 0) -> str:
+        """Create a standardized JSON error result."""
+        return json.dumps(
+            {
+                "status": "error",
+                "code": code,
+                "stdout": "",
+                "stderr": "",
+                "error": error_message,
+                "duration_ms": duration_ms,
+            },
+            indent=2,
+        )
+
+    async def _process_execution_result(
+        self,
+        runtime: SandboxRuntime,
+        result: Any,
+        code: str,
+        created_files: list[str],
+        duration_ms: int,
+    ) -> str | dict[str, Any]:
+        """Build the JSON result, and an artifacts response when files were created."""
+        result_dict = self._result_to_dict(result, code, duration_ms)
+
+        if not created_files:
+            result_dict["message"] = "No new files were created during execution"
+            return json.dumps(result_dict, indent=2)
+
+        downloaded_artifacts = await self._create_artifacts_from_files(runtime, created_files)
+        if downloaded_artifacts:
+            execution_summary = self._create_execution_summary(result_dict, created_files)
+            return create_multiple_artifacts_response(result=execution_summary, artifacts=downloaded_artifacts)
+
+        logger.warning("Files were detected but could not be downloaded")
+        return json.dumps(result_dict, indent=2)
+
+    async def _create_artifacts_from_files(self, runtime: SandboxRuntime, file_paths: list[str]) -> list[dict[str, Any]]:
+        """Download each created file and turn it into an artifact dict."""
+        artifacts: list[dict[str, Any]] = []
+        try:
+            for file_path in file_paths:
+                logger.debug(f"Processing newly created file: {file_path}")
+                artifact_dict = await self._download_and_create_artifact(runtime, file_path)
+                if artifact_dict:
+                    artifacts.append(artifact_dict)
+        except Exception as e:
+            logger.error(f"Error creating artifacts from files: {str(e)}")
+        return artifacts
+
+    async def _download_and_create_artifact(self, runtime: SandboxRuntime, file_path: str) -> dict[str, Any] | None:
+        """Download a single file via the runtime and create an artifact."""
+        try:
+            file_content = await runtime.download_file(file_path)
+            if file_content is None:
+                logger.warning(f"Failed to download file: {file_path}")
+                return None
+
+            filename = self._extract_filename(file_path)
+            artifact_response = ArtifactHandler().create_file_artifact(
+                result=f"File created during execution: {file_path}",
+                artifact_data=file_content,
+                artifact_name=filename,
+                artifact_description=f"File created during execution: {file_path}",
+                enable_deduplication=False,
+            )
+            logger.info(f"Successfully downloaded newly created file: {file_path} ({len(file_content)} bytes)")
+            return artifact_response["artifact"]
+        except Exception as e:
+            logger.warning(f"Failed to download file {file_path}: {str(e)}")
+            return None
+
+    def _create_execution_summary(self, result_dict: dict[str, Any], created_files: list[str]) -> str:
+        """Create a markdown execution summary listing the downloaded files."""
+        execution_summary = f"Code executed successfully (status: {result_dict['status']}).\n"
+        if result_dict["stdout"]:
+            execution_summary += f"Output: {result_dict['stdout']}\n"
+        execution_summary += f"Downloaded {len(created_files)} file(s) created during execution:\n\n"
+
+        for file_path in created_files:
+            filename = self._extract_filename(file_path)
+            mime_type = get_mime_type_from_filename(filename)
+            if mime_type.startswith("image/"):
+                execution_summary += f"- **{filename}** (Image): ![{filename}]({filename})\n"
+            else:
+                execution_summary += f"- **{filename}** ({mime_type}): `{file_path}`\n"
+        return execution_summary
+
+    @staticmethod
+    def _extract_filename(file_path: str) -> str:
+        """Extract a filename from a path, with a stable fallback."""
+        filename = os.path.basename(file_path)
+        if not filename:
+            filename = f"sandbox_file_{hash(file_path) % 10000}"
+        return filename

aip_agents/tools/code_sandbox/constant.py

@@ -0,0 +1,13 @@
+"""Constants for Code Sandbox tool.
+
+Authors:
+    Komang Elang Surya Prawira (komang.e.s.prawira@gdplabs.id)
+"""
+
+# The filename `data.csv` is used as a hard-coded value because we currently lack a mechanism
+# to dynamically determine the actual filename. This static filename is used for saving files
+# inside the sandbox, and any pre-population steps will need to read from this filename.
+# Code interacting with the sandbox can directly access the content of the file using the pre-defined variable,
+# eliminating the need to load the file again. Agents or LLMs do not need to be aware of this filename.
+DATA_FILE_NAME = "data.csv"
+DATA_FILE_PATH = f"/files/{DATA_FILE_NAME}"