|
| 1 | +import logging |
| 2 | + |
| 3 | +from griptape_nodes.node_library.advanced_node_library import AdvancedNodeLibrary |
| 4 | +from griptape_nodes.node_library.library_registry import Library, LibrarySchema |
| 5 | +from griptape_nodes.retained_mode.events.base_events import RequestPayload, ResultPayload |
| 6 | +from griptape_nodes.retained_mode.events.workflow_events import ( |
| 7 | + PublishWorkflowRegisteredEventData, |
| 8 | + PublishWorkflowRequest, |
| 9 | +) |
| 10 | +from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes |
| 11 | + |
| 12 | +logging.basicConfig(level=logging.INFO) |
| 13 | +logger = logging.getLogger("griptape_nodes") |
| 14 | + |
| 15 | + |
| 16 | +def _publish_workflow_request_handler(request: RequestPayload) -> ResultPayload: |
| 17 | + if not isinstance(request, PublishWorkflowRequest): |
| 18 | + msg = f"Expected PublishWorkflowRequest, got {type(request).__name__}" |
| 19 | + raise TypeError(msg) |
| 20 | + from griptape_nodes_library.publish_workflow.local_publisher import LocalPublisher |
| 21 | + |
| 22 | + publisher = LocalPublisher( |
| 23 | + workflow_name=request.workflow_name, |
| 24 | + pickle_control_flow_result=request.pickle_control_flow_result, |
| 25 | + ) |
| 26 | + return publisher.publish_workflow() |
| 27 | + |
| 28 | + |
| 29 | +class GriptapeNodesLibraryAdvanced(AdvancedNodeLibrary): |
| 30 | + """Advanced library implementation for the default Griptape Nodes Library.""" |
| 31 | + |
| 32 | + def before_library_nodes_loaded(self, library_data: LibrarySchema, library: Library) -> None: # noqa: ARG002 |
| 33 | + """Called before any nodes are loaded from the library.""" |
| 34 | + msg = f"Starting to load nodes for '{library_data.name}' library..." |
| 35 | + logger.info(msg) |
| 36 | + |
| 37 | + def after_library_nodes_loaded(self, library_data: LibrarySchema, library: Library) -> None: # noqa: ARG002 |
| 38 | + """Called after all nodes have been loaded from the library.""" |
| 39 | + GriptapeNodes.LibraryManager().on_register_event_handler( |
| 40 | + request_type=PublishWorkflowRequest, |
| 41 | + handler=_publish_workflow_request_handler, |
| 42 | + library_data=library_data, |
| 43 | + event_data=PublishWorkflowRegisteredEventData( |
| 44 | + start_flow_node_type="StartFlow", |
| 45 | + start_flow_node_library_name=library_data.name, |
| 46 | + end_flow_node_type="EndFlow", |
| 47 | + end_flow_node_library_name=library_data.name, |
| 48 | + ), |
| 49 | + ) |
0 commit comments