ms-enclave 0.0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

ms_enclave/sandbox/tools/sandbox_tools/file_operation.py

@@ -0,0 +1,215 @@
+"""File operation tool for reading and writing files."""
+
+import io
+import os
+import tarfile
+from typing import TYPE_CHECKING, Literal, Optional
+
+from ms_enclave.sandbox.model import ExecutionStatus, SandboxType, ToolResult
+from ms_enclave.sandbox.tools.base import register_tool
+from ms_enclave.sandbox.tools.sandbox_tool import SandboxTool
+from ms_enclave.sandbox.tools.tool_info import ToolParams
+
+if TYPE_CHECKING:
+    from ms_enclave.sandbox.boxes import Sandbox
+
+
+@register_tool('file_operation')
+class FileOperation(SandboxTool):
+
+    _name = 'file_operation'
+    _sandbox_type = SandboxType.DOCKER
+    _description = 'Perform file operations like read, write, delete, and list files'
+    _parameters = ToolParams(
+        type='object',
+        properties={
+            'operation': {
+                'type': 'string',
+                'description': 'Type of file operation to perform',
+                'enum': ['create', 'read', 'write', 'delete', 'list', 'exists']
+            },
+            'file_path': {
+                'type': 'string',
+                'description': 'Path to the file or directory'
+            },
+            'content': {
+                'type': 'string',
+                'description': 'Content to write to file (only for write operation)'
+            },
+            'encoding': {
+                'type': 'string',
+                'description': 'File encoding',
+                'default': 'utf-8'
+            }
+        },
+        required=['operation', 'file_path']
+    )
+
+    async def execute(
+        self,
+        sandbox_context: 'Sandbox',
+        operation: str,
+        file_path: str,
+        content: Optional[str] = None,
+        encoding: str = 'utf-8'
+    ) -> ToolResult:
+        """Perform file operations in the Docker container."""
+
+        if not file_path.strip():
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error='No file path provided'
+            )
+
+        try:
+            if operation == 'read':
+                return await self._read_file(sandbox_context, file_path, encoding)
+            elif operation == 'write':
+                if content is None:
+                    return ToolResult(
+                        tool_name=self.name,
+                        status=ExecutionStatus.ERROR,
+                        output='',
+                        error='Content is required for write operation'
+                    )
+                return await self._write_file(sandbox_context, file_path, content, encoding)
+            elif operation == 'delete':
+                return await self._delete_file(sandbox_context, file_path)
+            elif operation == 'list':
+                return await self._list_directory(sandbox_context, file_path)
+            elif operation == 'exists':
+                return await self._check_exists(sandbox_context, file_path)
+            elif operation == 'create':
+                if content is None:
+                    content = ''
+                return await self._write_file(sandbox_context, file_path, content, encoding)
+            else:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.ERROR,
+                    output='',
+                    error=f'Unknown operation: {operation}'
+                )
+
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Operation failed: {str(e)}'
+            )
+
+    async def _read_file(self, sandbox_context: 'Sandbox', file_path: str, encoding: str) -> ToolResult:
+        """Read file content from the container."""
+        try:
+            # Use cat command to read file content
+            result = await sandbox_context.execute_command(f'cat "{file_path}"')
+
+            if result.exit_code == 0:
+                return ToolResult(tool_name=self.name, status=ExecutionStatus.SUCCESS, output=result.stdout, error=None)
+            else:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.ERROR,
+                    output='',
+                    error=result.stderr or f'Failed to read file: {file_path}'
+                )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Read failed: {str(e)}'
+            )
+
+    async def _write_file(self, sandbox_context: 'Sandbox', file_path: str, content: str, encoding: str) -> ToolResult:
+        """Write content to a file in the container."""
+        try:
+            # Create directory if it doesn't exist
+            dir_path = os.path.dirname(file_path)
+            if dir_path:
+                await sandbox_context.execute_command(f'mkdir -p "{dir_path}"')
+
+            # Write file using tar archive (similar to python_executor)
+            await self._write_file_to_container(sandbox_context, file_path, content)
+
+            return ToolResult(
+                tool_name=self.name,
+                status=ExecutionStatus.SUCCESS,
+                output=f'File written successfully: {file_path}',
+                error=None
+            )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Write failed: {str(e)}'
+            )
+
+    async def _delete_file(self, sandbox_context: 'Sandbox', file_path: str) -> ToolResult:
+        """Delete a file or directory from the container."""
+        try:
+            result = await sandbox_context.execute_command(f'rm -rf "{file_path}"')
+
+            if result.exit_code == 0:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.SUCCESS,
+                    output=f'Successfully deleted: {file_path}',
+                    error=None
+                )
+            else:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.ERROR,
+                    output='',
+                    error=result.stderr or f'Failed to delete: {file_path}'
+                )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Delete failed: {str(e)}'
+            )
+
+    async def _list_directory(self, sandbox_context: 'Sandbox', dir_path: str) -> ToolResult:
+        """List contents of a directory."""
+        try:
+            result = await sandbox_context.execute_command(f'ls -la "{dir_path}"')
+
+            if result.exit_code == 0:
+                return ToolResult(tool_name=self.name, status=ExecutionStatus.SUCCESS, output=result.stdout, error=None)
+            else:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.ERROR,
+                    output='',
+                    error=result.stderr or f'Failed to list directory: {dir_path}'
+                )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'List failed: {str(e)}'
+            )
+
+    async def _check_exists(self, sandbox_context: 'Sandbox', file_path: str) -> ToolResult:
+        """Check if a file or directory exists."""
+        try:
+            result = await sandbox_context.execute_command(f'test -e "{file_path}"')
+
+            exists = result.exit_code == 0
+            return ToolResult(
+                tool_name=self.name,
+                status=ExecutionStatus.SUCCESS,
+                output=f'{"exists" if exists else "does not exist"}',
+                error=None
+            )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Exists check failed: {str(e)}'
+            )
+
+    async def _write_file_to_container(self, sandbox_context: 'Sandbox', file_path: str, content: str) -> None:
+        """Write content to a file in the container using tar archive."""
+        # Create a tar archive in memory
+        tar_stream = io.BytesIO()
+        tar = tarfile.TarFile(fileobj=tar_stream, mode='w')
+
+        # Add file to tar
+        file_data = content.encode('utf-8')
+        tarinfo = tarfile.TarInfo(name=os.path.basename(file_path))
+        tarinfo.size = len(file_data)
+        tar.addfile(tarinfo, io.BytesIO(file_data))
+        tar.close()
+
+        # Write to container
+        tar_stream.seek(0)
+        sandbox_context.container.put_archive(os.path.dirname(file_path), tar_stream.getvalue())

ms_enclave/sandbox/tools/sandbox_tools/notebook_executor.py

@@ -0,0 +1,167 @@
+"""Notebook code execution tool for Jupyter kernels."""
+
+import json
+import time
+import uuid
+from typing import TYPE_CHECKING, Optional
+
+from ms_enclave.sandbox.model import CommandResult, ExecutionStatus, SandboxType, ToolResult
+from ms_enclave.sandbox.tools.base import Tool, register_tool
+from ms_enclave.sandbox.tools.sandbox_tool import SandboxTool
+from ms_enclave.sandbox.tools.tool_info import ToolParams
+from ms_enclave.utils import get_logger
+
+if TYPE_CHECKING:
+    from ms_enclave.sandbox.boxes import Sandbox
+
+logger = get_logger()
+
+
+@register_tool('notebook_executor')
+class NotebookExecutor(SandboxTool):
+
+    _name = 'notebook_executor'
+    _sandbox_type = SandboxType.DOCKER_NOTEBOOK
+    _description = 'Execute Python code in a Jupyter kernel environment'
+    _parameters = ToolParams(
+        type='object',
+        properties={
+            'code': {
+                'type': 'string',
+                'description': 'Python code to execute in the notebook kernel'
+            },
+            'timeout': {
+                'type': 'integer',
+                'description': 'Execution timeout in seconds',
+                'default': 30
+            }
+        },
+        required=['code']
+    )
+
+    async def execute(self, sandbox_context: 'Sandbox', code: str, timeout: Optional[int] = 30) -> ToolResult:
+        """Execute Python code in the Jupyter kernel."""
+
+        if not code.strip():
+            return ToolResult(tool_name=self.name, status=ExecutionStatus.ERROR, output='', error='No code provided')
+
+        try:
+            # Execute code using the sandbox's Jupyter kernel
+            result = await self._execute_in_kernel(sandbox_context, code, timeout)
+
+            if result.exit_code == 0:
+                status = ExecutionStatus.SUCCESS
+            else:
+                status = ExecutionStatus.ERROR
+
+            return ToolResult(
+                tool_name=self.name,
+                status=status,
+                output=result.stdout,
+                error=result.stderr if result.stderr else None
+            )
+
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Execution failed: {str(e)}'
+            )
+
+    async def _execute_in_kernel(self, sandbox_context: 'Sandbox', code: str, timeout: Optional[int]):
+        """Execute code in the Jupyter kernel via websocket."""
+
+        # Check if sandbox has the required Jupyter components
+        if not hasattr(sandbox_context, 'ws') or not hasattr(sandbox_context, 'kernel_id'):
+            raise RuntimeError('Sandbox does not have Jupyter kernel setup')
+
+        if not sandbox_context.ws or not sandbox_context.kernel_id:
+            raise RuntimeError('Jupyter kernel is not ready')
+
+        # Send execute request
+        msg_id = self._send_execute_request(sandbox_context, code)
+
+        outputs = []
+        result = None
+        error_occurred = False
+        error_msg = ''
+        actual_timeout = timeout or 30
+        start_time = time.time()
+
+        while True:
+            elapsed = time.time() - start_time
+            if elapsed >= actual_timeout:
+                error_occurred = True
+                error_msg = f'Execution timed out after {actual_timeout} seconds'
+                logger.error(error_msg)
+                break
+
+            try:
+                # Set a short timeout for recv to allow periodic timeout check
+                sandbox_context.ws.settimeout(min(1, actual_timeout - elapsed))
+                msg = json.loads(sandbox_context.ws.recv())
+                parent_msg_id = msg.get('parent_header', {}).get('msg_id')
+
+                # Skip unrelated messages
+                if parent_msg_id != msg_id:
+                    continue
+
+                msg_type = msg.get('msg_type', '')
+                msg_content = msg.get('content', {})
+
+                if msg_type == 'stream':
+                    outputs.append(msg_content['text'])
+                elif msg_type == 'execute_result':
+                    result = msg_content['data'].get('text/plain', '')
+                elif msg_type == 'error':
+                    error_occurred = True
+                    error_msg = '\n'.join(msg_content.get('traceback', []))
+                    outputs.append(error_msg)
+                elif msg_type == 'status' and msg_content.get('execution_state') == 'idle':
+                    break
+
+            except Exception as e:
+                logger.error(f'Error receiving message: {e}')
+                error_occurred = True
+                error_msg = str(e)
+                break
+
+        output_text = ''.join(outputs)
+        if result:
+            output_text += f'\nResult: {result}'
+        if error_msg and error_msg not in output_text:
+            output_text += f'\n{error_msg}'
+
+        return CommandResult(
+            status=ExecutionStatus.SUCCESS if not error_occurred else ExecutionStatus.ERROR,
+            command=code,
+            exit_code=1 if error_occurred else 0,
+            stdout=output_text if not error_occurred else '',
+            stderr=output_text if error_occurred else ''
+        )
+
+    def _send_execute_request(self, sandbox_context: 'Sandbox', code: str) -> str:
+        """Send code execution request to kernel."""
+        # Generate a unique message ID
+        msg_id = str(uuid.uuid4())
+
+        # Create execute request
+        execute_request = {
+            'header': {
+                'msg_id': msg_id,
+                'username': 'anonymous',
+                'session': str(uuid.uuid4()),
+                'msg_type': 'execute_request',
+                'version': '5.0',
+            },
+            'parent_header': {},
+            'metadata': {},
+            'content': {
+                'code': code,
+                'silent': False,
+                'store_history': True,
+                'user_expressions': {},
+                'allow_stdin': False,
+            },
+        }
+
+        sandbox_context.ws.send(json.dumps(execute_request))
+        return msg_id

ms_enclave/sandbox/tools/sandbox_tools/python_executor.py

@@ -0,0 +1,87 @@
+"""Python code execution tool."""
+
+import os
+import uuid
+from typing import TYPE_CHECKING, Optional
+
+from ms_enclave.sandbox.model import ExecutionStatus, SandboxType, ToolResult
+from ms_enclave.sandbox.tools.base import Tool, register_tool
+from ms_enclave.sandbox.tools.sandbox_tool import SandboxTool
+from ms_enclave.sandbox.tools.tool_info import ToolParams
+
+if TYPE_CHECKING:
+    from ms_enclave.sandbox.boxes import DockerSandbox
+
+
+@register_tool('python_executor')
+class PythonExecutor(SandboxTool):
+
+    _name = 'python_executor'
+    _sandbox_type = SandboxType.DOCKER
+    _description = 'Execute Python code in an isolated environment using IPython'
+    _parameters = ToolParams(
+        type='object',
+        properties={
+            'code': {
+                'type': 'string',
+                'description': 'Python code to execute'
+            },
+            'timeout': {
+                'type': 'integer',
+                'description': 'Execution timeout in seconds',
+                'default': 30
+            }
+        },
+        required=['code']
+    )
+
+    async def execute(self, sandbox_context: 'DockerSandbox', code: str, timeout: Optional[int] = 30) -> ToolResult:
+        """Execute Python code by writing to a temporary file and executing it."""
+
+        script_basename = f'exec_script_{uuid.uuid4().hex}.py'
+        script_path = f'/tmp/{script_basename}'
+
+        if not code.strip():
+            return ToolResult(tool_name=self.name, status=ExecutionStatus.ERROR, output='', error='No code provided')
+
+        try:
+
+            # Write script to container to avoid long code errors
+            await self._write_file_to_container(sandbox_context, script_path, code)
+
+            # Execute using python
+            command = f'python {script_path}'
+            result = await sandbox_context.execute_command(command, timeout=timeout)
+
+            if result.exit_code == 0:
+                status = ExecutionStatus.SUCCESS
+            else:
+                status = ExecutionStatus.ERROR
+
+            return ToolResult(
+                tool_name=self.name,
+                status=status,
+                output=result.stdout,
+                error=result.stderr if result.stderr else None
+            )
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Execution failed: {str(e)}'
+            )
+
+    async def _write_file_to_container(self, sandbox_context: 'DockerSandbox', file_path: str, content: str) -> None:
+        """Write content to a file in the container."""
+        import io
+        import tarfile
+
+        # Create a tar archive in memory using context managers
+        with io.BytesIO() as tar_stream:
+            with tarfile.TarFile(fileobj=tar_stream, mode='w') as tar:
+                file_data = content.encode('utf-8')
+                tarinfo = tarfile.TarInfo(name=os.path.basename(file_path))
+                tarinfo.size = len(file_data)
+                tar.addfile(tarinfo, io.BytesIO(file_data))
+
+                # Reset stream position and put archive into container
+                tar_stream.seek(0)
+                sandbox_context.container.put_archive(os.path.dirname(file_path), tar_stream.getvalue())

ms_enclave/sandbox/tools/sandbox_tools/shell_executor.py

@@ -0,0 +1,63 @@
+"""Shell command execution tool."""
+
+from typing import TYPE_CHECKING, Optional
+
+from ms_enclave.sandbox.model import ExecutionStatus, SandboxType, ToolResult
+from ms_enclave.sandbox.tools.base import register_tool
+from ms_enclave.sandbox.tools.sandbox_tool import SandboxTool
+from ms_enclave.sandbox.tools.tool_info import ToolParams
+
+if TYPE_CHECKING:
+    from ms_enclave.sandbox.boxes import Sandbox
+
+
+@register_tool('shell_executor')
+class ShellExecutor(SandboxTool):
+
+    _name = 'shell_executor'
+    _sandbox_type = SandboxType.DOCKER
+    _description = 'Execute shell commands in an isolated environment'
+    _parameters = ToolParams(
+        type='object',
+        properties={
+            'command': {
+                'type': 'string',
+                'description': 'Shell command to execute'
+            },
+            'timeout': {
+                'type': 'integer',
+                'description': 'Execution timeout in seconds',
+                'default': 30
+            }
+        },
+        required=['command']
+    )
+
+    async def execute(self, sandbox_context: 'Sandbox', command: str, timeout: Optional[int] = 30) -> ToolResult:
+        """Execute shell command in the Docker container."""
+
+        if not command.strip():
+            return ToolResult(tool_name=self.name, status=ExecutionStatus.ERROR, output='', error='No command provided')
+
+        try:
+            result = await sandbox_context.execute_command(command, timeout=timeout)
+
+            if result.exit_code == 0:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.SUCCESS,
+                    output=result.stdout,
+                    error=result.stderr if result.stderr else None
+                )
+            else:
+                return ToolResult(
+                    tool_name=self.name,
+                    status=ExecutionStatus.ERROR,
+                    output=result.stdout,
+                    error=result.stderr if result.stderr else f'Command failed with exit code {result.exit_code}'
+                )
+
+        except Exception as e:
+            return ToolResult(
+                tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Execution failed: {str(e)}'
+            )

ms_enclave/sandbox/tools/tool_info.py

@@ -0,0 +1,141 @@
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Literal, Optional, TypeAlias, Union, get_args, get_type_hints
+
+from docstring_parser import Docstring, parse
+from pydantic import BaseModel, Field
+
+from ms_enclave.utils.json_schema import JSONSchema, JSONType, json_schema, python_type_to_json_type
+
+ToolParam: TypeAlias = JSONSchema
+"""Description of tool parameter in JSON Schema format."""
+
+
+class ToolParams(BaseModel):
+    """Description of tool parameters object in JSON Schema format."""
+
+    type: Literal['object'] = Field(default='object', description="Params type (always 'object')")
+    properties: Dict[str, ToolParam] = Field(default_factory=dict, description='Tool function parameters.')
+    required: List[str] = Field(default_factory=list, description='List of required fields.')
+    additionalProperties: bool = Field(
+        default=False, description='Are additional object properties allowed? (always `False`)'
+    )
+
+
+class ToolInfo(BaseModel):
+    """Specification of a tool (JSON Schema compatible)
+
+    If you are implementing a ModelAPI, most LLM libraries can
+    be passed this object (dumped to a dict) directly as a function
+    specification. For example, in the OpenAI provider:
+
+    ```python
+    ChatCompletionToolParam(
+        type="function",
+        function=tool.model_dump(exclude_none=True),
+    )
+    ```
+
+    In some cases the field names don't match up exactly. In that case
+    call `model_dump()` on the `parameters` field. For example, in the
+    Anthropic provider:
+
+    ```python
+    ToolParam(
+        name=tool.name,
+        description=tool.description,
+        input_schema=tool.parameters.model_dump(exclude_none=True),
+    )
+    ```
+    """
+
+    name: str = Field(description='Name of tool.')
+    description: str = Field(description='Short description of tool.')
+    parameters: ToolParams = Field(default_factory=ToolParams, description='JSON Schema of tool parameters object.')
+    options: Optional[Dict[str, object]] = Field(
+        default=None,
+        description=
+        'Optional property bag that can be used by the model provider to customize the implementation of the tool'
+    )
+
+
+def parse_tool_info(func: Callable[..., Any]) -> ToolInfo:
+    # tool may already have registry attributes w/ tool info
+
+    if (getattr(func, 'name', None) and getattr(func, 'description', None) and getattr(func, 'parameters', None)):
+        return ToolInfo(
+            name=func.name,
+            description=func.description,
+            parameters=func.parameters,
+        )
+
+    signature = inspect.signature(func)
+    type_hints = get_type_hints(func)
+    docstring = inspect.getdoc(func)
+    parsed_docstring: Optional[Docstring] = parse(docstring) if docstring else None
+
+    info = ToolInfo(name=func.__name__, description='')
+
+    for param_name, param in signature.parameters.items():
+        tool_param = ToolParam()
+
+        # Parse docstring
+        docstring_info = parse_docstring(docstring, param_name)
+
+        # Get type information from type annotations
+        if param_name in type_hints:
+            tool_param = json_schema(type_hints[param_name])
+        # as a fallback try to parse it from the docstring
+        # (this is minimally necessary for backwards compatiblity
+        #  with tools gen1 type parsing, which only used docstrings)
+        elif 'docstring_type' in docstring_info:
+            json_type = python_type_to_json_type(docstring_info['docstring_type'])
+            if json_type and (json_type in get_args(JSONType)):
+                tool_param = ToolParam(type=json_type)
+
+        # Get default value
+        if param.default is param.empty:
+            info.parameters.required.append(param_name)
+        else:
+            tool_param.default = param.default
+
+        # Add description from docstring
+        if 'description' in docstring_info:
+            tool_param.description = docstring_info['description']
+
+        # append the tool param
+        info.parameters.properties[param_name] = tool_param
+
+    # Add function description if available
+    if parsed_docstring:
+        if parsed_docstring.description:
+            info.description = parsed_docstring.description.strip()
+        elif parsed_docstring.long_description:
+            info.description = parsed_docstring.long_description.strip()
+        elif parsed_docstring.short_description:
+            info.description = parsed_docstring.short_description.strip()
+
+        # Add examples if available
+        if parsed_docstring.examples:
+            examples = '\n\n'.join([(example.description or '') for example in parsed_docstring.examples])
+            info.description = f'{info.description}\n\nExamples\n\n{examples}'
+
+    return info
+
+
+def parse_docstring(docstring: Optional[str], param_name: str) -> Dict[str, str]:
+    if not docstring:
+        return {}
+
+    parsed_docstring: Docstring = parse(docstring)
+
+    for param in parsed_docstring.params:
+        if param.arg_name == param_name:
+            schema: Dict[str, str] = {'description': param.description or ''}
+
+            if param.type_name:
+                schema['docstring_type'] = param.type_name
+
+            return schema
+
+    return {}

ms_enclave/utils/__init__.py

@@ -0,0 +1 @@
+from .logger import get_logger