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

griptape_nodes/bootstrap/workflow_executors/subprocess_workflow_executor.py

@@ -0,0 +1,326 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import tempfile
+import threading
+import uuid
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Self
+
+import anyio
+from websockets.exceptions import ConnectionClosed, ConnectionClosedError
+
+from griptape_nodes.app.app import _create_websocket_connection, _send_subscribe_command
+from griptape_nodes.bootstrap.utils.python_subprocess_executor import PythonSubprocessExecutor
+from griptape_nodes.bootstrap.workflow_executors.local_session_workflow_executor import LocalSessionWorkflowExecutor
+from griptape_nodes.drivers.storage import StorageBackend
+from griptape_nodes.retained_mode.events.base_events import (
+    EventResultFailure,
+    EventResultSuccess,
+    ExecutionEvent,
+    ResultPayload,
+)
+from griptape_nodes.retained_mode.events.execution_events import (
+    ControlFlowCancelledEvent,
+    ControlFlowResolvedEvent,
+    StartFlowRequest,
+)
+from griptape_nodes.retained_mode.events.payload_registry import PayloadRegistry
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from types import TracebackType
+
+logger = logging.getLogger(__name__)
+
+
+class SubprocessWorkflowExecutorError(Exception):
+    """Exception raised during subprocess workflow execution."""
+
+
+class SubprocessWorkflowExecutor(LocalSessionWorkflowExecutor, PythonSubprocessExecutor):
+    def __init__(
+        self,
+        workflow_path: str,
+        on_start_flow_result: Callable[[ResultPayload], None] | None = None,
+        session_id: str | None = None,
+    ) -> None:
+        PythonSubprocessExecutor.__init__(self)
+        self._workflow_path = workflow_path
+        self._on_start_flow_result = on_start_flow_result
+        # Generate a unique session ID for this execution
+        self._session_id = session_id or uuid.uuid4().hex
+        self._websocket_thread: threading.Thread | None = None
+        self._websocket_event_loop: asyncio.AbstractEventLoop | None = None
+        self._websocket_event_loop_ready = threading.Event()
+        self._event_handlers: dict[str, list] = {}
+        self._shutdown_event: asyncio.Event | None = None
+        self._stored_exception: SubprocessWorkflowExecutorError | None = None
+
+    async def __aenter__(self) -> Self:
+        """Async context manager entry: start WebSocket connection."""
+        logger.info("Starting WebSocket listener for session %s", self._session_id)
+        await self._start_websocket_listener()
+        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 WebSocket connection."""
+        logger.info("Stopping WebSocket listener for session %s", self._session_id)
+        self._stop_websocket_listener()
+
+    async def arun(
+        self,
+        workflow_name: str,  # noqa: ARG002
+        flow_input: Any,
+        storage_backend: StorageBackend = StorageBackend.LOCAL,
+        **kwargs: Any,  # noqa: ARG002
+    ) -> None:
+        """Execute a workflow in a subprocess and wait for completion."""
+        script_path = Path(__file__).parent / "utils" / "subprocess_script.py"
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_workflow_path = Path(tmpdir) / "workflow.py"
+            tmp_script_path = Path(tmpdir) / "subprocess_script.py"
+
+            try:
+                async with (
+                    await anyio.open_file(self._workflow_path, "rb") as src,
+                    await anyio.open_file(tmp_workflow_path, "wb") as dst,
+                ):
+                    await dst.write(await src.read())
+
+                async with (
+                    await anyio.open_file(script_path, "rb") as src,
+                    await anyio.open_file(tmp_script_path, "wb") as dst,
+                ):
+                    await dst.write(await src.read())
+            except Exception as e:
+                msg = f"Failed to copy workflow or script to temp directory: {e}"
+                logger.exception(msg)
+                raise SubprocessWorkflowExecutorError(msg) from e
+
+            args = [
+                "--json-input",
+                json.dumps(flow_input),
+                "--session-id",
+                self._session_id,
+                "--storage-backend",
+                storage_backend.value,
+                "--workflow-path",
+                str(tmp_workflow_path),
+            ]
+
+            try:
+                await self.execute_python_script(
+                    script_path=tmp_script_path,
+                    args=args,
+                    cwd=Path(tmpdir),
+                )
+            except Exception as e:
+                msg = f"Failed to execute subprocess script: {e}"
+                logger.exception(msg)
+                raise SubprocessWorkflowExecutorError(msg) from e
+            finally:
+                # Check if an exception was stored coming from the WebSocket
+                if self._stored_exception:
+                    raise self._stored_exception
+
+    async def _start_websocket_listener(self) -> None:
+        """Start WebSocket connection to listen for events from the subprocess."""
+        logger.info("Starting WebSocket listener for session %s", self._session_id)
+        api_key = self._get_api_key()
+        if api_key is None:
+            logger.warning("No API key found, WebSocket listener will not be started")
+            return
+
+        logger.info("API key found, starting WebSocket listener thread")
+        self._websocket_thread = threading.Thread(target=self._start_websocket_thread, args=(api_key,), daemon=True)
+        self._websocket_thread.start()
+
+        if self._websocket_event_loop_ready.wait(timeout=10):
+            logger.info("WebSocket listener thread ready")
+        else:
+            logger.error("Timeout waiting for WebSocket listener thread to start")
+
+    def _stop_websocket_listener(self) -> None:
+        """Stop the WebSocket listener thread."""
+        if self._websocket_thread is None or not self._websocket_thread.is_alive():
+            return
+
+        logger.info("Stopping WebSocket listener thread")
+        self._websocket_event_loop_ready.clear()
+
+        # Signal shutdown to the websocket tasks
+        if self._websocket_event_loop and self._websocket_event_loop.is_running() and self._shutdown_event:
+
+            def signal_shutdown() -> None:
+                if self._shutdown_event:
+                    self._shutdown_event.set()
+
+            self._websocket_event_loop.call_soon_threadsafe(signal_shutdown)
+
+        # Wait for thread to finish
+        self._websocket_thread.join(timeout=5.0)
+        if self._websocket_thread.is_alive():
+            logger.warning("WebSocket listener thread did not stop gracefully")
+        else:
+            logger.info("WebSocket listener thread stopped successfully")
+
+    def _start_websocket_thread(self, api_key: str) -> None:
+        """Run WebSocket tasks in a separate thread with its own async loop."""
+        try:
+            # Create a new event loop for this thread
+            loop = asyncio.new_event_loop()
+            self._websocket_event_loop = loop
+            asyncio.set_event_loop(loop)
+
+            # Create shutdown event
+            self._shutdown_event = asyncio.Event()
+
+            # Signal that websocket_event_loop is ready
+            self._websocket_event_loop_ready.set()
+            logger.info("WebSocket listener thread started and ready")
+
+            # Run the async WebSocket listener
+            loop.run_until_complete(self._run_websocket_listener(api_key))
+        except Exception as e:
+            logger.error("WebSocket listener thread error: %s", e)
+        finally:
+            self._websocket_event_loop = None
+            self._websocket_event_loop_ready.clear()
+            self._shutdown_event = None
+            logger.info("WebSocket listener thread ended")
+
+    async def _run_websocket_listener(self, api_key: str) -> None:
+        """Run WebSocket listener - establish connection and handle incoming messages."""
+        logger.info("Creating WebSocket connection stream for listening")
+        connection_stream = _create_websocket_connection(api_key)
+
+        async for ws_connection in connection_stream:
+            logger.info("WebSocket connection established for session %s", self._session_id)
+            try:
+                # Listen for incoming messages
+                await self._listen_for_messages(ws_connection)
+
+            except (ConnectionClosed, ConnectionClosedError):
+                logger.info("WebSocket connection closed, reconnecting...")
+                continue
+            except asyncio.CancelledError:
+                logger.info("WebSocket listener task cancelled, shutting down")
+                break
+            except Exception:
+                logger.exception("WebSocket listener failed")
+                await asyncio.sleep(2.0)
+                continue
+
+            # Check if shutdown was requested
+            if self._shutdown_event and self._shutdown_event.is_set():
+                logger.info("Shutdown requested, ending WebSocket listener connection loop")
+                break
+
+        logger.info("WebSocket listener connection loop ended")
+
+    async def _listen_for_messages(self, ws_connection: Any) -> None:
+        """Listen for incoming WebSocket messages from the subprocess."""
+        logger.info("Starting to listen for WebSocket messages")
+
+        # Subscribe to the session topic to receive messages
+        topic = f"sessions/{self._session_id}/response"
+
+        try:
+            await _send_subscribe_command(
+                ws_connection=ws_connection,
+                topic=topic,
+            )
+            async for message in ws_connection:
+                if self._shutdown_event and self._shutdown_event.is_set():
+                    logger.info("Shutdown requested, ending message listener")
+                    break
+
+                try:
+                    data = json.loads(message)
+                    logger.debug("Received WebSocket message: %s", data.get("type"))
+                    await self._process_event(data)
+
+                except json.JSONDecodeError:
+                    logger.warning("Failed to parse WebSocket message: %s", message)
+                except Exception:
+                    logger.exception("Error processing WebSocket message")
+
+        except Exception as e:
+            logger.error("Error in WebSocket message listener: %s", e)
+            raise
+
+    async def _process_event(self, event: dict) -> None:
+        """Process events received from the subprocess via WebSocket."""
+        event_type = event.get("type", "unknown")
+        if event_type == "execution_event":
+            await self._process_execution_event(event)
+        elif event_type in ["success_result", "failure_result"]:
+            await self._process_result_event(event)
+
+    async def _process_execution_event(self, event: dict) -> None:
+        payload = event.get("payload", {})
+        event_type = payload.get("event_type", "")
+        payload_type_name = payload.get("payload_type", "")
+        payload_type = PayloadRegistry.get_type(payload_type_name)
+
+        # Focusing on ExecutionEvent types for the workflow executor
+        if event_type not in ["ExecutionEvent", "EventResultSuccess", "EventResultFailure"]:
+            logger.debug("Ignoring event type: %s", event_type)
+            return
+
+        if payload_type is None:
+            logger.warning("Unknown payload type: %s", payload_type_name)
+            return
+
+        ex_event = ExecutionEvent.from_dict(data=payload, payload_type=payload_type)
+
+        if isinstance(ex_event.payload, ControlFlowResolvedEvent):
+            logger.info("Workflow execution completed successfully")
+            self.output = {ex_event.payload.end_node_name: ex_event.payload.parameter_output_values}
+
+        if isinstance(ex_event.payload, ControlFlowCancelledEvent):
+            logger.error("Workflow execution cancelled")
+
+            details = ex_event.payload.result_details or "No details provided"
+            msg = f"Workflow execution cancelled: {details}"
+
+            if ex_event.payload.exception:
+                msg = f"Exception running workflow: {ex_event.payload.exception}"
+                self._stored_exception = SubprocessWorkflowExecutorError(ex_event.payload.exception)
+            else:
+                self._stored_exception = SubprocessWorkflowExecutorError(msg)
+
+    async def _process_result_event(self, event: dict) -> None:
+        payload = event.get("payload", {})
+        request_type_name = payload.get("request_type", "")
+        response_type_name = payload.get("result_type", "")
+        request_payload_type = PayloadRegistry.get_type(request_type_name)
+        response_payload_type = PayloadRegistry.get_type(response_type_name)
+
+        if request_payload_type is None or response_payload_type is None:
+            logger.warning("Unknown payload types: %s, %s", request_type_name, response_type_name)
+            return
+        if payload.get("type", "unknown") == "success_result":
+            result_event = EventResultSuccess.from_dict(
+                data=payload, req_payload_type=request_payload_type, res_payload_type=response_payload_type
+            )
+        else:
+            result_event = EventResultFailure.from_dict(
+                data=payload, req_payload_type=request_payload_type, res_payload_type=response_payload_type
+            )
+
+        if isinstance(result_event.request, StartFlowRequest):
+            logger.info("Received StartFlowRequest result event")
+            if self._on_start_flow_result:
+                self._on_start_flow_result(result_event.result)
+        else:
+            logger.warning("Ignoring result event for request type: %s", request_type_name)

griptape_nodes/bootstrap/workflow_executors/utils/__init__.py

@@ -0,0 +1 @@
+"""Workflow executors utils package."""

griptape_nodes/bootstrap/workflow_executors/utils/subprocess_script.py

@@ -0,0 +1,51 @@
+"""Subprocess script to execute a Griptape Nodes workflow.
+
+This script is intended to be run as a subprocess by the SubprocessWorkflowExecutor.
+"""
+
+import json
+from argparse import ArgumentParser
+
+from workflow import execute_workflow  # type: ignore[attr-defined]
+
+from griptape_nodes.bootstrap.workflow_executors.local_session_workflow_executor import LocalSessionWorkflowExecutor
+from griptape_nodes.drivers.storage import StorageBackend
+
+
+def _main() -> None:
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--json-input",
+        default=json.dumps({}),
+        help="JSON string representing the flow input",
+    )
+    parser.add_argument(
+        "--session-id",
+        default=None,
+        help="ID of the session to use",
+    )
+    parser.add_argument(
+        "--storage-backend",
+        default="local",
+        help="Storage backend to use",
+    )
+    parser.add_argument(
+        "--workflow-path",
+        default=None,
+        help="Path to the Griptape Nodes workflow file",
+    )
+    args = parser.parse_args()
+    flow_input = json.loads(args.json_input)
+
+    local_session_workflow_executor = LocalSessionWorkflowExecutor(
+        session_id=args.session_id, storage_backend=StorageBackend(args.storage_backend)
+    )
+
+    execute_workflow(
+        input=flow_input,
+        workflow_executor=local_session_workflow_executor,
+    )
+
+
+if __name__ == "__main__":
+    _main()

griptape_nodes/bootstrap/workflow_publishers/__init__.py

@@ -0,0 +1 @@
+"""Workflow publishers package."""

griptape_nodes/bootstrap/workflow_publishers/local_workflow_publisher.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from griptape_nodes.bootstrap.workflow_executors.local_workflow_executor import LocalWorkflowExecutor
+from griptape_nodes.retained_mode.events.workflow_events import PublishWorkflowRequest, PublishWorkflowResultSuccess
+from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+
+logger = logging.getLogger(__name__)
+
+
+class LocalPublisherError(Exception):
+    """Exception raised during local workflow publish."""
+
+
+class LocalWorkflowPublisher(LocalWorkflowExecutor):
+    def __init__(self) -> None:
+        return
+
+    async def arun(
+        self,
+        workflow_name: str,
+        workflow_path: str,
+        publisher_name: str,
+        published_workflow_file_name: str,
+        **kwargs: Any,  # noqa: ARG002
+    ) -> None:
+        # Load the workflow into memory
+        await self.aprepare_workflow_for_run(workflow_name=workflow_name, flow_input={}, workflow_path=workflow_path)
+        publish_workflow_request = PublishWorkflowRequest(
+            workflow_name=workflow_name,
+            publisher_name=publisher_name,
+            execute_on_publish=False,
+            published_workflow_file_name=published_workflow_file_name,
+        )
+        publish_workflow_result = await GriptapeNodes.ahandle_request(publish_workflow_request)
+
+        if isinstance(publish_workflow_result, PublishWorkflowResultSuccess):
+            logger.info("Published workflow to %s", publish_workflow_result.published_workflow_file_path)
+        else:
+            msg = f"Failed to publish workflow: {publish_workflow_result}"
+            raise LocalPublisherError(msg)

griptape_nodes/bootstrap/workflow_publishers/subprocess_workflow_publisher.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import logging
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Self
+
+import anyio
+
+from griptape_nodes.bootstrap.utils.python_subprocess_executor import PythonSubprocessExecutor
+from griptape_nodes.bootstrap.workflow_publishers.local_workflow_publisher import LocalWorkflowPublisher
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+logger = logging.getLogger(__name__)
+
+
+class SubprocessWorkflowPublisherError(Exception):
+    """Exception raised during subprocess workflow publishing."""
+
+
+class SubprocessWorkflowPublisher(LocalWorkflowPublisher, PythonSubprocessExecutor):
+    def __init__(self) -> None:
+        PythonSubprocessExecutor.__init__(self)
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        return
+
+    async def arun(
+        self,
+        workflow_name: str,
+        workflow_path: str,
+        publisher_name: str,
+        published_workflow_file_name: str,
+        **kwargs: Any,  # noqa: ARG002
+    ) -> None:
+        """Publish a workflow in a subprocess and wait for completion."""
+        script_path = Path(__file__).parent / "utils" / "subprocess_script.py"
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_workflow_path = Path(tmpdir) / "workflow.py"
+            tmp_script_path = Path(tmpdir) / "subprocess_script.py"
+
+            try:
+                async with (
+                    await anyio.open_file(workflow_path, "rb") as src,
+                    await anyio.open_file(tmp_workflow_path, "wb") as dst,
+                ):
+                    await dst.write(await src.read())
+
+                async with (
+                    await anyio.open_file(script_path, "rb") as src,
+                    await anyio.open_file(tmp_script_path, "wb") as dst,
+                ):
+                    await dst.write(await src.read())
+            except Exception as e:
+                msg = f"Failed to copy workflow or script to temp directory: {e}"
+                logger.exception(msg)
+                raise SubprocessWorkflowPublisherError(msg) from e
+
+            args = [
+                "--workflow-name",
+                workflow_name,
+                "--workflow-path",
+                str(tmp_workflow_path),
+                "--publisher-name",
+                publisher_name,
+                "--published-workflow-file-name",
+                published_workflow_file_name,
+            ]
+            await self.execute_python_script(
+                script_path=tmp_script_path,
+                args=args,
+                cwd=Path(tmpdir),
+            )

griptape_nodes/bootstrap/workflow_publishers/utils/__init__.py

@@ -0,0 +1 @@
+"""Workflow publishers utils package."""

griptape_nodes/bootstrap/workflow_publishers/utils/subprocess_script.py

@@ -0,0 +1,54 @@
+import asyncio
+import logging
+from argparse import ArgumentParser
+
+from griptape_nodes.bootstrap.workflow_publishers.local_workflow_publisher import LocalWorkflowPublisher
+
+logging.basicConfig(level=logging.INFO)
+
+logger = logging.getLogger(__name__)
+
+
+async def _main(workflow_name: str, workflow_path: str, publisher_name: str, published_workflow_file_name: str) -> None:
+    local_publisher = LocalWorkflowPublisher()
+    async with local_publisher as publisher:
+        await publisher.arun(
+            workflow_name=workflow_name,
+            workflow_path=workflow_path,
+            publisher_name=publisher_name,
+            published_workflow_file_name=published_workflow_file_name,
+        )
+
+    msg = f"Published workflow to file: {published_workflow_file_name}"
+    logger.info(msg)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--workflow-name",
+        help="Name of the workflow to publish",
+        required=True,
+    )
+    parser.add_argument(
+        "--workflow-path",
+        help="Path to the workflow file to publish",
+        required=True,
+    )
+    parser.add_argument(
+        "--publisher-name",
+        help="Name of the publisher to use",
+        required=True,
+    )
+    parser.add_argument(
+        "--published-workflow-file-name", help="Name to use for the published workflow file", required=True
+    )
+    args = parser.parse_args()
+    asyncio.run(
+        _main(
+            workflow_name=args.workflow_name,
+            workflow_path=args.workflow_path,
+            publisher_name=args.publisher_name,
+            published_workflow_file_name=args.published_workflow_file_name,
+        )
+    )

griptape_nodes/cli/commands/engine.py

@@ -1,6 +1,5 @@
 """Engine command for Griptape Nodes CLI."""
 
-import typer
 from rich.prompt import Confirm
 
 from griptape_nodes.app import start_app
@@ -21,19 +20,13 @@ from griptape_nodes.cli.shared import (
 from griptape_nodes.utils.version_utils import get_current_version, get_install_source
 
 
-def engine_command(
-    no_update: bool = typer.Option(False, "--no-update", help="Skip the auto-update check."),  # noqa: FBT001
-) -> None:
+def engine_command() -> None:
     """Run the Griptape Nodes engine."""
-    _start_engine(no_update=no_update)
+    _start_engine()
 
 
-def _start_engine(*, no_update: bool = False) -> None:
-    """Starts the Griptape Nodes engine.
-
-    Args:
-        no_update (bool): If True, skips the auto-update check.
-    """
+def _start_engine() -> None:
+    """Starts the Griptape Nodes engine."""
     if not CONFIG_DIR.exists():
         # Default init flow if there is no config directory
         console.print("[bold green]Config directory not found. Initializing...[/bold green]")
@@ -51,10 +44,6 @@ def _start_engine(*, no_update: bool = False) -> None:
             )
         )
 
-    # Confusing double negation -- If `no_update` is set, we want to skip the update
-    if not no_update:
-        _auto_update_self()
-
     console.print("[bold green]Starting Griptape Nodes engine...[/bold green]")
     start_app()
 

griptape_nodes/cli/main.py

@@ -10,6 +10,7 @@ from rich.console import Console
 sys.path.append(str(Path.cwd()))
 
 from griptape_nodes.cli.commands import config, engine, init, libraries, models, self
+from griptape_nodes.cli.commands.engine import _auto_update_self
 from griptape_nodes.utils.version_utils import get_complete_version_string
 
 console = Console()
@@ -47,9 +48,13 @@ def main(
         console.print(f"[bold green]{version_string}[/bold green]")
         raise typer.Exit
 
+    # Run auto-update check for any command (unless disabled)
+    if not no_update:
+        _auto_update_self()
+
     if ctx.invoked_subcommand is None:
         # Default to engine command when no subcommand is specified
-        engine.engine_command(no_update=no_update)
+        engine.engine_command()
 
 
 if __name__ == "__main__":