vellum-ai 0.9.16__py3-none-any.whl → 0.9.16rc4__py3-none-any.whl

vellum/workflows/state/base.py

@@ -0,0 +1,327 @@
+from collections import defaultdict
+from copy import deepcopy
+from dataclasses import field
+from datetime import datetime
+from queue import Queue
+from threading import Lock
+from uuid import UUID, uuid4
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Sequence, Set, Tuple, Type, cast
+from typing_extensions import dataclass_transform
+
+from pydantic import GetCoreSchemaHandler, field_serializer
+from pydantic_core import core_schema
+
+from vellum.core.pydantic_utilities import UniversalBaseModel
+
+from vellum.workflows.constants import UNDEF
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.references import ExternalInputReference, OutputReference, StateValueReference
+from vellum.workflows.types.generics import StateType
+from vellum.workflows.types.stack import Stack
+from vellum.workflows.types.utils import datetime_now, deepcopy_with_exclusions, get_class_by_qualname, infer_types
+
+if TYPE_CHECKING:
+    from vellum.workflows.nodes.bases import BaseNode
+
+
+class _Snapshottable:
+    _snapshot_callback: Callable[[], None]
+
+    def __deepcopy__(self, memo: Any) -> "_Snapshottable":
+        return deepcopy_with_exclusions(
+            self,
+            memo=memo,
+            exclusions={
+                "_snapshot_callback": self._snapshot_callback,
+            },
+        )
+
+
+@dataclass_transform(kw_only_default=True)
+class _BaseStateMeta(type):
+    def __getattribute__(cls, name: str) -> Any:
+        if not name.startswith("_"):
+            instance = vars(cls).get(name)
+            types = infer_types(cls, name)
+            return StateValueReference(name=name, types=types, instance=instance)
+
+        return super().__getattribute__(name)
+
+
+class _SnapshottableDict(dict, _Snapshottable):
+    def __setitem__(self, key: Any, value: Any) -> None:
+        super().__setitem__(key, value)
+        self._snapshot_callback()
+
+
+def _make_snapshottable(value: Any, snapshot_callback: Callable[[], None]) -> Any:
+    """
+    Edits any value to make it snapshottable on edit. Made as a separate function from `BaseState` to
+    avoid namespace conflicts with subclasses.
+    """
+    if isinstance(value, _Snapshottable):
+        return value
+
+    if isinstance(value, dict):
+        snapshottable_dict = _SnapshottableDict(value)
+        snapshottable_dict._snapshot_callback = snapshot_callback
+        return snapshottable_dict
+
+    return value
+
+
+class NodeExecutionCache:
+    _node_execution_ids: Dict[Type["BaseNode"], Stack[UUID]]
+    _node_executions_initiated: Dict[Type["BaseNode"], Set[UUID]]
+    _dependencies_invoked: Dict[str, Set[str]]
+
+    def __init__(
+        self,
+        dependencies_invoked: Optional[Dict[str, Sequence[str]]] = None,
+        node_execution_ids: Optional[Dict[str, Sequence[str]]] = None,
+        node_executions_initiated: Optional[Dict[str, Sequence[str]]] = None,
+    ) -> None:
+        self._dependencies_invoked = defaultdict(set)
+        self._node_execution_ids = defaultdict(Stack[UUID])
+        self._node_executions_initiated = defaultdict(set)
+
+        for node, dependencies in (dependencies_invoked or {}).items():
+            self._dependencies_invoked[node].update(dependencies)
+
+        for node, execution_ids in (node_execution_ids or {}).items():
+            node_class = get_class_by_qualname(node)
+            self._node_execution_ids[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
+
+        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})
+
+    @property
+    def dependencies_invoked(self) -> Dict[str, Set[str]]:
+        return self._dependencies_invoked
+
+    def is_node_initiated(self, node: Type["BaseNode"]) -> bool:
+        return node in self._node_executions_initiated and len(self._node_executions_initiated[node]) > 0
+
+    def initiate_node_execution(self, node: Type["BaseNode"], execution_id: UUID) -> None:
+        self._node_executions_initiated[node].add(execution_id)
+
+    def fulfill_node_execution(self, node: Type["BaseNode"], execution_id: UUID) -> None:
+        self._node_executions_initiated[node].remove(execution_id)
+        self._node_execution_ids[node].push(execution_id)
+
+    def get_execution_count(self, node: Type["BaseNode"]) -> int:
+        return self._node_execution_ids[node].size()
+
+    def dump(self) -> Dict[str, Any]:
+        return {
+            "dependencies_invoked": {
+                node: list(dependencies) for node, dependencies in self._dependencies_invoked.items()
+            },
+            "node_executions_initiated": {
+                str(node): list(execution_ids) for node, execution_ids in self._node_executions_initiated.items()
+            },
+            "node_execution_ids": {
+                str(node): execution_ids.dump() for node, execution_ids in self._node_execution_ids.items()
+            },
+        }
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Type[Any], handler: GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        return core_schema.is_instance_schema(cls)
+
+
+def uuid4_default_factory() -> UUID:
+    """
+    Allows us to mock the uuid4 for testing.
+    """
+    return uuid4()
+
+
+def default_datetime_factory() -> datetime:
+    """
+    Makes it possible to mock the datetime factory for testing.
+    """
+
+    return datetime_now()
+
+
+class StateMeta(UniversalBaseModel):
+    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)
+    external_inputs: Dict[ExternalInputReference, Any] = field(default_factory=dict)
+    node_outputs: Dict[OutputReference, Any] = field(default_factory=dict)
+    node_execution_cache: NodeExecutionCache = field(default_factory=NodeExecutionCache)
+    parent: Optional["BaseState"] = None
+    is_terminated: Optional[bool] = None
+    __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:
+        self.node_outputs = _make_snapshottable(self.node_outputs, callback)
+        self.__snapshot_callback__ = callback
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name.startswith("__") or name == "updated_ts":
+            super().__setattr__(name, value)
+            return
+
+        super().__setattr__(name, value)
+        if callable(self.__snapshot_callback__):
+            self.__snapshot_callback__()
+
+    @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()}
+
+    @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()}
+
+    def __deepcopy__(self, memo: Optional[Dict[int, Any]] = None) -> "StateMeta":
+        if not memo:
+            memo = {}
+
+        new_node_outputs = {
+            descriptor: value if isinstance(value, Queue) else deepcopy(value, memo)
+            for descriptor, value in self.node_outputs.items()
+        }
+
+        memo[id(self.node_outputs)] = new_node_outputs
+        memo[id(self.__snapshot_callback__)] = None
+
+        return super().__deepcopy__(memo)
+
+
+class BaseState(metaclass=_BaseStateMeta):
+    meta: StateMeta = field(init=False)
+
+    __lock__: Lock = field(init=False)
+    __is_initializing__: bool = field(init=False)
+    __snapshot_callback__: Callable[["BaseState"], None] = field(init=False)
+
+    def __init__(self, meta: Optional[StateMeta] = None, **kwargs: Any) -> None:
+        self.__is_initializing__ = True
+        self.__snapshot_callback__ = lambda state: None
+        self.__lock__ = Lock()
+
+        self.meta = meta or StateMeta()
+        self.meta.add_snapshot_callback(self.__snapshot__)
+
+        # Make all class attribute values snapshottable
+        for name, value in self.__class__.__dict__.items():
+            if not name.startswith("_") and name != "meta":
+                # Bypass __is_initializing__ instead of `setattr`
+                snapshottable_value = _make_snapshottable(value, self.__snapshot__)
+                super().__setattr__(name, snapshottable_value)
+
+        for name, value in kwargs.items():
+            setattr(self, name, value)
+
+        self.__is_initializing__ = False
+
+    def __deepcopy__(self, memo: Any) -> "BaseState":
+        new_state = deepcopy_with_exclusions(
+            self,
+            exclusions={
+                "__lock__": Lock(),
+            },
+            memo=memo,
+        )
+        new_state.meta.add_snapshot_callback(new_state.__snapshot__)
+        return new_state
+
+    def __repr__(self) -> str:
+        values = "\n".join(
+            [f"    {key}={value}" for key, value in vars(self).items() if not key.startswith("_") and key != "meta"]
+        )
+        node_outputs = "\n".join([f"            {key}={value}" for key, value in self.meta.node_outputs.items()])
+        return f"""\
+{self.__class__.__name__}:
+{values}
+    meta:
+        id={self.meta.id}
+        is_terminated={self.meta.is_terminated}
+        updated_ts={self.meta.updated_ts}
+        node_outputs:{' Empty' if not node_outputs else ''}
+{node_outputs}
+"""
+
+    def __iter__(self) -> Iterator[Tuple[Any, Any]]:
+        """
+        Returns an iterator treating all state keys as (key, value) items, allowing consumers to call `dict()`
+        on an instance of this class.
+        """
+
+        # If the user sets a default value on state (e.g. something = "foo"), it's not on `instance_attributes` below.
+        # So we need to include class_attributes here just in case
+        class_attributes = {key: value for key, value in self.__class__.__dict__.items() if not key.startswith("_")}
+        instance_attributes = {key: value for key, value in self.__dict__.items() if not key.startswith("__")}
+
+        all_attributes = {**class_attributes, **instance_attributes}
+        items = [(key, value) for key, value in all_attributes.items() if key not in ["_lock"]]
+        return iter(items)
+
+    def __getitem__(self, key: str) -> Any:
+        return self.__dict__[key]
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name.startswith("_") or self.__is_initializing__:
+            super().__setattr__(name, value)
+            return
+
+        snapshottable_value = _make_snapshottable(value, self.__snapshot__)
+        super().__setattr__(name, snapshottable_value)
+        self.meta.updated_ts = datetime_now()
+        self.__snapshot__()
+
+    def __add__(self, other: StateType) -> StateType:
+        """
+        Handles merging two states together, preferring the latest state by updated_ts for any given node output.
+        """
+
+        if not isinstance(other, type(self)):
+            raise TypeError(f"Cannot add {type(other).__name__} to {type(self).__name__}]")
+
+        latest_state = self if self.meta.updated_ts >= other.meta.updated_ts else other
+        oldest_state = other if latest_state == self else self
+
+        for descriptor, value in oldest_state.meta.node_outputs.items():
+            if descriptor not in latest_state.meta.node_outputs:
+                latest_state.meta.node_outputs[descriptor] = value
+
+        for key, value in oldest_state:
+            if not isinstance(key, str):
+                continue
+
+            if key.startswith("_"):
+                continue
+
+            if getattr(latest_state, key, UNDEF) == UNDEF:
+                setattr(latest_state, key, value)
+
+        return cast(StateType, latest_state)
+
+    def __snapshot__(self) -> None:
+        """
+        Snapshots the current state to the workflow emitter. The invoked callback is overridden by the
+        workflow runner.
+        """
+        self.__snapshot_callback__(deepcopy(self))
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Type[Any], handler: GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        return core_schema.is_instance_schema(cls)

vellum/workflows/state/context.py

@@ -0,0 +1,18 @@
+from functools import cached_property
+from typing import Optional
+
+from vellum import Vellum
+
+from vellum.workflows.vellum_client import create_vellum_client
+
+
+class WorkflowContext:
+    def __init__(self, _vellum_client: Optional[Vellum] = None):
+        self._vellum_client = _vellum_client
+
+    @cached_property
+    def vellum_client(self) -> Vellum:
+        if self._vellum_client:
+            return self._vellum_client
+
+        return create_vellum_client()

vellum/workflows/state/encoder.py

@@ -0,0 +1,57 @@
+from dataclasses import asdict, is_dataclass
+from datetime import datetime
+import enum
+from json import JSONEncoder
+from uuid import UUID
+from typing import Any, Callable, Dict, Type
+
+from pydantic import BaseModel
+
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
+from vellum.workflows.state.base import BaseState, NodeExecutionCache
+
+
+class DefaultStateEncoder(JSONEncoder):
+    encoders: Dict[Type, Callable] = {}
+
+    def default(self, obj: Any) -> Any:
+        if isinstance(obj, BaseState):
+            return dict(obj)
+
+        if isinstance(obj, (BaseInputs, BaseOutputs)):
+            return {descriptor.name: value for descriptor, value in obj}
+
+        if isinstance(obj, BaseOutput):
+            return obj.serialize()
+
+        if isinstance(obj, NodeExecutionCache):
+            return obj.dump()
+
+        if isinstance(obj, UUID):
+            return str(obj)
+
+        if isinstance(obj, set):
+            return list(obj)
+
+        if isinstance(obj, BaseModel):
+            return obj.model_dump()
+
+        if isinstance(obj, datetime):
+            return obj.isoformat()
+
+        if isinstance(obj, enum.Enum):
+            return obj.value
+
+        if is_dataclass(obj):
+            # Technically, obj is DataclassInstance | type[DataclassInstance], but asdict expects a DataclassInstance
+            # in practice, we only ever pass the former
+            return asdict(obj)  # type: ignore[call-overload]
+
+        if isinstance(obj, type):
+            return str(obj)
+
+        if obj.__class__ in self.encoders:
+            return self.encoders[obj.__class__](obj)
+
+        return super().default(obj)

vellum/workflows/state/store.py

@@ -0,0 +1,28 @@
+from typing import Iterator, List
+
+from vellum.workflows.events.workflow import WorkflowEvent
+from vellum.workflows.state.base import BaseState
+
+
+class Store:
+    def __init__(self) -> None:
+        self._events: List[WorkflowEvent] = []
+        self._state_snapshots: List[BaseState] = []
+
+    def append_event(self, event: WorkflowEvent) -> None:
+        self._events.append(event)
+
+    def append_state_snapshot(self, state: BaseState) -> None:
+        self._state_snapshots.append(state)
+
+    def clear(self) -> None:
+        self._events = []
+        self._state_snapshots = []
+
+    @property
+    def events(self) -> Iterator[WorkflowEvent]:
+        return iter(self._events)
+
+    @property
+    def state_snapshots(self) -> Iterator[BaseState]:
+        return iter(self._state_snapshots)

vellum/workflows/state/tests/test_state.py

@@ -0,0 +1,113 @@
+from collections import defaultdict
+from copy import deepcopy
+import json
+from typing import Dict
+
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.state.encoder import DefaultStateEncoder
+
+snapshot_count: Dict[int, int] = defaultdict(int)
+
+
+class MockState(BaseState):
+    foo: str
+    nested_dict: Dict[str, int] = {}
+
+    def __snapshot__(self) -> None:
+        global snapshot_count
+        snapshot_count[id(self)] += 1
+
+
+class MockNode(BaseNode):
+    class Outputs(BaseOutputs):
+        baz: str
+
+
+def test_state_snapshot__node_attribute_edit():
+    # GIVEN an initial state instance
+    state = MockState(foo="bar")
+    assert snapshot_count[id(state)] == 0
+
+    # WHEN we edit an attribute
+    state.foo = "baz"
+
+    # THEN the snapshot is emitted
+    assert snapshot_count[id(state)] == 1
+
+
+def test_state_snapshot__node_output_edit():
+    # GIVEN an initial state instance
+    state = MockState(foo="bar")
+    assert snapshot_count[id(state)] == 0
+
+    # WHEN we add a Node Output to state
+    for output in MockNode.Outputs:
+        state.meta.node_outputs[output] = "hello"
+
+    # THEN the snapshot is emitted
+    assert snapshot_count[id(state)] == 1
+
+
+def test_state_snapshot__nested_dictionary_edit():
+    # GIVEN an initial state instance
+    state = MockState(foo="bar")
+    assert snapshot_count[id(state)] == 0
+
+    # WHEN we edit a nested dictionary
+    state.nested_dict["hello"] = 1
+
+    # THEN the snapshot is emitted
+    assert snapshot_count[id(state)] == 1
+
+
+def test_state_deepcopy():
+    # GIVEN an initial state instance
+    state = MockState(foo="bar")
+
+    # AND we add a Node Output to state
+    state.meta.node_outputs[MockNode.Outputs.baz] = "hello"
+
+    # WHEN we deepcopy the state
+    deepcopied_state = deepcopy(state)
+
+    # THEN node outputs are deepcopied
+    assert deepcopied_state.meta.node_outputs == state.meta.node_outputs
+
+
+def test_state_deepcopy__with_node_output_updates():
+    # GIVEN an initial state instance
+    state = MockState(foo="bar")
+
+    # AND we add a Node Output to state
+    state.meta.node_outputs[MockNode.Outputs.baz] = "hello"
+
+    # AND we deepcopy the state
+    deepcopied_state = deepcopy(state)
+
+    # AND we update the original state
+    state.meta.node_outputs[MockNode.Outputs.baz] = "world"
+
+    # THEN the copied state is not updated
+    assert deepcopied_state.meta.node_outputs[MockNode.Outputs.baz] == "hello"
+
+    # AND the original state has had the correct number of snapshots
+    assert snapshot_count[id(state)] == 2
+
+    # AND the copied state has had the correct number of snapshots
+    assert snapshot_count[id(deepcopied_state)] == 0
+
+
+def test_state_json_serialization__with_node_output_updates():
+    # GIVEN an initial state instance
+    state = MockState(foo="bar")
+
+    # AND we add a Node Output to state
+    state.meta.node_outputs[MockNode.Outputs.baz] = "hello"
+
+    # WHEN we serialize the state
+    json_state = json.loads(json.dumps(state, cls=DefaultStateEncoder))
+
+    # THEN the state is serialized correctly
+    assert json_state["meta"]["node_outputs"] == {"MockNode.Outputs.baz": "hello"}

vellum/workflows/types/core.py

@@ -0,0 +1,91 @@
+from enum import Enum
+from typing import (  # type: ignore[attr-defined]
+    Dict,
+    List,
+    Union,
+    _GenericAlias,
+    _SpecialGenericAlias,
+    _UnionGenericAlias,
+)
+
+from vellum import (
+    ChatMessage,
+    FunctionCall,
+    FunctionCallRequest,
+    SearchResult,
+    SearchResultRequest,
+    VellumAudio,
+    VellumAudioRequest,
+    VellumError,
+    VellumErrorRequest,
+    VellumImage,
+    VellumImageRequest,
+    VellumValue,
+    VellumValueRequest,
+)
+
+JsonArray = List["Json"]
+JsonObject = Dict[str, "Json"]
+Json = Union[None, bool, int, float, str, JsonArray, JsonObject]
+
+# Unions and Generics inherit from `_GenericAlias` instead of `type`
+# In future versions of python, we'll see `_UnionGenericAlias`
+UnderGenericAlias = _GenericAlias
+SpecialGenericAlias = _SpecialGenericAlias
+UnionGenericAlias = _UnionGenericAlias
+
+
+class VellumSecret:
+    name: str
+
+    def __init__(self, name: str):
+        self.name = name
+
+
+VellumValuePrimitive = Union[
+    # String inputs
+    str,
+    # Chat history inputs
+    List[ChatMessage],
+    List[ChatMessage],
+    # Search results inputs
+    List[SearchResultRequest],
+    List[SearchResult],
+    # JSON inputs
+    Json,
+    # Number inputs
+    float,
+    # Function Call Inputs
+    FunctionCall,
+    FunctionCallRequest,
+    # Error Inputs
+    VellumError,
+    VellumErrorRequest,
+    # Array Inputs
+    List[VellumValueRequest],
+    List[VellumValue],
+    # Image Inputs
+    VellumImage,
+    VellumImageRequest,
+    # Audio Inputs
+    VellumAudio,
+    VellumAudioRequest,
+    # Vellum Secrets
+    VellumSecret,
+]
+
+EntityInputsInterface = Dict[
+    str,
+    VellumValuePrimitive,
+]
+
+
+class MergeBehavior(Enum):
+    AWAIT_ALL = "AWAIT_ALL"
+    AWAIT_ANY = "AWAIT_ANY"
+
+
+class ConditionType(Enum):
+    IF = "IF"
+    ELIF = "ELIF"
+    ELSE = "ELSE"

vellum/workflows/types/generics.py

@@ -0,0 +1,14 @@
+from typing import TYPE_CHECKING, TypeVar
+
+if TYPE_CHECKING:
+    from vellum.workflows import BaseWorkflow
+    from vellum.workflows.inputs import BaseInputs
+    from vellum.workflows.nodes import BaseNode
+    from vellum.workflows.outputs import BaseOutputs
+    from vellum.workflows.state import BaseState
+
+NodeType = TypeVar("NodeType", bound="BaseNode")
+StateType = TypeVar("StateType", bound="BaseState")
+WorkflowType = TypeVar("WorkflowType", bound="BaseWorkflow")
+WorkflowInputsType = TypeVar("WorkflowInputsType", bound="BaseInputs")
+OutputsType = TypeVar("OutputsType", bound="BaseOutputs")

vellum/workflows/types/stack.py

@@ -0,0 +1,39 @@
+from collections import deque
+from typing import Deque, Generic, Iterable, List, TypeVar
+
+_T = TypeVar("_T")
+
+
+class Stack(Generic[_T]):
+    def __init__(self) -> None:
+        self._items: Deque[_T] = deque()
+
+    def push(self, item: _T) -> None:
+        self._items.append(item)
+
+    def extend(self, items: Iterable[_T]) -> None:
+        item_list = list(items)
+        for item in item_list[::-1]:
+            self._items.append(item)
+
+    def pop(self) -> _T:
+        if not self.is_empty():
+            return self._items.pop()
+        raise IndexError("pop from empty stack")
+
+    def peek(self) -> _T:
+        if not self.is_empty():
+            return self._items[-1]
+        raise IndexError("peek from empty stack")
+
+    def is_empty(self) -> bool:
+        return len(self._items) == 0
+
+    def size(self) -> int:
+        return len(self._items)
+
+    def __repr__(self) -> str:
+        return f"Stack({self.dump()})"
+
+    def dump(self) -> List[_T]:
+        return [item for item in self._items][::-1]