reactor-sdk 0.1.0__py3-none-any.whl

reactor_sdk/types.py

@@ -0,0 +1,255 @@
+"""
+Type definitions for the Reactor SDK.
+
+This module contains all the type definitions, enums, and data classes
+used throughout the SDK.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Callable, Literal, Optional, TypedDict
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+# =============================================================================
+# Status Enums
+# =============================================================================
+
+
+class ReactorStatus(Enum):
+    """Status of the Reactor connection."""
+
+    DISCONNECTED = "disconnected"  # Not connected to anything
+    CONNECTING = "connecting"  # Establishing connection to coordinator
+    WAITING = "waiting"  # Connected to coordinator, waiting for GPU assignment
+    READY = "ready"  # Connected to GPU machine, can send/receive messages
+
+
+class GPUMachineStatus(Enum):
+    """Status of the GPU machine WebRTC connection."""
+
+    DISCONNECTED = "disconnected"
+    CONNECTING = "connecting"
+    CONNECTED = "connected"
+    ERROR = "error"
+
+
+# =============================================================================
+# Event Types
+# =============================================================================
+
+# Reactor event types
+ReactorEvent = Literal[
+    "status_changed",  # Updates on the reactor status
+    "session_id_changed",  # Updates on the session ID
+    "new_message",  # New messages from the machine
+    "stream_changed",  # Video stream has changed
+    "error",  # Error events with ReactorError details
+    "session_expiration_changed",  # Session expiration has changed
+]
+
+# GPU Machine event types
+GPUMachineEvent = Literal[
+    "status_changed",  # Connection state changes
+    "track_received",  # Remote track received
+    "track_removed",  # Remote track removed
+    "application",  # Data channel messages
+]
+
+
+# =============================================================================
+# Error Types
+# =============================================================================
+
+
+@dataclass
+class ReactorError:
+    """Information about an error that occurred in the Reactor."""
+
+    code: str
+    message: str
+    timestamp: float
+    recoverable: bool
+    component: Literal["coordinator", "gpu"]
+    retry_after: Optional[float] = None
+
+    def __str__(self) -> str:
+        return f"[{self.component}:{self.code}] {self.message}"
+
+
+class ConflictError(Exception):
+    """Raised when a connection conflict occurs (e.g., superseded by newer request)."""
+
+    pass
+
+
+# =============================================================================
+# State Types
+# =============================================================================
+
+
+@dataclass
+class ReactorState:
+    """Current state of the Reactor including status and error info."""
+
+    status: ReactorStatus
+    last_error: Optional[ReactorError] = None
+
+
+# =============================================================================
+# Callback Types
+# =============================================================================
+
+# Type for frame callback function - receives numpy RGB frame (H, W, 3)
+FrameCallback = Callable[[NDArray[np.uint8]], None]
+
+# Type for event handler function
+EventHandler = Callable[..., None]
+
+
+# =============================================================================
+# API Response Types
+# =============================================================================
+
+
+class IceServerCredentials(TypedDict, total=False):
+    """Credentials for an ICE server."""
+
+    username: str
+    password: str
+
+
+class IceServerConfig(TypedDict):
+    """Configuration for a single ICE server."""
+
+    uris: list[str]
+    credentials: Optional[IceServerCredentials]
+
+
+class IceServersResponse(TypedDict):
+    """Response from the ICE servers endpoint."""
+
+    ice_servers: list[IceServerConfig]
+
+
+class ModelConfig(TypedDict):
+    """Model configuration in session requests."""
+
+    name: str
+
+
+class CreateSessionRequest(TypedDict):
+    """Request body for creating a session."""
+
+    model: ModelConfig
+    sdp_offer: str
+    extra_args: dict[str, Any]
+
+
+class CreateSessionResponse(TypedDict):
+    """Response from creating a session."""
+
+    session_id: str
+
+
+class SDPParamsRequest(TypedDict):
+    """Request body for SDP params endpoint."""
+
+    sdp_offer: str
+    extra_args: dict[str, Any]
+
+
+class SDPParamsResponse(TypedDict):
+    """Response from SDP params endpoint."""
+
+    sdp_answer: str
+    extra_args: dict[str, Any]
+
+
+class SessionState(Enum):
+    """State of a session on the coordinator."""
+
+    CREATED = "CREATED"
+    PENDING = "PENDING"
+    SUSPENDED = "SUSPENDED"
+    WAITING = "WAITING"
+    ACTIVE = "ACTIVE"
+    INACTIVE = "INACTIVE"
+    CLOSED = "CLOSED"
+
+
+class SessionInfoResponse(TypedDict):
+    """Response from session info endpoint."""
+
+    session_id: str
+    state: str
+
+
+# =============================================================================
+# WebRTC Types
+# =============================================================================
+
+
+@dataclass
+class RTCIceServer:
+    """ICE server configuration for WebRTC."""
+
+    urls: list[str]
+    username: Optional[str] = None
+    credential: Optional[str] = None
+
+
+@dataclass
+class WebRTCConfig:
+    """Configuration for WebRTC peer connection."""
+
+    ice_servers: list[RTCIceServer]
+    data_channel_label: str = "data"
+
+
+# =============================================================================
+# Command Schema Types (for capabilities)
+# =============================================================================
+
+
+class ParameterSchema(TypedDict, total=False):
+    """Schema for a command parameter."""
+
+    description: str
+    type: str  # "number", "integer", "string", "boolean"
+    minimum: float
+    maximum: float
+    required: bool
+    enum: list[str]
+
+
+class CommandSchema(TypedDict):
+    """Schema for a command."""
+
+    description: str
+    schema: dict[str, ParameterSchema]
+
+
+class CapabilitiesMessage(TypedDict):
+    """Message containing model capabilities/commands."""
+
+    commands: dict[str, CommandSchema]
+
+
+# =============================================================================
+# Video Frame Types
+# =============================================================================
+
+
+@dataclass
+class VideoFrameInfo:
+    """Information about a video frame."""
+
+    width: int
+    height: int
+    format: str = "rgb24"
+    timestamp: float = field(default=0.0)

reactor_sdk/utils/__init__.py

@@ -0,0 +1,25 @@
+"""Utility modules for the Reactor SDK."""
+
+from reactor_sdk.utils.tokens import fetch_jwt_token
+from reactor_sdk.utils.webrtc import (
+    WebRTCConfig,
+    create_peer_connection,
+    create_data_channel,
+    create_offer,
+    set_remote_description,
+    transform_ice_servers,
+    wait_for_ice_gathering,
+)
+
+__all__ = [
+    # Token utilities
+    "fetch_jwt_token",
+    # WebRTC utilities
+    "WebRTCConfig",
+    "create_peer_connection",
+    "create_data_channel",
+    "create_offer",
+    "set_remote_description",
+    "transform_ice_servers",
+    "wait_for_ice_gathering",
+]

reactor_sdk/utils/tokens.py

@@ -0,0 +1,64 @@
+"""
+Token utilities for the Reactor SDK.
+
+This module provides functions for fetching JWT tokens from the coordinator.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import aiohttp
+
+logger = logging.getLogger(__name__)
+
+
+# Default coordinator URL
+PROD_COORDINATOR_URL = "https://api.reactor.inc"
+
+
+async def fetch_jwt_token(
+    api_key: str,
+    coordinator_url: str = PROD_COORDINATOR_URL,
+) -> str:
+    """
+    Fetch a JWT token from the coordinator using an API key.
+
+    This is safe to use in Python applications (CLI tools, scripts, servers)
+    since the API key is not exposed to end users like it would be in
+    browser-based JavaScript applications.
+
+    Args:
+        api_key: Your Reactor API key.
+        coordinator_url: Optional coordinator URL, defaults to production.
+
+    Returns:
+        The JWT token string.
+
+    Raises:
+        RuntimeError: If the token fetch fails.
+
+    Example:
+        >>> token = await fetch_jwt_token("your-api-key")
+        >>> reactor = Reactor(model_name="my-model")
+        >>> await reactor.connect(jwt_token=token)
+    """
+    url = f"{coordinator_url.rstrip('/')}/tokens"
+
+    async with aiohttp.ClientSession() as session:
+        async with session.get(
+            url,
+            headers={"X-API-Key": api_key},
+        ) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise RuntimeError(
+                    f"Failed to fetch JWT token: {response.status} {error_text}"
+                )
+
+            data = await response.json()
+            jwt_token: str = data["jwt"]
+
+            logger.debug("Successfully fetched JWT token")
+            return jwt_token

reactor_sdk/utils/webrtc.py

@@ -0,0 +1,315 @@
+"""
+WebRTC utility functions for the Reactor SDK.
+
+This module provides stateless utility functions for WebRTC operations
+using aiortc.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from aiortc import RTCConfiguration, RTCDataChannel, RTCIceServer, RTCPeerConnection
+from aiortc import RTCSessionDescription
+
+from reactor_sdk.types import IceServersResponse
+
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# Configuration
+# =============================================================================
+
+
+@dataclass
+class WebRTCConfig:
+    """Configuration for WebRTC peer connection."""
+
+    ice_servers: list[RTCIceServer]
+    data_channel_label: str = "data"
+
+
+DEFAULT_DATA_CHANNEL_LABEL = "data"
+DEFAULT_ICE_GATHERING_TIMEOUT = 5.0  # seconds
+
+
+# =============================================================================
+# Peer Connection Creation
+# =============================================================================
+
+
+def create_peer_connection(config: WebRTCConfig) -> RTCPeerConnection:
+    """
+    Create a new RTCPeerConnection with the specified configuration.
+
+    Args:
+        config: WebRTC configuration with ICE servers.
+
+    Returns:
+        A new RTCPeerConnection instance.
+    """
+    rtc_config = RTCConfiguration(iceServers=config.ice_servers)
+    return RTCPeerConnection(configuration=rtc_config)
+
+
+def create_data_channel(
+    pc: RTCPeerConnection,
+    label: Optional[str] = None,
+) -> RTCDataChannel:
+    """
+    Create a data channel on the peer connection.
+
+    Args:
+        pc: The RTCPeerConnection.
+        label: Label for the data channel (defaults to "data").
+
+    Returns:
+        The created RTCDataChannel.
+    """
+    return pc.createDataChannel(label or DEFAULT_DATA_CHANNEL_LABEL)
+
+
+# =============================================================================
+# SDP Offer/Answer
+# =============================================================================
+
+
+async def create_offer(pc: RTCPeerConnection) -> str:
+    """
+    Create an SDP offer on the peer connection.
+
+    Waits for ICE gathering to complete before returning.
+
+    Args:
+        pc: The RTCPeerConnection.
+
+    Returns:
+        The SDP offer string.
+
+    Raises:
+        RuntimeError: If local description creation fails.
+    """
+    offer = await pc.createOffer()
+    await pc.setLocalDescription(offer)
+
+    await wait_for_ice_gathering(pc)
+
+    local_description = pc.localDescription
+    if local_description is None:
+        raise RuntimeError("Failed to create local description")
+
+    return local_description.sdp
+
+
+async def create_answer(pc: RTCPeerConnection, offer_sdp: str) -> str:
+    """
+    Create an SDP answer in response to a received offer.
+
+    Waits for ICE gathering to complete before returning.
+
+    Args:
+        pc: The RTCPeerConnection.
+        offer_sdp: The SDP offer string from the remote peer.
+
+    Returns:
+        The SDP answer string.
+
+    Raises:
+        RuntimeError: If local description creation fails.
+    """
+    await set_remote_description(pc, offer_sdp, "offer")
+
+    answer = await pc.createAnswer()
+    await pc.setLocalDescription(answer)
+
+    await wait_for_ice_gathering(pc)
+
+    local_description = pc.localDescription
+    if local_description is None:
+        raise RuntimeError("Failed to create local description")
+
+    return local_description.sdp
+
+
+async def set_remote_description(
+    pc: RTCPeerConnection,
+    sdp: str,
+    sdp_type: str = "answer",
+) -> None:
+    """
+    Set the remote description on the peer connection.
+
+    Args:
+        pc: The RTCPeerConnection.
+        sdp: The SDP string.
+        sdp_type: The type of SDP ("offer" or "answer").
+    """
+    session_description = RTCSessionDescription(sdp=sdp, type=sdp_type)
+    await pc.setRemoteDescription(session_description)
+
+
+def get_local_description(pc: RTCPeerConnection) -> Optional[str]:
+    """
+    Get the local SDP description from the peer connection.
+
+    Args:
+        pc: The RTCPeerConnection.
+
+    Returns:
+        The local SDP string, or None if not set.
+    """
+    desc = pc.localDescription
+    if desc is None:
+        return None
+    return desc.sdp
+
+
+# =============================================================================
+# ICE Handling
+# =============================================================================
+
+
+def transform_ice_servers(response: IceServersResponse) -> list[RTCIceServer]:
+    """
+    Transform ICE servers from the coordinator API format to RTCIceServer format.
+
+    Args:
+        response: The parsed IceServersResponse from the coordinator.
+
+    Returns:
+        List of RTCIceServer objects for WebRTC peer connection configuration.
+    """
+    ice_servers: list[RTCIceServer] = []
+
+    for server in response["ice_servers"]:
+        if server.get("credentials"):
+            creds = server["credentials"]
+            ice_server = RTCIceServer(
+                urls=server["uris"],
+                username=creds.get("username"),
+                credential=creds.get("password"),
+            )
+        else:
+            ice_server = RTCIceServer(urls=server["uris"])
+
+        ice_servers.append(ice_server)
+
+    return ice_servers
+
+
+async def wait_for_ice_gathering(
+    pc: RTCPeerConnection,
+    timeout: float = DEFAULT_ICE_GATHERING_TIMEOUT,
+) -> None:
+    """
+    Wait for ICE gathering to complete with a timeout.
+
+    Args:
+        pc: The RTCPeerConnection.
+        timeout: Maximum time to wait in seconds.
+    """
+    if pc.iceGatheringState == "complete":
+        return
+
+    gathering_complete = asyncio.Event()
+
+    @pc.on("icegatheringstatechange")
+    def on_ice_gathering_state_change() -> None:
+        if pc.iceGatheringState == "complete":
+            gathering_complete.set()
+
+    try:
+        await asyncio.wait_for(gathering_complete.wait(), timeout=timeout)
+    except asyncio.TimeoutError:
+        logger.warning(
+            f"ICE gathering timed out after {timeout}s, proceeding with current candidates"
+        )
+
+
+# =============================================================================
+# Data Channel Messaging
+# =============================================================================
+
+
+def send_message(channel: RTCDataChannel, command: str, data: Any) -> None:
+    """
+    Send a message through a data channel.
+
+    Args:
+        channel: The RTCDataChannel.
+        command: The command type.
+        data: The data to send with the command.
+
+    Raises:
+        RuntimeError: If the data channel is not open.
+    """
+    if channel.readyState != "open":
+        raise RuntimeError(f"Data channel not open: {channel.readyState}")
+
+    json_data = data if isinstance(data, dict) else json.loads(data) if isinstance(data, str) else data
+    payload = {"type": command, "data": json_data}
+    channel.send(json.dumps(payload))
+
+
+def parse_message(data: Any) -> Any:
+    """
+    Parse a received data channel message, attempting JSON parse.
+
+    Args:
+        data: The raw message data.
+
+    Returns:
+        The parsed message (dict if JSON, otherwise original data).
+    """
+    if isinstance(data, str):
+        try:
+            return json.loads(data)
+        except json.JSONDecodeError:
+            return data
+    return data
+
+
+# =============================================================================
+# Connection State
+# =============================================================================
+
+
+def is_connected(pc: RTCPeerConnection) -> bool:
+    """
+    Check if the peer connection is in a connected state.
+
+    Args:
+        pc: The RTCPeerConnection.
+
+    Returns:
+        True if connected.
+    """
+    return pc.connectionState == "connected"
+
+
+def is_closed(pc: RTCPeerConnection) -> bool:
+    """
+    Check if the peer connection is closed or failed.
+
+    Args:
+        pc: The RTCPeerConnection.
+
+    Returns:
+        True if closed or failed.
+    """
+    return pc.connectionState in ("closed", "failed")
+
+
+async def close_peer_connection(pc: RTCPeerConnection) -> None:
+    """
+    Close the peer connection and clean up.
+
+    Args:
+        pc: The RTCPeerConnection.
+    """
+    await pc.close()