agentrun-inner-test 0.0.46__py3-none-any.whl

agentrun/sandbox/code_interpreter_sandbox.py

@@ -0,0 +1,823 @@
+"""
+This file is auto generated by the code generation script.
+Do not modify this file manually.
+Use the `make codegen` command to regenerate.
+
+当前文件为自动生成的控制 API 客户端代码。请勿手动修改此文件。
+使用 `make codegen` 命令重新生成。
+
+source: agentrun/sandbox/__code_interpreter_sandbox_async_template.py
+
+代码解释器沙箱高层API模板 / Code Interpreter Sandbox High-Level API Template
+
+此模板用于生成代码解释器沙箱资源的高级API代码。
+This template is used to generate high-level API code for code interpreter sandbox resources.
+"""
+
+import asyncio
+import time
+from typing import Optional
+
+from agentrun.sandbox.api.code_interpreter_data import CodeInterpreterDataAPI
+from agentrun.sandbox.model import CodeLanguage, TemplateType
+from agentrun.utils.exception import ServerError
+from agentrun.utils.log import logger
+
+from .sandbox import Sandbox
+
+
+class FileOperations:
+    """File upload/download operations."""
+
+    def __init__(self, sandbox: "CodeInterpreterSandbox"):
+        self._sandbox = sandbox
+
+    async def read_async(self, path: str):
+        """Read a file from the code interpreter (async).
+        Args:
+            path: Remote file path in the code interpreter
+        Returns:
+            File content
+        """
+        return await self._sandbox.data_api.read_file_async(path=path)
+
+    def read(self, path: str):
+        """Read a file from the code interpreter (async).
+        Args:
+            path: Remote file path in the code interpreter
+        Returns:
+            File content
+        """
+        return self._sandbox.data_api.read_file(path=path)
+
+    async def write_async(
+        self,
+        path: str,
+        content: str,
+        mode: str = "644",
+        encoding: str = "utf-8",
+        create_dir=True,
+    ):
+        """Write a file to the code interpreter (async).
+        Args:
+            path: Remote file path in the code interpreter
+            content: File content
+        """
+        return await self._sandbox.data_api.write_file_async(
+            path=path,
+            content=content,
+            mode=mode,
+            encoding=encoding,
+            create_dir=create_dir,
+        )
+
+    def write(
+        self,
+        path: str,
+        content: str,
+        mode: str = "644",
+        encoding: str = "utf-8",
+        create_dir=True,
+    ):
+        """Write a file to the code interpreter (async).
+        Args:
+            path: Remote file path in the code interpreter
+            content: File content
+        """
+        return self._sandbox.data_api.write_file(
+            path=path,
+            content=content,
+            mode=mode,
+            encoding=encoding,
+            create_dir=create_dir,
+        )
+
+
+class FileSystemOperations:
+    """File system operations (list, move, remove, stat, mkdir)."""
+
+    def __init__(self, sandbox: "CodeInterpreterSandbox"):
+        self._sandbox = sandbox
+
+    async def list_async(
+        self, path: Optional[str] = None, depth: Optional[int] = None
+    ):
+        """List directory contents (async).
+
+        Args:
+            path: Directory path (optional)
+            depth: Traversal depth (optional)
+
+        Returns:
+            Directory contents
+        """
+        return await self._sandbox.data_api.list_directory_async(
+            path=path, depth=depth
+        )
+
+    def list(self, path: Optional[str] = None, depth: Optional[int] = None):
+        """List directory contents (async).
+
+        Args:
+            path: Directory path (optional)
+            depth: Traversal depth (optional)
+
+        Returns:
+            Directory contents
+        """
+        return self._sandbox.data_api.list_directory(path=path, depth=depth)
+
+    async def move_async(self, source: str, destination: str):
+        """Move a file or directory (async).
+
+        Args:
+            source: Source file or directory path
+            destination: Target file or directory path
+
+        Returns:
+            Move operation result
+        """
+        return await self._sandbox.data_api.move_file_async(
+            source=source, destination=destination
+        )
+
+    def move(self, source: str, destination: str):
+        """Move a file or directory (async).
+
+        Args:
+            source: Source file or directory path
+            destination: Target file or directory path
+
+        Returns:
+            Move operation result
+        """
+        return self._sandbox.data_api.move_file(
+            source=source, destination=destination
+        )
+
+    async def remove_async(self, path: str):
+        """Remove a file or directory (async).
+
+        Args:
+            path: File or directory path to remove
+
+        Returns:
+            Remove operation result
+        """
+        return await self._sandbox.data_api.remove_file_async(path=path)
+
+    def remove(self, path: str):
+        """Remove a file or directory (async).
+
+        Args:
+            path: File or directory path to remove
+
+        Returns:
+            Remove operation result
+        """
+        return self._sandbox.data_api.remove_file(path=path)
+
+    async def stat_async(self, path: str):
+        """Get file or directory statistics (async).
+
+        Args:
+            path: File or directory path
+
+        Returns:
+            File/directory statistics
+        """
+        return await self._sandbox.data_api.stat_async(path=path)
+
+    def stat(self, path: str):
+        """Get file or directory statistics (async).
+
+        Args:
+            path: File or directory path
+
+        Returns:
+            File/directory statistics
+        """
+        return self._sandbox.data_api.stat(path=path)
+
+    async def mkdir_async(
+        self,
+        path: str,
+        parents: Optional[bool] = True,
+        mode: Optional[str] = "0755",
+    ):
+        """Create a directory (async).
+
+        Args:
+            path: Directory path to create
+            parents: Whether to create parent directories (default: True)
+            mode: Directory permissions mode (default: "0755")
+
+        Returns:
+            Mkdir operation result
+        """
+        return await self._sandbox.data_api.mkdir_async(
+            path=path, parents=parents, mode=mode
+        )
+
+    def mkdir(
+        self,
+        path: str,
+        parents: Optional[bool] = True,
+        mode: Optional[str] = "0755",
+    ):
+        """Create a directory (async).
+
+        Args:
+            path: Directory path to create
+            parents: Whether to create parent directories (default: True)
+            mode: Directory permissions mode (default: "0755")
+
+        Returns:
+            Mkdir operation result
+        """
+        return self._sandbox.data_api.mkdir(
+            path=path, parents=parents, mode=mode
+        )
+
+    async def upload_async(
+        self,
+        local_file_path: str,
+        target_file_path: str,
+    ):
+        """Upload a file to the code interpreter (async).
+
+        Args:
+            local_file_path: Local file path to upload
+            target_file_path: Target file path in code interpreter
+
+        Returns:
+            Upload result
+        """
+        return await self._sandbox.data_api.upload_file_async(
+            local_file_path=local_file_path, target_file_path=target_file_path
+        )
+
+    def upload(
+        self,
+        local_file_path: str,
+        target_file_path: str,
+    ):
+        """Upload a file to the code interpreter (async).
+
+        Args:
+            local_file_path: Local file path to upload
+            target_file_path: Target file path in code interpreter
+
+        Returns:
+            Upload result
+        """
+        return self._sandbox.data_api.upload_file(
+            local_file_path=local_file_path, target_file_path=target_file_path
+        )
+
+    async def download_async(self, path: str, save_path: str):
+        """Download a file from the code interpreter (async).
+
+        Args:
+            path: Remote file path in the code interpreter
+            save_path: Local file path to save the downloaded file
+
+        Returns:
+            Download result with 'saved_path' and 'size'
+        """
+        return await self._sandbox.data_api.download_file_async(
+            path=path, save_path=save_path
+        )
+
+    def download(self, path: str, save_path: str):
+        """Download a file from the code interpreter (async).
+
+        Args:
+            path: Remote file path in the code interpreter
+            save_path: Local file path to save the downloaded file
+
+        Returns:
+            Download result with 'saved_path' and 'size'
+        """
+        return self._sandbox.data_api.download_file(
+            path=path, save_path=save_path
+        )
+
+
+class ProcessOperations:
+    """Process management operations."""
+
+    def __init__(self, sandbox: "CodeInterpreterSandbox"):
+        self._sandbox = sandbox
+
+    async def cmd_async(
+        self, command: str, cwd: str, timeout: Optional[int] = 30
+    ):
+        """Execute a command in the code interpreter (async).
+
+        Args:
+            command: Command to execute
+            cwd: Working directory
+            timeout: Execution timeout in seconds (default: 30)
+
+        Returns:
+            Command execution result
+        """
+        return await self._sandbox.data_api.cmd_async(
+            command=command, cwd=cwd, timeout=timeout
+        )
+
+    def cmd(self, command: str, cwd: str, timeout: Optional[int] = 30):
+        """Execute a command in the code interpreter (async).
+
+        Args:
+            command: Command to execute
+            cwd: Working directory
+            timeout: Execution timeout in seconds (default: 30)
+
+        Returns:
+            Command execution result
+        """
+        return self._sandbox.data_api.cmd(
+            command=command, cwd=cwd, timeout=timeout
+        )
+
+    async def list_async(self):
+        """List all processes (async).
+
+        Returns:
+            List of processes
+        """
+        return await self._sandbox.data_api.list_processes_async()
+
+    def list(self):
+        """List all processes (async).
+
+        Returns:
+            List of processes
+        """
+        return self._sandbox.data_api.list_processes()
+
+    async def get_async(self, pid: str):
+        """Get a specific process by PID (async).
+
+        Args:
+            pid: Process ID
+
+        Returns:
+            Process information
+        """
+        return await self._sandbox.data_api.get_process_async(pid=pid)
+
+    def get(self, pid: str):
+        """Get a specific process by PID (async).
+
+        Args:
+            pid: Process ID
+
+        Returns:
+            Process information
+        """
+        return self._sandbox.data_api.get_process(pid=pid)
+
+    async def kill_async(self, pid: str):
+        """Kill a specific process by PID (async).
+
+        Args:
+            pid: Process ID
+
+        Returns:
+            Kill operation result
+        """
+        return await self._sandbox.data_api.kill_process_async(pid=pid)
+
+    def kill(self, pid: str):
+        """Kill a specific process by PID (async).
+
+        Args:
+            pid: Process ID
+
+        Returns:
+            Kill operation result
+        """
+        return self._sandbox.data_api.kill_process(pid=pid)
+
+
+class ContextOperations:
+    """Context management operations."""
+
+    def __init__(self, sandbox: "CodeInterpreterSandbox"):
+        self._sandbox = sandbox
+        self._context_id: Optional[str] = None
+        self._language: Optional[str] = None
+        self._cwd: Optional[str] = None
+
+    @property
+    def context_id(self) -> Optional[str]:
+        """Get the current context ID."""
+        return self._context_id
+
+    async def list_async(self):
+        """List all contexts (async)."""
+        return await self._sandbox.data_api.list_contexts_async()
+
+    def list(self):
+        """List all contexts (async)."""
+        return self._sandbox.data_api.list_contexts()
+
+    async def create_async(
+        self,
+        language: Optional[CodeLanguage] = CodeLanguage.PYTHON,
+        cwd: str = "/home/user",
+    ) -> "ContextOperations":
+        """Create a new context and save its ID (async).
+
+        Args:
+            name: Context name
+            language: Programming language (default: "python")
+            config: Context configuration (optional)
+
+        Returns:
+            ContextOperations: Returns self for chaining and context manager support
+        """
+        result = await self._sandbox.data_api.create_context_async(
+            language=language, cwd=cwd
+        )
+        if all(result.get(key) for key in ("id", "cwd", "language")):
+            self._context_id = result["id"]
+            self._language = result["language"]
+            self._cwd = result["cwd"]
+            return self
+        raise ServerError(500, "Failed to create context")
+
+    def create(
+        self,
+        language: Optional[CodeLanguage] = CodeLanguage.PYTHON,
+        cwd: str = "/home/user",
+    ) -> "ContextOperations":
+        """Create a new context and save its ID (async).
+
+        Args:
+            name: Context name
+            language: Programming language (default: "python")
+            config: Context configuration (optional)
+
+        Returns:
+            ContextOperations: Returns self for chaining and context manager support
+        """
+        result = self._sandbox.data_api.create_context(
+            language=language, cwd=cwd
+        )
+        if all(result.get(key) for key in ("id", "cwd", "language")):
+            self._context_id = result["id"]
+            self._language = result["language"]
+            self._cwd = result["cwd"]
+            return self
+        raise ServerError(500, "Failed to create context")
+
+    async def get_async(
+        self, context_id: Optional[str] = None
+    ) -> "ContextOperations":
+        """Get a specific context by ID (async).
+        Args:
+            context_id: Context ID
+        Returns:
+            ContextOperations: Returns self for chaining and context manager support
+        """
+        if context_id is None:
+            context_id = self._context_id
+        if context_id is None:
+            logger.error(f"context id is not set")
+            raise ValueError("context id is not set,")
+        result = await self._sandbox.data_api.get_context_async(
+            context_id=context_id
+        )
+        if all(result.get(key) for key in ("id", "cwd", "language")):
+            self._context_id = result["id"]
+            self._language = result["language"]
+            self._cwd = result["cwd"]
+            return self
+        raise ServerError(500, "Failed to create context")
+
+    def get(self, context_id: Optional[str] = None) -> "ContextOperations":
+        """Get a specific context by ID (async).
+        Args:
+            context_id: Context ID
+        Returns:
+            ContextOperations: Returns self for chaining and context manager support
+        """
+        if context_id is None:
+            context_id = self._context_id
+        if context_id is None:
+            logger.error(f"context id is not set")
+            raise ValueError("context id is not set,")
+        result = self._sandbox.data_api.get_context(context_id=context_id)
+        if all(result.get(key) for key in ("id", "cwd", "language")):
+            self._context_id = result["id"]
+            self._language = result["language"]
+            self._cwd = result["cwd"]
+            return self
+        raise ServerError(500, "Failed to create context")
+
+    async def execute_async(
+        self,
+        code: str,
+        language: Optional[CodeLanguage] = None,
+        context_id: Optional[str] = None,
+        timeout: Optional[int] = 30,
+    ):
+        """Execute code in a context (async).
+
+        Args:
+            code: Code to execute
+            context_id: Context ID (optional, uses saved context_id if not provided)
+            timeout: Execution timeout in seconds (default: 30)
+
+        Returns:
+            Execution result
+
+        Raises:
+            ValueError: If no context_id is provided and none is saved
+        """
+        if context_id is None:
+            context_id = self._context_id
+        if context_id is None and language is None:
+            logger.debug("context id is not set, use default language: python")
+            language = CodeLanguage.PYTHON
+        return await self._sandbox.data_api.execute_code_async(
+            context_id=context_id, language=language, code=code, timeout=timeout
+        )
+
+    def execute(
+        self,
+        code: str,
+        language: Optional[CodeLanguage] = None,
+        context_id: Optional[str] = None,
+        timeout: Optional[int] = 30,
+    ):
+        """Execute code in a context (async).
+
+        Args:
+            code: Code to execute
+            context_id: Context ID (optional, uses saved context_id if not provided)
+            timeout: Execution timeout in seconds (default: 30)
+
+        Returns:
+            Execution result
+
+        Raises:
+            ValueError: If no context_id is provided and none is saved
+        """
+        if context_id is None:
+            context_id = self._context_id
+        if context_id is None and language is None:
+            logger.debug("context id is not set, use default language: python")
+            language = CodeLanguage.PYTHON
+        return self._sandbox.data_api.execute_code(
+            context_id=context_id, language=language, code=code, timeout=timeout
+        )
+
+    async def delete_async(self, context_id: Optional[str] = None):
+        """Delete a context (async).
+
+        Args:
+            context_id: Context ID (optional, uses saved context_id if not provided)
+
+        Returns:
+            Delete result
+
+        Raises:
+            ValueError: If no context_id is provided and none is saved
+        """
+        if context_id is None:
+            context_id = self._context_id
+        if context_id is None:
+            raise ValueError(
+                "context_id is required. Either pass it as parameter or create"
+                " a context first."
+            )
+        result = await self._sandbox.data_api.delete_context_async(
+            context_id=context_id
+        )
+        # Clear the saved context_id after deletion
+        self._context_id = None
+        return result
+
+    def delete(self, context_id: Optional[str] = None):
+        """Delete a context (async).
+
+        Args:
+            context_id: Context ID (optional, uses saved context_id if not provided)
+
+        Returns:
+            Delete result
+
+        Raises:
+            ValueError: If no context_id is provided and none is saved
+        """
+        if context_id is None:
+            context_id = self._context_id
+        if context_id is None:
+            raise ValueError(
+                "context_id is required. Either pass it as parameter or create"
+                " a context first."
+            )
+        result = self._sandbox.data_api.delete_context(context_id=context_id)
+        # Clear the saved context_id after deletion
+        self._context_id = None
+        return result
+
+    async def __aenter__(self):
+        """Asynchronous context manager entry."""
+        if self._context_id is None:
+            raise ValueError(
+                "No context has been created. Call create() first or use: "
+                "async with await sandbox.context.create_async(...) as ctx:"
+            )
+        return self
+
+    def __enter__(self):
+        """Synchronous context manager entry."""
+        if self._context_id is None:
+            raise ValueError(
+                "No context has been created. Call create() first or use: "
+                "with sandbox.context.create(...) as ctx:"
+            )
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Asynchronous context manager exit - deletes the context."""
+        if self._context_id is not None:
+            try:
+                await self._sandbox.data_api.delete_context_async(
+                    self._context_id
+                )
+            except Exception as e:
+                logger.error(
+                    f"Warning: Failed to delete context {self._context_id}: {e}"
+                )
+        return False
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Synchronous context manager exit - deletes the context."""
+        if self._context_id is not None:
+            try:
+                self._sandbox.data_api.delete_context(self._context_id)
+            except Exception as e:
+                logger.error(
+                    f"Warning: Failed to delete context {self._context_id}: {e}"
+                )
+        return False
+
+
+class CodeInterpreterSandbox(Sandbox):
+    _template_type = TemplateType.CODE_INTERPRETER
+
+    _data_api: Optional["CodeInterpreterDataAPI"] = None
+    _file: Optional[FileOperations] = None
+    _file_system: Optional[FileSystemOperations] = None
+    _context: Optional[ContextOperations] = None
+    _process: Optional[ProcessOperations] = None
+
+    @property
+    def data_api(self) -> "CodeInterpreterDataAPI":
+        """Get data client."""
+        if self._data_api is None:
+            self._data_api = CodeInterpreterDataAPI(
+                sandbox_id=self.sandbox_id or "", config=self._config
+            )
+        return self._data_api
+
+    @property
+    def file(self) -> FileOperations:
+        """Access file upload/download operations."""
+        if self._file is None:
+            self._file = FileOperations(self)
+        return self._file
+
+    @property
+    def file_system(self) -> FileSystemOperations:
+        """Access file system operations."""
+        if self._file_system is None:
+            self._file_system = FileSystemOperations(self)
+        return self._file_system
+
+    @property
+    def context(self) -> ContextOperations:
+        """Access context management operations."""
+        if self._context is None:
+            self._context = ContextOperations(self)
+        return self._context
+
+    @property
+    def process(self) -> ProcessOperations:
+        """Access process management operations."""
+        if self._process is None:
+            self._process = ProcessOperations(self)
+        return self._process
+
+    async def check_health_async(self):
+        return await self.data_api.check_health_async()
+
+    def check_health(self):
+        return self.data_api.check_health()
+
+    async def __aenter__(self):
+        """Asynchronous context manager entry."""
+        # Poll health check asynchronously
+        max_retries = 60  # Maximum 60 seconds
+        retry_count = 0
+
+        logger.debug("Waiting for code interpreter to be ready...")
+
+        while retry_count < max_retries:
+            retry_count += 1
+
+            try:
+                health = await self.check_health_async()
+
+                if health["status"] == "ok":
+                    logger.debug(
+                        "Code Interpreter is ready! (took"
+                        f" {retry_count} seconds)"
+                    )
+                    return self
+
+                logger.debug(
+                    f"[{retry_count}/{max_retries}] Health status: not ready"
+                )
+
+                logger.debug(
+                    f"[{retry_count}/{max_retries}] Health status:"
+                    f" {health.get('code')} { health.get('message')}",
+                )
+
+            except Exception as e:
+                logger.error(
+                    f"[{retry_count}/{max_retries}] Health check failed: {e}"
+                )
+
+            if retry_count < max_retries:
+                await asyncio.sleep(1)
+
+        raise RuntimeError(
+            f"Health check timeout after {max_retries} seconds. "
+            "Code interpreter did not become ready in time."
+        )
+
+    def __enter__(self):
+        """Synchronous context manager entry."""
+        # Poll health check asynchronously
+        max_retries = 60  # Maximum 60 seconds
+        retry_count = 0
+
+        logger.debug("Waiting for code interpreter to be ready...")
+
+        while retry_count < max_retries:
+            retry_count += 1
+
+            try:
+                health = self.check_health()
+
+                if health["status"] == "ok":
+                    logger.debug(
+                        "Code Interpreter is ready! (took"
+                        f" {retry_count} seconds)"
+                    )
+                    return self
+
+                logger.debug(
+                    f"[{retry_count}/{max_retries}] Health status: not ready"
+                )
+
+                logger.debug(
+                    f"[{retry_count}/{max_retries}] Health status:"
+                    f" {health.get('code')} { health.get('message')}",
+                )
+
+            except Exception as e:
+                logger.error(
+                    f"[{retry_count}/{max_retries}] Health check failed: {e}"
+                )
+
+            if retry_count < max_retries:
+                time.sleep(1)
+
+        raise RuntimeError(
+            f"Health check timeout after {max_retries} seconds. "
+            "Code interpreter did not become ready in time."
+        )
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Asynchronous context manager exit."""
+        if self.sandbox_id is None:
+            raise ValueError("Sandbox ID is not set")
+        logger.debug(f"Deleting code interpreter sandbox {self.sandbox_id}...")
+        await self.delete_async()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Synchronous context manager exit."""
+        if self.sandbox_id is None:
+            raise ValueError("Sandbox ID is not set")
+        logger.debug(f"Deleting code interpreter sandbox {self.sandbox_id}...")
+        self.delete()