vellum-ai 0.14.37__py3-none-any.whl → 0.14.39__py3-none-any.whl

vellum/workflows/workflows/base.py

@@ -22,6 +22,7 @@ from typing import (
     Union,
     cast,
     get_args,
+    overload,
 )
 
 from vellum.workflows.edges import Edge
@@ -145,7 +146,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
 
     WorkflowEvent = Union[  # type: ignore
         GenericWorkflowEvent,
-        WorkflowExecutionInitiatedEvent[InputsType],  # type: ignore[valid-type]
+        WorkflowExecutionInitiatedEvent[InputsType, StateType],  # type: ignore[valid-type]
         WorkflowExecutionFulfilledEvent[Outputs],
         WorkflowExecutionSnapshottedEvent[StateType],  # type: ignore[valid-type]
     ]
@@ -334,7 +335,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
 
         if not last_event:
             return WorkflowExecutionRejectedEvent(
-                trace_id=uuid4(),
+                trace_id=self._execution_context.trace_id,
                 span_id=uuid4(),
                 body=WorkflowExecutionRejectedBody(
                     error=WorkflowError(
@@ -347,7 +348,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
 
         if not first_event:
             return WorkflowExecutionRejectedEvent(
-                trace_id=uuid4(),
+                trace_id=self._execution_context.trace_id,
                 span_id=uuid4(),
                 body=WorkflowExecutionRejectedBody(
                     error=WorkflowError(
@@ -366,7 +367,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
             return last_event
 
         return WorkflowExecutionRejectedEvent(
-            trace_id=first_event.trace_id,
+            trace_id=self._execution_context.trace_id,
             span_id=first_event.span_id,
             body=WorkflowExecutionRejectedBody(
                 workflow_definition=self.__class__,
@@ -481,19 +482,11 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         return self.get_inputs_class()()
 
     def get_default_state(self, workflow_inputs: Optional[InputsType] = None) -> StateType:
-        execution_context = self._execution_context
         return self.get_state_class()(
-            meta=(
-                StateMeta(
-                    parent=self._parent_state,
-                    workflow_inputs=workflow_inputs or self.get_default_inputs(),
-                    trace_id=execution_context.trace_id,
-                )
-                if execution_context and int(execution_context.trace_id)
-                else 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(),
+                workflow_definition=self.__class__,
             )
         )
 
@@ -530,18 +523,30 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
 
         return most_recent_state_snapshot
 
+    @overload
+    @classmethod
+    def deserialize_state(cls, state: dict, workflow_inputs: Optional[InputsType] = None) -> StateType: ...
+
+    @overload
     @classmethod
-    def deserialize_state(cls, state: dict, workflow_inputs: Optional[InputsType] = None) -> StateType:
+    def deserialize_state(cls, state: None, workflow_inputs: Optional[InputsType] = None) -> None: ...
+
+    @classmethod
+    def deserialize_state(
+        cls, state: Optional[dict], workflow_inputs: Optional[InputsType] = None
+    ) -> Optional[StateType]:
+        if state is None:
+            return None
+
         state_class = cls.get_state_class()
         if "meta" in state:
-            nodes = list(cls.get_nodes())
             state["meta"] = StateMeta.model_validate(
                 {
                     **state["meta"],
                     "workflow_inputs": workflow_inputs,
                 },
                 context={
-                    "nodes": nodes,
+                    "workflow_definition": cls,
                 },
             )
 
@@ -601,3 +606,5 @@ NodeExecutionRejectedEvent.model_rebuild()
 NodeExecutionPausedEvent.model_rebuild()
 NodeExecutionResumedEvent.model_rebuild()
 NodeExecutionStreamingEvent.model_rebuild()
+
+StateMeta.model_rebuild()

vellum/workflows/workflows/tests/test_base_workflow.py

@@ -1,4 +1,5 @@
 import pytest
+from uuid import uuid4
 
 from vellum.workflows.edges.edge import Edge
 from vellum.workflows.inputs.base import BaseInputs
@@ -9,6 +10,14 @@ from vellum.workflows.state.base import BaseState
 from vellum.workflows.workflows.base import BaseWorkflow
 
 
+class GlobalTestWorkflow(BaseWorkflow):
+    """
+    Used as part of `test_base_workflow__deserialize_state_with_invalid_workflow_definition` below.
+    """
+
+    pass
+
+
 def test_base_workflow__inherit_base_outputs():
     class MyNode(BaseNode):
         class Outputs(BaseNode.Outputs):
@@ -336,13 +345,34 @@ def test_base_workflow__deserialize_state():
         graph = NodeA
 
     # WHEN we deserialize the state
+    last_span_id = str(uuid4())
     state = TestWorkflow.deserialize_state(
         {
             "bar": "My state bar",
             "meta": {
+                "id": "b70a5a4f-8253-4a38-aeaf-0700b4783a78",
+                "trace_id": "9dcfb309-81e9-4b75-9b21-edd31cf9685f",
+                "span_id": "1c2e310c-3624-4f4f-b7ac-1e429de29bbf",
+                "updated_ts": "2025-04-14T19:22:18.504902",
+                "external_inputs": {},
                 "node_outputs": {
                     "test_base_workflow__deserialize_state.<locals>.NodeA.Outputs.foo": "My node A output foo"
-                }
+                },
+                "node_execution_cache": {
+                    "dependencies_invoked": {
+                        last_span_id: ["test_base_workflow__deserialize_state.<locals>.NodeA"],
+                    },
+                    "node_executions_initiated": {
+                        "test_base_workflow__deserialize_state.<locals>.NodeA": [last_span_id],
+                    },
+                    "node_executions_fulfilled": {
+                        "test_base_workflow__deserialize_state.<locals>.NodeA": [last_span_id],
+                    },
+                    "node_executions_queued": {
+                        "test_base_workflow__deserialize_state.<locals>.NodeA": [],
+                    },
+                },
+                "parent": None,
             },
         },
         workflow_inputs=Inputs(baz="My input baz"),
@@ -353,3 +383,204 @@ def test_base_workflow__deserialize_state():
     assert state.meta.node_outputs == {NodeA.Outputs.foo: "My node A output foo"}
     assert isinstance(state.meta.workflow_inputs, Inputs)
     assert state.meta.workflow_inputs.baz == "My input baz"
+
+    # AND the node execution cache should deserialize
+    assert state.meta.node_execution_cache
+
+
+def test_base_workflow__deserialize_state_with_optional_inputs():
+
+    # GIVEN a state definition
+    class State(BaseState):
+        bar: str
+
+    # AND an inputs definition with an optional field
+    class Inputs(BaseInputs):
+        baz: str = "My default baz"
+
+    # AND a node
+    class NodeA(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            foo: str
+
+    # AND a workflow that uses all three
+    class TestWorkflow(BaseWorkflow[Inputs, State]):
+        graph = NodeA
+
+    # WHEN we deserialize the state
+    state = TestWorkflow.deserialize_state(
+        {
+            "bar": "My state bar",
+            "meta": {
+                "node_outputs": {},
+                "parent": None,
+            },
+        },
+    )
+
+    # THEN the state should be correct
+    assert state.bar == "My state bar"
+    assert isinstance(state.meta.workflow_inputs, Inputs)
+    assert state.meta.workflow_inputs.baz == "My default baz"
+
+
+def test_base_workflow__deserialize_nested_state():
+    # GIVEN a state definition
+    class State(BaseState):
+        foo: str
+
+    # AND a nested state definition
+    class NestedState(BaseState):
+        bar: str
+
+    # AND an inner node
+    class InnerNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            baz: str
+
+    # AND a subworkflow definition with the nested state and inner node
+    class InnerWorkflow(BaseWorkflow[BaseInputs, NestedState]):
+        graph = InnerNode
+
+    # AND a subworkflow node
+    class OuterNode(InlineSubworkflowNode[State, BaseInputs, NestedState]):
+        subworkflow = InnerWorkflow
+
+        class Outputs(InlineSubworkflowNode.Outputs):
+            qux: str
+
+    # AND a workflow that uses the outer state and subworkflow node
+    class TestWorkflow(BaseWorkflow[BaseInputs, State]):
+        graph = OuterNode
+
+    # WHEN we deserialize the nested state
+    inner_node_output_definition = "test_base_workflow__deserialize_nested_state.<locals>.InnerNode.Outputs.baz"
+    outer_node_output_definition = "test_base_workflow__deserialize_nested_state.<locals>.OuterNode.Outputs.qux"
+    state = InnerWorkflow.deserialize_state(
+        {
+            "bar": "My nested state bar",
+            "meta": {
+                "workflow_definition": {
+                    "id": str(InnerWorkflow.__id__),
+                    "name": "test_base_workflow__deserialize_nested_state.<locals>.InnerWorkflow",
+                    "module": ["vellum", "workflows", "workflows", "tests", "test_base_workflow"],
+                },
+                "node_outputs": {inner_node_output_definition: "My inner node output baz"},
+                "parent": {
+                    "foo": "My outer state foo",
+                    "meta": {
+                        "node_outputs": {outer_node_output_definition: "My outer node output qux"},
+                        "workflow_definition": {
+                            "id": str(TestWorkflow.__id__),
+                            "name": "test_base_workflow__deserialize_nested_state.<locals>.TestWorkflow",
+                            "module": ["vellum", "workflows", "workflows", "tests", "test_base_workflow"],
+                        },
+                    },
+                },
+            },
+        },
+    )
+
+    # THEN the state should be correct
+    assert state.bar == "My nested state bar"
+    assert state.meta.node_outputs == {InnerNode.Outputs.baz: "My inner node output baz"}
+    assert state.meta.workflow_definition == InnerWorkflow
+
+    # AND the parent state should be correct
+    assert isinstance(state.meta.parent, State)
+    assert state.meta.parent.foo == "My outer state foo"
+    assert state.meta.parent.meta.node_outputs == {OuterNode.Outputs.qux: "My outer node output qux"}
+    assert state.meta.parent.meta.workflow_definition == TestWorkflow
+
+
+def test_base_workflow__deserialize_state_with_none():
+
+    # GIVEN a state definition
+    class State(BaseState):
+        bar: str
+
+    # AND an inputs definition
+    class Inputs(BaseInputs):
+        baz: str
+
+    # AND a node
+    class NodeA(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            foo: str
+
+    # AND a workflow that uses all three
+    class TestWorkflow(BaseWorkflow[Inputs, State]):
+        graph = NodeA
+
+    # WHEN we deserialize None
+    state = TestWorkflow.deserialize_state(None)
+
+    # THEN we should get back None
+    assert state is None
+
+
+@pytest.mark.parametrize(
+    "raw_workflow_definition",
+    [
+        {
+            "name": "test_base_workflow__deserialize_state_with_invalid_workflow_definition.<locals>.TestWorkflow",
+            "module": ["invalid", "module", "path"],
+        },
+        {
+            "name": "test_base_workflow__deserialize_state_with_invalid_workflow_definition.<locals>.InvalidWorkflow",
+            "module": ["vellum", "workflows", "workflows", "tests", "test_base_workflow"],
+        },
+        {
+            "name": "invalid_local_function_name.<locals>.TestWorkflow",
+            "module": ["vellum", "workflows", "workflows", "tests", "test_base_workflow"],
+        },
+        {
+            "name": "GlobalTestWorkflow",
+            "module": ["invalid", "module", "path"],
+        },
+        {
+            "name": "InlineSubworkflowNode",
+            "module": ["vellum", "workflows", "workflows", "tests", "test_base_workflow"],
+        },
+        {
+            "name": "InvalidWorkflow",
+            "module": ["vellum", "workflows", "workflows", "tests", "test_base_workflow"],
+        },
+    ],
+    ids=[
+        "invalid_module_path_with_locals",
+        "invalid_local_name",
+        "invalid_local_container",
+        "invalid_global_module",
+        "invalid_global_name",
+        "missing_global_name",
+    ],
+)
+def test_base_workflow__deserialize_state_with_invalid_workflow_definition(raw_workflow_definition):
+
+    # GIVEN a state definition
+    class State(BaseState):
+        bar = "My default bar"
+
+    # AND a workflow
+    class TestWorkflow(BaseWorkflow[BaseInputs, State]):
+        pass
+
+    # WHEN we deserialize a state with an invalid module path
+    state = TestWorkflow.deserialize_state(
+        {
+            "meta": {
+                "workflow_definition": {
+                    "id": str(TestWorkflow.__id__),
+                    **raw_workflow_definition,
+                },
+            },
+        },
+    )
+
+    # THEN we should get back the default state
+    assert isinstance(state, State)
+    assert state.bar == "My default bar"
+
+    # AND the workflow definition should be BaseWorkflow
+    assert state.meta.workflow_definition == BaseWorkflow

{vellum_ai-0.14.37.dist-info → vellum_ai-0.14.39.dist-info}/METADATA

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