scalebox-sdk 0.1.15__py3-none-any.whl → 0.1.17__py3-none-any.whl

scalebox/code_interpreter/code_interpreter_async.py

@@ -1,369 +1,370 @@
-import logging
-from typing import Dict, Literal, Optional, Union, overload
-
-import aiohttp
-import httpx
-from aiohttp import TCPConnector
-from httpx import AsyncClient
-
-from ..connection_config import ConnectionConfig
-from ..exceptions import InvalidArgumentException
-from ..generated import api_pb2, api_pb2_connect
-from ..sandbox_async.main import AsyncSandbox as BaseAsyncSandbox
-from .constants import DEFAULT_TEMPLATE, DEFAULT_TIMEOUT, JUPYTER_PORT
-from .exceptions import format_execution_timeout_error, format_request_timeout_error
-from .models import (
-    Context,
-    Execution,
-    ExecutionError,
-    OutputHandler,
-    OutputMessage,
-    Result,
-    aextract_exception,
-    parse_output,
-)
-
-logger = logging.getLogger(__name__)
-
-
-class AsyncSandbox(BaseAsyncSandbox):
-    """
-    Scalebox cloud sandbox is a secure and isolated cloud environment.
-
-    The sandbox allows you to:
-    - Access Linux OS
-    - Create, list, and delete files and directories
-    - Run commands
-    - Run isolated code
-    - Access the internet
-
-    Check docs for more information.
-
-    Use the `AsyncSandbox.create()` to create a new sandbox.
-
-    Example:
-    ```python
-    from scalebox.code_interpreter import AsyncSandbox
-    sandbox = await AsyncSandbox.create()
-    ```
-    """
-
-    default_template = DEFAULT_TEMPLATE
-
-    @property
-    def _jupyter_url(self) -> str:
-        return f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(JUPYTER_PORT)}"
-
-    @property
-    def _client(self) -> AsyncClient:
-        return AsyncClient(transport=self._transport)
-
-    @overload
-    async def run_code(
-        self,
-        code: str,
-        language: Union[Literal["python"], None] = None,
-        on_stdout: Optional[OutputHandler[OutputMessage]] = None,
-        on_stderr: Optional[OutputHandler[OutputMessage]] = None,
-        on_result: Optional[OutputHandler[Result]] = None,
-        on_error: Optional[OutputHandler[ExecutionError]] = None,
-        envs: Optional[Dict[str, str]] = None,
-        timeout: Optional[float] = None,
-        request_timeout: Optional[float] = None,
-    ) -> Execution:
-        """
-        Runs the code as Python.
-
-        Specify the `language` or `context` option to run the code as a different language or in a different `Context`.
-
-        You can reference previously defined variables, imports, and functions in the code.
-
-        :param code: Code to execute
-        :param language: Language to use for code execution. If not defined, the default Python context is used.
-        :param on_stdout: Callback for stdout messages
-        :param on_stderr: Callback for stderr messages
-        :param on_result: Callback for the `Result` object
-        :param on_error: Callback for the `ExecutionError` object
-        :param envs: Custom environment variables
-        :param timeout: Timeout for the code execution in **seconds**
-        :param request_timeout: Timeout for the request in **seconds**
-
-        :return: `Execution` result object
-        """
-        ...
-
-    @overload
-    async def run_code(
-        self,
-        code: str,
-        language: Optional[str] = "python",
-        on_stdout: Optional[OutputHandler[OutputMessage]] = None,
-        on_stderr: Optional[OutputHandler[OutputMessage]] = None,
-        on_result: Optional[OutputHandler[Result]] = None,
-        on_error: Optional[OutputHandler[ExecutionError]] = None,
-        envs: Optional[Dict[str, str]] = None,
-        timeout: Optional[float] = None,
-        request_timeout: Optional[float] = None,
-    ) -> Execution:
-        """
-        Runs the code for the specified language.
-
-        Specify the `language` or `context` option to run the code as a different language or in a different `Context`.
-        If no language is specified, Python is used.
-
-        You can reference previously defined variables, imports, and functions in the code.
-
-        :param code: Code to execute
-        :param language: Language to use for code execution. If not defined, the default Python context is used.
-        :param on_stdout: Callback for stdout messages
-        :param on_stderr: Callback for stderr messages
-        :param on_result: Callback for the `Result` object
-        :param on_error: Callback for the `ExecutionError` object
-        :param envs: Custom environment variables
-        :param timeout: Timeout for the code execution in **seconds**
-        :param request_timeout: Timeout for the request in **seconds**
-
-        :return: `Execution` result object
-        """
-        ...
-
-    @overload
-    async def run_code(
-        self,
-        code: str,
-        context: Optional[Context] = None,
-        on_stdout: Optional[OutputHandler[OutputMessage]] = None,
-        on_stderr: Optional[OutputHandler[OutputMessage]] = None,
-        on_result: Optional[OutputHandler[Result]] = None,
-        on_error: Optional[OutputHandler[ExecutionError]] = None,
-        envs: Optional[Dict[str, str]] = None,
-        timeout: Optional[float] = None,
-        request_timeout: Optional[float] = None,
-    ) -> Execution:
-        """
-        Runs the code in the specified context, if not specified, the default context is used.
-
-        Specify the `language` or `context` option to run the code as a different language or in a different `Context`.
-
-        You can reference previously defined variables, imports, and functions in the code.
-
-        :param code: Code to execute
-        :param context: Concrete context to run the code in. If not specified, the default context for the language is used. It's mutually exclusive with the language.
-        :param on_stdout: Callback for stdout messages
-        :param on_stderr: Callback for stderr messages
-        :param on_result: Callback for the `Result` object
-        :param on_error: Callback for the `ExecutionError` object
-        :param envs: Custom environment variables
-        :param timeout: Timeout for the code execution in **seconds**
-        :param request_timeout: Timeout for the request in **seconds**
-
-        :return: `Execution` result object
-        """
-        ...
-
-    async def run_code(
-        self,
-        code: str,
-        language: Optional[str] = None,
-        context: Optional[Context] = None,
-        on_stdout: Optional[OutputHandler[OutputMessage]] = None,
-        on_stderr: Optional[OutputHandler[OutputMessage]] = None,
-        on_result: Optional[OutputHandler[Result]] = None,
-        on_error: Optional[OutputHandler[ExecutionError]] = None,
-        envs: Optional[Dict[str, str]] = None,
-        timeout: Optional[float] = None,
-        request_timeout: Optional[float] = None,
-    ) -> Execution:
-        logger.debug(f"Executing code {code}")
-
-        """
-            Execute code in the sandbox and return the execution result.
-            """
-        logger.debug(f"Executing code: {code}")
-
-        if context and language:
-            raise Exception(
-                "You can provide context or language, but not both at the same time."
-            )
-
-        # timeout = None if timeout == 0 else (timeout or 30)
-        request_timeout = request_timeout or self._connection_config.request_timeout
-        context_id = context.id if context else None
-
-        # Ensure session exists
-        if not self._session or self._session.closed:
-            connector = aiohttp.TCPConnector(
-                limit=100,
-                limit_per_host=20,
-                keepalive_timeout=30,
-                enable_cleanup_closed=True,
-                ssl=False,
-            )
-            self._session = aiohttp.ClientSession(
-                connector=connector,
-                timeout=aiohttp.ClientTimeout(total=request_timeout),
-            )
-
-        # Set headers
-        # headers = {
-        #     "Authorization": "Bearer root",
-        # }
-
-        try:
-            # Create client and execute request
-            client = api_pb2_connect.AsyncExecutionServiceClient(
-                http_client=self._session,
-                base_url=self.envd_api_url,
-            )
-
-            # Build request
-            request = api_pb2.ExecuteRequest(
-                code=code,
-                language=language or "",
-                context_id=context_id or "",
-                env_vars=envs or {},
-            )
-
-            # Execute request and get response stream
-            responses = client.execute(
-                req=request,
-                extra_headers=self.connection_config.headers,
-                timeout_seconds=request_timeout,
-            )
-
-            execution = Execution()
-
-            # Process response stream
-            async for response in responses:
-                parse_output(
-                    execution,
-                    response,
-                    on_stdout=on_stdout,
-                    on_stderr=on_stderr,
-                    on_result=on_result,
-                    on_error=on_error,
-                )
-
-            return execution
-
-        except Exception as e:
-            # Handle exception
-            logger.error(f"Error executing code: {e}")
-            raise
-        finally:
-            # Don't close session here, let context manager handle it
-            pass
-
-    async def create_code_context(
-        self,
-        cwd: Optional[str] = None,
-        language: Optional[str] = None,
-        request_timeout: Optional[float] = None,
-    ) -> Context:
-        """
-        Creates a new context to run code in.
-
-        :param cwd: Set the current working directory for the context, defaults to `/home/user`
-        :param language: Language of the context. If not specified, defaults to Python
-        :param request_timeout: Timeout for the request in **milliseconds**
-
-        :return: Context object
-        """
-        logger.debug(f"Creating new {language} context")
-
-        request_timeout = request_timeout or self._connection_config.request_timeout
-
-        # Ensure session exists
-        if not self._session or self._session.closed:
-            connector = aiohttp.TCPConnector(
-                limit=100,
-                limit_per_host=20,
-                keepalive_timeout=30,
-                enable_cleanup_closed=True,
-                ssl=False,
-            )
-            self._session = aiohttp.ClientSession(
-                connector=connector,
-                timeout=aiohttp.ClientTimeout(total=request_timeout),
-            )
-
-        data = {}
-        if language:
-            data["language"] = language
-        if cwd:
-            data["cwd"] = cwd
-
-        try:
-            client = api_pb2_connect.AsyncContextServiceClient(
-                http_client=self._session,
-                base_url=self.envd_api_url,
-            )
-            # headers = {
-            #     "Authorization": "Bearer root",
-            # }
-
-            # Build request
-            request = api_pb2.CreateContextRequest(
-                language=language or "",
-                cwd=cwd or "",
-            )
-
-            # Execute request and get response stream
-            response = await client.create_context(
-                req=request,
-                extra_headers=self.connection_config.headers,
-            )
-            return Context.from_json(
-                {
-                    "id": response.id,
-                    "language": response.language,
-                    "cwd": response.cwd,
-                }
-            )
-        except Exception as e:
-            logger.error(f"Error create_code_context: {e}")
-            raise e
-        finally:
-            # Don't close session here, let context manager handle it
-            pass
-
-    async def destroy_context(self, context: Context) -> None:
-        """
-        Destroys a context.
-
-        :param context: Context to destroy
-        """
-        logger.debug(f"Destroying context {context.id}")
-
-        request_timeout = self._connection_config.request_timeout
-
-        # Ensure session exists
-        if not self._session or self._session.closed:
-            connector = aiohttp.TCPConnector(
-                limit=100,
-                limit_per_host=20,
-                keepalive_timeout=30,
-                enable_cleanup_closed=True,
-                ssl=False,
-            )
-            self._session = aiohttp.ClientSession(
-                connector=connector,
-                timeout=aiohttp.ClientTimeout(total=request_timeout),
-            )
-
-        # Create destroy context request
-        destroy_context_request = api_pb2.DestroyContextRequest(
-            context_id=context.id,
-        )
-
-        try:
-            client = api_pb2_connect.AsyncContextServiceClient(
-                base_url=self.envd_api_url,
-                http_client=self._session,
-            )
-            # headers = {
-            #     "Authorization": "Bearer root",
-            # }
-            await client.destroy_context(destroy_context_request,extra_headers=self.connection_config.headers)
-        except Exception as e:
-            logger.warning(f"Failed to destroy context {context.id}: {e}")
+import logging
+from typing import Dict, Literal, Optional, Union, overload
+
+import aiohttp
+import httpx
+from aiohttp import TCPConnector
+from httpx import AsyncClient
+
+from ..connection_config import ConnectionConfig
+from ..exceptions import InvalidArgumentException
+from ..generated import api_pb2, api_pb2_connect
+from ..sandbox_async.main import AsyncSandbox as BaseAsyncSandbox
+from .constants import DEFAULT_TEMPLATE, DEFAULT_TIMEOUT, JUPYTER_PORT
+from .exceptions import format_execution_timeout_error, format_request_timeout_error
+from .models import (
+    Context,
+    Execution,
+    ExecutionError,
+    OutputHandler,
+    OutputMessage,
+    Result,
+    aextract_exception,
+    parse_output,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncSandbox(BaseAsyncSandbox):
+    """
+    Scalebox cloud sandbox is a secure and isolated cloud environment.
+
+    The sandbox allows you to:
+    - Access Linux OS
+    - Create, list, and delete files and directories
+    - Run commands
+    - Run isolated code
+    - Access the internet
+
+    Check docs for more information.
+
+    Use the `AsyncSandbox.create()` to create a new sandbox.
+
+    Example:
+    ```python
+    from scalebox.code_interpreter import AsyncSandbox
+    sandbox = await AsyncSandbox.create()
+    ```
+    """
+
+    default_template = DEFAULT_TEMPLATE
+
+    @property
+    def _jupyter_url(self) -> str:
+        return f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(JUPYTER_PORT)}"
+
+    @property
+    def _client(self) -> AsyncClient:
+        return AsyncClient(transport=self._transport)
+
+    @overload
+    async def run_code(
+        self,
+        code: str,
+        language: Union[Literal["python"], None] = None,
+        on_stdout: Optional[OutputHandler[OutputMessage]] = None,
+        on_stderr: Optional[OutputHandler[OutputMessage]] = None,
+        on_result: Optional[OutputHandler[Result]] = None,
+        on_error: Optional[OutputHandler[ExecutionError]] = None,
+        envs: Optional[Dict[str, str]] = None,
+        timeout: Optional[float] = None,
+        request_timeout: Optional[float] = None,
+    ) -> Execution:
+        """
+        Runs the code as Python.
+
+        Specify the `language` or `context` option to run the code as a different language or in a different `Context`.
+
+        You can reference previously defined variables, imports, and functions in the code.
+
+        :param code: Code to execute
+        :param language: Language to use for code execution. If not defined, the default Python context is used.
+        :param on_stdout: Callback for stdout messages
+        :param on_stderr: Callback for stderr messages
+        :param on_result: Callback for the `Result` object
+        :param on_error: Callback for the `ExecutionError` object
+        :param envs: Custom environment variables
+        :param timeout: Timeout for the code execution in **seconds**
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: `Execution` result object
+        """
+        ...
+
+    @overload
+    async def run_code(
+        self,
+        code: str,
+        language: Optional[str] = "python",
+        on_stdout: Optional[OutputHandler[OutputMessage]] = None,
+        on_stderr: Optional[OutputHandler[OutputMessage]] = None,
+        on_result: Optional[OutputHandler[Result]] = None,
+        on_error: Optional[OutputHandler[ExecutionError]] = None,
+        envs: Optional[Dict[str, str]] = None,
+        timeout: Optional[float] = None,
+        request_timeout: Optional[float] = None,
+    ) -> Execution:
+        """
+        Runs the code for the specified language.
+
+        Specify the `language` or `context` option to run the code as a different language or in a different `Context`.
+        If no language is specified, Python is used.
+
+        You can reference previously defined variables, imports, and functions in the code.
+
+        :param code: Code to execute
+        :param language: Language to use for code execution. If not defined, the default Python context is used.
+        :param on_stdout: Callback for stdout messages
+        :param on_stderr: Callback for stderr messages
+        :param on_result: Callback for the `Result` object
+        :param on_error: Callback for the `ExecutionError` object
+        :param envs: Custom environment variables
+        :param timeout: Timeout for the code execution in **seconds**
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: `Execution` result object
+        """
+        ...
+
+    @overload
+    async def run_code(
+        self,
+        code: str,
+        context: Optional[Context] = None,
+        on_stdout: Optional[OutputHandler[OutputMessage]] = None,
+        on_stderr: Optional[OutputHandler[OutputMessage]] = None,
+        on_result: Optional[OutputHandler[Result]] = None,
+        on_error: Optional[OutputHandler[ExecutionError]] = None,
+        envs: Optional[Dict[str, str]] = None,
+        timeout: Optional[float] = None,
+        request_timeout: Optional[float] = None,
+    ) -> Execution:
+        """
+        Runs the code in the specified context, if not specified, the default context is used.
+
+        Specify the `language` or `context` option to run the code as a different language or in a different `Context`.
+
+        You can reference previously defined variables, imports, and functions in the code.
+
+        :param code: Code to execute
+        :param context: Concrete context to run the code in. If not specified, the default context for the language is used. It's mutually exclusive with the language.
+        :param on_stdout: Callback for stdout messages
+        :param on_stderr: Callback for stderr messages
+        :param on_result: Callback for the `Result` object
+        :param on_error: Callback for the `ExecutionError` object
+        :param envs: Custom environment variables
+        :param timeout: Timeout for the code execution in **seconds**
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: `Execution` result object
+        """
+        ...
+
+    async def run_code(
+        self,
+        code: str,
+        language: Optional[str] = None,
+        context: Optional[Context] = None,
+        on_stdout: Optional[OutputHandler[OutputMessage]] = None,
+        on_stderr: Optional[OutputHandler[OutputMessage]] = None,
+        on_result: Optional[OutputHandler[Result]] = None,
+        on_error: Optional[OutputHandler[ExecutionError]] = None,
+        envs: Optional[Dict[str, str]] = None,
+        timeout: Optional[float] = None,
+        request_timeout: Optional[float] = None,
+    ) -> Execution:
+        logger.debug(f"Executing code {code}")
+
+        """
+            Execute code in the sandbox and return the execution result.
+            """
+        logger.debug(f"Executing code: {code}")
+
+        if context and language:
+            raise Exception(
+                "You can provide context or language, but not both at the same time."
+            )
+        if language is None and context is None:
+            language = "python"
+        # timeout = None if timeout == 0 else (timeout or 30)
+        request_timeout = request_timeout or self._connection_config.request_timeout
+        context_id = context.id if context else None
+
+        # Ensure session exists
+        if not self._session or self._session.closed:
+            connector = aiohttp.TCPConnector(
+                limit=100,
+                limit_per_host=20,
+                keepalive_timeout=30,
+                enable_cleanup_closed=True,
+                ssl=False,
+            )
+            self._session = aiohttp.ClientSession(
+                connector=connector,
+                timeout=aiohttp.ClientTimeout(total=request_timeout),
+            )
+
+        # Set headers
+        # headers = {
+        #     "Authorization": "Bearer root",
+        # }
+
+        try:
+            # Create client and execute request
+            client = api_pb2_connect.AsyncExecutionServiceClient(
+                http_client=self._session,
+                base_url=self.envd_api_url,
+            )
+
+            # Build request
+            request = api_pb2.ExecuteRequest(
+                code=code,
+                language=language or "",
+                context_id=context_id or "",
+                env_vars=envs or {},
+            )
+
+            # Execute request and get response stream
+            responses = client.execute(
+                req=request,
+                extra_headers=self.connection_config.headers,
+                timeout_seconds=request_timeout,
+            )
+
+            execution = Execution()
+
+            # Process response stream
+            async for response in responses:
+                parse_output(
+                    execution,
+                    response,
+                    on_stdout=on_stdout,
+                    on_stderr=on_stderr,
+                    on_result=on_result,
+                    on_error=on_error,
+                )
+
+            return execution
+
+        except Exception as e:
+            # Handle exception
+            logger.error(f"Error executing code: {e}")
+            raise
+        finally:
+            # Don't close session here, let context manager handle it
+            pass
+
+    async def create_code_context(
+        self,
+        cwd: Optional[str] = None,
+        language: Optional[str] = None,
+        request_timeout: Optional[float] = None,
+    ) -> Context:
+        """
+        Creates a new context to run code in.
+
+        :param cwd: Set the current working directory for the context, defaults to `/home/user`
+        :param language: Language of the context. If not specified, defaults to Python
+        :param request_timeout: Timeout for the request in **milliseconds**
+
+        :return: Context object
+        """
+        logger.debug(f"Creating new {language} context")
+
+        request_timeout = request_timeout or self._connection_config.request_timeout
+
+        # Ensure session exists
+        if not self._session or self._session.closed:
+            connector = aiohttp.TCPConnector(
+                limit=100,
+                limit_per_host=20,
+                keepalive_timeout=30,
+                enable_cleanup_closed=True,
+                ssl=False,
+            )
+            self._session = aiohttp.ClientSession(
+                connector=connector,
+                timeout=aiohttp.ClientTimeout(total=request_timeout),
+            )
+
+        data = {}
+        if language:
+            data["language"] = language
+        if cwd:
+            data["cwd"] = cwd
+
+        try:
+            client = api_pb2_connect.AsyncContextServiceClient(
+                http_client=self._session,
+                base_url=self.envd_api_url,
+            )
+            # headers = {
+            #     "Authorization": "Bearer root",
+            # }
+
+            # Build request
+            request = api_pb2.CreateContextRequest(
+                language=language or "",
+                cwd=cwd or "",
+            )
+
+            # Execute request and get response stream
+            response = await client.create_context(
+                req=request,
+                extra_headers=self.connection_config.headers,
+            )
+            return Context.from_json(
+                {
+                    "id": response.id,
+                    "language": response.language,
+                    "cwd": response.cwd,
+                }
+            )
+        except Exception as e:
+            logger.error(f"Error create_code_context: {e}")
+            raise e
+        finally:
+            # Don't close session here, let context manager handle it
+            pass
+
+    async def destroy_context(self, context: Context) -> None:
+        """
+        Destroys a context.
+
+        :param context: Context to destroy
+        """
+        logger.debug(f"Destroying context {context.id}")
+
+        request_timeout = self._connection_config.request_timeout
+
+        # Ensure session exists
+        if not self._session or self._session.closed:
+            connector = aiohttp.TCPConnector(
+                limit=100,
+                limit_per_host=20,
+                keepalive_timeout=30,
+                enable_cleanup_closed=True,
+                ssl=False,
+            )
+            self._session = aiohttp.ClientSession(
+                connector=connector,
+                timeout=aiohttp.ClientTimeout(total=request_timeout),
+            )
+
+        # Create destroy context request
+        destroy_context_request = api_pb2.DestroyContextRequest(
+            context_id=context.id,
+        )
+
+        try:
+            client = api_pb2_connect.AsyncContextServiceClient(
+                base_url=self.envd_api_url,
+                http_client=self._session,
+            )
+            # headers = {
+            #     "Authorization": "Bearer root",
+            # }
+            await client.destroy_context(destroy_context_request,extra_headers=self.connection_config.headers)
+        except Exception as e:
+            logger.warning(f"Failed to destroy context {context.id}: {e}")