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

vellum/workflows/state/base.py

@@ -6,7 +6,7 @@ import logging
 from queue import Queue
 from threading import Lock
 from uuid import UUID, uuid4
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Type, cast
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, cast
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler, ValidationInfo, field_serializer, field_validator
@@ -16,19 +16,14 @@ from vellum.core.pydantic_utilities import UniversalBaseModel
 from vellum.workflows.constants import undefined
 from vellum.workflows.edges.edge import Edge
 from vellum.workflows.inputs.base import BaseInputs
-from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references import ExternalInputReference, OutputReference, StateValueReference
-from vellum.workflows.types.generics import StateType
+from vellum.workflows.types.definition import CodeResourceDefinition, serialize_type_encoder_with_id
+from vellum.workflows.types.generics import StateType, import_workflow_class, is_workflow_class
 from vellum.workflows.types.stack import Stack
-from vellum.workflows.types.utils import (
-    datetime_now,
-    deepcopy_with_exclusions,
-    get_class_attr_names,
-    get_class_by_qualname,
-    infer_types,
-)
+from vellum.workflows.types.utils import datetime_now, deepcopy_with_exclusions, get_class_attr_names, infer_types
 
 if TYPE_CHECKING:
+    from vellum.workflows import BaseWorkflow
     from vellum.workflows.nodes.bases import BaseNode
 
 logger = logging.getLogger(__name__)
@@ -94,38 +89,65 @@ def _make_snapshottable(value: Any, snapshot_callback: Callable[[], None]) -> An
     return value
 
 
+NodeExecutionsFulfilled = Dict[Type["BaseNode"], Stack[UUID]]
+NodeExecutionsInitiated = Dict[Type["BaseNode"], Set[UUID]]
+NodeExecutionsQueued = Dict[Type["BaseNode"], List[UUID]]
+DependenciesInvoked = Dict[UUID, Set[Type["BaseNode"]]]
+
+
 class NodeExecutionCache:
-    _node_executions_fulfilled: Dict[Type["BaseNode"], Stack[UUID]]
-    _node_executions_initiated: Dict[Type["BaseNode"], Set[UUID]]
-    _node_executions_queued: Dict[Type["BaseNode"], List[UUID]]
-    _dependencies_invoked: Dict[UUID, Set[Type["BaseNode"]]]
+    _node_executions_fulfilled: NodeExecutionsFulfilled
+    _node_executions_initiated: NodeExecutionsInitiated
+    _node_executions_queued: NodeExecutionsQueued
+    _dependencies_invoked: DependenciesInvoked
 
-    def __init__(
-        self,
-        dependencies_invoked: Optional[Dict[str, Sequence[str]]] = None,
-        node_executions_fulfilled: Optional[Dict[str, Sequence[str]]] = None,
-        node_executions_initiated: Optional[Dict[str, Sequence[str]]] = None,
-        node_executions_queued: Optional[Dict[str, Sequence[str]]] = None,
-    ) -> None:
+    def __init__(self) -> None:
         self._dependencies_invoked = defaultdict(set)
         self._node_executions_fulfilled = defaultdict(Stack[UUID])
         self._node_executions_initiated = defaultdict(set)
         self._node_executions_queued = defaultdict(list)
 
-        for execution_id, dependencies in (dependencies_invoked or {}).items():
-            self._dependencies_invoked[UUID(execution_id)] = {get_class_by_qualname(dep) for dep in dependencies}
+    @classmethod
+    def deserialize(cls, raw_data: dict, nodes: Dict[str, Type["BaseNode"]]):
+        cache = cls()
+
+        dependencies_invoked = raw_data.get("dependencies_invoked")
+        if isinstance(dependencies_invoked, dict):
+            for execution_id, dependencies in dependencies_invoked.items():
+                cache._dependencies_invoked[UUID(execution_id)] = {nodes[dep] for dep in dependencies if dep in nodes}
+
+        node_executions_fulfilled = raw_data.get("node_executions_fulfilled")
+        if isinstance(node_executions_fulfilled, dict):
+            for node, execution_ids in node_executions_fulfilled.items():
+                node_class = nodes.get(node)
+                if not node_class:
+                    continue
+
+                cache._node_executions_fulfilled[node_class].extend(
+                    UUID(execution_id) for execution_id in execution_ids
+                )
+
+        node_executions_initiated = raw_data.get("node_executions_initiated")
+        if isinstance(node_executions_initiated, dict):
+            for node, execution_ids in node_executions_initiated.items():
+                node_class = nodes.get(node)
+                if not node_class:
+                    continue
+
+                cache._node_executions_initiated[node_class].update(
+                    {UUID(execution_id) for execution_id in execution_ids}
+                )
 
-        for node, execution_ids in (node_executions_fulfilled or {}).items():
-            node_class = get_class_by_qualname(node)
-            self._node_executions_fulfilled[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
+        node_executions_queued = raw_data.get("node_executions_queued")
+        if isinstance(node_executions_queued, dict):
+            for node, execution_ids in node_executions_queued.items():
+                node_class = nodes.get(node)
+                if not node_class:
+                    continue
 
-        for node, execution_ids in (node_executions_initiated or {}).items():
-            node_class = get_class_by_qualname(node)
-            self._node_executions_initiated[node_class].update({UUID(execution_id) for execution_id in execution_ids})
+                cache._node_executions_queued[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
 
-        for node, execution_ids in (node_executions_queued or {}).items():
-            node_class = get_class_by_qualname(node)
-            self._node_executions_queued[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
+        return cache
 
     def _invoke_dependency(
         self,
@@ -207,8 +229,8 @@ def default_datetime_factory() -> datetime:
 
 
 class StateMeta(UniversalBaseModel):
+    workflow_definition: Type["BaseWorkflow"] = field(default_factory=import_workflow_class)
     id: UUID = field(default_factory=uuid4_default_factory)
-    trace_id: UUID = field(default_factory=uuid4_default_factory)
     span_id: UUID = field(default_factory=uuid4_default_factory)
     updated_ts: datetime = field(default_factory=default_datetime_factory)
     workflow_inputs: BaseInputs = field(default_factory=BaseInputs)
@@ -219,8 +241,6 @@ class StateMeta(UniversalBaseModel):
     __snapshot_callback__: Optional[Callable[[], None]] = field(init=False, default=None)
 
     def model_post_init(self, context: Any) -> None:
-        if self.parent:
-            self.trace_id = self.parent.meta.trace_id
         self.__snapshot_callback__ = None
 
     def add_snapshot_callback(self, callback: Callable[[], None]) -> None:
@@ -237,6 +257,25 @@ class StateMeta(UniversalBaseModel):
         if callable(self.__snapshot_callback__):
             self.__snapshot_callback__()
 
+    @field_serializer("workflow_definition")
+    def serialize_workflow_definition(self, workflow_definition: Type["BaseWorkflow"], _info: Any) -> Dict[str, Any]:
+        return serialize_type_encoder_with_id(workflow_definition)
+
+    @field_validator("workflow_definition", mode="before")
+    @classmethod
+    def deserialize_workflow_definition(cls, workflow_definition: Any, info: ValidationInfo):
+        if isinstance(workflow_definition, dict):
+            deserialized_workflow_definition = CodeResourceDefinition.model_validate(workflow_definition).decode()
+            if not is_workflow_class(deserialized_workflow_definition):
+                return import_workflow_class()
+
+            return deserialized_workflow_definition
+
+        if is_workflow_class(workflow_definition):
+            return workflow_definition
+
+        return import_workflow_class()
+
     @field_serializer("node_outputs")
     def serialize_node_outputs(self, node_outputs: Dict[OutputReference, Any], _info: Any) -> Dict[str, Any]:
         return {str(descriptor): value for descriptor, value in node_outputs.items()}
@@ -244,17 +283,16 @@ class StateMeta(UniversalBaseModel):
     @field_validator("node_outputs", mode="before")
     @classmethod
     def deserialize_node_outputs(cls, node_outputs: Any, info: ValidationInfo):
-        if isinstance(node_outputs, dict) and isinstance(info.context, dict):
-            raw_workflow_nodes = info.context.get("nodes")
-            workflow_node_outputs = {}
-            if isinstance(raw_workflow_nodes, list):
-                for node in raw_workflow_nodes:
-                    Outputs = getattr(node, "Outputs", None)
-                    if not isinstance(Outputs, type) or not issubclass(Outputs, BaseOutputs):
-                        continue
+        if isinstance(node_outputs, dict):
+            workflow_definition = cls._get_workflow(info)
+            if not workflow_definition:
+                return node_outputs
 
-                    for output in Outputs:
-                        workflow_node_outputs[str(output)] = output
+            raw_workflow_nodes = workflow_definition.get_nodes()
+            workflow_node_outputs = {}
+            for node in raw_workflow_nodes:
+                for output in node.Outputs:
+                    workflow_node_outputs[str(output)] = output
 
             node_output_keys = list(node_outputs.keys())
             deserialized_node_outputs = {}
@@ -269,12 +307,64 @@ class StateMeta(UniversalBaseModel):
 
         return node_outputs
 
+    @field_validator("node_execution_cache", mode="before")
+    @classmethod
+    def deserialize_node_execution_cache(cls, node_execution_cache: Any, info: ValidationInfo):
+        if isinstance(node_execution_cache, dict):
+            workflow_definition = cls._get_workflow(info)
+            if not workflow_definition:
+                return node_execution_cache
+
+            nodes_cache: Dict[str, Type["BaseNode"]] = {}
+            raw_workflow_nodes = workflow_definition.get_nodes()
+            for node in raw_workflow_nodes:
+                nodes_cache[str(node)] = node
+
+            return NodeExecutionCache.deserialize(node_execution_cache, nodes_cache)
+
+        return node_execution_cache
+
+    @field_validator("workflow_inputs", mode="before")
+    @classmethod
+    def deserialize_workflow_inputs(cls, workflow_inputs: Any, info: ValidationInfo):
+        workflow_definition = cls._get_workflow(info)
+
+        if workflow_definition:
+            if workflow_inputs is None:
+                return workflow_definition.get_inputs_class()()
+            if isinstance(workflow_inputs, dict):
+                return workflow_definition.get_inputs_class()(**workflow_inputs)
+
+        return workflow_inputs
+
     @field_serializer("external_inputs")
     def serialize_external_inputs(
         self, external_inputs: Dict[ExternalInputReference, Any], _info: Any
     ) -> Dict[str, Any]:
         return {str(descriptor): value for descriptor, value in external_inputs.items()}
 
+    @field_validator("parent", mode="before")
+    @classmethod
+    def deserialize_parent(cls, parent: Any, info: ValidationInfo):
+        if isinstance(parent, dict):
+            workflow_definition = cls._get_workflow(info)
+            if not workflow_definition:
+                return parent
+
+            parent_meta = parent.get("meta")
+            if not isinstance(parent_meta, dict):
+                return parent
+
+            parent_workflow_definition = cls.deserialize_workflow_definition(
+                parent_meta.get("workflow_definition"), info
+            )
+            if not is_workflow_class(parent_workflow_definition):
+                return parent
+
+            return parent_workflow_definition.deserialize_state(parent)
+
+        return parent
+
     def __deepcopy__(self, memo: Optional[Dict[int, Any]] = None) -> "StateMeta":
         if not memo:
             memo = {}
@@ -295,6 +385,20 @@ class StateMeta(UniversalBaseModel):
 
         return super().__deepcopy__(memo)
 
+    @classmethod
+    def _get_workflow(cls, info: ValidationInfo) -> Optional[Type["BaseWorkflow"]]:
+        if not info.context:
+            return None
+
+        if not isinstance(info.context, dict):
+            return None
+
+        workflow_definition = info.context.get("workflow_definition")
+        if not is_workflow_class(workflow_definition):
+            return None
+
+        return workflow_definition
+
 
 class BaseState(metaclass=_BaseStateMeta):
     meta: StateMeta = field(init=False)

vellum/workflows/types/definition.py

@@ -0,0 +1,71 @@
+import importlib
+import inspect
+from types import FrameType
+from uuid import UUID
+from typing import Annotated, Any, Dict, Optional, Union
+
+from pydantic import BeforeValidator
+
+from vellum.client.types.code_resource_definition import CodeResourceDefinition as ClientCodeResourceDefinition
+
+
+def serialize_type_encoder(obj: type) -> Dict[str, Any]:
+    return {
+        "name": obj.__name__,
+        "module": obj.__module__.split("."),
+    }
+
+
+def serialize_type_encoder_with_id(obj: Union[type, "CodeResourceDefinition"]) -> Dict[str, Any]:
+    if hasattr(obj, "__id__") and isinstance(obj, type):
+        return {
+            "id": getattr(obj, "__id__"),
+            **serialize_type_encoder(obj),
+        }
+    elif isinstance(obj, CodeResourceDefinition):
+        return obj.model_dump(mode="json")
+
+    raise AttributeError(f"The object of type '{type(obj).__name__}' must have an '__id__' attribute.")
+
+
+class CodeResourceDefinition(ClientCodeResourceDefinition):
+    id: UUID
+
+    @staticmethod
+    def encode(obj: type) -> "CodeResourceDefinition":
+        return CodeResourceDefinition(**serialize_type_encoder_with_id(obj))
+
+    def decode(self) -> Any:
+        if ".<locals>." in self.name:
+            # We are decoding a local class that should already be loaded in our stack frame. So
+            # we climb up to look for it.
+            frame = inspect.currentframe()
+            return self._resolve_local(frame)
+
+        try:
+            imported_module = importlib.import_module(".".join(self.module))
+        except ImportError:
+            return None
+
+        return getattr(imported_module, self.name, None)
+
+    def _resolve_local(self, frame: Optional[FrameType]) -> Any:
+        if not frame:
+            return None
+
+        frame_module = frame.f_globals.get("__name__")
+        if not isinstance(frame_module, str) or frame_module.split(".") != self.module:
+            return self._resolve_local(frame.f_back)
+
+        outer, inner = self.name.split(".<locals>.")
+        frame_outer = frame.f_code.co_name
+        if frame_outer != outer:
+            return self._resolve_local(frame.f_back)
+
+        return frame.f_locals.get(inner)
+
+
+VellumCodeResourceDefinition = Annotated[
+    CodeResourceDefinition,
+    BeforeValidator(lambda d: (d if type(d) is dict else serialize_type_encoder_with_id(d))),
+]

vellum/workflows/types/generics.py

@@ -1,4 +1,6 @@
-from typing import TYPE_CHECKING, TypeVar
+from functools import cache
+from typing import TYPE_CHECKING, Any, Type, TypeVar
+from typing_extensions import TypeGuard
 
 if TYPE_CHECKING:
     from vellum.workflows import BaseWorkflow
@@ -12,3 +14,34 @@ StateType = TypeVar("StateType", bound="BaseState")
 WorkflowType = TypeVar("WorkflowType", bound="BaseWorkflow")
 InputsType = TypeVar("InputsType", bound="BaseInputs")
 OutputsType = TypeVar("OutputsType", bound="BaseOutputs")
+
+
+@cache
+def _import_node_class() -> Type["BaseNode"]:
+    """
+    Helper function to help avoid circular imports.
+    """
+
+    from vellum.workflows.nodes import BaseNode
+
+    return BaseNode
+
+
+def import_workflow_class() -> Type["BaseWorkflow"]:
+    """
+    Helper function to help avoid circular imports.
+    """
+
+    from vellum.workflows.workflows import BaseWorkflow
+
+    return BaseWorkflow
+
+
+def is_node_class(obj: Any) -> TypeGuard[Type["BaseNode"]]:
+    base_node_class = _import_node_class()
+    return isinstance(obj, type) and issubclass(obj, base_node_class)
+
+
+def is_workflow_class(obj: Any) -> TypeGuard[Type["BaseWorkflow"]]:
+    base_workflow_class = import_workflow_class()
+    return isinstance(obj, type) and issubclass(obj, base_workflow_class)

vellum/workflows/workflows/base.py

@@ -22,6 +22,7 @@ from typing import (
     Union,
     cast,
     get_args,
+    overload,
 )
 
 from vellum.workflows.edges import Edge
@@ -488,11 +489,13 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
                     parent=self._parent_state,
                     workflow_inputs=workflow_inputs or self.get_default_inputs(),
                     trace_id=execution_context.trace_id,
+                    workflow_definition=self.__class__,
                 )
                 if execution_context and int(execution_context.trace_id)
                 else StateMeta(
                     parent=self._parent_state,
                     workflow_inputs=workflow_inputs or self.get_default_inputs(),
+                    workflow_definition=self.__class__,
                 )
             )
         )
@@ -530,18 +533,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 +616,5 @@ NodeExecutionRejectedEvent.model_rebuild()
 NodeExecutionPausedEvent.model_rebuild()
 NodeExecutionResumedEvent.model_rebuild()
 NodeExecutionStreamingEvent.model_rebuild()
+
+StateMeta.model_rebuild()