nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/tool/code_execution/code_sandbox.py

@@ -0,0 +1,267 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import json
+import logging
+import textwrap
+from typing import Any
+from urllib.parse import urljoin
+
+import requests
+import requests.adapters
+from pydantic import HttpUrl
+
+from nat.utils.type_utils import override
+
+logger = logging.getLogger(__file__)
+
+
+class Sandbox(abc.ABC):
+    """Code execution sandbox.
+
+    Args:
+        host: Optional[str] = '127.0.0.1' - Host of the sandbox server.
+            Can also be specified through NEMO_SKILLS_SANDBOX_HOST env var.
+        port: Optional[str] = '5000' - Port of the sandbox server.
+            Can also be specified through NEMO_SKILLS_SANDBOX_PORT env var.
+        ssh_server: Optional[str] = None - SSH server for tunneling requests.
+            Useful if server is running on slurm cluster to which there is an ssh access.
+            Can also be specified through NEMO_SKILLS_SSH_SERVER env var.
+        ssh_key_path: Optional[str] = None - Path to the ssh key for tunneling.
+            Can also be specified through NEMO_SKILLS_SSH_KEY_PATH env var.
+    """
+
+    def __init__(
+        self,
+        *,
+        uri: HttpUrl,
+    ):
+        self.url: str = self._get_execute_url(uri)
+        session = requests.Session()
+        adapter = requests.adapters.HTTPAdapter(pool_maxsize=1500, pool_connections=1500, max_retries=3)
+        session.mount('http://', adapter)
+        session.mount('https://', adapter)
+        self.http_session: requests.Session = session
+
+    def _send_request(self, request: dict[str, Any], timeout_seconds: float) -> dict[str, str]:
+        output = self.http_session.post(
+            url=self.url,
+            data=json.dumps(request),
+            timeout=timeout_seconds,
+            headers={"Content-Type": "application/json"},
+        )
+        # retrying 502 errors
+        if output.status_code == 502:
+            raise requests.exceptions.Timeout
+
+        return self._parse_request_output(output)
+
+    @abc.abstractmethod
+    def _parse_request_output(self, output: requests.Response) -> dict[str, str]:
+        pass
+
+    @abc.abstractmethod
+    def _get_execute_url(self, uri: HttpUrl) -> str:
+        pass
+
+    @abc.abstractmethod
+    def _prepare_request(self, generated_code: str, timeout_seconds: float) -> dict[str, Any]:
+        pass
+
+    async def execute_code(
+        self,
+        generated_code: str,
+        timeout_seconds: float = 10.0,
+        language: str = "python",
+        max_output_characters: int = 1000,
+    ) -> dict[str, str]:
+
+        if language != "python":
+            raise ValueError(f"Language {language} not supported")
+
+        generated_code = generated_code.strip().strip("`")
+        code_to_execute = textwrap.dedent("""
+            import traceback
+            import json
+            import os
+            import warnings
+            import contextlib
+            import io
+            warnings.filterwarnings('ignore')
+            os.environ['OPENBLAS_NUM_THREADS'] = '16'
+        """).strip()
+
+        # Use json.dumps to properly escape the generated_code instead of repr()
+        escaped_code = json.dumps(generated_code)
+        code_to_execute += textwrap.dedent(f"""
+
+            generated_code = {escaped_code}
+
+            stdout = io.StringIO()
+            stderr = io.StringIO()
+
+            with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
+                try:
+                    exec(generated_code)
+                    status = "completed"
+                except Exception:
+                    status = "error"
+                    stderr.write(traceback.format_exc())
+            stdout = stdout.getvalue()
+            stderr = stderr.getvalue()
+            if len(stdout) > {max_output_characters}:
+                stdout = stdout[:{max_output_characters}] + "<output cut>"
+            if len(stderr) > {max_output_characters}:
+                stderr = stderr[:{max_output_characters}] + "<output cut>"
+            if stdout:
+                stdout += "\\n"
+            if stderr:
+                stderr += "\\n"
+            output = {{"process_status": status, "stdout": stdout, "stderr": stderr}}
+            print(json.dumps(output))
+        """).strip()
+        request = self._prepare_request(code_to_execute, timeout_seconds)
+        try:
+            return self._send_request(request, timeout_seconds)
+        except requests.exceptions.Timeout:
+            return {"process_status": "timeout", "stdout": "", "stderr": "Timed out\n"}
+
+
+class LocalSandbox(Sandbox):
+    """Locally hosted sandbox."""
+
+    def __init__(self, *, uri: HttpUrl):
+        super().__init__(uri=uri)
+
+    @override
+    def _get_execute_url(self, uri: HttpUrl) -> str:
+        return urljoin(str(uri), "execute")
+
+    @override
+    def _parse_request_output(self, output: requests.Response) -> dict[str, str]:
+        try:
+            output_json = output.json()
+            assert isinstance(output_json, dict)
+            return output_json
+        except json.JSONDecodeError as e:
+            logger.exception("Error parsing output: %s. %s", output.text, e)
+            return {'process_status': 'error', 'stdout': '', 'stderr': f'Unknown error: {e} \"{output.text}\"'}
+
+    @override
+    def _prepare_request(self,
+                         generated_code: str,
+                         timeout_seconds: float,
+                         language: str = "python",
+                         **kwargs) -> dict[str, Any]:
+        request = {
+            "generated_code": generated_code,
+            "timeout": timeout_seconds,
+            "language": language,
+        }
+        return request
+
+    @override
+    async def execute_code(
+        self,
+        generated_code: str,
+        timeout_seconds: float = 10.0,
+        language: str = "python",
+        max_output_characters: int = 1000,
+    ) -> dict[str, str]:
+        """Override execute_code to bypass the wrapper logic and send user code directly to our server."""
+
+        logger.debug("Raw input generated_code: %s", generated_code)
+
+        # The input appears to be a string representation of a dictionary
+        # We need to parse it and extract the actual code
+        try:
+            # Try to evaluate the string as a Python literal (dictionary)
+            import ast
+            parsed_dict = ast.literal_eval(generated_code)
+            if isinstance(parsed_dict, dict) and 'generated_code' in parsed_dict:
+                actual_code = parsed_dict['generated_code']
+                assert isinstance(actual_code, str)
+                logger.debug("Extracted code from dict: %s...", actual_code[:100])
+            else:
+                # If it's not a dict or doesn't have the expected key, use as-is
+                actual_code = generated_code
+                logger.debug("Using code as-is: %s...", actual_code[:100])
+        except (ValueError, SyntaxError):
+            # If parsing fails, use the input as-is
+            actual_code = generated_code
+            logger.debug("Failed to parse, using as-is: %s...", actual_code[:100])
+
+        # Clean the actual code more carefully to avoid removing backticks that are part of Python code
+        # remove all leading/trailing whitespace -- strip()
+        # remove all leading/trailing backticks -- strip("`")
+        # may potentially start with python, so just trim from the front.
+        POTENTIAL_PREFIXES = ["python"]
+        actual_code = actual_code.strip().strip("`")
+        for prefix in POTENTIAL_PREFIXES:
+            if actual_code.startswith(prefix):
+                actual_code = actual_code[len(prefix):]
+                break
+
+        # Send the user's code directly to our server without any wrapper logic
+        # Our server already handles stdout/stderr capture and error handling
+        request = self._prepare_request(actual_code, timeout_seconds, language)
+        try:
+            return self._send_request(request, timeout_seconds)
+        except requests.exceptions.Timeout:
+            return {"process_status": "timeout", "stdout": "", "stderr": "Timed out\n"}
+
+
+class PistonSandbox(Sandbox):
+    """Piston sandbox (https://github.com/engineer-man/piston)"""
+
+    @override
+    def _get_execute_url(self, uri: HttpUrl) -> str:
+        return urljoin(str(uri), "execute")
+
+    @override
+    def _parse_request_output(self, output: requests.Response) -> dict[str, str]:
+        output_json = output.json()
+        assert isinstance(output_json, dict)
+        assert 'run' in output_json
+        run_json = output_json['run']
+        assert isinstance(run_json, dict)
+        if run_json["code"] != 0:
+            return {'process_status': "error", 'stdout': run_json['stdout'], 'stderr': run_json['stderr']}
+        return {'process_status': "completed", 'stdout': run_json['stdout'], 'stderr': run_json['stderr']}
+
+    @override
+    def _prepare_request(self, generated_code: str, timeout_seconds: float, **kwargs) -> dict[str, Any]:
+        return {
+            "language": "py",
+            "version": "3.10.0",
+            "files": [{
+                "content": generated_code,
+            }],
+            "stdin": "",
+            "args": [],
+            "run_timeout": timeout_seconds * 1000.0,  # milliseconds
+            "compile_memory_limit": -1,
+            "run_memory_limit": -1,
+        }
+
+
+def get_sandbox(sandbox_type: str = "local", **kwargs):
+    """A helper function to make it easier to set sandbox through cmd."""
+    sandboxes = {
+        'local': LocalSandbox,
+        'piston': PistonSandbox,
+    }
+    sandbox_class = sandboxes[sandbox_type.lower()]
+    return sandbox_class(**kwargs)

nat/tool/code_execution/local_sandbox/.gitignore

@@ -0,0 +1 @@
+persistence_test.*

nat/tool/code_execution/local_sandbox/Dockerfile.sandbox

@@ -0,0 +1,60 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Use the base image with Python 3.10 and Flask
+FROM tiangolo/uwsgi-nginx-flask:python3.10
+
+# Install dependencies required for Lean 4 and other tools
+RUN apt-get update && \
+    apt-get install -y curl git && \
+    curl https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -sSf | sh -s -- -y && \
+    /root/.elan/bin/elan toolchain install leanprover/lean4:v4.12.0 && \
+    /root/.elan/bin/elan default leanprover/lean4:v4.12.0 && \
+    /root/.elan/bin/elan self update
+
+# Set environment variables to include Lean and elan/lake in the PATH
+ENV PATH="/root/.elan/bin:$PATH"
+
+# Create Lean project directory and initialize a new Lean project with Mathlib4
+RUN mkdir -p /lean4 && cd /lean4 && \
+    /root/.elan/bin/lake new my_project && \
+    cd my_project && \
+    echo 'leanprover/lean4:v4.12.0' > lean-toolchain && \
+    echo 'require mathlib from git "https://github.com/leanprover-community/mathlib4" @ "v4.12.0"' >> lakefile.lean
+
+# Download and cache Mathlib4 to avoid recompiling, then build the project
+RUN cd /lean4/my_project && \
+    /root/.elan/bin/lake exe cache get && \
+    /root/.elan/bin/lake build
+
+# Set environment variables to include Lean project path
+ENV LEAN_PATH="/lean4/my_project"
+ENV PATH="/lean4/my_project:$PATH"
+
+# Set up application code and install Python dependencies
+COPY sandbox.requirements.txt /app/requirements.txt
+RUN pip install --no-cache-dir -r /app/requirements.txt
+COPY local_sandbox_server.py /app/main.py
+
+# Set the working directory to /app
+WORKDIR /app
+
+# Set Flask app environment variables and ports
+ARG UWSGI_CHEAPER
+ENV UWSGI_CHEAPER=$UWSGI_CHEAPER
+
+ARG UWSGI_PROCESSES
+ENV UWSGI_PROCESSES=$UWSGI_PROCESSES
+
+ENV LISTEN_PORT=6000

nat/tool/code_execution/local_sandbox/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nat/tool/code_execution/local_sandbox/local_sandbox_server.py

@@ -0,0 +1,198 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import multiprocessing
+import os
+import resource
+from enum import Enum
+from io import StringIO
+
+from flask import Flask
+from flask import Request
+from flask import Response
+from flask import request
+from pydantic import BaseModel
+from pydantic import Field
+
+app = Flask(__name__)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+
+class CodeExecutionStatus(str, Enum):
+    """
+    Status of code execution.
+    """
+    COMPLETED = "completed"
+    ERROR = "error"
+    TIMEOUT = "timeout"
+
+
+class CodeExecutionResult(BaseModel):
+    """
+    Result of code execution.
+    """
+    process_status: CodeExecutionStatus = Field(default=CodeExecutionStatus.COMPLETED,
+                                                description="Status of the process")
+    stdout: str = Field(description="Standard output of the process")
+    stderr: str = Field(description="Standard error of the process")
+
+
+class CodeExecutionResponse(Response):
+    """
+    Response class that returns a JSON response with the given status code and result.
+    """
+
+    def __init__(self, status_code: int, result: CodeExecutionResult):
+        super().__init__(status=status_code, mimetype="application/json", response=result.model_dump_json())
+
+    @classmethod
+    def with_error(cls, status_code: int, error_message: str) -> CodeExecutionResponse:
+        return cls(status_code,
+                   CodeExecutionResult(process_status=CodeExecutionStatus.ERROR, stdout="", stderr=error_message))
+
+
+@app.after_request
+def add_hsts_header(response):
+    response.headers['Strict-Transport-Security'] = 'max-age=31536000; includeSubDomains'
+    response.headers['X-Content-Type-Options'] = 'nosniff'
+    response.headers['X-Frame-Options'] = 'SAMEORIGIN'
+    response.headers['X-XSS-Protection'] = '1; mode=block'
+
+    return response
+
+
+def execute_python(generated_code: str, timeout: float) -> CodeExecutionResult:
+    """
+    Execute Python code in a subprocess.
+
+    Args:
+        generated_code: The code to execute
+        timeout: The timeout for the execution
+
+    Returns:
+        CodeExecutionResult object containing the execution result
+    """
+
+    # running in a separate process to ensure any kind of crashes are properly handled
+    queue = multiprocessing.Queue()
+    process = multiprocessing.Process(target=execute_code_subprocess, args=(generated_code, queue))
+
+    process.start()
+    # wait until the process finishes or the timeout expires
+    process.join(timeout=timeout)
+    if process.exitcode is None:
+        process.kill()
+        return CodeExecutionResult(process_status=CodeExecutionStatus.TIMEOUT, stdout="", stderr="Timed out\n")
+
+    return queue.get()
+
+
+# need to memory-limit to avoid common errors of allocating too much
+# but this has to be done in a subprocess to not crush server itself
+def execute_code_subprocess(generated_code: str, queue):
+    """
+    Execute code in a subprocess.
+
+    Args:
+        generated_code: The code to execute
+        queue: The queue to put the result in
+    """
+
+    logger.debug("execute_code_subprocess started, PID: %s", os.getpid())
+
+    try:
+        limit = 1024 * 1024 * 1024 * 10  # 10gb - somehow with a smaller limit the server dies when numpy is used
+        resource.setrlimit(resource.RLIMIT_AS, (limit, limit))
+        resource.setrlimit(resource.RLIMIT_DATA, (limit, limit))
+    except Exception as e:
+        logger.exception("Failed to set resource limits, PID: %s, error: %s", os.getpid(), e)
+
+    stdout_capture = StringIO()
+    stderr_capture = StringIO()
+    try:
+        with contextlib.redirect_stdout(stdout_capture), contextlib.redirect_stderr(stderr_capture):
+            exec(generated_code, {})
+        logger.debug("execute_code_subprocess finished, PID: %s", os.getpid())
+        queue.put(CodeExecutionResult(stdout=stdout_capture.getvalue(), stderr=stderr_capture.getvalue()))
+    except Exception as e:
+        import traceback
+        with contextlib.redirect_stderr(stderr_capture):
+            traceback.print_exc()
+        logger.debug("execute_code_subprocess failed, PID: %s, error: %s", os.getpid(), e)
+        queue.put(
+            CodeExecutionResult(process_status=CodeExecutionStatus.ERROR,
+                                stdout=stdout_capture.getvalue(),
+                                stderr=stderr_capture.getvalue()))
+
+
+def do_execute(request: Request) -> CodeExecutionResponse:
+    """
+    Main function to handle execution requests.
+
+    Args:
+        request: Request object containing the execution request
+
+    Returns:
+        CodeExecutionResponse object containing the execution result
+    """
+    try:
+        # Check if request has JSON data
+        if not request.is_json:
+            return CodeExecutionResponse.with_error(400, "Request must be JSON")
+
+        # Get JSON data safely
+        json_data = request.get_json(silent=True)
+
+        if json_data is None:
+            return CodeExecutionResponse.with_error(400, "Invalid JSON data")
+
+        # Check for required fields
+        if 'generated_code' not in json_data:
+            return CodeExecutionResponse.with_error(400, "Missing required field: generated_code")
+
+        if 'timeout' not in json_data:
+            return CodeExecutionResponse.with_error(400, "Missing required field: timeout")
+
+        if 'language' not in json_data:
+            return CodeExecutionResponse.with_error(400, "Missing required field: language")
+
+        generated_code: str | None = json_data.get('generated_code', None)
+        assert generated_code is not None
+        timeout: float | None = json_data.get('timeout', None)
+        assert timeout is not None
+        language: str | None = json_data.get('language', None)
+        assert language is not None
+
+        if language != 'python':
+            return CodeExecutionResponse.with_error(400, "Only python execution is supported")
+
+        return CodeExecutionResponse(200, execute_python(generated_code, timeout))
+
+    except Exception as e:
+        return CodeExecutionResponse.with_error(500, f"Server error: {str(e)}")
+
+
+# Main Flask endpoint to handle execution requests
+@app.route("/execute", methods=["POST"])
+def execute():
+    return do_execute(request)
+
+
+if __name__ == '__main__':
+    app.run(port=6000)

nat/tool/code_execution/local_sandbox/sandbox.requirements.txt

@@ -0,0 +1,6 @@
+numpy
+pandas
+scipy
+ipython
+plotly
+pydantic

nat/tool/code_execution/local_sandbox/start_local_sandbox.sh

@@ -0,0 +1,50 @@
+#!/bin/bash
+
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Usage: ./start_local_sandbox.sh [SANDBOX_NAME] [OUTPUT_DATA_PATH]
+# NOTE: needs to run from the root of the repo!
+
+DOCKER_COMMAND=${DOCKER_COMMAND:-"docker"}
+SANDBOX_NAME=${1:-'local-sandbox'}
+NUM_THREADS=10
+
+# Get the output_data directory path for mounting
+# Priority: command line argument > environment variable > default path (current directory)
+OUTPUT_DATA_PATH=${2:-${OUTPUT_DATA_PATH:-$(pwd)}}
+
+echo "Starting sandbox with container name: ${SANDBOX_NAME}"
+echo "Mounting output_data directory: ${OUTPUT_DATA_PATH}"
+
+# Verify the path exists before mounting, create if it doesn't
+if [ ! -d "${OUTPUT_DATA_PATH}" ]; then
+    echo "Output data directory does not exist, creating: ${OUTPUT_DATA_PATH}"
+    mkdir -p "${OUTPUT_DATA_PATH}"
+fi
+
+# Check if the Docker image already exists
+if ! ${DOCKER_COMMAND} images ${SANDBOX_NAME} | grep -q "${SANDBOX_NAME}"; then
+    echo "Docker image not found locally. Building ${SANDBOX_NAME}..."
+    ${DOCKER_COMMAND} build --tag=${SANDBOX_NAME} --build-arg="UWSGI_PROCESSES=$((${NUM_THREADS} * 10))" --build-arg="UWSGI_CHEAPER=${NUM_THREADS}" -f Dockerfile.sandbox .
+else
+    echo "Using existing Docker image: ${SANDBOX_NAME}"
+fi
+
+# Mount the output_data directory directly so files created in container appear in the local directory
+${DOCKER_COMMAND} run --rm --name=local-sandbox \
+  --network=host \
+  -v "${OUTPUT_DATA_PATH}:/workspace" \
+  -w /workspace \
+  ${SANDBOX_NAME}

nat/tool/code_execution/register.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from typing import Literal
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import HttpUrl
+
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class CodeExecutionToolConfig(FunctionBaseConfig, name="code_execution"):
+    """
+    Tool for executing python code in a remotely hosted sandbox environment.
+    """
+    uri: HttpUrl = Field(default=HttpUrl("http://127.0.0.1:6000"),
+                         description="URI for the code execution sandbox server")
+    sandbox_type: Literal["local", "piston"] = Field(default="local", description="The type of code execution sandbox")
+    timeout: float = Field(default=10.0, description="Number of seconds to wait for a code execution request")
+    max_output_characters: int = Field(default=1000, description="Maximum number of characters that can be returned")
+
+
+@register_function(config_type=CodeExecutionToolConfig)
+async def code_execution_tool(config: CodeExecutionToolConfig, builder: Builder):
+    from nat.tool.code_execution.code_sandbox import get_sandbox
+
+    class CodeExecutionInputSchema(BaseModel):
+        generated_code: str = Field(description="String containing the code to be executed")
+
+    # Create sandbox without working_directory
+    sandbox_kwargs = {"uri": config.uri}
+
+    sandbox = get_sandbox(sandbox_type=config.sandbox_type, **sandbox_kwargs)
+    logger.info(f"[DEBUG] Created sandbox of type: {config.sandbox_type}")
+
+    async def _execute_code(generated_code: str) -> dict:
+        logger.info("Executing code in the sandbox at %s", config.uri)
+        try:
+            output = await sandbox.execute_code(
+                generated_code=generated_code,
+                language="python",
+                timeout_seconds=config.timeout,
+                max_output_characters=config.max_output_characters,
+            )
+        except Exception as e:
+            logger.exception("Error when executing code in the sandbox, %s", e)
+            return {"process_status": "error", "stdout": "", "stderr": str(e)}
+        return output
+
+    yield FunctionInfo.from_fn(
+        fn=_execute_code,
+        input_schema=CodeExecutionInputSchema,
+        description="""Executes the provied 'generated_code' in a python sandbox environment and returns
+        a dictionary containing stdout, stderr, and the execution status, as well as a session_id. The
+        session_id can be used to append to code that was previously executed.""")