rocket-welder-sdk 1.1.42__py3-none-any.whl → 1.1.44__py3-none-any.whl

rocket_welder_sdk/session_id.py

@@ -1,4 +1,4 @@
-"""SessionId parsing utilities for NNG URL generation.
+"""SessionId parsing utilities.
 
 SessionId format: ps-{guid} (e.g., ps-a1b2c3d4-e5f6-7890-abcd-ef1234567890)
 Prefix "ps" = PipelineSession.
@@ -6,33 +6,20 @@ 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
+## URL Configuration
 
-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).
+Output URLs are configured via environment variables set by rocket-welder2:
+- SEGMENTATION_SINK_URL: URL for segmentation output (e.g., socket:///tmp/seg.sock)
+- KEYPOINTS_SINK_URL: URL for keypoints output (e.g., socket:///tmp/kp.sock)
+- ACTIONS_SINK_URL: URL for actions output
 """
 
 from __future__ import annotations
 
-import logging
 import os
 import uuid
 
-logger = logging.getLogger(__name__)
-
 SESSION_ID_PREFIX = "ps-"
 SESSION_ID_ENV_VAR = "SessionId"
 
@@ -73,166 +60,3 @@ def get_session_id_from_env() -> str | None:
         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

@@ -6,7 +6,6 @@ Provides transport-agnostic frame sink/source abstractions for protocols.
 
 from .frame_sink import IFrameSink, NullFrameSink
 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 (
@@ -14,12 +13,16 @@ from .unix_socket_transport import (
     UnixSocketFrameSource,
     UnixSocketServer,
 )
+from .websocket_transport import (
+    WebSocketFrameSink,
+    WebSocketFrameSource,
+    connect_websocket_sink,
+    connect_websocket_source,
+)
 
 __all__ = [
     "IFrameSink",
     "IFrameSource",
-    "NngFrameSink",
-    "NngFrameSource",
     "NullFrameSink",
     "StreamFrameSink",
     "StreamFrameSource",
@@ -28,4 +31,8 @@ __all__ = [
     "UnixSocketFrameSink",
     "UnixSocketFrameSource",
     "UnixSocketServer",
+    "WebSocketFrameSink",
+    "WebSocketFrameSource",
+    "connect_websocket_sink",
+    "connect_websocket_source",
 ]

rocket_welder_sdk/transport/frame_sink.py

@@ -9,7 +9,7 @@ class IFrameSink(ABC):
 
     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.
+    transport mechanisms (File, TCP, WebSocket, Unix Socket). Each frame is written atomically.
     """
 
     @abstractmethod
@@ -37,7 +37,7 @@ class IFrameSink(ABC):
         """
         Flush any buffered data to the transport synchronously.
 
-        For message-based transports (NNG, WebSocket), this may be a no-op.
+        For message-based transports (WebSocket), this may be a no-op.
         """
         pass
 
@@ -46,7 +46,7 @@ class IFrameSink(ABC):
         """
         Flush any buffered data to the transport asynchronously.
 
-        For message-based transports (NNG, WebSocket), this may be a no-op.
+        For message-based transports (WebSocket), this may be a no-op.
         """
         pass
 

rocket_welder_sdk/transport/frame_source.py

@@ -10,7 +10,7 @@ class IFrameSource(ABC):
 
     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.
+    transport mechanisms (File, TCP, WebSocket, Unix Socket). Each frame is read atomically.
     """
 
     @abstractmethod
@@ -40,7 +40,7 @@ class IFrameSource(ABC):
         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.
+        For message-based transports, this may always return True until disconnection.
 
         Returns:
             True if more frames are available, False otherwise

rocket_welder_sdk/transport/websocket_transport.py

@@ -0,0 +1,316 @@
+"""
+WebSocket transport for reading/writing frames.
+Matches C# WebSocketFrameSink and WebSocketFrameSource from RocketWelder.SDK.Transport.
+
+Uses the websockets library for async WebSocket support.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Optional, Protocol
+
+from .frame_sink import IFrameSink
+from .frame_source import IFrameSource
+
+
+class WebSocketProtocol(Protocol):
+    """Protocol for WebSocket-like objects (for type checking)."""
+
+    @property
+    def closed(self) -> bool: ...
+
+    @property
+    def state(self) -> Any: ...
+
+    async def send(self, message: bytes) -> None: ...
+
+    async def recv(self) -> bytes: ...
+
+    async def close(self) -> None: ...
+
+
+class WebSocketFrameSink(IFrameSink):
+    """
+    Frame sink that writes to a WebSocket connection.
+
+    Each frame is sent as a single binary WebSocket message.
+
+    Attributes:
+        leave_open: If True, doesn't close WebSocket on disposal
+    """
+
+    __slots__ = ("_closed", "_leave_open", "_websocket")
+
+    def __init__(self, websocket: WebSocketProtocol, leave_open: bool = False) -> None:
+        """
+        Create a WebSocket frame sink.
+
+        Args:
+            websocket: WebSocket connection to write to
+            leave_open: If True, doesn't close WebSocket on disposal
+        """
+        if websocket is None:
+            raise ValueError("websocket cannot be None")
+        self._websocket = websocket
+        self._leave_open = leave_open
+        self._closed = False
+
+    def write_frame(self, frame_data: bytes) -> None:
+        """
+        Write a complete frame to the WebSocket synchronously.
+
+        Note: WebSocket is inherently async, so this runs the async
+        version in a new event loop. Prefer write_frame_async for better performance.
+
+        Args:
+            frame_data: Complete frame data to write
+
+        Raises:
+            RuntimeError: If sink is closed or WebSocket is not open
+        """
+        if self._closed:
+            raise RuntimeError("WebSocketFrameSink is closed")
+
+        # Use asyncio.run for simplicity (creates new event loop)
+        try:
+            loop = asyncio.get_running_loop()
+            # If we're already in an async context, use run_coroutine_threadsafe
+            future = asyncio.run_coroutine_threadsafe(self.write_frame_async(frame_data), loop)
+            future.result()
+        except RuntimeError:
+            # No running event loop, create a new one
+            asyncio.run(self.write_frame_async(frame_data))
+
+    async def write_frame_async(self, frame_data: bytes) -> None:
+        """
+        Write a complete frame to the WebSocket asynchronously.
+
+        Args:
+            frame_data: Complete frame data to write
+
+        Raises:
+            RuntimeError: If sink is closed or WebSocket is not open
+        """
+        if self._closed:
+            raise RuntimeError("WebSocketFrameSink is closed")
+
+        if self._websocket.closed:
+            raise RuntimeError(f"WebSocket is not open: {self._websocket.state.name}")
+
+        # Send as single binary message
+        await self._websocket.send(frame_data)
+
+    def flush(self) -> None:
+        """
+        Flush any buffered data.
+
+        WebSocket sends immediately, so this is a no-op.
+        """
+        pass
+
+    async def flush_async(self) -> None:
+        """
+        Flush any buffered data asynchronously.
+
+        WebSocket sends immediately, so this is a no-op.
+        """
+        pass
+
+    def close(self) -> None:
+        """Close the sink and release resources."""
+        if self._closed:
+            return
+        self._closed = True
+
+        if not self._leave_open and not self._websocket.closed:
+            try:
+                loop = asyncio.get_running_loop()
+                future = asyncio.run_coroutine_threadsafe(self.close_async(), loop)
+                future.result()
+            except RuntimeError:
+                asyncio.run(self._close_websocket())
+
+    async def close_async(self) -> None:
+        """Close the sink and release resources asynchronously."""
+        if self._closed:
+            return
+        self._closed = True
+
+        await self._close_websocket()
+
+    async def _close_websocket(self) -> None:
+        """Internal async close helper."""
+        if not self._leave_open and not self._websocket.closed:
+            try:  # noqa: SIM105
+                await self._websocket.close()
+            except Exception:
+                # Best effort close
+                pass
+
+
+class WebSocketFrameSource(IFrameSource):
+    """
+    Frame source that reads from a WebSocket connection.
+
+    Each WebSocket binary message is treated as a complete frame.
+
+    Attributes:
+        leave_open: If True, doesn't close WebSocket on disposal
+    """
+
+    __slots__ = ("_closed", "_leave_open", "_websocket")
+
+    def __init__(self, websocket: WebSocketProtocol, leave_open: bool = False) -> None:
+        """
+        Create a WebSocket frame source.
+
+        Args:
+            websocket: WebSocket connection to read from
+            leave_open: If True, doesn't close WebSocket on disposal
+        """
+        if websocket is None:
+            raise ValueError("websocket cannot be None")
+        self._websocket = websocket
+        self._leave_open = leave_open
+        self._closed = False
+
+    @property
+    def has_more_frames(self) -> bool:
+        """
+        Check if more frames are available.
+
+        Returns True if the WebSocket is open or in CloseSent state.
+        """
+        return not self._websocket.closed
+
+    def read_frame(self) -> Optional[bytes]:
+        """
+        Read a complete frame from the WebSocket synchronously.
+
+        Note: WebSocket is inherently async, so this runs the async
+        version. Prefer read_frame_async for better performance.
+
+        Returns:
+            Complete frame data, or None if connection closed
+
+        Raises:
+            RuntimeError: If source is closed
+        """
+        if self._closed:
+            raise RuntimeError("WebSocketFrameSource is closed")
+
+        try:
+            loop = asyncio.get_running_loop()
+            future = asyncio.run_coroutine_threadsafe(self.read_frame_async(), loop)
+            return future.result()
+        except RuntimeError:
+            return asyncio.run(self.read_frame_async())
+
+    async def read_frame_async(self) -> Optional[bytes]:
+        """
+        Read a complete frame from the WebSocket asynchronously.
+
+        Returns:
+            Complete frame data, or None if connection closed
+
+        Raises:
+            RuntimeError: If source is closed
+            ValueError: If received non-binary message
+        """
+        if self._closed:
+            raise RuntimeError("WebSocketFrameSource is closed")
+
+        if not self.has_more_frames:
+            return None
+
+        try:
+            # websockets library handles message framing automatically
+            message = await self._websocket.recv()
+
+            # Ensure we got binary data
+            if isinstance(message, str):
+                raise ValueError("Expected binary message, got text")
+
+            return message
+
+        except Exception:
+            # Connection closed or error
+            return None
+
+    def close(self) -> None:
+        """Close the source and release resources."""
+        if self._closed:
+            return
+        self._closed = True
+
+        if not self._leave_open and not self._websocket.closed:
+            try:
+                loop = asyncio.get_running_loop()
+                future = asyncio.run_coroutine_threadsafe(self.close_async(), loop)
+                future.result()
+            except RuntimeError:
+                asyncio.run(self._close_websocket())
+
+    async def close_async(self) -> None:
+        """Close the source and release resources asynchronously."""
+        if self._closed:
+            return
+        self._closed = True
+
+        await self._close_websocket()
+
+    async def _close_websocket(self) -> None:
+        """Internal async close helper."""
+        if not self._leave_open and not self._websocket.closed:
+            try:  # noqa: SIM105
+                await self._websocket.close()
+            except Exception:
+                # Best effort close
+                pass
+
+
+async def connect_websocket_sink(url: str, leave_open: bool = False) -> WebSocketFrameSink:
+    """
+    Connect to a WebSocket server and return a frame sink.
+
+    Args:
+        url: WebSocket URL (ws:// or wss://)
+        leave_open: If True, doesn't close WebSocket on disposal
+
+    Returns:
+        Connected WebSocketFrameSink
+    """
+    try:
+        import websockets
+    except ImportError as e:
+        raise ImportError(
+            "websockets package is required for WebSocket transport. "
+            "Install with: pip install websockets"
+        ) from e
+
+    websocket = await websockets.connect(url)
+    return WebSocketFrameSink(websocket, leave_open)
+
+
+async def connect_websocket_source(url: str, leave_open: bool = False) -> WebSocketFrameSource:
+    """
+    Connect to a WebSocket server and return a frame source.
+
+    Args:
+        url: WebSocket URL (ws:// or wss://)
+        leave_open: If True, doesn't close WebSocket on disposal
+
+    Returns:
+        Connected WebSocketFrameSource
+    """
+    try:
+        import websockets
+    except ImportError as e:
+        raise ImportError(
+            "websockets package is required for WebSocket transport. "
+            "Install with: pip install websockets"
+        ) from e
+
+    websocket = await websockets.connect(url)
+    return WebSocketFrameSource(websocket, leave_open)