griptape-nodes 0.57.0__py3-none-any.whl → 0.58.0__py3-none-any.whl

griptape_nodes/api_client/__init__.py

@@ -0,0 +1,9 @@
+"""API client for Nodes API communication."""
+
+from griptape_nodes.api_client.client import Client
+from griptape_nodes.api_client.request_client import RequestClient
+
+__all__ = [
+    "Client",
+    "RequestClient",
+]

griptape_nodes/api_client/client.py

@@ -0,0 +1,279 @@
+"""Unified WebSocket client for Nodes API communication."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+import logging
+import os
+from typing import TYPE_CHECKING, Any, Self
+from urllib.parse import urljoin
+
+from websockets.asyncio.client import connect
+from websockets.exceptions import ConnectionClosed
+
+from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+    from types import TracebackType
+
+logger = logging.getLogger("griptape_nodes_client")
+
+
+def get_default_websocket_url() -> str:
+    """Get the default WebSocket endpoint URL for connecting to Nodes API.
+
+    Returns:
+        WebSocket URL for Nodes API events endpoint
+    """
+    return urljoin(
+        os.getenv("GRIPTAPE_NODES_API_BASE_URL", "https://api.nodes.griptape.ai").replace("http", "ws"),
+        "/ws/engines/events?version=v2",
+    )
+
+
+class Client:
+    """WebSocket client for Nodes API pub/sub communication.
+
+    Provides connection management, topic-based pub/sub, and message routing.
+    Handles WebSocket reconnection and async event streaming.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        url: str | None = None,
+    ):
+        """Initialize Nodes API client.
+
+        Args:
+            api_key: API key for authentication (defaults to GT_CLOUD_API_KEY from SecretsManager)
+            url: WebSocket URL to connect to (defaults to Nodes API endpoint)
+        """
+        self.url = url if url is not None else get_default_websocket_url()
+
+        # Get API key from SecretsManager if not provided
+        if api_key is None:
+            api_key = GriptapeNodes.SecretsManager().get_secret("GT_CLOUD_API_KEY")
+
+        self.api_key = api_key
+
+        self.headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
+
+        # Event streaming management
+        self._message_queue: asyncio.Queue = asyncio.Queue()
+        self._subscribed_topics: set[str] = set()
+        self._receiving_task: asyncio.Task | None = None
+        self._sending_task: asyncio.Task | None = None
+        self._websocket: Any = None
+        self._connection_ready = asyncio.Event()
+        self._reconnect_delay = 2.0
+
+    async def __aenter__(self) -> Self:
+        """Async context manager entry: connect to WebSocket server."""
+        await self._connect()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Async context manager exit: disconnect from WebSocket server."""
+        await self._disconnect()
+
+    def __aiter__(self) -> AsyncIterator[dict[str, Any]]:
+        """Return self as async iterator."""
+        return self
+
+    async def __anext__(self) -> dict[str, Any]:
+        """Get next message from the message queue.
+
+        Returns:
+            Next message dictionary from subscribed topics
+
+        Raises:
+            StopAsyncIteration: When iteration is cancelled
+        """
+        try:
+            return await self._message_queue.get()
+        except asyncio.CancelledError:
+            raise StopAsyncIteration from None
+
+    @property
+    def messages(self) -> AsyncIterator[dict[str, Any]]:
+        """Async iterator for receiving messages from subscribed topics.
+
+        Returns:
+            Async iterator yielding message dictionaries
+
+        Example:
+            async with Client(...) as client:
+                await client.subscribe("topic")
+                async for message in client.messages:
+                    print(message)
+        """
+        return self
+
+    async def subscribe(self, topic: str) -> None:
+        """Subscribe to a topic by sending subscribe command to server.
+
+        Args:
+            topic: Topic name to subscribe to
+
+        Example:
+            await client.subscribe("sessions/123/response")
+        """
+        self._subscribed_topics.add(topic)
+        await self._send_subscribe_command(topic)
+
+    async def unsubscribe(self, topic: str) -> None:
+        """Unsubscribe from a topic.
+
+        Args:
+            topic: Topic name to unsubscribe from
+        """
+        self._subscribed_topics.discard(topic)
+        await self._send_unsubscribe_command(topic)
+
+    async def publish(self, event_type: str, payload: dict[str, Any], topic: str) -> None:
+        """Publish an event to the server.
+
+        Args:
+            event_type: Type of event to publish
+            payload: Event payload data
+            topic: Topic to publish to
+        """
+        message = {"type": event_type, "payload": payload, "topic": topic}
+        await self._send_message(message)
+
+    async def _connect(self) -> None:
+        """Connect to the WebSocket server and start receiving messages.
+
+        This method starts the connection manager task.
+        It returns once the initial connection is established.
+
+        Raises:
+            ConnectionError: If connection fails
+        """
+        # Start connection manager task
+        self._receiving_task = asyncio.create_task(self._manage_connection())
+
+        # Wait for initial connection to be established
+        try:
+            await asyncio.wait_for(self._connection_ready.wait(), timeout=10.0)
+            logger.debug("WebSocket client connected")
+        except TimeoutError as e:
+            logger.error("Failed to connect WebSocket client: timeout")
+            msg = "Connection timeout"
+            raise ConnectionError(msg) from e
+
+    async def _disconnect(self) -> None:
+        """Disconnect from the WebSocket server and clean up tasks."""
+        # Cancel tasks
+        if self._receiving_task:
+            self._receiving_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._receiving_task
+
+        if self._sending_task:
+            self._sending_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._sending_task
+
+        # Close websocket connection
+        if self._websocket:
+            await self._websocket.close()
+        logger.info("WebSocket client disconnected")
+
+    async def _manage_connection(self) -> None:
+        """Manage WebSocket connection lifecycle with automatic reconnection.
+
+        This method establishes and maintains the WebSocket connection,
+        automatically reconnecting on failures.
+        """
+        try:
+            async for websocket in connect(self.url, additional_headers=self.headers):
+                self._websocket = websocket
+                self._connection_ready.set()
+                logger.debug("WebSocket connection established: %s", self.url)
+
+                # Resubscribe to all topics after reconnection
+                if self._subscribed_topics:
+                    logger.debug("Resubscribing to %d topics after reconnection", len(self._subscribed_topics))
+                    for topic in self._subscribed_topics:
+                        await self._send_subscribe_command(topic)
+
+                try:
+                    await self._receive_messages(websocket)
+                except ConnectionClosed:
+                    logger.info("WebSocket connection closed, reconnecting...")
+                    self._connection_ready.clear()
+                    continue
+
+        except asyncio.CancelledError:
+            logger.debug("Connection manager task cancelled")
+
+    async def _receive_messages(self, websocket: Any) -> None:
+        """Receive messages from WebSocket and put them in message queue.
+
+        Args:
+            websocket: WebSocket connection to receive messages from
+
+        Raises:
+            ConnectionClosed: When the WebSocket connection is closed
+        """
+        try:
+            async for message in websocket:
+                try:
+                    data = json.loads(message)
+                    await self._message_queue.put(data)
+                except json.JSONDecodeError:
+                    logger.error("Failed to parse message: %s", message)
+                except Exception as e:
+                    logger.error("Error receiving message: %s", e)
+        except asyncio.CancelledError:
+            logger.debug("Receive messages task cancelled")
+            raise
+
+    async def _send_message(self, message: dict[str, Any]) -> None:
+        """Send a message through the WebSocket connection.
+
+        Args:
+            message: Message dictionary to send
+
+        Raises:
+            ConnectionError: If not connected
+        """
+        if not self._websocket:
+            msg = "Not connected to WebSocket"
+            raise ConnectionError(msg)
+
+        try:
+            await self._websocket.send(json.dumps(message))
+            logger.debug("Sent message type: %s", message.get("type"))
+        except Exception as e:
+            logger.error("Failed to send message: %s", e)
+
+    async def _send_subscribe_command(self, topic: str) -> None:
+        """Send subscribe command to server.
+
+        Args:
+            topic: Topic to subscribe to
+        """
+        message = {"type": "subscribe", "topic": topic, "payload": {}}
+        await self._send_message(message)
+        logger.debug("Sent subscribe command for topic: %s", topic)
+
+    async def _send_unsubscribe_command(self, topic: str) -> None:
+        """Send unsubscribe command to server.
+
+        Args:
+            topic: Topic to unsubscribe from
+        """
+        message = {"type": "unsubscribe", "topic": topic, "payload": {}}
+        await self._send_message(message)
+        logger.debug("Sent unsubscribe command for topic: %s", topic)

griptape_nodes/api_client/request_client.py

@@ -0,0 +1,273 @@
+"""Request/response tracking with futures and timeouts."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+import uuid
+from typing import TYPE_CHECKING, Any, Self, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from types import TracebackType
+
+    from griptape_nodes.api_client.client import Client
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+class RequestClient:
+    """Request/response client built on top of Client.
+
+    Wraps a Client to provide request/response semantics on top of
+    pub/sub messaging. Tracks pending requests by request_id and resolves/rejects
+    futures when responses arrive. Supports timeouts for requests that don't
+    receive responses.
+    """
+
+    def __init__(
+        self,
+        client: Client,
+        request_topic_fn: Callable[[], str] | None = None,
+        response_topic_fn: Callable[[], str] | None = None,
+    ) -> None:
+        """Initialize request/response client.
+
+        Args:
+            client: Client instance to use for communication
+            request_topic_fn: Function to determine request topic (defaults to "request")
+            response_topic_fn: Function to determine response topic (defaults to "response")
+        """
+        self.client = client
+        self.request_topic_fn = request_topic_fn or (lambda: "request")
+        self.response_topic_fn = response_topic_fn or (lambda: "response")
+
+        # Map of request_id -> Future that will be resolved when response arrives
+        self._pending_requests: dict[str, asyncio.Future] = {}
+        self._lock = asyncio.Lock()
+
+        # Track subscribed response topics
+        self._subscribed_response_topics: set[str] = set()
+
+        # Background task for listening to responses
+        self._response_listener_task: asyncio.Task | None = None
+
+    async def __aenter__(self) -> Self:
+        """Async context manager entry: start response listener."""
+        self._response_listener_task = asyncio.create_task(self._listen_for_responses())
+        logger.debug("RequestClient started")
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Async context manager exit: stop response listener."""
+        if self._response_listener_task:
+            self._response_listener_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._response_listener_task
+        logger.debug("RequestClient stopped")
+
+    async def request(
+        self, request_type: str, payload: dict[str, Any], timeout_ms: int | None = None
+    ) -> dict[str, Any]:
+        """Send a request and wait for its response.
+
+        This method automatically:
+        - Generates a request_id
+        - Determines request and response topics
+        - Subscribes to response topic if needed
+        - Sends the request
+        - Waits for and returns the response
+
+        Args:
+            request_type: Type of request to send
+            payload: Request payload data
+            timeout_ms: Optional timeout in milliseconds
+
+        Returns:
+            Response data from the server
+
+        Raises:
+            TimeoutError: If request times out
+            Exception: If request fails
+        """
+        # Generate request ID and track it
+        request_id = str(uuid.uuid4())
+        payload["request_id"] = request_id
+
+        response_future = await self._track_request(request_id)
+
+        # Determine topics
+        request_topic = self.request_topic_fn()
+        response_topic = self.response_topic_fn()
+
+        # Subscribe to response topic if not already subscribed
+        if response_topic not in self._subscribed_response_topics:
+            await self.client.subscribe(response_topic)
+            self._subscribed_response_topics.add(response_topic)
+
+        # Send the request as an EventRequest
+        event_payload = {"event_type": "EventRequest", "request_type": request_type, "request": payload}
+
+        logger.debug("Sending request %s: %s", request_id, request_type)
+
+        try:
+            await self.client.publish("EventRequest", event_payload, request_topic)
+
+            # Wait for response with optional timeout
+            if timeout_ms:
+                timeout_sec = timeout_ms / 1000
+                result = await asyncio.wait_for(response_future, timeout=timeout_sec)
+            else:
+                result = await response_future
+
+        except TimeoutError:
+            logger.error("Request %s timed out", request_id)
+            await self._cancel_request(request_id)
+            raise
+
+        except Exception as e:
+            logger.error("Request %s failed: %s", request_id, e)
+            await self._cancel_request(request_id)
+            raise
+        else:
+            logger.debug("Request %s completed successfully", request_id)
+            return result
+
+    async def _track_request(self, request_id: str) -> asyncio.Future:
+        """Start tracking a request and return a future that will be resolved on response.
+
+        Args:
+            request_id: Unique identifier for this request
+
+        Returns:
+            Future that will be resolved when response arrives
+
+        Raises:
+            ValueError: If request_id is already being tracked
+        """
+        async with self._lock:
+            if request_id in self._pending_requests:
+                msg = f"Request ID already exists: {request_id}"
+                raise ValueError(msg)
+
+            future: asyncio.Future = asyncio.Future()
+            self._pending_requests[request_id] = future
+            logger.debug("Tracking request: %s", request_id)
+            return future
+
+    async def _resolve_request(self, request_id: str, result: Any) -> None:
+        """Mark a request as successful and resolve its future with a result.
+
+        Args:
+            request_id: Request identifier
+            result: Result data to return to the requester
+        """
+        async with self._lock:
+            future = self._pending_requests.pop(request_id, None)
+
+            if future is None:
+                logger.warning("Received response for unknown request: %s", request_id)
+                return
+
+            if not future.done():
+                future.set_result(result)
+                logger.debug("Resolved request: %s", request_id)
+
+    async def _reject_request(self, request_id: str, error: Exception) -> None:
+        """Mark a request as failed and reject its future with an exception.
+
+        Args:
+            request_id: Request identifier
+            error: Exception to raise for the requester
+        """
+        async with self._lock:
+            future = self._pending_requests.pop(request_id, None)
+
+            if future is None:
+                logger.warning("Received error for unknown request: %s", request_id)
+                return
+
+            if not future.done():
+                future.set_exception(error)
+                logger.debug("Rejected request: %s with error: %s", request_id, error)
+
+    async def _cancel_request(self, request_id: str) -> None:
+        """Cancel a pending request and clean up its tracking.
+
+        Args:
+            request_id: Request identifier
+        """
+        async with self._lock:
+            future = self._pending_requests.pop(request_id, None)
+
+            if future is None:
+                logger.debug("Request already completed or unknown: %s", request_id)
+                return
+
+            if not future.done():
+                future.cancel()
+                logger.debug("Cancelled request: %s", request_id)
+
+    @property
+    def pending_count(self) -> int:
+        """Get number of currently pending requests.
+
+        Returns:
+            Count of pending requests
+        """
+        return len(self._pending_requests)
+
+    @property
+    def pending_request_ids(self) -> list[str]:
+        """Get list of all pending request IDs.
+
+        Returns:
+            List of request_id strings
+        """
+        return list(self._pending_requests.keys())
+
+    async def _listen_for_responses(self) -> None:
+        """Listen for response messages from subscribed topics."""
+        try:
+            async for message in self.client.messages:
+                try:
+                    await self._handle_response(message)
+                except Exception as e:
+                    logger.error("Error handling response message: %s", e)
+        except asyncio.CancelledError:
+            logger.debug("Response listener cancelled")
+            raise
+
+    async def _handle_response(self, message: dict[str, Any]) -> None:
+        """Handle response messages by resolving tracked requests.
+
+        Args:
+            message: WebSocket message containing response
+        """
+        message_type = message.get("type")
+
+        # Only handle success/failure result messages
+        if message_type not in ("success_result", "failure_result"):
+            return
+
+        payload = message.get("payload", {})
+        request_id = payload.get("request", {}).get("request_id")
+
+        if not request_id:
+            logger.debug("Response message has no request_id")
+            return
+
+        if message_type == "success_result":
+            result = payload.get("result", "Success")
+            await self._resolve_request(request_id, result)
+        else:
+            error_msg = payload.get("result", {}).get("exception", "Unknown error") or "Unknown error"
+            await self._reject_request(request_id, Exception(error_msg))