vellum-ai 0.10.8__py3-none-any.whl → 0.11.0__py3-none-any.whl

vellum/workflows/workflows/base.py

@@ -4,6 +4,7 @@ import importlib
 import inspect
 
 from vellum.plugins.utils import load_runtime_plugins
+from vellum.workflows.workflows.event_filters import workflow_event_filter
 
 load_runtime_plugins()
 
@@ -13,6 +14,7 @@ from threading import Event as ThreadingEvent
 from uuid import uuid4
 from typing import (
     Any,
+    Callable,
     ClassVar,
     Dict,
     Generator,
@@ -47,7 +49,6 @@ from vellum.workflows.events.node import (
     NodeExecutionStreamingBody,
     NodeExecutionStreamingEvent,
 )
-from vellum.workflows.events.types import WorkflowEventType
 from vellum.workflows.events.workflow import (
     GenericWorkflowEvent,
     WorkflowExecutionFulfilledBody,
@@ -60,6 +61,8 @@ from vellum.workflows.events.workflow import (
     WorkflowExecutionRejectedEvent,
     WorkflowExecutionResumedBody,
     WorkflowExecutionResumedEvent,
+    WorkflowExecutionSnapshottedBody,
+    WorkflowExecutionSnapshottedEvent,
     WorkflowExecutionStreamingBody,
     WorkflowExecutionStreamingEvent,
 )
@@ -80,13 +83,16 @@ from vellum.workflows.types.utils import get_original_base
 class _BaseWorkflowMeta(type):
     def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
         if "graph" not in dct:
-            dct["graph"] = []
+            dct["graph"] = set()
 
         return super().__new__(mcs, name, bases, dct)
 
 
+GraphAttribute = Union[Type[BaseNode], Graph, Set[Type[BaseNode]], Set[Graph]]
+
+
 class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkflowMeta):
-    graph: ClassVar[Union[Type[BaseNode], Graph, Set[Type[BaseNode]], Set[Graph]]]
+    graph: ClassVar[GraphAttribute]
     emitters: List[BaseWorkflowEmitter]
     resolvers: List[BaseWorkflowResolver]
 
@@ -97,6 +103,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
         GenericWorkflowEvent,
         WorkflowExecutionInitiatedEvent[WorkflowInputsType],  # type: ignore[valid-type]
         WorkflowExecutionFulfilledEvent[Outputs],
+        WorkflowExecutionSnapshottedEvent[StateType],  # type: ignore[valid-type]
     ]
 
     TerminalWorkflowEvent = Union[
@@ -181,6 +188,31 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
         external_inputs: Optional[ExternalInputsArg] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
     ) -> TerminalWorkflowEvent:
+        """
+        Invoke a Workflow, returning the last event emitted, which should be one of:
+        - `WorkflowExecutionFulfilledEvent` if the Workflow Execution was successful
+        - `WorkflowExecutionRejectedEvent` if the Workflow Execution was rejected
+        - `WorkflowExecutionPausedEvent` if the Workflow Execution was paused
+
+        Parameters
+        ----------
+        inputs: Optional[WorkflowInputsType] = None
+            The Inputs instance used to initiate the Workflow Execution.
+
+        state: Optional[StateType] = None
+            The State instance to run the Workflow with. Workflows maintain a global state that can be used to
+            deterministically resume execution from any point.
+
+        entrypoint_nodes: Optional[RunFromNodeArg] = None
+            The entrypoint nodes to run the Workflow with. Useful for resuming execution from a specific node.
+
+        external_inputs: Optional[ExternalInputsArg] = None
+            External inputs to pass to the Workflow. Useful for providing human-in-the-loop behavior to the Workflow.
+
+        cancel_signal: Optional[ThreadingEvent] = None
+            A threading event that can be used to cancel the Workflow Execution.
+        """
+
         events = WorkflowRunner(
             self,
             inputs=inputs,
@@ -188,6 +220,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
             entrypoint_nodes=entrypoint_nodes,
             external_inputs=external_inputs,
             cancel_signal=cancel_signal,
+            parent_context=self._context.parent_context,
         ).stream()
         first_event: Optional[Union[WorkflowExecutionInitiatedEvent, WorkflowExecutionResumedEvent]] = None
         last_event = None
@@ -201,7 +234,10 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
                 trace_id=uuid4(),
                 span_id=uuid4(),
                 body=WorkflowExecutionRejectedBody(
-                    error=VellumError(code=VellumErrorCode.INTERNAL_ERROR, message="No events were emitted"),
+                    error=VellumError(
+                        code=VellumErrorCode.INTERNAL_ERROR,
+                        message="No events were emitted",
+                    ),
                     workflow_definition=self.__class__,
                 ),
             )
@@ -212,7 +248,8 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
                 span_id=uuid4(),
                 body=WorkflowExecutionRejectedBody(
                     error=VellumError(
-                        code=VellumErrorCode.INTERNAL_ERROR, message="Initiated event was never emitted"
+                        code=VellumErrorCode.INTERNAL_ERROR,
+                        message="Initiated event was never emitted",
                     ),
                     workflow_definition=self.__class__,
                 ),
@@ -239,14 +276,40 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
 
     def stream(
         self,
-        event_types: Optional[Set[WorkflowEventType]] = None,
+        event_filter: Optional[Callable[[Type["BaseWorkflow"], WorkflowEvent], bool]] = None,
         inputs: Optional[WorkflowInputsType] = None,
         state: Optional[StateType] = None,
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
     ) -> WorkflowEventStream:
-        event_types = event_types or {WorkflowEventType.WORKFLOW}
+        """
+        Invoke a Workflow, yielding events as they are emitted.
+
+        Parameters
+        ----------
+        event_filter: Optional[Callable[[Type["BaseWorkflow"], WorkflowEvent], bool]] = None
+            A filter that can be used to filter events based on the Workflow Class and the event itself. If the method
+            returns `False`, the event will not be yielded.
+
+        inputs: Optional[WorkflowInputsType] = None
+            The Inputs instance used to initiate the Workflow Execution.
+
+        state: Optional[StateType] = None
+            The State instance to run the Workflow with. Workflows maintain a global state that can be used to
+            deterministically resume execution from any point.
+
+        entrypoint_nodes: Optional[RunFromNodeArg] = None
+            The entrypoint nodes to run the Workflow with. Useful for resuming execution from a specific node.
+
+        external_inputs: Optional[ExternalInputsArg] = None
+            External inputs to pass to the Workflow. Useful for providing human-in-the-loop behavior to the Workflow.
+
+        cancel_signal: Optional[ThreadingEvent] = None
+            A threading event that can be used to cancel the Workflow Execution.
+        """
+
+        should_yield = event_filter or workflow_event_filter
         for event in WorkflowRunner(
             self,
             inputs=inputs,
@@ -254,8 +317,9 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
             entrypoint_nodes=entrypoint_nodes,
             external_inputs=external_inputs,
             cancel_signal=cancel_signal,
+            parent_context=self.context.parent_context,
         ).stream():
-            if WorkflowEventType(event.name.split(".")[0].upper()) in event_types:
+            if should_yield(self.__class__, event):
                 yield event
 
     def validate(self) -> None:
@@ -268,7 +332,9 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
 
     @classmethod
     @lru_cache
-    def _get_parameterized_classes(cls) -> Tuple[Type[WorkflowInputsType], Type[StateType]]:
+    def _get_parameterized_classes(
+        cls,
+    ) -> Tuple[Type[WorkflowInputsType], Type[StateType]]:
         original_base = get_original_base(cls)
 
         inputs_type, state_type = get_args(original_base)
@@ -299,7 +365,10 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
 
     def get_default_state(self, workflow_inputs: Optional[WorkflowInputsType] = None) -> StateType:
         return self.get_state_class()(
-            meta=StateMeta(parent=self._parent_state, workflow_inputs=workflow_inputs or self.get_default_inputs())
+            meta=StateMeta(
+                parent=self._parent_state,
+                workflow_inputs=workflow_inputs or self.get_default_inputs(),
+            )
         )
 
     def get_state_at_node(self, node: Type[BaseNode]) -> StateType:
@@ -342,6 +411,9 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
 
         workflows: List[Type[BaseWorkflow]] = []
         for name in dir(module):
+            if name.startswith("__"):
+                continue
+
             attr = getattr(module, name)
             if (
                 inspect.isclass(attr)
@@ -365,6 +437,7 @@ WorkflowExecutionRejectedBody.model_rebuild()
 WorkflowExecutionPausedBody.model_rebuild()
 WorkflowExecutionResumedBody.model_rebuild()
 WorkflowExecutionStreamingBody.model_rebuild()
+WorkflowExecutionSnapshottedBody.model_rebuild()
 
 NodeExecutionInitiatedBody.model_rebuild()
 NodeExecutionFulfilledBody.model_rebuild()
@@ -379,6 +452,7 @@ WorkflowExecutionRejectedEvent.model_rebuild()
 WorkflowExecutionPausedEvent.model_rebuild()
 WorkflowExecutionResumedEvent.model_rebuild()
 WorkflowExecutionStreamingEvent.model_rebuild()
+WorkflowExecutionSnapshottedEvent.model_rebuild()
 
 NodeExecutionInitiatedEvent.model_rebuild()
 NodeExecutionFulfilledEvent.model_rebuild()

vellum/workflows/workflows/event_filters.py

@@ -0,0 +1,53 @@
+from typing import TYPE_CHECKING, Type
+
+from vellum.workflows.events.types import CodeResourceDefinition
+
+if TYPE_CHECKING:
+    from vellum.workflows.events.workflow import WorkflowEvent
+    from vellum.workflows.workflows.base import BaseWorkflow
+
+
+def workflow_event_filter(workflow_definition: Type["BaseWorkflow"], event: "WorkflowEvent") -> bool:
+    """
+    Filters for only Workflow events that were emitted by the `workflow_definition` parameter.
+    """
+
+    if (
+        event.name == "workflow.execution.initiated"
+        or event.name == "workflow.execution.resumed"
+        or event.name == "workflow.execution.fulfilled"
+        or event.name == "workflow.execution.rejected"
+        or event.name == "workflow.execution.paused"
+        or event.name == "workflow.execution.streaming"
+    ):
+        return event.workflow_definition == workflow_definition
+
+    return False
+
+
+def root_workflow_event_filter(workflow_definition: Type["BaseWorkflow"], event: "WorkflowEvent") -> bool:
+    """
+    Filters for Workflow and Node events that were emitted by the `workflow_definition` parameter.
+    """
+
+    if (
+        event.name == "workflow.execution.initiated"
+        or event.name == "workflow.execution.resumed"
+        or event.name == "workflow.execution.fulfilled"
+        or event.name == "workflow.execution.rejected"
+        or event.name == "workflow.execution.paused"
+        or event.name == "workflow.execution.streaming"
+    ):
+        return event.workflow_definition == workflow_definition
+
+    if not event.parent:
+        return False
+
+    if event.parent.type != "WORKFLOW":
+        return False
+
+    return event.parent.workflow_definition == CodeResourceDefinition.encode(workflow_definition)
+
+
+def all_workflow_event_filter(workflow_definition: Type["BaseWorkflow"], event: "WorkflowEvent") -> bool:
+    return True

{vellum_ai-0.10.8.dist-info → vellum_ai-0.11.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.10.8
+Version: 0.11.0
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0