scalebox-sdk 0.1.25__py3-none-any.whl → 1.0.1__py3-none-any.whl

scalebox/code_interpreter/models.py

@@ -1,485 +1,485 @@
-import asyncio
-import json
-import logging
-from dataclasses import dataclass
-from datetime import datetime
-from typing import (
-    Any,
-    Awaitable,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    Optional,
-    TypeVar,
-    Union,
-)
-
-from google.protobuf.timestamp_pb2 import Timestamp
-from httpx import Response
-
-from ..exceptions import NotFoundException, SandboxException, TimeoutException
-from ..generated import api_pb2
-
-T = TypeVar("T")
-OutputHandler = Union[
-    Callable[[T], Any],
-    Callable[[T], Awaitable[Any]],
-]
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class OutputMessage:
-    """
-    Represents an output message from the sandbox code execution.
-    """
-
-    content: str
-    """
-    The output content.
-    """
-    timestamp: int
-    """
-    Unix epoch in nanoseconds
-    """
-    error: bool = False
-    """
-    Whether the output is an error.
-    """
-
-    def __str__(self):
-        return self.content
-
-
-@dataclass
-class ExecutionError:
-    """
-    Represents an error that occurred during the execution of a cell.
-    """
-
-    name: str
-    """
-    Name of the error.
-    """
-    value: str
-    """
-    Value of the error.
-    """
-    traceback: str
-    """
-    The raw traceback of the error.
-    """
-
-    def __init__(self, name: str, value: str, traceback: str, **kwargs):
-        self.name = name
-        self.value = value
-        self.traceback = traceback
-
-    def to_json(self) -> str:
-        """
-        Returns the JSON representation of the Error object.
-        """
-        data = {"name": self.name, "value": self.value, "traceback": self.traceback}
-        return json.dumps(data)
-
-
-class MIMEType(str):
-    """
-    Represents a MIME type.
-    """
-
-
-@dataclass
-class Result:
-    """
-    Represents the result of code execution.
-    """
-
-    exit_code: int = 0
-    """Process exit code."""
-    started_at: Optional[datetime] = None
-    """Start time."""
-    finished_at: Optional[datetime] = None
-    """End time."""
-    text: Optional[str] = None
-    """Text representation."""
-    html: Optional[str] = None
-    """HTML representation."""
-    markdown: Optional[str] = None
-    """Markdown representation."""
-    svg: Optional[str] = None
-    """SVG representation."""
-    png: Optional[str] = None
-    """PNG representation."""
-    jpeg: Optional[str] = None
-    """JPEG representation."""
-    pdf: Optional[str] = None
-    """PDF representation."""
-    latex: Optional[str] = None
-    """LaTeX representation."""
-    json_data: Optional[dict] = None
-    """JSON data."""
-    javascript: Optional[str] = None
-    """JavaScript representation."""
-    data: Optional[dict] = None
-    """Raw data."""
-    chart: Optional[api_pb2.Chart] = None
-    """Chart data."""
-    execution_count: Optional[int] = None
-    """Execution count."""
-    is_main_result: bool = False
-    """Whether this is the main result."""
-    extra: Optional[dict] = None
-    """Extra data."""
-
-    def __init__(
-        self,
-        exit_code: int = 0,
-        started_at: Optional[Timestamp] = None,
-        finished_at: Optional[Timestamp] = None,
-        text: Optional[str] = None,
-        html: Optional[str] = None,
-        markdown: Optional[str] = None,
-        svg: Optional[str] = None,
-        png: Optional[str] = None,
-        jpeg: Optional[str] = None,
-        pdf: Optional[str] = None,
-        latex: Optional[str] = None,
-        json: Optional[str] = None,
-        javascript: Optional[str] = None,
-        data: Optional[str] = None,
-        chart: Optional[api_pb2.Chart] = None,
-        execution_count: Optional[int] = None,
-        is_main_result: bool = False,
-        extra: Optional[Dict[str, str]] = None,
-        **kwargs,
-    ):
-        self.exit_code = exit_code
-        self.started_at = self._convert_timestamp(started_at) if started_at else None
-        self.finished_at = self._convert_timestamp(finished_at) if finished_at else None
-        self.text = text
-        self.html = html
-        self.markdown = markdown
-        self.svg = svg
-        self.png = png
-        self.jpeg = jpeg
-        self.pdf = pdf
-        self.latex = latex
-        self.json_data = json.loads(json) if json else None
-        self.javascript = javascript
-        self.data = json.loads(data) if data else None
-        self.chart = chart
-        self.execution_count = execution_count
-        self.is_main_result = is_main_result
-        self.extra = dict(extra) if extra else {}
-
-    def _convert_timestamp(self, timestamp: Timestamp) -> datetime:
-        """Convert protobuf timestamp to datetime."""
-        return datetime.fromtimestamp(timestamp.seconds + timestamp.nanos / 1e9)
-
-    def formats(self) -> Iterable[str]:
-        """
-        Returns all available formats of the result.
-
-        :return: All available formats of the result.
-        """
-        formats = []
-        if self.text:
-            formats.append("text")
-        if self.html:
-            formats.append("html")
-        if self.markdown:
-            formats.append("markdown")
-        if self.svg:
-            formats.append("svg")
-        if self.png:
-            formats.append("png")
-        if self.jpeg:
-            formats.append("jpeg")
-        if self.pdf:
-            formats.append("pdf")
-        if self.latex:
-            formats.append("latex")
-        if self.json_data:
-            formats.append("json")
-        if self.javascript:
-            formats.append("javascript")
-        if self.data:
-            formats.append("data")
-        if self.chart:
-            formats.append("chart")
-
-        if self.extra:
-            for key in self.extra:
-                formats.append(key)
-
-        return formats
-
-    def __str__(self) -> str:
-        """
-        Returns the text representation of the data.
-
-        :return: The text representation of the data.
-        """
-        if self.text:
-            return f"Result(exit_code={self.exit_code}, text={self.text})"
-        else:
-            return f"Result(exit_code={self.exit_code}, formats={list(self.formats())})"
-
-
-@dataclass
-class Logs:
-    """
-    Data printed to stdout and stderr during execution.
-    """
-
-    stdout: List[str]
-    """List of strings printed to stdout."""
-    stderr: List[str]
-    """List of strings printed to stderr."""
-
-    def __init__(self, stdout: List[str] = None, stderr: List[str] = None, **kwargs):
-        self.stdout = stdout or []
-        self.stderr = stderr or []
-
-    def __repr__(self):
-        return f"Logs(stdout={self.stdout}, stderr={self.stderr})"
-
-    def to_json(self) -> str:
-        """
-        Returns the JSON representation of the Logs object.
-        """
-        data = {"stdout": self.stdout, "stderr": self.stderr}
-        return json.dumps(data)
-
-
-@dataclass
-class Execution:
-    """
-    Represents the result of a code execution.
-    """
-
-    results: List[Result]
-    """List of execution results."""
-    logs: Logs
-    """Logs printed during execution."""
-    error: Optional[ExecutionError]
-    """Error object if an error occurred."""
-    execution_count: Optional[int]
-    """Execution count."""
-
-    def __init__(
-        self,
-        results: List[Result] = None,
-        logs: Logs = None,
-        error: Optional[ExecutionError] = None,
-        execution_count: Optional[int] = None,
-        **kwargs,
-    ):
-        self.results = results or []
-        self.logs = logs or Logs()
-        self.error = error
-        self.execution_count = execution_count
-
-    def __repr__(self):
-        return f"Execution(results={self.results}, logs={self.logs}, error={self.error}, execution_count={self.execution_count})"
-
-    @property
-    def text(self) -> Optional[str]:
-        """
-        Returns the text representation of the main result.
-
-        :return: The text representation of the main result.
-        """
-        for result in self.results:
-            if result.is_main_result:
-                return result.text
-        return None
-
-    def to_json(self) -> str:
-        """
-        Returns the JSON representation of the Execution object.
-        """
-        data = {
-            "results": [self._serialize_result(result) for result in self.results],
-            "logs": self.logs.to_json(),
-            "error": self.error.to_json() if self.error else None,
-            "execution_count": self.execution_count,
-        }
-        return json.dumps(data)
-
-    def _serialize_result(self, result: Result) -> Dict[str, Any]:
-        """Serialize a single result to JSON-serializable format."""
-        serialized = {
-            "exit_code": result.exit_code,
-            "started_at": result.started_at.isoformat() if result.started_at else None,
-            "finished_at": (
-                result.finished_at.isoformat() if result.finished_at else None
-            ),
-            "text": result.text,
-            "html": result.html,
-            "markdown": result.markdown,
-            "svg": result.svg,
-            "png": result.png,
-            "jpeg": result.jpeg,
-            "pdf": result.pdf,
-            "latex": result.latex,
-            "json": result.json_data,
-            "javascript": result.javascript,
-            "data": result.data,
-            "execution_count": result.execution_count,
-            "is_main_result": result.is_main_result,
-            "extra": result.extra,
-        }
-
-        # Remove None values
-        return {k: v for k, v in serialized.items() if v is not None}
-
-
-@dataclass
-class Context:
-    """
-    Represents a context for code execution.
-    """
-
-    id: str
-    """The ID of the context."""
-    language: str
-    """The language of the context."""
-    cwd: str
-    """The working directory of the context."""
-
-    def __init__(self, context_id: str, language: str, cwd: str, **kwargs):
-        self.id = context_id
-        self.language = language
-        self.cwd = cwd
-
-    @classmethod
-    def from_json(cls, data: Dict[str, str]):
-        return cls(
-            context_id=data.get("id"),
-            language=data.get("language"),
-            cwd=data.get("cwd"),
-        )
-
-
-def parse_output(
-    execution: Execution,
-    output: api_pb2.ExecuteResponse,
-    on_stdout: Optional[OutputHandler[OutputMessage]] = None,
-    on_stderr: Optional[OutputHandler[OutputMessage]] = None,
-    on_result: Optional[OutputHandler[Result]] = None,
-    on_error: Optional[OutputHandler[ExecutionError]] = None,
-):
-    """
-    Parse the output from the execution service and update the execution object.
-    """
-    if output.HasField("stdout"):
-        content = output.stdout.content
-        execution.logs.stdout.append(content)
-        if on_stdout:
-            message = OutputMessage(
-                content=content,
-                timestamp=int(datetime.now().timestamp() * 1e9),
-                error=False,
-            )
-            if asyncio.iscoroutinefunction(on_stdout):
-                asyncio.create_task(on_stdout(message))
-            else:
-                on_stdout(message)
-
-    elif output.HasField("stderr"):
-        content = output.stderr.content
-        execution.logs.stderr.append(content)
-        if on_stderr:
-            message = OutputMessage(
-                content=content,
-                timestamp=int(datetime.now().timestamp() * 1e9),
-                error=True,
-            )
-            if asyncio.iscoroutinefunction(on_stderr):
-                asyncio.create_task(on_stderr(message))
-            else:
-                on_stderr(message)
-
-    elif output.HasField("result"):
-        result_msg = output.result
-        result = Result(
-            exit_code=result_msg.exit_code,
-            started_at=result_msg.started_at,
-            finished_at=result_msg.finished_at,
-            text=result_msg.text,
-            html=result_msg.html,
-            markdown=result_msg.markdown,
-            svg=result_msg.svg,
-            png=result_msg.png,
-            jpeg=result_msg.jpeg,
-            pdf=result_msg.pdf,
-            latex=result_msg.latex,
-            json=result_msg.json,
-            javascript=result_msg.javascript,
-            data=result_msg.data,
-            chart=result_msg.chart,
-            execution_count=result_msg.execution_count,
-            is_main_result=result_msg.is_main_result,
-            extra=dict(result_msg.extra),
-        )
-        execution.results.append(result)
-        if on_result:
-            if asyncio.iscoroutinefunction(on_result):
-                asyncio.create_task(on_result(result))
-            else:
-                on_result(result)
-
-        # Update execution count if this is the main result
-        if result.is_main_result and result.execution_count is not None:
-            execution.execution_count = result.execution_count
-
-    elif output.HasField("error"):
-        error_msg = output.error
-        error = ExecutionError(
-            name=error_msg.name,
-            value=error_msg.value,
-            traceback=error_msg.traceback,
-        )
-        execution.error = error
-        if on_error:
-            if asyncio.iscoroutinefunction(on_error):
-                asyncio.create_task(on_error(error))
-            else:
-                on_error(error)
-
-
-async def aextract_exception(res: Response):
-    """Asynchronously extract exception from response."""
-    if res.is_success:
-        return None
-
-    await res.aread()
-    return extract_exception(res)
-
-
-def extract_exception(res: Response):
-    """Extract exception from response."""
-    if res.is_success:
-        return None
-
-    res.read()
-    return format_exception(res)
-
-
-def format_exception(res: Response):
-    """Format exception from response."""
-    if res.is_success:
-        return None
-
-    if res.status_code == 404:
-        return NotFoundException(res.text)
-    elif res.status_code == 502:
-        return TimeoutException(
-            f"{res.text}: This error is likely due to sandbox timeout. You can modify the sandbox timeout by passing 'timeout' when starting the sandbox or calling '.set_timeout' on the sandbox with the desired timeout."
-        )
-    else:
-        return SandboxException(f"{res.status_code}: {res.text}")
+import asyncio
+import json
+import logging
+from dataclasses import dataclass
+from datetime import datetime
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    TypeVar,
+    Union,
+)
+
+from google.protobuf.timestamp_pb2 import Timestamp
+from httpx import Response
+
+from ..exceptions import NotFoundException, SandboxException, TimeoutException
+from ..generated import api_pb2
+
+T = TypeVar("T")
+OutputHandler = Union[
+    Callable[[T], Any],
+    Callable[[T], Awaitable[Any]],
+]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OutputMessage:
+    """
+    Represents an output message from the sandbox code execution.
+    """
+
+    content: str
+    """
+    The output content.
+    """
+    timestamp: int
+    """
+    Unix epoch in nanoseconds
+    """
+    error: bool = False
+    """
+    Whether the output is an error.
+    """
+
+    def __str__(self):
+        return self.content
+
+
+@dataclass
+class ExecutionError:
+    """
+    Represents an error that occurred during the execution of a cell.
+    """
+
+    name: str
+    """
+    Name of the error.
+    """
+    value: str
+    """
+    Value of the error.
+    """
+    traceback: str
+    """
+    The raw traceback of the error.
+    """
+
+    def __init__(self, name: str, value: str, traceback: str, **kwargs):
+        self.name = name
+        self.value = value
+        self.traceback = traceback
+
+    def to_json(self) -> str:
+        """
+        Returns the JSON representation of the Error object.
+        """
+        data = {"name": self.name, "value": self.value, "traceback": self.traceback}
+        return json.dumps(data)
+
+
+class MIMEType(str):
+    """
+    Represents a MIME type.
+    """
+
+
+@dataclass
+class Result:
+    """
+    Represents the result of code execution.
+    """
+
+    exit_code: int = 0
+    """Process exit code."""
+    started_at: Optional[datetime] = None
+    """Start time."""
+    finished_at: Optional[datetime] = None
+    """End time."""
+    text: Optional[str] = None
+    """Text representation."""
+    html: Optional[str] = None
+    """HTML representation."""
+    markdown: Optional[str] = None
+    """Markdown representation."""
+    svg: Optional[str] = None
+    """SVG representation."""
+    png: Optional[str] = None
+    """PNG representation."""
+    jpeg: Optional[str] = None
+    """JPEG representation."""
+    pdf: Optional[str] = None
+    """PDF representation."""
+    latex: Optional[str] = None
+    """LaTeX representation."""
+    json_data: Optional[dict] = None
+    """JSON data."""
+    javascript: Optional[str] = None
+    """JavaScript representation."""
+    data: Optional[dict] = None
+    """Raw data."""
+    chart: Optional[api_pb2.Chart] = None
+    """Chart data."""
+    execution_count: Optional[int] = None
+    """Execution count."""
+    is_main_result: bool = False
+    """Whether this is the main result."""
+    extra: Optional[dict] = None
+    """Extra data."""
+
+    def __init__(
+        self,
+        exit_code: int = 0,
+        started_at: Optional[Timestamp] = None,
+        finished_at: Optional[Timestamp] = None,
+        text: Optional[str] = None,
+        html: Optional[str] = None,
+        markdown: Optional[str] = None,
+        svg: Optional[str] = None,
+        png: Optional[str] = None,
+        jpeg: Optional[str] = None,
+        pdf: Optional[str] = None,
+        latex: Optional[str] = None,
+        json: Optional[str] = None,
+        javascript: Optional[str] = None,
+        data: Optional[str] = None,
+        chart: Optional[api_pb2.Chart] = None,
+        execution_count: Optional[int] = None,
+        is_main_result: bool = False,
+        extra: Optional[Dict[str, str]] = None,
+        **kwargs,
+    ):
+        self.exit_code = exit_code
+        self.started_at = self._convert_timestamp(started_at) if started_at else None
+        self.finished_at = self._convert_timestamp(finished_at) if finished_at else None
+        self.text = text
+        self.html = html
+        self.markdown = markdown
+        self.svg = svg
+        self.png = png
+        self.jpeg = jpeg
+        self.pdf = pdf
+        self.latex = latex
+        self.json_data = json.loads(json) if json else None
+        self.javascript = javascript
+        self.data = json.loads(data) if data else None
+        self.chart = chart
+        self.execution_count = execution_count
+        self.is_main_result = is_main_result
+        self.extra = dict(extra) if extra else {}
+
+    def _convert_timestamp(self, timestamp: Timestamp) -> datetime:
+        """Convert protobuf timestamp to datetime."""
+        return datetime.fromtimestamp(timestamp.seconds + timestamp.nanos / 1e9)
+
+    def formats(self) -> Iterable[str]:
+        """
+        Returns all available formats of the result.
+
+        :return: All available formats of the result.
+        """
+        formats = []
+        if self.text:
+            formats.append("text")
+        if self.html:
+            formats.append("html")
+        if self.markdown:
+            formats.append("markdown")
+        if self.svg:
+            formats.append("svg")
+        if self.png:
+            formats.append("png")
+        if self.jpeg:
+            formats.append("jpeg")
+        if self.pdf:
+            formats.append("pdf")
+        if self.latex:
+            formats.append("latex")
+        if self.json_data:
+            formats.append("json")
+        if self.javascript:
+            formats.append("javascript")
+        if self.data:
+            formats.append("data")
+        if self.chart:
+            formats.append("chart")
+
+        if self.extra:
+            for key in self.extra:
+                formats.append(key)
+
+        return formats
+
+    def __str__(self) -> str:
+        """
+        Returns the text representation of the data.
+
+        :return: The text representation of the data.
+        """
+        if self.text:
+            return f"Result(exit_code={self.exit_code}, text={self.text})"
+        else:
+            return f"Result(exit_code={self.exit_code}, formats={list(self.formats())})"
+
+
+@dataclass
+class Logs:
+    """
+    Data printed to stdout and stderr during execution.
+    """
+
+    stdout: List[str]
+    """List of strings printed to stdout."""
+    stderr: List[str]
+    """List of strings printed to stderr."""
+
+    def __init__(self, stdout: List[str] = None, stderr: List[str] = None, **kwargs):
+        self.stdout = stdout or []
+        self.stderr = stderr or []
+
+    def __repr__(self):
+        return f"Logs(stdout={self.stdout}, stderr={self.stderr})"
+
+    def to_json(self) -> str:
+        """
+        Returns the JSON representation of the Logs object.
+        """
+        data = {"stdout": self.stdout, "stderr": self.stderr}
+        return json.dumps(data)
+
+
+@dataclass
+class Execution:
+    """
+    Represents the result of a code execution.
+    """
+
+    results: List[Result]
+    """List of execution results."""
+    logs: Logs
+    """Logs printed during execution."""
+    error: Optional[ExecutionError]
+    """Error object if an error occurred."""
+    execution_count: Optional[int]
+    """Execution count."""
+
+    def __init__(
+        self,
+        results: List[Result] = None,
+        logs: Logs = None,
+        error: Optional[ExecutionError] = None,
+        execution_count: Optional[int] = None,
+        **kwargs,
+    ):
+        self.results = results or []
+        self.logs = logs or Logs()
+        self.error = error
+        self.execution_count = execution_count
+
+    def __repr__(self):
+        return f"Execution(results={self.results}, logs={self.logs}, error={self.error}, execution_count={self.execution_count})"
+
+    @property
+    def text(self) -> Optional[str]:
+        """
+        Returns the text representation of the main result.
+
+        :return: The text representation of the main result.
+        """
+        for result in self.results:
+            if result.is_main_result:
+                return result.text
+        return None
+
+    def to_json(self) -> str:
+        """
+        Returns the JSON representation of the Execution object.
+        """
+        data = {
+            "results": [self._serialize_result(result) for result in self.results],
+            "logs": self.logs.to_json(),
+            "error": self.error.to_json() if self.error else None,
+            "execution_count": self.execution_count,
+        }
+        return json.dumps(data)
+
+    def _serialize_result(self, result: Result) -> Dict[str, Any]:
+        """Serialize a single result to JSON-serializable format."""
+        serialized = {
+            "exit_code": result.exit_code,
+            "started_at": result.started_at.isoformat() if result.started_at else None,
+            "finished_at": (
+                result.finished_at.isoformat() if result.finished_at else None
+            ),
+            "text": result.text,
+            "html": result.html,
+            "markdown": result.markdown,
+            "svg": result.svg,
+            "png": result.png,
+            "jpeg": result.jpeg,
+            "pdf": result.pdf,
+            "latex": result.latex,
+            "json": result.json_data,
+            "javascript": result.javascript,
+            "data": result.data,
+            "execution_count": result.execution_count,
+            "is_main_result": result.is_main_result,
+            "extra": result.extra,
+        }
+
+        # Remove None values
+        return {k: v for k, v in serialized.items() if v is not None}
+
+
+@dataclass
+class Context:
+    """
+    Represents a context for code execution.
+    """
+
+    id: str
+    """The ID of the context."""
+    language: str
+    """The language of the context."""
+    cwd: str
+    """The working directory of the context."""
+
+    def __init__(self, context_id: str, language: str, cwd: str, **kwargs):
+        self.id = context_id
+        self.language = language
+        self.cwd = cwd
+
+    @classmethod
+    def from_json(cls, data: Dict[str, str]):
+        return cls(
+            context_id=data.get("id"),
+            language=data.get("language"),
+            cwd=data.get("cwd"),
+        )
+
+
+def parse_output(
+    execution: Execution,
+    output: api_pb2.ExecuteResponse,
+    on_stdout: Optional[OutputHandler[OutputMessage]] = None,
+    on_stderr: Optional[OutputHandler[OutputMessage]] = None,
+    on_result: Optional[OutputHandler[Result]] = None,
+    on_error: Optional[OutputHandler[ExecutionError]] = None,
+):
+    """
+    Parse the output from the execution service and update the execution object.
+    """
+    if output.HasField("stdout"):
+        content = output.stdout.content
+        execution.logs.stdout.append(content)
+        if on_stdout:
+            message = OutputMessage(
+                content=content,
+                timestamp=int(datetime.now().timestamp() * 1e9),
+                error=False,
+            )
+            if asyncio.iscoroutinefunction(on_stdout):
+                asyncio.create_task(on_stdout(message))
+            else:
+                on_stdout(message)
+
+    elif output.HasField("stderr"):
+        content = output.stderr.content
+        execution.logs.stderr.append(content)
+        if on_stderr:
+            message = OutputMessage(
+                content=content,
+                timestamp=int(datetime.now().timestamp() * 1e9),
+                error=True,
+            )
+            if asyncio.iscoroutinefunction(on_stderr):
+                asyncio.create_task(on_stderr(message))
+            else:
+                on_stderr(message)
+
+    elif output.HasField("result"):
+        result_msg = output.result
+        result = Result(
+            exit_code=result_msg.exit_code,
+            started_at=result_msg.started_at,
+            finished_at=result_msg.finished_at,
+            text=result_msg.text,
+            html=result_msg.html,
+            markdown=result_msg.markdown,
+            svg=result_msg.svg,
+            png=result_msg.png,
+            jpeg=result_msg.jpeg,
+            pdf=result_msg.pdf,
+            latex=result_msg.latex,
+            json=result_msg.json,
+            javascript=result_msg.javascript,
+            data=result_msg.data,
+            chart=result_msg.chart,
+            execution_count=result_msg.execution_count,
+            is_main_result=result_msg.is_main_result,
+            extra=dict(result_msg.extra),
+        )
+        execution.results.append(result)
+        if on_result:
+            if asyncio.iscoroutinefunction(on_result):
+                asyncio.create_task(on_result(result))
+            else:
+                on_result(result)
+
+        # Update execution count if this is the main result
+        if result.is_main_result and result.execution_count is not None:
+            execution.execution_count = result.execution_count
+
+    elif output.HasField("error"):
+        error_msg = output.error
+        error = ExecutionError(
+            name=error_msg.name,
+            value=error_msg.value,
+            traceback=error_msg.traceback,
+        )
+        execution.error = error
+        if on_error:
+            if asyncio.iscoroutinefunction(on_error):
+                asyncio.create_task(on_error(error))
+            else:
+                on_error(error)
+
+
+async def aextract_exception(res: Response):
+    """Asynchronously extract exception from response."""
+    if res.is_success:
+        return None
+
+    await res.aread()
+    return extract_exception(res)
+
+
+def extract_exception(res: Response):
+    """Extract exception from response."""
+    if res.is_success:
+        return None
+
+    res.read()
+    return format_exception(res)
+
+
+def format_exception(res: Response):
+    """Format exception from response."""
+    if res.is_success:
+        return None
+
+    if res.status_code == 404:
+        return NotFoundException(res.text)
+    elif res.status_code == 502:
+        return TimeoutException(
+            f"{res.text}: This error is likely due to sandbox timeout. You can modify the sandbox timeout by passing 'timeout' when starting the sandbox or calling '.set_timeout' on the sandbox with the desired timeout."
+        )
+    else:
+        return SandboxException(f"{res.status_code}: {res.text}")