griptape-nodes 0.57.1__py3-none-any.whl → 0.58.1__py3-none-any.whl

griptape_nodes/servers/ws_request_manager.py

@@ -1,268 +0,0 @@
-import asyncio
-import contextlib
-import json
-import logging
-import os
-import uuid
-from collections.abc import Callable
-from typing import Any, Generic, TypeVar
-from urllib.parse import urljoin
-
-import websockets
-
-logger = logging.getLogger("griptape_nodes_mcp_server")
-
-T = TypeVar("T")
-
-
-class WebSocketConnectionManager:
-    """Python equivalent of WebSocketConnectionManager in TypeScript."""
-
-    def __init__(
-        self,
-        websocket_url: str = urljoin(
-            os.getenv("GRIPTAPE_NODES_API_BASE_URL", "https://api.nodes.griptape.ai").replace("http", "ws"),
-            "/ws/engines/events?version=v2",
-        ),
-    ):
-        self.websocket_url = websocket_url
-        self.websocket: Any = None
-        self.connected = False
-        self.event_handlers: dict[str, list[Callable]] = {}
-        self.request_handlers: dict[str, tuple[Callable, Callable]] = {}
-        self._process_task: asyncio.Task | None = None
-
-    async def send(self, data: dict[str, Any]) -> None:
-        """Send a message to the WebSocket server."""
-        if not self.websocket:
-            msg = "Not connected to WebSocket server"
-            raise ConnectionError(msg)
-
-        try:
-            message = json.dumps(data)
-            await self.websocket.send(message)
-            logger.debug("Sent message: %s", message)
-        except Exception as e:
-            logger.error("Failed to send message: %s", e)
-            raise
-
-    async def _process_messages(self) -> None:
-        """Process incoming WebSocket messages."""
-        if not self.websocket:
-            logger.warning("WebSocket is not connected, cannot process messages")
-            return
-
-        try:
-            async for message in self.websocket:
-                try:
-                    data = json.loads(message)
-                    logger.debug("📥 Received message: %s", message)
-                    await self._handle_message(data)
-                except json.JSONDecodeError:
-                    logger.error("Failed to parse message: %s", message)
-                except Exception as e:
-                    logger.error("Error processing message: %s", e)
-        except websockets.exceptions.ConnectionClosed:
-            logger.warning("WebSocket connection closed")
-            self.connected = False
-        except asyncio.CancelledError:
-            # Task was cancelled, just exit
-            pass
-        except Exception as e:
-            logger.error("Error in message processing loop: %s", e)
-            self.connected = False
-
-    async def _handle_message(self, data: dict[str, Any]) -> None:
-        request = data.get("payload", {}).get("request", {})
-        request_id = request.get("request_id")
-
-        if request_id and request_id in self.request_handlers:
-            success_handler, failure_handler = self.request_handlers[request_id]
-            try:
-                if data.get("type") == "success_result":
-                    success_handler(data, request)
-                else:
-                    failure_handler(data, request)
-            except Exception as e:
-                logger.error("Error in request handler: %s", e)
-
-    def subscribe_to_request_event(
-        self, success_handler: Callable[[Any, Any], None], failure_handler: Callable[[Any, Any], None]
-    ) -> str:
-        """Subscribe to a request-response event."""
-        request_id = str(uuid.uuid4())
-        self.request_handlers[request_id] = (success_handler, failure_handler)
-        return request_id
-
-    def unsubscribe_from_request_event(self, request_id: str) -> None:
-        """Unsubscribe from a request-response event."""
-        if request_id in self.request_handlers:
-            del self.request_handlers[request_id]
-
-
-class AsyncRequestManager(Generic[T]):  # noqa: UP046
-    def __init__(self, connection_manager: WebSocketConnectionManager, api_key: str):
-        self.connection_manager = connection_manager
-        self.api_key = api_key
-
-    async def _subscribe_to_topic(self, topic: str) -> None:
-        """Subscribe to a specific topic in the message bus."""
-        if self.connection_manager.websocket is None:
-            logger.warning("WebSocket connection not available for subscribing to topic")
-            return
-
-        try:
-            body = {"type": "subscribe", "topic": topic, "payload": {}}
-            await self.connection_manager.websocket.send(json.dumps(body))
-            logger.debug("Subscribed to topic: %s", topic)
-        except websockets.exceptions.WebSocketException as e:
-            logger.error("Error subscribing to topic %s: %s", topic, e)
-        except Exception as e:
-            logger.error("Unexpected error while subscribing to topic %s: %s", topic, e)
-
-    async def _subscribe_to_topics(self) -> None:
-        from griptape_nodes.retained_mode.managers.session_manager import SessionManager
-        from griptape_nodes.retained_mode.utils.engine_identity import EngineIdentity
-
-        # Subscribe to response topic (engine discovery)
-        await self._subscribe_to_topic("response")
-
-        # Get engine ID and subscribe to engine_id/response
-        engine_id = EngineIdentity.get_engine_id()
-        if engine_id:
-            await self._subscribe_to_topic(f"engines/{engine_id}/response")
-        else:
-            logger.warning("Engine ID not available for subscription")
-
-        # Get session ID and subscribe to session_id/response if available
-        session_id = SessionManager.get_saved_session_id()
-        if session_id:
-            topic = f"sessions/{session_id}/response"
-            await self._subscribe_to_topic(topic)
-        else:
-            logger.info("No session ID available for subscription")
-
-    async def connect(self) -> None:
-        """Connect to the WebSocket server."""
-        from griptape_nodes.app.app import _create_websocket_connection
-
-        try:
-            self.connection_manager.websocket = await _create_websocket_connection(self.api_key)
-            logger.debug("🟢 WebSocket connection established: %s", self.connection_manager.websocket)
-
-            await self._subscribe_to_topics()
-
-            # Start processing messages
-            self.connection_manager._process_task = asyncio.create_task(self.connection_manager._process_messages())
-
-        except Exception as e:
-            self.connection_manager.connected = False
-            logger.error("[red]X[/red] WebSocket connection failed: %s", str(e))
-            msg = f"Failed to connect to WebSocket: {e!s}"
-            raise ConnectionError(msg) from e
-
-    async def disconnect(self) -> None:
-        """Disconnect from the WebSocket server."""
-        if self.connection_manager.websocket:
-            await self.connection_manager.websocket.close()
-            self.connection_manager.websocket = None
-        self.connection_manager.connected = False
-
-        # Cancel processing task if it's running
-        if self.connection_manager._process_task:
-            self.connection_manager._process_task.cancel()
-            with contextlib.suppress(asyncio.CancelledError):
-                await self.connection_manager._process_task
-            self.connection_manager._process_task = None
-
-        logger.debug("WebSocket disconnected")
-
-    async def send_api_message(self, data: dict[str, Any]) -> None:
-        """Send a message to the API via WebSocket."""
-        try:
-            await self.connection_manager.send(data)
-        except ConnectionError as e:
-            logger.error("Failed to send API message: %s", e)
-            raise
-        except Exception as e:
-            logger.error("Unexpected error sending API message: %s", e)
-            raise
-
-    async def create_event(self, request_type: str, payload: dict[str, Any]) -> None:
-        """Send an event to the API without waiting for a response."""
-        from griptape_nodes.app.app import determine_request_topic
-
-        logger.debug("Creating Event: %s - %s", request_type, json.dumps(payload))
-
-        data = {"event_type": "EventRequest", "request_type": request_type, "request": payload}
-        topic = determine_request_topic()
-
-        request_data = {"payload": data, "type": data["event_type"], "topic": topic}
-
-        if not request_data["payload"]["request"].get("request_id"):
-            request_data["payload"]["request"]["request_id"] = ""
-
-        await self.send_api_message(request_data)
-
-    async def create_request_event(
-        self, request_type: str, payload: dict[str, Any], timeout_ms: int | None = None
-    ) -> T:
-        """Send a request and wait for its response.
-
-        Args:
-            request_type: Type of request to send
-            payload: Data to send with the request
-            timeout_ms: Optional timeout in milliseconds
-
-        Returns:
-            The response data
-
-        Raises:
-            asyncio.TimeoutError: If the request times out
-            Exception: If the request fails
-        """
-        # Create a future that will be resolved when the response arrives
-        response_future = asyncio.Future()
-
-        # Convert timeout from milliseconds to seconds for asyncio
-        timeout_sec = timeout_ms / 1000 if timeout_ms else None
-
-        # Define handlers that will resolve/reject the future
-        def success_handler(response: Any, _: Any) -> None:
-            if not response_future.done():
-                result = response.get("payload", {}).get("result", "Success")
-                logger.debug("[green]OK[/green] Request succeeded: %s", result)
-                response_future.set_result(result)
-
-        def failure_handler(response: Any, _: Any) -> None:
-            if not response_future.done():
-                error = (
-                    response.get("payload", {}).get("result", {}).get("exception", "Unknown error") or "Unknown error"
-                )
-                logger.error("[red]X[/red] Request failed: %s", error)
-                response_future.set_exception(Exception(error))
-
-        # Generate request ID and subscribe
-        request_id = self.connection_manager.subscribe_to_request_event(success_handler, failure_handler)
-        payload["request_id"] = request_id
-
-        logger.debug("Request (%s): %s %s", request_id, request_type, json.dumps(payload))
-
-        try:
-            # Send the event
-            await self.create_event(request_type, payload)
-
-            # Wait for the response with optional timeout
-            if timeout_sec:
-                return await asyncio.wait_for(response_future, timeout=timeout_sec)
-            return await response_future
-
-        except TimeoutError:
-            logger.error("Request timed out after %s ms: %s", timeout_ms, request_id)
-            self.connection_manager.unsubscribe_from_request_event(request_id)
-            raise
-
-        except Exception as e:
-            logger.error("Request failed: %s - %s", request_id, e)
-            self.connection_manager.unsubscribe_from_request_event(request_id)
-            raise