initial commit
This commit is contained in:
445
venv/Lib/site-packages/langsmith/sandbox/_sandbox.py
Normal file
445
venv/Lib/site-packages/langsmith/sandbox/_sandbox.py
Normal file
@@ -0,0 +1,445 @@
|
||||
"""Sandbox class for interacting with a specific sandbox instance."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union, overload
|
||||
|
||||
import httpx
|
||||
|
||||
from langsmith.sandbox._exceptions import (
|
||||
DataplaneNotConfiguredError,
|
||||
ResourceNotFoundError,
|
||||
SandboxConnectionError,
|
||||
SandboxNotReadyError,
|
||||
)
|
||||
from langsmith.sandbox._helpers import handle_sandbox_http_error
|
||||
from langsmith.sandbox._models import (
|
||||
CommandHandle,
|
||||
ExecutionResult,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langsmith.sandbox._client import SandboxClient
|
||||
|
||||
|
||||
@dataclass
|
||||
class Sandbox:
|
||||
"""Represents an active sandbox for running commands and file operations.
|
||||
|
||||
This class is typically obtained from SandboxClient.sandbox() and supports
|
||||
the context manager protocol for automatic cleanup.
|
||||
|
||||
Attributes:
|
||||
name: Display name (can be updated).
|
||||
template_name: Name of the template used to create this sandbox.
|
||||
dataplane_url: URL for data plane operations (file I/O, command execution).
|
||||
Only functional when status is "ready".
|
||||
id: Unique identifier (UUID). Remains constant even if name changes.
|
||||
May be None for resources created before ID support was added.
|
||||
status: Sandbox lifecycle status. One of "provisioning", "ready", "failed".
|
||||
status_message: Human-readable details when status is "failed", None otherwise.
|
||||
created_at: Timestamp when the sandbox was created.
|
||||
updated_at: Timestamp when the sandbox was last updated.
|
||||
|
||||
Example:
|
||||
with client.sandbox(template_name="python-sandbox") as sandbox:
|
||||
result = sandbox.run("python --version")
|
||||
print(result.stdout)
|
||||
"""
|
||||
|
||||
# Data fields (from API response)
|
||||
name: str
|
||||
template_name: str
|
||||
dataplane_url: Optional[str] = None
|
||||
id: Optional[str] = None
|
||||
status: str = "ready"
|
||||
status_message: Optional[str] = None
|
||||
created_at: Optional[str] = None
|
||||
updated_at: Optional[str] = None
|
||||
|
||||
# Internal fields (not from API)
|
||||
_client: SandboxClient = field(repr=False, default=None) # type: ignore
|
||||
_auto_delete: bool = field(repr=False, default=True)
|
||||
|
||||
@classmethod
|
||||
def from_dict(
|
||||
cls,
|
||||
data: dict[str, Any],
|
||||
client: SandboxClient,
|
||||
auto_delete: bool = True,
|
||||
) -> Sandbox:
|
||||
"""Create a Sandbox from API response dict.
|
||||
|
||||
Args:
|
||||
data: API response dictionary containing sandbox data.
|
||||
client: Parent SandboxClient for operations.
|
||||
auto_delete: Whether to delete the sandbox on context exit.
|
||||
|
||||
Returns:
|
||||
Sandbox instance.
|
||||
"""
|
||||
return cls(
|
||||
name=data.get("name", ""),
|
||||
template_name=data.get("template_name", ""),
|
||||
dataplane_url=data.get("dataplane_url"),
|
||||
id=data.get("id"),
|
||||
status=data.get("status", "ready"),
|
||||
status_message=data.get("status_message"),
|
||||
created_at=data.get("created_at"),
|
||||
updated_at=data.get("updated_at"),
|
||||
_client=client,
|
||||
_auto_delete=auto_delete,
|
||||
)
|
||||
|
||||
def __enter__(self) -> Sandbox:
|
||||
"""Enter context manager."""
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: Optional[type],
|
||||
exc_val: Optional[BaseException],
|
||||
exc_tb: Optional[Any],
|
||||
) -> None:
|
||||
"""Exit context manager, optionally deleting the sandbox."""
|
||||
if self._auto_delete:
|
||||
try:
|
||||
self._client.delete_sandbox(self.name)
|
||||
except Exception:
|
||||
# Don't raise on cleanup errors
|
||||
pass
|
||||
|
||||
def _require_dataplane_url(self) -> str:
|
||||
"""Validate and return the dataplane URL.
|
||||
|
||||
Returns:
|
||||
The dataplane URL.
|
||||
|
||||
Raises:
|
||||
SandboxNotReadyError: If sandbox status is not "ready".
|
||||
DataplaneNotConfiguredError: If dataplane_url is not configured.
|
||||
"""
|
||||
if self.status != "ready":
|
||||
raise SandboxNotReadyError(
|
||||
f"Sandbox '{self.name}' is not ready (status: {self.status}). "
|
||||
"Wait for status 'ready' before running operations."
|
||||
)
|
||||
if not self.dataplane_url:
|
||||
raise DataplaneNotConfiguredError(
|
||||
f"Sandbox '{self.name}' does not have a dataplane_url configured. "
|
||||
"Runtime operations require a dataplane URL."
|
||||
)
|
||||
return self.dataplane_url
|
||||
|
||||
@overload
|
||||
def run(
|
||||
self,
|
||||
command: str,
|
||||
*,
|
||||
timeout: int = ...,
|
||||
env: Optional[dict[str, str]] = ...,
|
||||
cwd: Optional[str] = ...,
|
||||
shell: str = ...,
|
||||
on_stdout: Optional[Callable[[str], Any]] = ...,
|
||||
on_stderr: Optional[Callable[[str], Any]] = ...,
|
||||
wait: Literal[True] = ...,
|
||||
) -> ExecutionResult: ...
|
||||
|
||||
@overload
|
||||
def run(
|
||||
self,
|
||||
command: str,
|
||||
*,
|
||||
timeout: int = ...,
|
||||
env: Optional[dict[str, str]] = ...,
|
||||
cwd: Optional[str] = ...,
|
||||
shell: str = ...,
|
||||
on_stdout: Optional[Callable[[str], Any]] = ...,
|
||||
on_stderr: Optional[Callable[[str], Any]] = ...,
|
||||
wait: Literal[False],
|
||||
) -> CommandHandle: ...
|
||||
|
||||
def run(
|
||||
self,
|
||||
command: str,
|
||||
*,
|
||||
timeout: int = 60,
|
||||
env: Optional[dict[str, str]] = None,
|
||||
cwd: Optional[str] = None,
|
||||
shell: str = "/bin/bash",
|
||||
on_stdout: Optional[Callable[[str], Any]] = None,
|
||||
on_stderr: Optional[Callable[[str], Any]] = None,
|
||||
wait: bool = True,
|
||||
) -> Union[ExecutionResult, CommandHandle]:
|
||||
"""Execute a command in the sandbox.
|
||||
|
||||
Args:
|
||||
command: Shell command to execute.
|
||||
timeout: Command timeout in seconds.
|
||||
env: Environment variables to set for the command.
|
||||
cwd: Working directory for command execution. If None, uses sandbox default.
|
||||
shell: Shell to use for command execution. Defaults to "/bin/bash".
|
||||
on_stdout: Callback invoked with each stdout chunk as it arrives.
|
||||
Blocks until the command completes and returns ExecutionResult.
|
||||
Cannot be combined with wait=False.
|
||||
on_stderr: Callback invoked with each stderr chunk as it arrives.
|
||||
Blocks until the command completes and returns ExecutionResult.
|
||||
Cannot be combined with wait=False.
|
||||
wait: If True (default), block until the command completes and
|
||||
return ExecutionResult. If False, return a
|
||||
CommandHandle immediately for streaming output,
|
||||
kill, stdin input, and reconnection. Cannot be combined with
|
||||
on_stdout/on_stderr callbacks.
|
||||
|
||||
Returns:
|
||||
ExecutionResult when wait=True (default).
|
||||
CommandHandle when wait=False.
|
||||
|
||||
Raises:
|
||||
ValueError: If wait=False is combined with callbacks.
|
||||
DataplaneNotConfiguredError: If dataplane_url is not configured.
|
||||
SandboxOperationError: If command execution fails.
|
||||
CommandTimeoutError: If command exceeds its timeout.
|
||||
SandboxConnectionError: If connection to sandbox fails after retries.
|
||||
SandboxNotReadyError: If sandbox is not ready.
|
||||
SandboxClientError: For other errors.
|
||||
"""
|
||||
if not wait and (on_stdout or on_stderr):
|
||||
raise ValueError(
|
||||
"Cannot combine wait=False with on_stdout/on_stderr callbacks. "
|
||||
"Use wait=False and iterate the CommandHandle, or use callbacks."
|
||||
)
|
||||
|
||||
self._require_dataplane_url()
|
||||
|
||||
# When not waiting or callbacks are requested, WS is required
|
||||
use_ws = not wait or on_stdout or on_stderr
|
||||
if use_ws:
|
||||
return self._run_ws(
|
||||
command,
|
||||
timeout=timeout,
|
||||
env=env,
|
||||
cwd=cwd,
|
||||
shell=shell,
|
||||
wait=wait,
|
||||
on_stdout=on_stdout,
|
||||
on_stderr=on_stderr,
|
||||
)
|
||||
|
||||
# Default (wait=True, no callbacks): try WS, fall back to HTTP.
|
||||
# Catch broad exceptions so that unexpected WS failures (e.g. version
|
||||
# incompatibilities) don't break users who don't need WS features.
|
||||
try:
|
||||
return self._run_ws(
|
||||
command,
|
||||
timeout=timeout,
|
||||
env=env,
|
||||
cwd=cwd,
|
||||
shell=shell,
|
||||
wait=True,
|
||||
on_stdout=None,
|
||||
on_stderr=None,
|
||||
)
|
||||
except (SandboxConnectionError, ImportError, OSError, TypeError):
|
||||
return self._run_http(
|
||||
command,
|
||||
timeout=timeout,
|
||||
env=env,
|
||||
cwd=cwd,
|
||||
shell=shell,
|
||||
)
|
||||
|
||||
def _run_ws(
|
||||
self,
|
||||
command: str,
|
||||
*,
|
||||
timeout: int,
|
||||
env: Optional[dict[str, str]],
|
||||
cwd: Optional[str],
|
||||
shell: str,
|
||||
wait: bool,
|
||||
on_stdout: Optional[Callable[[str], Any]],
|
||||
on_stderr: Optional[Callable[[str], Any]],
|
||||
) -> Union[ExecutionResult, CommandHandle]:
|
||||
"""Execute via WebSocket /execute/ws."""
|
||||
from langsmith.sandbox._ws_execute import run_ws_stream
|
||||
|
||||
dataplane_url = self._require_dataplane_url()
|
||||
api_key = self._client._api_key
|
||||
|
||||
msg_stream, control = run_ws_stream(
|
||||
dataplane_url,
|
||||
api_key,
|
||||
command,
|
||||
timeout=timeout,
|
||||
env=env,
|
||||
cwd=cwd,
|
||||
shell=shell,
|
||||
on_stdout=on_stdout,
|
||||
on_stderr=on_stderr,
|
||||
)
|
||||
|
||||
handle = CommandHandle(msg_stream, control, self)
|
||||
|
||||
if not wait:
|
||||
return handle
|
||||
|
||||
return handle.result # blocks until command completes
|
||||
|
||||
def _run_http(
|
||||
self,
|
||||
command: str,
|
||||
*,
|
||||
timeout: int,
|
||||
env: Optional[dict[str, str]],
|
||||
cwd: Optional[str],
|
||||
shell: str,
|
||||
) -> ExecutionResult:
|
||||
"""Execute via HTTP POST /execute (existing implementation)."""
|
||||
dataplane_url = self._require_dataplane_url()
|
||||
url = f"{dataplane_url}/execute"
|
||||
payload: dict[str, Any] = {
|
||||
"command": command,
|
||||
"timeout": timeout,
|
||||
"shell": shell,
|
||||
}
|
||||
if env is not None:
|
||||
payload["env"] = env
|
||||
if cwd is not None:
|
||||
payload["cwd"] = cwd
|
||||
|
||||
try:
|
||||
response = self._client._http.post(url, json=payload, timeout=timeout + 10)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return ExecutionResult(
|
||||
stdout=data.get("stdout", ""),
|
||||
stderr=data.get("stderr", ""),
|
||||
exit_code=data.get("exit_code", -1),
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
handle_sandbox_http_error(e)
|
||||
raise # pragma: no cover
|
||||
|
||||
def reconnect(
|
||||
self,
|
||||
command_id: str,
|
||||
*,
|
||||
stdout_offset: int = 0,
|
||||
stderr_offset: int = 0,
|
||||
) -> CommandHandle:
|
||||
"""Reconnect to a running or recently-finished command.
|
||||
|
||||
Resumes output from the given byte offsets. Any output produced while
|
||||
the client was disconnected is replayed from the server's ring buffer.
|
||||
|
||||
Args:
|
||||
command_id: The command ID from handle.command_id.
|
||||
stdout_offset: Byte offset to resume stdout from (default: 0).
|
||||
stderr_offset: Byte offset to resume stderr from (default: 0).
|
||||
|
||||
Returns:
|
||||
A CommandHandle for the command.
|
||||
|
||||
Raises:
|
||||
SandboxOperationError: If command_id is not found or session expired.
|
||||
SandboxConnectionError: If connection to sandbox fails after retries.
|
||||
"""
|
||||
from langsmith.sandbox._ws_execute import reconnect_ws_stream
|
||||
|
||||
dataplane_url = self._require_dataplane_url()
|
||||
api_key = self._client._api_key
|
||||
|
||||
msg_stream, control = reconnect_ws_stream(
|
||||
dataplane_url,
|
||||
api_key,
|
||||
command_id,
|
||||
stdout_offset=stdout_offset,
|
||||
stderr_offset=stderr_offset,
|
||||
)
|
||||
|
||||
return CommandHandle(
|
||||
msg_stream,
|
||||
control,
|
||||
self,
|
||||
command_id=command_id,
|
||||
stdout_offset=stdout_offset,
|
||||
stderr_offset=stderr_offset,
|
||||
)
|
||||
|
||||
def write(
|
||||
self,
|
||||
path: str,
|
||||
content: Union[str, bytes],
|
||||
*,
|
||||
timeout: int = 60,
|
||||
) -> None:
|
||||
"""Write content to a file in the sandbox.
|
||||
|
||||
Args:
|
||||
path: Target file path in the sandbox.
|
||||
content: File content (str or bytes).
|
||||
timeout: Request timeout in seconds.
|
||||
|
||||
Raises:
|
||||
DataplaneNotConfiguredError: If dataplane_url is not configured.
|
||||
SandboxOperationError: If file write fails.
|
||||
SandboxConnectionError: If connection to sandbox fails after retries.
|
||||
SandboxNotReadyError: If sandbox is not ready.
|
||||
SandboxClientError: For other errors.
|
||||
"""
|
||||
dataplane_url = self._require_dataplane_url()
|
||||
url = f"{dataplane_url}/upload"
|
||||
|
||||
# Ensure content is bytes for multipart upload
|
||||
if isinstance(content, str):
|
||||
content = content.encode("utf-8")
|
||||
|
||||
files = {"file": ("file", content)}
|
||||
|
||||
try:
|
||||
response = self._client._http.post(
|
||||
url, params={"path": path}, files=files, timeout=timeout
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
handle_sandbox_http_error(e)
|
||||
|
||||
def read(self, path: str, *, timeout: int = 60) -> bytes:
|
||||
"""Read a file from the sandbox.
|
||||
|
||||
Args:
|
||||
path: File path to read. Supports both absolute paths (e.g., /tmp/file.txt)
|
||||
and relative paths (resolved from /home/user/).
|
||||
timeout: Request timeout in seconds.
|
||||
|
||||
Returns:
|
||||
File contents as bytes.
|
||||
|
||||
Raises:
|
||||
DataplaneNotConfiguredError: If dataplane_url is not configured.
|
||||
ResourceNotFoundError: If the file doesn't exist.
|
||||
SandboxOperationError: If file read fails.
|
||||
SandboxConnectionError: If connection to sandbox fails after retries.
|
||||
SandboxNotReadyError: If sandbox is not ready.
|
||||
SandboxClientError: For other errors.
|
||||
"""
|
||||
dataplane_url = self._require_dataplane_url()
|
||||
url = f"{dataplane_url}/download"
|
||||
|
||||
try:
|
||||
response = self._client._http.get(
|
||||
url, params={"path": path}, timeout=timeout
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
except httpx.HTTPStatusError as e:
|
||||
if e.response.status_code == 404:
|
||||
raise ResourceNotFoundError(
|
||||
f"File '{path}' not found in sandbox '{self.name}'",
|
||||
resource_type="file",
|
||||
) from e
|
||||
handle_sandbox_http_error(e)
|
||||
# This line should never be reached but satisfies type checker
|
||||
raise # pragma: no cover
|
||||
Reference in New Issue
Block a user