griptape-nodes 0.41.0__py3-none-any.whl → 0.43.0__py3-none-any.whl

griptape_nodes/machines/fsm.py

@@ -1,4 +1,4 @@
-from typing import Any, Generic, TypeVar
+from typing import Any, TypeVar
 
 T = TypeVar("T")
 
@@ -19,8 +19,13 @@ class State:
         """Called when exiting the state."""
         return
 
+    @staticmethod
+    def on_event(context: Any, event: Any) -> type["State"] | None:  # noqa: ARG004
+        """Called on an event, which may trigger a State transition."""
+        return None
+
 
-class FSM(Generic[T]):
+class FSM[T]:
     def __init__(self, context: T) -> None:
         self._context = context
         self._current_state = None
@@ -53,3 +58,12 @@ class FSM(Generic[T]):
 
         if new_state is not None:
             self.transition_state(new_state)
+
+    def handle_event(self, event: Any) -> None:
+        if self._current_state is None:
+            new_state = None
+        else:
+            new_state = self._current_state.on_event(self._context, event)
+
+        if new_state is not None:
+            self.transition_state(new_state)

griptape_nodes/machines/node_resolution.py

@@ -4,7 +4,7 @@ import logging
 from collections.abc import Generator
 from concurrent.futures import Future, ThreadPoolExecutor
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from griptape.events import EventBus
 from griptape.utils import with_contextvars
@@ -31,10 +31,6 @@ from griptape_nodes.retained_mode.events.parameter_events import (
     SetParameterValueRequest,
 )
 
-if TYPE_CHECKING:
-    from griptape_nodes.exe_types.flow import ControlFlow
-
-
 logger = logging.getLogger("griptape_nodes")
 
 
@@ -47,12 +43,10 @@ class Focus:
 
 # This is on a per-node basis
 class ResolutionContext:
-    flow: ControlFlow
     focus_stack: list[Focus]
     paused: bool
 
-    def __init__(self, flow: ControlFlow) -> None:
-        self.flow = flow
+    def __init__(self) -> None:
         self.focus_stack = []
         self.paused = False
 
@@ -90,7 +84,9 @@ class InitializeSpotlightState(State):
         if current_node.state == NodeResolutionState.UNRESOLVED:
             # Mark all future nodes unresolved.
             # TODO: https://github.com/griptape-ai/griptape-nodes/issues/862
-            context.flow.connections.unresolve_future_nodes(current_node)
+            from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+
+            GriptapeNodes.FlowManager().get_connections().unresolve_future_nodes(current_node)
             current_node.initialize_spotlight()
         # Set node to resolving - we are now resolving this node.
         current_node.state = NodeResolutionState.RESOLVING
@@ -132,7 +128,9 @@ class EvaluateParameterState(State):
     def on_update(context: ResolutionContext) -> type[State] | None:
         current_node = context.focus_stack[-1].node
         current_parameter = current_node.get_current_parameter()
-        connections = context.flow.connections
+        from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+
+        connections = GriptapeNodes.FlowManager().get_connections()
         if current_parameter is None:
             msg = "No current parameter set."
             raise ValueError(msg)
@@ -270,7 +268,10 @@ class ExecuteNodeState(State):
             )
             current_focus.process_generator = None
             current_focus.scheduled_value = None
-            context.flow.cancel_flow_run()
+
+            from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+
+            GriptapeNodes.FlowManager().cancel_flow_run()
 
             EventBus.publish_event(
                 ExecutionGriptapeNodeEvent(
@@ -320,7 +321,7 @@ class ExecuteNodeState(State):
                 )
             )
             # Pass the value through to the new nodes.
-            conn_output_nodes = context.flow.get_connected_output_parameters(current_node, parameter)
+            conn_output_nodes = GriptapeNodes.FlowManager().get_connected_output_parameters(current_node, parameter)
             for target_node, target_parameter in conn_output_nodes:
                 GriptapeNodes.get_instance().handle_request(
                     SetParameterValueRequest(
@@ -416,6 +417,7 @@ class ExecuteNodeState(State):
 
                 # Once we've passed on the scheduled value, we should clear it out just in case
                 current_focus.scheduled_value = None
+
                 future = ExecuteNodeState.executor.submit(with_contextvars(func))
                 future.add_done_callback(with_contextvars(on_future_done))
             except StopIteration:
@@ -445,8 +447,8 @@ class CompleteState(State):
 class NodeResolutionMachine(FSM[ResolutionContext]):
     """State machine for resolving node dependencies."""
 
-    def __init__(self, flow: ControlFlow) -> None:
-        resolution_context = ResolutionContext(flow)  # Gets the flow
+    def __init__(self) -> None:
+        resolution_context = ResolutionContext()
         super().__init__(resolution_context)
 
     def resolve_node(self, node: BaseNode) -> None:

griptape_nodes/mcp_server/__init__.py

@@ -0,0 +1 @@
+"""Runs the Griptape Nodes MCP server."""

griptape_nodes/mcp_server/server.py

@@ -0,0 +1,126 @@
+import contextlib
+import json
+import logging
+import os
+from collections.abc import AsyncIterator
+
+import uvicorn
+from fastapi import FastAPI
+from mcp.server.lowlevel import Server
+from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
+from mcp.types import (
+    TextContent,
+    Tool,
+)
+from pydantic import TypeAdapter
+from rich.logging import RichHandler
+from starlette.types import Receive, Scope, Send
+
+from griptape_nodes.mcp_server.ws_request_manager import AsyncRequestManager, WebSocketConnectionManager
+from griptape_nodes.retained_mode.events.base_events import RequestPayload
+from griptape_nodes.retained_mode.events.connection_events import (
+    CreateConnectionRequest,
+    DeleteConnectionRequest,
+    ListConnectionsForNodeRequest,
+)
+from griptape_nodes.retained_mode.events.flow_events import ListNodesInFlowRequest
+from griptape_nodes.retained_mode.events.node_events import (
+    CreateNodeRequest,
+    DeleteNodeRequest,
+    GetNodeMetadataRequest,
+    GetNodeResolutionStateRequest,
+    ListParametersOnNodeRequest,
+    SetNodeMetadataRequest,
+)
+from griptape_nodes.retained_mode.events.parameter_events import (
+    GetParameterValueRequest,
+    SetParameterValueRequest,
+)
+from griptape_nodes.retained_mode.managers.config_manager import ConfigManager
+from griptape_nodes.retained_mode.managers.secrets_manager import SecretsManager
+
+SUPPORTED_REQUEST_EVENTS: dict[str, type[RequestPayload]] = {
+    # Nodes
+    "CreateNodeRequest": CreateNodeRequest,
+    "DeleteNodeRequest": DeleteNodeRequest,
+    "ListNodesInFlowRequest": ListNodesInFlowRequest,
+    "GetNodeResolutionStateRequest": GetNodeResolutionStateRequest,
+    "GetNodeMetadataRequest": GetNodeMetadataRequest,
+    "SetNodeMetadataRequest": SetNodeMetadataRequest,
+    # Connections
+    "CreateConnectionRequest": CreateConnectionRequest,
+    "DeleteConnectionRequest": DeleteConnectionRequest,
+    "ListConnectionsForNodeRequest": ListConnectionsForNodeRequest,
+    # Parameters
+    "ListParametersOnNodeRequest": ListParametersOnNodeRequest,
+    "GetParameterValueRequest": GetParameterValueRequest,
+    "SetParameterValueRequest": SetParameterValueRequest,
+}
+
+GTN_MCP_SERVER_PORT = int(os.getenv("GTN_MCP_SERVER_PORT", "9927"))
+
+config_manager = ConfigManager()
+secrets_manager = SecretsManager(config_manager)
+
+
+def main(api_key: str) -> None:
+    """Main entry point for the Griptape Nodes MCP server."""
+    mcp_server_logger = logging.getLogger("griptape_nodes_mcp_server")
+    mcp_server_logger.addHandler(RichHandler(show_time=True, show_path=False, markup=True, rich_tracebacks=True))
+    mcp_server_logger.setLevel(logging.INFO)
+    mcp_server_logger.info("Starting MCP GTN server...")
+
+    # Give these a session ID
+    connection_manager = WebSocketConnectionManager()
+    request_manager = AsyncRequestManager(connection_manager, api_key)
+
+    app = Server("mcp-gtn")
+
+    @app.list_tools()
+    async def list_tools() -> list[Tool]:
+        return [
+            Tool(name=event.__name__, description=event.__doc__, inputSchema=TypeAdapter(event).json_schema())
+            for (name, event) in SUPPORTED_REQUEST_EVENTS.items()
+        ]
+
+    @app.call_tool()
+    async def call_tool(name: str, arguments: dict) -> list[TextContent]:
+        if name not in SUPPORTED_REQUEST_EVENTS:
+            msg = f"Unsupported tool: {name}"
+            raise ValueError(msg)
+
+        request_payload = SUPPORTED_REQUEST_EVENTS[name](**arguments)
+
+        await request_manager.connect()
+        result = await request_manager.create_request_event(
+            request_payload.__class__.__name__, request_payload.__dict__, timeout_ms=5000
+        )
+        mcp_server_logger.debug("Got result: %s", result)
+
+        return [TextContent(type="text", text=json.dumps(result))]
+
+    # Create the session manager with our app and event store
+    session_manager = StreamableHTTPSessionManager(
+        app=app,
+    )
+
+    @contextlib.asynccontextmanager
+    async def lifespan(_: FastAPI) -> AsyncIterator[None]:
+        """Context manager for managing session manager lifecycle."""
+        async with session_manager.run():
+            mcp_server_logger.info("GTN MCP server started with StreamableHTTP session manager!")
+            try:
+                yield
+            finally:
+                mcp_server_logger.info("GTN MCP server shutting down...")
+
+    # Create an ASGI application using the transport
+    mcp_server_app = FastAPI(lifespan=lifespan)
+
+    # ASGI handler for streamable HTTP connections
+    async def handle_streamable_http(scope: Scope, receive: Receive, send: Send) -> None:
+        await session_manager.handle_request(scope, receive, send)
+
+    mcp_server_app.mount("/mcp", app=handle_streamable_http)
+
+    uvicorn.run(mcp_server_app, host="127.0.0.1", port=GTN_MCP_SERVER_PORT)

griptape_nodes/mcp_server/ws_request_manager.py

@@ -0,0 +1,268 @@
+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("🔴 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("✅ 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("❌ 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

griptape_nodes/node_library/workflow_registry.py

@@ -13,19 +13,19 @@ from griptape_nodes.utils.metaclasses import SingletonMeta
 
 
 class WorkflowMetadata(BaseModel):
-    LATEST_SCHEMA_VERSION: ClassVar[str] = "0.4.0"
+    LATEST_SCHEMA_VERSION: ClassVar[str] = "0.6.1"
 
     name: str
     schema_version: str
     engine_version_created_with: str
     node_libraries_referenced: list[LibraryNameAndVersion]
+    workflows_referenced: list[str] | None = None
     description: str | None = None
     image: str | None = None
     is_griptape_provided: bool | None = False
     is_template: bool | None = False
     creation_date: datetime | None = Field(default=None)
     last_modified_date: datetime | None = Field(default=None)
-    published_workflow_id: str | None = Field(default=None)
 
 
 class WorkflowRegistry(metaclass=SingletonMeta):