vellum-ai 0.12.13__py3-none-any.whl → 0.12.15__py3-none-any.whl

vellum/workflows/nodes/core/inline_subworkflow_node/node.py

@@ -10,6 +10,7 @@ from vellum.workflows.state.base import BaseState
 from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.types.core import EntityInputsInterface
 from vellum.workflows.types.generics import StateType, WorkflowInputsType
+from vellum.workflows.workflows.event_filters import all_workflow_event_filter
 
 if TYPE_CHECKING:
     from vellum.workflows.workflows.base import BaseWorkflow
@@ -36,6 +37,7 @@ class InlineSubworkflowNode(BaseNode[StateType], Generic[StateType, WorkflowInpu
             )
             subworkflow_stream = subworkflow.stream(
                 inputs=self._compile_subworkflow_inputs(),
+                event_filter=all_workflow_event_filter,
             )
 
         outputs: Optional[BaseOutputs] = None

vellum/workflows/nodes/core/map_node/node.py

@@ -9,21 +9,21 @@ from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.events.types import ParentContext
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.inputs.base import BaseInputs
-from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
+from vellum.workflows.nodes.utils import create_adornment
 from vellum.workflows.outputs import BaseOutputs
-from vellum.workflows.state.base import BaseState
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.state.context import WorkflowContext
-from vellum.workflows.types.generics import NodeType, StateType
+from vellum.workflows.types.generics import StateType
 from vellum.workflows.workflows.event_filters import all_workflow_event_filter
 
 if TYPE_CHECKING:
-    from vellum.workflows import BaseWorkflow
     from vellum.workflows.events.workflow import WorkflowEvent
 
 MapNodeItemType = TypeVar("MapNodeItemType")
 
 
-class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
+class MapNode(BaseAdornmentNode[StateType], Generic[StateType, MapNodeItemType]):
     """
     Used to map over a list of items and execute a Subworkflow on each iteration.
 
@@ -33,11 +33,10 @@ class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
     """
 
     items: List[MapNodeItemType]
-    subworkflow: Type["BaseWorkflow"]
     concurrency: Optional[int] = None
 
-    class Outputs(BaseOutputs):
-        mapped_items: list
+    class Outputs(BaseAdornmentNode.Outputs):
+        pass
 
     class SubworkflowInputs(BaseInputs):
         # TODO: Both type: ignore's below are believed to be incorrect and both have the following error:
@@ -54,6 +53,7 @@ class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
             mapped_items[output_descripter.name] = [None] * len(self.items)
 
         self._event_queue: Queue[Tuple[int, WorkflowEvent]] = Queue()
+        self._concurrency_queue: Queue[Thread] = Queue()
         fulfilled_iterations: List[bool] = []
         for index, item in enumerate(self.items):
             fulfilled_iterations.append(False)
@@ -66,11 +66,21 @@ class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
                     "parent_context": parent_context,
                 },
             )
-            thread.start()
+            if self.concurrency is None:
+                thread.start()
+            else:
+                self._concurrency_queue.put(thread)
+
+        if self.concurrency is not None:
+            concurrency_count = 0
+            while concurrency_count < self.concurrency:
+                is_empty = self._start_thread()
+                if is_empty:
+                    break
+
+                concurrency_count += 1
 
         try:
-            # We should consolidate this logic with the logic workflow runner uses
-            # https://app.shortcut.com/vellum/story/4736
             while map_node_event := self._event_queue.get():
                 index = map_node_event[0]
                 terminal_event = map_node_event[1]
@@ -86,6 +96,9 @@ class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
                     fulfilled_iterations[index] = True
                     if all(fulfilled_iterations):
                         break
+
+                    if self.concurrency is not None:
+                        self._start_thread()
                 elif terminal_event.name == "workflow.execution.paused":
                     raise NodeException(
                         code=WorkflowErrorCode.INVALID_OUTPUTS,
@@ -98,7 +111,12 @@ class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
                     )
         except Empty:
             pass
-        return self.Outputs(**mapped_items)
+
+        outputs = self.Outputs()
+        for output_name, output_list in mapped_items.items():
+            setattr(outputs, output_name, output_list)
+
+        return outputs
 
     def _context_run_subworkflow(
         self, *, item: MapNodeItemType, index: int, parent_context: Optional[ParentContext] = None
@@ -109,7 +127,10 @@ class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
 
     def _run_subworkflow(self, *, item: MapNodeItemType, index: int) -> None:
         context = WorkflowContext(vellum_client=self._context.vellum_client)
-        subworkflow = self.subworkflow(parent_state=self.state, context=context)
+        subworkflow = self.subworkflow(
+            parent_state=self.state,
+            context=context,
+        )
         events = subworkflow.stream(
             inputs=self.SubworkflowInputs(index=index, item=item, all_items=self.items),
             event_filter=all_workflow_event_filter,
@@ -118,6 +139,14 @@ class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
         for event in events:
             self._event_queue.put((index, event))
 
+    def _start_thread(self) -> bool:
+        if self._concurrency_queue.empty():
+            return False
+
+        thread = self._concurrency_queue.get()
+        thread.start()
+        return True
+
     @overload
     @classmethod
     def wrap(cls, items: List[MapNodeItemType]) -> Callable[..., Type["MapNode[StateType, MapNodeItemType]"]]: ...
@@ -134,37 +163,12 @@ class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
     def wrap(
         cls, items: Union[List[MapNodeItemType], BaseDescriptor[List[MapNodeItemType]]]
     ) -> Callable[..., Type["MapNode[StateType, MapNodeItemType]"]]:
-        _items = items
+        return create_adornment(cls, attributes={"items": items})
 
-        def decorator(inner_cls: Type[NodeType]) -> Type["MapNode[StateType, MapNodeItemType]"]:
-            # Investigate how to use dependency injection to avoid circular imports
-            # https://app.shortcut.com/vellum/story/4116
-            from vellum.workflows import BaseWorkflow
-
-            class Subworkflow(BaseWorkflow[MapNode.SubworkflowInputs, BaseState]):
-                graph = inner_cls
-
-                # mypy is wrong here, this works and is defined
-                class Outputs(inner_cls.Outputs):  # type: ignore[name-defined]
-                    pass
-
-            class WrappedNodeOutputs(BaseOutputs):
-                pass
-
-            WrappedNodeOutputs.__annotations__ = {
-                # TODO: We'll need to infer the type T of Subworkflow.Outputs[name] so we could do List[T] here
-                # https://app.shortcut.com/vellum/story/4119
-                descriptor.name: List
-                for descriptor in inner_cls.Outputs
-            }
-
-            class WrappedNode(MapNode[StateType, MapNodeItemType]):
-                items = _items
-                subworkflow = Subworkflow
-
-                class Outputs(WrappedNodeOutputs):
-                    pass
-
-            return WrappedNode
+    @classmethod
+    def __annotate_outputs_class__(cls, outputs_class: Type[BaseOutputs], reference: OutputReference) -> None:
+        parameter_type = reference.types[0]
+        annotation = List[parameter_type]  # type: ignore[valid-type]
 
-        return decorator
+        previous_annotations = {prev: annotation for prev in outputs_class.__annotations__ if not prev.startswith("_")}
+        outputs_class.__annotations__ = {**previous_annotations, reference.name: annotation}

vellum/workflows/nodes/core/retry_node/node.py

@@ -1,27 +1,16 @@
-from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Optional, Type
+from typing import Callable, Generic, Optional, Type
 
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.nodes.bases.base import BaseNodeMeta
-from vellum.workflows.state.base import BaseState
+from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
+from vellum.workflows.nodes.utils import create_adornment
+from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.types.generics import StateType
 
-if TYPE_CHECKING:
-    from vellum.workflows import BaseWorkflow
 
-
-class _RetryNodeMeta(BaseNodeMeta):
-    @property
-    def _localns(cls) -> Dict[str, Any]:
-        return {
-            **super()._localns,
-            "SubworkflowInputs": getattr(cls, "SubworkflowInputs"),
-        }
-
-
-class RetryNode(BaseNode[StateType], Generic[StateType], metaclass=_RetryNodeMeta):
+class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
     """
     Used to retry a Subworkflow a specified number of times.
 
@@ -32,7 +21,6 @@ class RetryNode(BaseNode[StateType], Generic[StateType], metaclass=_RetryNodeMet
 
     max_attempts: int
     retry_on_error_code: Optional[WorkflowErrorCode] = None
-    subworkflow: Type["BaseWorkflow[SubworkflowInputs, BaseState]"]
 
     class SubworkflowInputs(BaseInputs):
         attempt_number: int
@@ -41,9 +29,10 @@ class RetryNode(BaseNode[StateType], Generic[StateType], metaclass=_RetryNodeMet
         last_exception = Exception("max_attempts must be greater than 0")
         for index in range(self.max_attempts):
             attempt_number = index + 1
+            context = WorkflowContext(vellum_client=self._context.vellum_client)
             subworkflow = self.subworkflow(
                 parent_state=self.state,
-                context=self._context,
+                context=context,
             )
             terminal_event = subworkflow.run(
                 inputs=self.SubworkflowInputs(attempt_number=attempt_number),
@@ -78,30 +67,6 @@ Message: {terminal_event.error.message}""",
     def wrap(
         cls, max_attempts: int, retry_on_error_code: Optional[WorkflowErrorCode] = None
     ) -> Callable[..., Type["RetryNode"]]:
-        _max_attempts = max_attempts
-        _retry_on_error_code = retry_on_error_code
-
-        def decorator(inner_cls: Type[BaseNode]) -> Type["RetryNode"]:
-            # Investigate how to use dependency injection to avoid circular imports
-            # https://app.shortcut.com/vellum/story/4116
-            from vellum.workflows import BaseWorkflow
-
-            class Subworkflow(BaseWorkflow[RetryNode.SubworkflowInputs, BaseState]):
-                graph = inner_cls
-
-                # mypy is wrong here, this works and is defined
-                class Outputs(inner_cls.Outputs):  # type: ignore[name-defined]
-                    pass
-
-            class WrappedNode(RetryNode[StateType]):
-                max_attempts = _max_attempts
-                retry_on_error_code = _retry_on_error_code
-
-                subworkflow = Subworkflow
-
-                class Outputs(Subworkflow.Outputs):
-                    pass
-
-            return WrappedNode
-
-        return decorator
+        return create_adornment(
+            cls, attributes={"max_attempts": max_attempts, "retry_on_error_code": retry_on_error_code}
+        )

vellum/workflows/nodes/core/try_node/node.py

@@ -1,61 +1,18 @@
-import sys
-from types import ModuleType
-from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Iterator, Optional, Set, Tuple, Type, TypeVar
+from typing import Callable, Generic, Iterator, Optional, Set, Type
 
 from vellum.workflows.errors.types import WorkflowError, WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.nodes.bases.base import BaseNodeMeta
-from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
+from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
+from vellum.workflows.nodes.utils import create_adornment
 from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.types.generics import StateType
 from vellum.workflows.workflows.event_filters import all_workflow_event_filter
 
-if TYPE_CHECKING:
-    from vellum.workflows import BaseWorkflow
 
-Subworkflow = Type["BaseWorkflow"]
-_T = TypeVar("_T", bound=BaseOutputs)
-
-
-class _TryNodeMeta(BaseNodeMeta):
-    def __new__(cls, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
-        node_class = super().__new__(cls, name, bases, dct)
-
-        subworkflow_attribute = dct.get("subworkflow")
-        if not subworkflow_attribute:
-            return node_class
-
-        subworkflow_outputs = getattr(subworkflow_attribute, "Outputs")
-        if not issubclass(subworkflow_outputs, BaseOutputs):
-            raise ValueError("subworkflow.Outputs must be a subclass of BaseOutputs")
-
-        outputs_class = dct.get("Outputs")
-        if not outputs_class:
-            raise ValueError("Outputs class not found in base classes")
-
-        if not issubclass(outputs_class, BaseNode.Outputs):
-            raise ValueError("Outputs class must be a subclass of BaseNode.Outputs")
-
-        for descriptor in subworkflow_outputs:
-            if descriptor.name == "error":
-                raise ValueError("`error` is a reserved name for TryNode.Outputs")
-
-            setattr(outputs_class, descriptor.name, descriptor)
-
-        return node_class
-
-    def __getattribute__(cls, name: str) -> Any:
-        try:
-            return super().__getattribute__(name)
-        except AttributeError:
-            if name != "__wrapped_node__" and issubclass(cls, TryNode):
-                return getattr(cls.__wrapped_node__, name)
-            raise
-
-
-class TryNode(BaseNode[StateType], Generic[StateType], metaclass=_TryNodeMeta):
+class TryNode(BaseAdornmentNode[StateType], Generic[StateType]):
     """
     Used to execute a Subworkflow and handle errors.
 
@@ -63,9 +20,7 @@ class TryNode(BaseNode[StateType], Generic[StateType], metaclass=_TryNodeMeta):
     subworkflow: Type["BaseWorkflow"] - The Subworkflow to execute
     """
 
-    __wrapped_node__: Optional[Type["BaseNode"]] = None
     on_error_code: Optional[WorkflowErrorCode] = None
-    subworkflow: Type["BaseWorkflow"]
 
     class Outputs(BaseNode.Outputs):
         error: Optional[WorkflowError] = None
@@ -129,38 +84,11 @@ Message: {event.error.message}""",
 
     @classmethod
     def wrap(cls, on_error_code: Optional[WorkflowErrorCode] = None) -> Callable[..., Type["TryNode"]]:
-        _on_error_code = on_error_code
-
-        def decorator(inner_cls: Type[BaseNode]) -> Type["TryNode"]:
-            # Investigate how to use dependency injection to avoid circular imports
-            # https://app.shortcut.com/vellum/story/4116
-            from vellum.workflows import BaseWorkflow
-
-            inner_cls._is_wrapped_node = True
-
-            class Subworkflow(BaseWorkflow):
-                graph = inner_cls
-
-                # mypy is wrong here, this works and is defined
-                class Outputs(inner_cls.Outputs):  # type: ignore[name-defined]
-                    pass
-
-            dynamic_module = f"{inner_cls.__module__}.{inner_cls.__name__}.{ADORNMENT_MODULE_NAME}"
-            # This dynamic module allows calls to `type_hints` to work
-            sys.modules[dynamic_module] = ModuleType(dynamic_module)
-
-            # We use a dynamic wrapped node class to be uniquely tied to this `inner_cls` node during serialization
-            WrappedNode = type(
-                cls.__name__,
-                (TryNode,),
-                {
-                    "__wrapped_node__": inner_cls,
-                    "__module__": dynamic_module,
-                    "on_error_code": _on_error_code,
-                    "subworkflow": Subworkflow,
-                    "Ports": type("Ports", (TryNode.Ports,), {port.name: port.copy() for port in inner_cls.Ports}),
-                },
-            )
-            return WrappedNode
+        return create_adornment(cls, attributes={"on_error_code": on_error_code})
+
+    @classmethod
+    def __annotate_outputs_class__(cls, outputs_class: Type[BaseOutputs], reference: OutputReference) -> None:
+        if reference.name == "error":
+            raise ValueError("`error` is a reserved name for TryNode.Outputs")
 
-        return decorator
+        setattr(outputs_class, reference.name, reference)

vellum/workflows/nodes/utils.py

@@ -1,5 +1,7 @@
 from functools import cache
-from typing import Type
+import sys
+from types import ModuleType
+from typing import Any, Callable, Optional, Type, TypeVar
 
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.ports.port import Port
@@ -42,3 +44,44 @@ def has_wrapped_node(node: Type[NodeType]) -> bool:
         return False
 
     return True
+
+
+AdornableNode = TypeVar("AdornableNode", bound=BaseNode)
+
+
+def create_adornment(
+    adornable_cls: Type[AdornableNode], attributes: Optional[dict[str, Any]] = None
+) -> Callable[..., Type["AdornableNode"]]:
+    def decorator(inner_cls: Type[BaseNode]) -> Type["AdornableNode"]:
+        # Investigate how to use dependency injection to avoid circular imports
+        # https://app.shortcut.com/vellum/story/4116
+        from vellum.workflows import BaseWorkflow
+
+        inner_cls._is_wrapped_node = True
+
+        class Subworkflow(BaseWorkflow):
+            graph = inner_cls
+
+            # mypy is wrong here, this works and is defined
+            class Outputs(inner_cls.Outputs):  # type: ignore[name-defined]
+                pass
+
+        dynamic_module = f"{inner_cls.__module__}.{inner_cls.__name__}.{ADORNMENT_MODULE_NAME}"
+        # This dynamic module allows calls to `type_hints` to work
+        sys.modules[dynamic_module] = ModuleType(dynamic_module)
+
+        # We use a dynamic wrapped node class to be uniquely tied to this `inner_cls` node during serialization
+        WrappedNode = type(
+            adornable_cls.__name__,
+            (adornable_cls,),
+            {
+                "__wrapped_node__": inner_cls,
+                "__module__": dynamic_module,
+                "subworkflow": Subworkflow,
+                "Ports": type("Ports", (adornable_cls.Ports,), {port.name: port.copy() for port in inner_cls.Ports}),
+                **(attributes or {}),
+            },
+        )
+        return WrappedNode
+
+    return decorator

vellum/workflows/references/constant.py

@@ -0,0 +1,21 @@
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+
+if TYPE_CHECKING:
+    from vellum.workflows.state.base import BaseState
+
+_T = TypeVar("_T")
+
+
+class ConstantValueReference(BaseDescriptor[_T], Generic[_T]):
+    def __init__(
+        self,
+        value: _T,
+    ) -> None:
+        self._value = value
+        types = (type(self._value),)
+        super().__init__(name=str(self._value), types=types)
+
+    def resolve(self, state: "BaseState") -> _T:
+        return self._value

vellum/workflows/runner/runner.py

@@ -44,11 +44,13 @@ from vellum.workflows.events.workflow import (
 )
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.bases.base import NodeRunResponse
 from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.ports.port import Port
 from vellum.workflows.references import ExternalInputReference, OutputReference
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.cycle_map import CycleMap
 from vellum.workflows.types.generics import OutputsType, StateType, WorkflowInputsType
 
 if TYPE_CHECKING:
@@ -124,9 +126,7 @@ class WorkflowRunner(Generic[StateType]):
 
         self._dependencies: Dict[Type[BaseNode], Set[Type[BaseNode]]] = defaultdict(set)
         self._state_forks: Set[StateType] = {self._initial_state}
-        self._mocks_by_node_outputs_class = (
-            {mock.__class__: mock for mock in node_output_mocks} if node_output_mocks else {}
-        )
+        self._mocks_by_node_outputs_class = CycleMap(items=node_output_mocks or [], key_by=lambda mock: mock.__class__)
 
         self._active_nodes_by_execution_id: Dict[UUID, BaseNode[StateType]] = {}
         self._cancel_signal = cancel_signal
@@ -182,6 +182,7 @@ class WorkflowRunner(Generic[StateType]):
                 node_definition=node.__class__,
                 parent=parent_context,
             )
+            node_run_response: NodeRunResponse
             if node.Outputs not in self._mocks_by_node_outputs_class:
                 with execution_context(parent_context=updated_parent_context):
                     node_run_response = node.run()

vellum/workflows/types/cycle_map.py

@@ -0,0 +1,34 @@
+from typing import Callable, Dict, Generic, List, TypeVar
+
+_K = TypeVar("_K")
+_T = TypeVar("_T")
+
+
+class CycleMap(Generic[_K, _T]):
+    """
+    A map that cycles through a list of items for each key.
+    """
+
+    def __init__(self, items: List[_T], key_by: Callable[[_T], _K]):
+        self._items: Dict[_K, List[_T]] = {}
+        for item in items:
+            self._add_item(key_by(item), item)
+
+    def _add_item(self, key: _K, item: _T):
+        if key not in self._items:
+            self._items[key] = []
+        self._items[key].append(item)
+
+    def _get_item(self, key: _K) -> _T:
+        item = self._items[key].pop(0)
+        self._items[key].append(item)
+        return item
+
+    def __getitem__(self, key: _K) -> _T:
+        return self._get_item(key)
+
+    def __setitem__(self, key: _K, value: _T):
+        self._add_item(key, value)
+
+    def __contains__(self, key: _K) -> bool:
+        return key in self._items

vellum/workflows/workflows/base.py

@@ -1,16 +1,7 @@
-# flake8: noqa: E402
-
-import importlib
-import inspect
-
-from vellum.plugins.utils import load_runtime_plugins
-from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum.workflows.workflows.event_filters import workflow_event_filter
-
-load_runtime_plugins()
-
 from datetime import datetime
 from functools import lru_cache
+import importlib
+import inspect
 from threading import Event as ThreadingEvent
 from uuid import UUID, uuid4
 from typing import (
@@ -79,6 +70,8 @@ from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.state.store import Store
 from vellum.workflows.types.generics import StateType, WorkflowInputsType
 from vellum.workflows.types.utils import get_original_base
+from vellum.workflows.utils.uuids import uuid4_from_hash
+from vellum.workflows.workflows.event_filters import workflow_event_filter
 
 
 class _BaseWorkflowMeta(type):

{vellum_ai-0.12.13.dist-info → vellum_ai-0.12.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.12.13
+Version: 0.12.15
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -95,7 +95,7 @@ more below](#workflows-sdk).
 
 ## Client SDK
 
-The Vellum Client SDK, found within `src/client` is a low-level client used to interact directly with the Vellum API.
+The Vellum Client SDK, found within `src/vellum/client` is a low-level client used to interact directly with the Vellum API.
 Learn more and get started by visiting the [Vellum Client SDK README](/src/vellum/client/README.md).
 
 ## Workflows SDK