rocket-welder-sdk 1.1.32__py3-none-any.whl → 1.1.34__py3-none-any.whl

rocket_welder_sdk/session_id.py

@@ -0,0 +1,238 @@
+"""SessionId parsing utilities for NNG URL generation.
+
+SessionId format: ps-{guid} (e.g., ps-a1b2c3d4-e5f6-7890-abcd-ef1234567890)
+Prefix "ps" = PipelineSession.
+
+This module provides utilities to:
+1. Parse SessionId from environment variable
+2. Extract the Guid portion
+3. Generate NNG IPC URLs for streaming results
+4. Read explicit NNG URLs from environment variables (preferred)
+
+## URL Configuration Priority
+
+The SDK supports two ways to configure NNG URLs:
+
+1. **Explicit URLs (PREFERRED)** - Set by rocket-welder2:
+   - SEGMENTATION_SINK_URL
+   - KEYPOINTS_SINK_URL
+   - ACTIONS_SINK_URL
+
+2. **Derived from SessionId (FALLBACK)** - For backwards compatibility:
+   - SessionId env var → parse GUID → generate URLs
+
+Use `get_nng_urls_from_env()` for explicit URLs (preferred).
+Use `get_nng_urls(session_id)` for SessionId-derived URLs (fallback).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import uuid
+
+logger = logging.getLogger(__name__)
+
+SESSION_ID_PREFIX = "ps-"
+SESSION_ID_ENV_VAR = "SessionId"
+
+# Explicit URL environment variables (set by rocket-welder2)
+SEGMENTATION_SINK_URL_ENV = "SEGMENTATION_SINK_URL"
+KEYPOINTS_SINK_URL_ENV = "KEYPOINTS_SINK_URL"
+ACTIONS_SINK_URL_ENV = "ACTIONS_SINK_URL"
+
+
+def parse_session_id(session_id: str) -> uuid.UUID:
+    """Parse SessionId (ps-{guid}) to extract Guid.
+
+    Args:
+        session_id: SessionId string (e.g., "ps-a1b2c3d4-...")
+
+    Returns:
+        UUID extracted from SessionId
+
+    Raises:
+        ValueError: If session_id format is invalid
+
+    Examples:
+        >>> parse_session_id("ps-a1b2c3d4-e5f6-7890-abcd-ef1234567890")
+        UUID('a1b2c3d4-e5f6-7890-abcd-ef1234567890')
+        >>> parse_session_id("a1b2c3d4-e5f6-7890-abcd-ef1234567890")  # backwards compat
+        UUID('a1b2c3d4-e5f6-7890-abcd-ef1234567890')
+    """
+    if session_id.startswith(SESSION_ID_PREFIX):
+        return uuid.UUID(session_id[len(SESSION_ID_PREFIX) :])
+    # Fallback: try parsing as raw guid for backwards compatibility
+    return uuid.UUID(session_id)
+
+
+def get_session_id_from_env() -> str | None:
+    """Get SessionId from environment variable.
+
+    Returns:
+        SessionId string or None if not set
+    """
+    return os.environ.get(SESSION_ID_ENV_VAR)
+
+
+def get_nng_urls(session_id: str) -> dict[str, str]:
+    """Generate NNG IPC URLs from SessionId.
+
+    Args:
+        session_id: SessionId string (e.g., "ps-a1b2c3d4-...")
+
+    Returns:
+        Dictionary with 'segmentation', 'keypoints', 'actions' URLs
+
+    Examples:
+        >>> urls = get_nng_urls("ps-a1b2c3d4-e5f6-7890-abcd-ef1234567890")
+        >>> urls["segmentation"]
+        'ipc:///tmp/rw-a1b2c3d4-e5f6-7890-abcd-ef1234567890-seg.sock'
+    """
+    guid = parse_session_id(session_id)
+    return {
+        "segmentation": f"ipc:///tmp/rw-{guid}-seg.sock",
+        "keypoints": f"ipc:///tmp/rw-{guid}-kp.sock",
+        "actions": f"ipc:///tmp/rw-{guid}-actions.sock",
+    }
+
+
+def get_segmentation_url(session_id: str) -> str:
+    """Get NNG URL for segmentation stream.
+
+    Args:
+        session_id: SessionId string (e.g., "ps-a1b2c3d4-...")
+
+    Returns:
+        IPC URL for segmentation stream
+    """
+    guid = parse_session_id(session_id)
+    return f"ipc:///tmp/rw-{guid}-seg.sock"
+
+
+def get_keypoints_url(session_id: str) -> str:
+    """Get NNG URL for keypoints stream.
+
+    Args:
+        session_id: SessionId string (e.g., "ps-a1b2c3d4-...")
+
+    Returns:
+        IPC URL for keypoints stream
+    """
+    guid = parse_session_id(session_id)
+    return f"ipc:///tmp/rw-{guid}-kp.sock"
+
+
+def get_actions_url(session_id: str) -> str:
+    """Get NNG URL for actions stream.
+
+    Args:
+        session_id: SessionId string (e.g., "ps-a1b2c3d4-...")
+
+    Returns:
+        IPC URL for actions stream
+    """
+    guid = parse_session_id(session_id)
+    return f"ipc:///tmp/rw-{guid}-actions.sock"
+
+
+# ============================================================================
+# Explicit URL functions (PREFERRED - URLs set by rocket-welder2)
+# ============================================================================
+
+
+def get_nng_urls_from_env() -> dict[str, str | None]:
+    """Get NNG URLs from explicit environment variables.
+
+    This is the PREFERRED method for getting NNG URLs. rocket-welder2
+    sets these environment variables when starting containers.
+
+    Returns:
+        Dictionary with 'segmentation', 'keypoints', 'actions' URLs.
+        Values are None if not configured.
+
+    Examples:
+        >>> os.environ["SEGMENTATION_SINK_URL"] = "ipc:///tmp/rw-abc-seg.sock"
+        >>> urls = get_nng_urls_from_env()
+        >>> urls["segmentation"]
+        'ipc:///tmp/rw-abc-seg.sock'
+    """
+    return {
+        "segmentation": os.environ.get(SEGMENTATION_SINK_URL_ENV),
+        "keypoints": os.environ.get(KEYPOINTS_SINK_URL_ENV),
+        "actions": os.environ.get(ACTIONS_SINK_URL_ENV),
+    }
+
+
+def get_segmentation_url_from_env() -> str | None:
+    """Get segmentation NNG URL from environment variable.
+
+    Returns:
+        IPC URL for segmentation stream, or None if not configured.
+    """
+    return os.environ.get(SEGMENTATION_SINK_URL_ENV)
+
+
+def get_keypoints_url_from_env() -> str | None:
+    """Get keypoints NNG URL from environment variable.
+
+    Returns:
+        IPC URL for keypoints stream, or None if not configured.
+    """
+    return os.environ.get(KEYPOINTS_SINK_URL_ENV)
+
+
+def get_actions_url_from_env() -> str | None:
+    """Get actions NNG URL from environment variable.
+
+    Returns:
+        IPC URL for actions stream, or None if not configured.
+    """
+    return os.environ.get(ACTIONS_SINK_URL_ENV)
+
+
+def has_explicit_nng_urls() -> bool:
+    """Check if explicit NNG URLs are configured.
+
+    Returns:
+        True if at least segmentation OR keypoints URL is configured.
+    """
+    urls = get_nng_urls_from_env()
+    return bool(urls["segmentation"] or urls["keypoints"])
+
+
+def get_configured_nng_urls() -> dict[str, str]:
+    """Get all configured NNG URLs (explicit or derived from SessionId).
+
+    Priority:
+    1. Explicit URLs from environment (SEGMENTATION_SINK_URL, etc.)
+    2. Derived from SessionId environment variable (fallback)
+
+    Returns:
+        Dictionary with 'segmentation', 'keypoints', 'actions' URLs.
+        Only includes URLs that are actually configured.
+
+    Raises:
+        ValueError: If no NNG URLs are configured (neither explicit nor SessionId).
+    """
+    # Try explicit URLs first (preferred)
+    explicit_urls = get_nng_urls_from_env()
+    result: dict[str, str] = {}
+
+    for name, url in explicit_urls.items():
+        if url:
+            result[name] = url
+
+    # If we have at least one explicit URL, return what we have
+    if result:
+        return result
+
+    # Fallback: derive from SessionId
+    session_id = get_session_id_from_env()
+    if session_id:
+        return get_nng_urls(session_id)
+
+    raise ValueError(
+        "No NNG URLs configured. Set SEGMENTATION_SINK_URL/KEYPOINTS_SINK_URL "
+        "environment variables, or set SessionId for URL derivation."
+    )

rocket_welder_sdk/transport/__init__.py

@@ -0,0 +1,30 @@
+"""
+Transport layer for RocketWelder SDK.
+
+Provides transport-agnostic frame sink/source abstractions for protocols.
+"""
+
+from .frame_sink import IFrameSink
+from .frame_source import IFrameSource
+from .nng_transport import NngFrameSink, NngFrameSource
+from .stream_transport import StreamFrameSink, StreamFrameSource
+from .tcp_transport import TcpFrameSink, TcpFrameSource
+from .unix_socket_transport import (
+    UnixSocketFrameSink,
+    UnixSocketFrameSource,
+    UnixSocketServer,
+)
+
+__all__ = [
+    "IFrameSink",
+    "IFrameSource",
+    "NngFrameSink",
+    "NngFrameSource",
+    "StreamFrameSink",
+    "StreamFrameSource",
+    "TcpFrameSink",
+    "TcpFrameSource",
+    "UnixSocketFrameSink",
+    "UnixSocketFrameSource",
+    "UnixSocketServer",
+]

rocket_welder_sdk/transport/frame_sink.py

@@ -0,0 +1,77 @@
+"""Frame sink abstraction for writing frames to any transport."""
+
+from abc import ABC, abstractmethod
+
+
+class IFrameSink(ABC):
+    """
+    Low-level abstraction for writing discrete frames to any transport.
+
+    Transport-agnostic interface that handles the question: "where do frames go?"
+    This abstraction decouples protocol logic (KeyPoints, SegmentationResults) from
+    transport mechanisms (File, TCP, WebSocket, NNG). Each frame is written atomically.
+    """
+
+    @abstractmethod
+    def write_frame(self, frame_data: bytes) -> None:
+        """
+        Write a complete frame to the underlying transport synchronously.
+
+        Args:
+            frame_data: Complete frame data to write
+        """
+        pass
+
+    @abstractmethod
+    async def write_frame_async(self, frame_data: bytes) -> None:
+        """
+        Write a complete frame to the underlying transport asynchronously.
+
+        Args:
+            frame_data: Complete frame data to write
+        """
+        pass
+
+    @abstractmethod
+    def flush(self) -> None:
+        """
+        Flush any buffered data to the transport synchronously.
+
+        For message-based transports (NNG, WebSocket), this may be a no-op.
+        """
+        pass
+
+    @abstractmethod
+    async def flush_async(self) -> None:
+        """
+        Flush any buffered data to the transport asynchronously.
+
+        For message-based transports (NNG, WebSocket), this may be a no-op.
+        """
+        pass
+
+    def __enter__(self) -> "IFrameSink":
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        """Context manager exit."""
+        self.close()
+
+    async def __aenter__(self) -> "IFrameSink":
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        """Async context manager exit."""
+        await self.close_async()
+
+    @abstractmethod
+    def close(self) -> None:
+        """Close the sink and release resources."""
+        pass
+
+    @abstractmethod
+    async def close_async(self) -> None:
+        """Close the sink and release resources asynchronously."""
+        pass

rocket_welder_sdk/transport/frame_source.py

@@ -0,0 +1,74 @@
+"""Frame source abstraction for reading frames from any transport."""
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+
+class IFrameSource(ABC):
+    """
+    Low-level abstraction for reading discrete frames from any transport.
+
+    Transport-agnostic interface that handles the question: "where do frames come from?"
+    This abstraction decouples protocol logic (KeyPoints, SegmentationResults) from
+    transport mechanisms (File, TCP, WebSocket, NNG). Each frame is read atomically.
+    """
+
+    @abstractmethod
+    def read_frame(self) -> Optional[bytes]:
+        """
+        Read a complete frame from the underlying transport synchronously.
+
+        Returns:
+            Complete frame data, or None if end of stream/no more messages
+        """
+        pass
+
+    @abstractmethod
+    async def read_frame_async(self) -> Optional[bytes]:
+        """
+        Read a complete frame from the underlying transport asynchronously.
+
+        Returns:
+            Complete frame data, or None if end of stream/no more messages
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def has_more_frames(self) -> bool:
+        """
+        Check if more frames are available.
+
+        For streaming transports (file), this checks for EOF.
+        For message-based transports (NNG), this may always return True until disconnection.
+
+        Returns:
+            True if more frames are available, False otherwise
+        """
+        pass
+
+    def __enter__(self) -> "IFrameSource":
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        """Context manager exit."""
+        self.close()
+
+    async def __aenter__(self) -> "IFrameSource":
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        """Async context manager exit."""
+        await self.close_async()
+
+    @abstractmethod
+    def close(self) -> None:
+        """Close the source and release resources."""
+        pass
+
+    @abstractmethod
+    async def close_async(self) -> None:
+        """Close the source and release resources asynchronously."""
+        pass

rocket_welder_sdk/transport/nng_transport.py

@@ -0,0 +1,197 @@
+"""NNG transport using pynng library.
+
+NNG (nanomsg next generation) provides high-performance, scalable messaging patterns.
+Supported patterns:
+- Pub/Sub: One publisher to many subscribers
+- Push/Pull: Load-balanced distribution to workers
+"""
+
+from typing import Any, Optional, cast
+
+import pynng
+
+from .frame_sink import IFrameSink
+from .frame_source import IFrameSource
+
+
+class NngFrameSink(IFrameSink):
+    """
+    Frame sink that publishes to NNG Pub/Sub or Push/Pull pattern.
+
+    Each frame is sent as a single NNG message (no framing needed - NNG handles message boundaries).
+    """
+
+    def __init__(self, socket: Any, leave_open: bool = False):
+        """
+        Create an NNG frame sink from a socket.
+
+        Args:
+            socket: pynng socket (Publisher or Pusher)
+            leave_open: If True, doesn't close socket on close
+        """
+        self._socket: Any = socket
+        self._leave_open = leave_open
+        self._closed = False
+
+    @classmethod
+    def create_publisher(cls, url: str) -> "NngFrameSink":
+        """
+        Create an NNG Publisher frame sink bound to the specified URL.
+
+        Args:
+            url: NNG URL (e.g., "tcp://127.0.0.1:5555", "ipc:///tmp/mysocket")
+
+        Returns:
+            Frame sink ready to publish messages
+        """
+        socket = pynng.Pub0()
+        socket.listen(url)
+        return cls(socket, leave_open=False)
+
+    @classmethod
+    def create_pusher(cls, url: str, bind_mode: bool = True) -> "NngFrameSink":
+        """
+        Create an NNG Pusher frame sink.
+
+        Args:
+            url: NNG URL (e.g., "tcp://127.0.0.1:5555", "ipc:///tmp/mysocket")
+            bind_mode: If True, listens (bind); if False, dials (connect)
+
+        Returns:
+            Frame sink ready to push messages
+        """
+        socket = pynng.Push0()
+        if bind_mode:
+            socket.listen(url)
+        else:
+            socket.dial(url)
+        return cls(socket, leave_open=False)
+
+    def write_frame(self, frame_data: bytes) -> None:
+        """Write frame to NNG socket (no length prefix - NNG handles message boundaries)."""
+        if self._closed:
+            raise ValueError("Cannot write to closed sink")
+
+        self._socket.send(frame_data)
+
+    async def write_frame_async(self, frame_data: bytes) -> None:
+        """Write frame asynchronously."""
+        if self._closed:
+            raise ValueError("Cannot write to closed sink")
+
+        await self._socket.asend(frame_data)
+
+    def flush(self) -> None:
+        """Flush is a no-op for NNG (data sent immediately)."""
+        pass
+
+    async def flush_async(self) -> None:
+        """Flush asynchronously is a no-op for NNG."""
+        pass
+
+    def close(self) -> None:
+        """Close the NNG sink."""
+        if self._closed:
+            return
+        self._closed = True
+        if not self._leave_open:
+            self._socket.close()
+
+    async def close_async(self) -> None:
+        """Close the NNG sink asynchronously."""
+        self.close()
+
+
+class NngFrameSource(IFrameSource):
+    """
+    Frame source that subscribes to NNG Pub/Sub or Pull pattern.
+
+    Each NNG message is treated as a complete frame (no framing needed - NNG handles message boundaries).
+    """
+
+    def __init__(self, socket: Any, leave_open: bool = False):
+        """
+        Create an NNG frame source from a socket.
+
+        Args:
+            socket: pynng socket (Subscriber or Puller)
+            leave_open: If True, doesn't close socket on close
+        """
+        self._socket: Any = socket
+        self._leave_open = leave_open
+        self._closed = False
+
+    @classmethod
+    def create_subscriber(cls, url: str, topic: bytes = b"") -> "NngFrameSource":
+        """
+        Create an NNG Subscriber frame source connected to the specified URL.
+
+        Args:
+            url: NNG URL (e.g., "tcp://127.0.0.1:5555", "ipc:///tmp/mysocket")
+            topic: Optional topic filter (empty for all messages)
+
+        Returns:
+            Frame source ready to receive messages
+        """
+        socket = pynng.Sub0()
+        socket.subscribe(topic)
+        socket.dial(url)
+        return cls(socket, leave_open=False)
+
+    @classmethod
+    def create_puller(cls, url: str, bind_mode: bool = True) -> "NngFrameSource":
+        """
+        Create an NNG Puller frame source.
+
+        Args:
+            url: NNG URL (e.g., "tcp://127.0.0.1:5555", "ipc:///tmp/mysocket")
+            bind_mode: If True, listens (bind); if False, dials (connect)
+
+        Returns:
+            Frame source ready to pull messages
+        """
+        socket = pynng.Pull0()
+        if bind_mode:
+            socket.listen(url)
+        else:
+            socket.dial(url)
+        return cls(socket, leave_open=False)
+
+    @property
+    def has_more_frames(self) -> bool:
+        """Check if more frames available (NNG blocks waiting for messages)."""
+        return not self._closed
+
+    def read_frame(self) -> Optional[bytes]:
+        """Read frame from NNG socket (blocking)."""
+        if self._closed:
+            return None
+
+        try:
+            return cast("bytes", self._socket.recv())
+        except pynng.Closed:
+            self._closed = True
+            return None
+
+    async def read_frame_async(self) -> Optional[bytes]:
+        """Read frame asynchronously."""
+        if self._closed:
+            return None
+
+        try:
+            return cast("bytes", await self._socket.arecv())
+        except pynng.Closed:
+            self._closed = True
+            return None
+
+    def close(self) -> None:
+        """Close the NNG source."""
+        if self._closed:
+            return
+        self._closed = True
+        if not self._leave_open:
+            self._socket.close()
+
+    async def close_async(self) -> None:
+        """Close the NNG source asynchronously."""
+        self.close()