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

vellum/workflows/nodes/bases/base.py

@@ -0,0 +1,332 @@
+from functools import cached_property, reduce
+import inspect
+from types import MappingProxyType
+from typing import Any, Dict, Generic, Iterator, Optional, Set, Tuple, Type, TypeVar, Union, cast, get_args
+
+from vellum.workflows.constants import UNDEF
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.descriptors.utils import resolve_value
+from vellum.workflows.edges.edge import Edge
+from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.graph import Graph
+from vellum.workflows.graph.graph import GraphTarget
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.outputs import BaseOutput, BaseOutputs
+from vellum.workflows.ports.node_ports import NodePorts
+from vellum.workflows.ports.port import Port
+from vellum.workflows.references import ExternalInputReference
+from vellum.workflows.references.execution_count import ExecutionCountReference
+from vellum.workflows.references.node import NodeReference
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.state.context import WorkflowContext
+from vellum.workflows.types.core import MergeBehavior
+from vellum.workflows.types.generics import StateType
+from vellum.workflows.types.utils import get_class_attr_names, get_original_base, infer_types
+
+
+def is_nested_class(nested: Any, parent: Type) -> bool:
+    return (
+        inspect.isclass(nested)
+        # If a class is defined within a function, we don't consider it nested in the class defining that function
+        # The example of this is a Subworkflow defined within TryNode.wrap()
+        and (len(nested.__qualname__.split(".")) < 2 or nested.__qualname__.split(".")[-2] != "<locals>")
+        and nested.__module__ == parent.__module__
+        and (nested.__qualname__.startswith(parent.__name__) or nested.__qualname__.startswith(parent.__qualname__))
+    ) or any(is_nested_class(nested, base) for base in parent.__bases__)
+
+
+class BaseNodeMeta(type):
+    def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
+        # TODO: Inherit the inner Output classes from every base class.
+        #   https://app.shortcut.com/vellum/story/4007/support-auto-inheriting-parent-node-outputs
+
+        if "Outputs" not in dct:
+            for base in reversed(bases):
+                if hasattr(base, "Outputs"):
+                    dct["Outputs"] = type(f"{name}.Outputs", (base.Outputs,), {"__module__": dct["__module__"]})
+                    break
+            else:
+                raise ValueError("Outputs class not found in base classes")
+
+        if "Ports" in dct:
+            dct["Ports"] = type(
+                f"{name}.Ports",
+                (NodePorts,),
+                dict(dct["Ports"].__dict__),
+            )
+        else:
+            for base in reversed(bases):
+                if issubclass(base, BaseNode):
+                    ports_dct = {p.name: Port(default=p.default) for p in base.Ports}
+                    ports_dct["__module__"] = dct["__module__"]
+                    dct["Ports"] = type(f"{name}.Ports", (NodePorts,), ports_dct)
+                    break
+
+        if "Execution" not in dct:
+            for base in reversed(bases):
+                if issubclass(base, BaseNode):
+                    dct["Execution"] = type(f"{name}.Execution", (base.Execution,), {"__module__": dct["__module__"]})
+                    break
+
+        if "Trigger" not in dct:
+            for base in reversed(bases):
+                if issubclass(base, BaseNode):
+                    trigger_dct = {**base.Trigger.__dict__, "__module__": dct["__module__"]}
+                    dct["Trigger"] = type(f"{name}.Trigger", (base.Trigger,), trigger_dct)
+                    break
+
+        cls = super().__new__(mcs, name, bases, dct)
+        node_class = cast(Type["BaseNode"], cls)
+
+        node_class.Outputs._node_class = node_class
+
+        # Add cls to relevant nested classes, since python should've been doing this by default
+        for port in node_class.Ports:
+            port.node_class = node_class
+
+        node_class.Execution.node_class = node_class
+        node_class.Trigger.node_class = node_class
+        node_class.ExternalInputs.__parent_class__ = node_class
+        return node_class
+
+    @property
+    def _localns(cls) -> Dict[str, Any]:
+        from vellum.workflows.workflows.base import BaseWorkflow
+
+        return {"BaseWorkflow": BaseWorkflow}
+
+    def __getattribute__(cls, name: str) -> Any:
+        attribute = super().__getattribute__(name)
+        if (
+            name.startswith("_")
+            or inspect.isfunction(attribute)
+            or inspect.ismethod(attribute)
+            or is_nested_class(attribute, cls)
+            or isinstance(attribute, (property, cached_property))
+            or not issubclass(cls, BaseNode)
+        ):
+            return attribute
+
+        types = infer_types(cls, name, cls._localns)
+        return NodeReference(
+            name=name,
+            types=types,
+            instance=attribute,
+            node_class=cls,
+        )
+
+    def __rshift__(cls, other_cls: GraphTarget) -> Graph:
+        if not issubclass(cls, BaseNode):
+            raise ValueError("BaseNodeMeta can only be extended from subclasses of BaseNode")
+
+        if not cls.Ports._default_port:
+            raise ValueError("No default port found on node")
+
+        if isinstance(other_cls, Graph) or isinstance(other_cls, set):
+            return Graph.from_node(cls) >> other_cls
+
+        return cls.Ports._default_port >> other_cls
+
+    def __rrshift__(cls, other_cls: GraphTarget) -> Graph:
+        if not issubclass(cls, BaseNode):
+            raise ValueError("BaseNodeMeta can only be extended from subclasses of BaseNode")
+
+        if not isinstance(other_cls, set):
+            other_cls = {other_cls}
+
+        return Graph.from_set(other_cls) >> cls
+
+    def __repr__(self) -> str:
+        return f"{self.__module__}.{self.__qualname__}"
+
+    def __iter__(cls) -> Iterator[NodeReference]:
+        # We iterate through the inheritance hierarchy to find all the OutputDescriptors attached to this Outputs class.
+        # __mro__ is the method resolution order, which is the order in which base classes are resolved.
+        yielded_attr_names: Set[str] = {"state"}
+
+        for resolved_cls in cls.__mro__:
+            attr_names = get_class_attr_names(resolved_cls)
+            for attr_name in attr_names:
+                if attr_name in yielded_attr_names:
+                    continue
+
+                attr_value = getattr(resolved_cls, attr_name, UNDEF)
+                if not isinstance(attr_value, NodeReference):
+                    continue
+
+                yield attr_value
+                yielded_attr_names.add(attr_name)
+
+
+class _BaseNodeTriggerMeta(type):
+    def __eq__(self, other: Any) -> bool:
+        """
+        We need to include custom eq logic to prevent infinite loops during ipython reloading.
+        """
+
+        if not isinstance(other, _BaseNodeTriggerMeta):
+            return False
+
+        if not self.__name__.endswith(".Trigger") or not other.__name__.endswith(".Trigger"):
+            return super().__eq__(other)
+
+        self_trigger_class = cast(Type["BaseNode.Trigger"], self)
+        other_trigger_class = cast(Type["BaseNode.Trigger"], other)
+
+        return self_trigger_class.node_class.__name__ == other_trigger_class.node_class.__name__
+
+
+class _BaseNodeExecutionMeta(type):
+    def __getattribute__(cls, name: str) -> Any:
+        if name.startswith("count") and issubclass(cls, BaseNode.Execution):
+            return ExecutionCountReference(cls.node_class)
+
+        return super().__getattribute__(name)
+
+    def __eq__(self, other: Any) -> bool:
+        """
+        We need to include custom eq logic to prevent infinite loops during ipython reloading.
+        """
+
+        if not isinstance(other, _BaseNodeExecutionMeta):
+            return False
+
+        if not self.__name__.endswith(".Execution") or not other.__name__.endswith(".Execution"):
+            return super().__eq__(other)
+
+        self_execution_class = cast(Type["BaseNode.Execution"], self)
+        other_execution_class = cast(Type["BaseNode.Execution"], other)
+
+        return self_execution_class.node_class.__name__ == other_execution_class.node_class.__name__
+
+
+class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
+    state: StateType
+    _context: WorkflowContext
+    _inputs: MappingProxyType[NodeReference, Any]
+    _is_wrapped_node: bool = False
+
+    class ExternalInputs(BaseInputs):
+        __descriptor_class__ = ExternalInputReference
+
+    # TODO: Consider using metaclasses to prevent the need for users to specify that the
+    #   "Outputs" class inherits from "BaseOutputs" and do so automatically.
+    #   https://app.shortcut.com/vellum/story/4008/auto-inherit-basenodeoutputs-in-outputs-classes
+    class Outputs(BaseOutputs):
+        _node_class: Optional[Type["BaseNode"]] = None
+        pass
+
+    class Ports(NodePorts):
+        default = Port(default=True)
+
+    class Trigger(metaclass=_BaseNodeTriggerMeta):
+        node_class: Type["BaseNode"]
+        merge_behavior = MergeBehavior.AWAIT_ANY
+
+        @classmethod
+        def should_initiate(
+            cls, state: StateType, dependencies: Set["Type[BaseNode]"], invoked_by: "Optional[Edge]" = None
+        ) -> bool:
+            """
+            Determines whether a Node's execution should be initiated. Override this method to define custom
+            trigger criteria.
+            """
+
+            if cls.merge_behavior == MergeBehavior.AWAIT_ANY:
+                if not invoked_by:
+                    return True
+
+                is_ready = not state.meta.node_execution_cache.is_node_initiated(cls.node_class)
+
+                invoked_identifier = str(invoked_by.from_port.node_class)
+                node_identifier = str(cls.node_class)
+
+                dependencies_invoked = state.meta.node_execution_cache.dependencies_invoked[node_identifier]
+                dependencies_invoked.add(invoked_identifier)
+                if all(str(dep) in dependencies_invoked for dep in dependencies):
+                    del state.meta.node_execution_cache.dependencies_invoked[node_identifier]
+
+                return is_ready
+
+            if cls.merge_behavior == MergeBehavior.AWAIT_ALL:
+                if not invoked_by:
+                    return True
+
+                if state.meta.node_execution_cache.is_node_initiated(cls.node_class):
+                    return False
+
+                """
+                A node utilizing an AWAIT_ALL merge strategy will only be considered ready for the Nth time
+                when all of its dependencies have been executed N times.
+                """
+                current_node_execution_count = state.meta.node_execution_cache.get_execution_count(cls.node_class)
+                is_ready_outcome = all(
+                    state.meta.node_execution_cache.get_execution_count(dep) == current_node_execution_count + 1
+                    for dep in dependencies
+                )
+
+                return is_ready_outcome
+
+            raise NodeException(message="Invalid Trigger Node Specification", code=VellumErrorCode.INVALID_INPUTS)
+
+    class Execution(metaclass=_BaseNodeExecutionMeta):
+        node_class: Type["BaseNode"]
+        count: int
+
+    def __init__(self, *, state: Optional[StateType] = None, context: Optional[WorkflowContext] = None):
+        if state:
+            self.state = state
+        else:
+            # Instantiate a default state class if one is not provided, for ease of testing
+
+            original_base = get_original_base(self.__class__)
+
+            args = get_args(original_base)
+            state_type = args[0]
+
+            if isinstance(state_type, TypeVar):
+                state_type = BaseState
+
+            self.state = state_type()
+
+        self._context = context or WorkflowContext()
+        inputs: Dict[str, Any] = {}
+        for descriptor in self.__class__:
+            if not descriptor.instance:
+                continue
+
+            resolved_value = resolve_value(descriptor.instance, self.state, path=descriptor.name, memo=inputs)
+            setattr(self, descriptor.name, resolved_value)
+
+        # Resolve descriptors set as defaults to the outputs class
+        def _outputs_post_init(outputs_self: "BaseNode.Outputs", **kwargs: Any) -> None:
+            for node_output_descriptor in self.Outputs:
+                if node_output_descriptor.name in kwargs:
+                    continue
+
+                if isinstance(node_output_descriptor.instance, BaseDescriptor):
+                    setattr(
+                        outputs_self,
+                        node_output_descriptor.name,
+                        node_output_descriptor.instance.resolve(self.state),
+                    )
+            delattr(self.Outputs, "_outputs_post_init")
+
+        setattr(self.Outputs, "_outputs_post_init", _outputs_post_init)
+
+        # We only want to store the attributes that were actually set as inputs, not every attribute that exists.
+        all_inputs = {}
+        for key, value in inputs.items():
+            path_parts = key.split(".")
+            node_attribute_discriptor = getattr(self.__class__, path_parts[0])
+            inputs_key = reduce(lambda acc, part: acc[part], path_parts[1:], node_attribute_discriptor)
+            all_inputs[inputs_key] = value
+
+        self._inputs = MappingProxyType(all_inputs)
+
+    def run(self) -> Union[BaseOutputs, Iterator[BaseOutput]]:
+        return self.Outputs()
+
+    def __repr__(self) -> str:
+        return str(self.__class__)

vellum/workflows/nodes/bases/base_subworkflow_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import BaseSubworkflowNode
+
+__all__ = [
+    "BaseSubworkflowNode",
+]

vellum/workflows/nodes/bases/base_subworkflow_node/node.py

@@ -0,0 +1,10 @@
+from typing import ClassVar, Generic
+
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.types.core import EntityInputsInterface
+from vellum.workflows.types.generics import StateType
+
+
+class BaseSubworkflowNode(BaseNode[StateType], Generic[StateType]):
+    # Inputs that are passed to the Subworkflow
+    subworkflow_inputs: ClassVar[EntityInputsInterface] = {}

vellum/workflows/nodes/bases/tests/test_base_node.py

@@ -0,0 +1,125 @@
+from typing import Optional
+
+from vellum.core.pydantic_utilities import UniversalBaseModel
+
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.state.base import BaseState, StateMeta
+
+
+def test_base_node__node_resolution__unset_pydantic_fields():
+    # GIVEN a pydantic class with an optional field
+    class Data(UniversalBaseModel):
+        hello: str
+        world: Optional[str] = None
+
+    # AND a node that uses the pydantic class only setting one field
+    my_data = Data(hello="hi")
+
+    class MyNode(BaseNode):
+        data = my_data
+
+    # WHEN the node is initialized
+    node = MyNode()
+
+    # THEN the node is initialized with the correct data
+    assert node.data.dict() == my_data.dict()
+
+
+def test_base_node__node_resolution__descriptor_in_dict():
+    # GIVEN an Input and State class
+    class Inputs(BaseInputs):
+        hello: str
+
+    class State(BaseState):
+        pass
+
+    # AND a node referencing a descriptor in a dict
+    class MyNode(BaseNode):
+        data = {"world": Inputs.hello}
+
+    # WHEN the node is initialized
+    node = MyNode(
+        state=State(
+            meta=StateMeta(
+                workflow_inputs=Inputs(hello="hi"),
+            )
+        )
+    )
+
+    # THEN the node is initialized with the correct data
+    assert node.data["world"] == "hi"
+
+    # AND there are inputs compiled
+    assert node._inputs == {
+        MyNode.data["world"]: "hi",
+    }
+
+
+def test_base_node__node_resolution__descriptor_in_pydantic():
+    # GIVEN an Input and State class
+    class Inputs(BaseInputs):
+        hello: str
+
+    class State(BaseState):
+        pass
+
+    class Data(UniversalBaseModel):
+        world: str
+
+    class MyNode(BaseNode):
+        data = Data(world=Inputs.hello)
+
+    # WHEN the node is initialized
+    node = MyNode(
+        state=State(
+            meta=StateMeta(
+                workflow_inputs=Inputs(hello="hi"),
+            )
+        )
+    )
+
+    # THEN the node is initialized with the correct data
+    assert node.data.world == "hi"
+
+    # AND there are inputs compiled
+    assert node._inputs == {
+        MyNode.data["world"]: "hi",
+    }
+
+
+def test_base_node__node_resolution__no_inputs():
+    # GIVEN a node that defines some attributes
+    class MyNode(BaseNode):
+        foo = "bar"
+        baz = 1
+
+    # WHEN the node is initialized
+    node = MyNode()
+
+    # THEN the node is initialized with the correct data
+    assert node.foo == "bar"
+    assert node.baz == 1
+
+    # AND there are no inputs
+    assert node._inputs == {}
+
+
+def test_base_node__node_resolution__coalesce_constants():
+    # GIVEN a State class
+    class State(BaseState):
+        pass
+
+    # AND a node that uses the coalesce operator with constants
+    class FirstNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            empty: str
+
+    class MyNode(BaseNode):
+        value = FirstNode.Outputs.empty.coalesce("world")
+
+    # WHEN the node is initialized
+    node = MyNode(state=State())
+
+    # THEN the node is initialized with the correct data
+    assert node.value == "world"

vellum/workflows/nodes/core/__init__.py

@@ -0,0 +1,16 @@
+from vellum.workflows.nodes.displayable.bases.api_node import BaseAPINode
+
+from .error_node import ErrorNode
+from .inline_subworkflow_node import InlineSubworkflowNode
+from .map_node import MapNode
+from .retry_node import RetryNode
+from .templating_node import TemplatingNode
+
+__all__ = [
+    "BaseAPINode",
+    "ErrorNode",
+    "InlineSubworkflowNode",
+    "MapNode",
+    "RetryNode",
+    "TemplatingNode",
+]

vellum/workflows/nodes/core/error_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import ErrorNode
+
+__all__ = [
+    "ErrorNode",
+]

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

@@ -0,0 +1,26 @@
+from typing import ClassVar, Union
+
+from vellum.workflows.errors.types import VellumError, VellumErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.bases.base import BaseNode
+
+
+class ErrorNode(BaseNode):
+    """
+    Used to raise an error to reject the surrounding Workflow.
+
+    error: Union[str, VellumError] - The error to raise.
+    """
+
+    error: ClassVar[Union[str, VellumError]]
+
+    def run(self) -> BaseNode.Outputs:
+        if isinstance(self.error, str):
+            raise NodeException(message=self.error, code=VellumErrorCode.USER_DEFINED_ERROR)
+        elif isinstance(self.error, VellumError):
+            raise NodeException(message=self.error.message, code=self.error.code)
+        else:
+            raise NodeException(
+                message=f"Error node received an unexpected input type: {self.error.__class__}",
+                code=VellumErrorCode.INVALID_INPUTS,
+            )

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

@@ -0,0 +1,5 @@
+from .node import InlineSubworkflowNode
+
+__all__ = [
+    "InlineSubworkflowNode",
+]

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

@@ -0,0 +1,73 @@
+from typing import TYPE_CHECKING, Generic, Iterator, Optional, Set, Type, TypeVar
+
+from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.bases.base_subworkflow_node import BaseSubworkflowNode
+from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.generics import StateType, WorkflowInputsType
+
+if TYPE_CHECKING:
+    from vellum.workflows.workflows.base import BaseWorkflow
+
+InnerStateType = TypeVar("InnerStateType", bound=BaseState)
+
+
+class InlineSubworkflowNode(BaseSubworkflowNode[StateType], Generic[StateType, WorkflowInputsType, InnerStateType]):
+    """
+    Used to execute a Subworkflow defined inline.
+
+    subworkflow: Type["BaseWorkflow[WorkflowInputsType, InnerStateType]"] - The Subworkflow to execute
+    """
+
+    subworkflow: Type["BaseWorkflow[WorkflowInputsType, InnerStateType]"]
+
+    def run(self) -> Iterator[BaseOutput]:
+        subworkflow = self.subworkflow(
+            parent_state=self.state,
+        )
+        subworkflow_stream = subworkflow.stream(
+            inputs=self._compile_subworkflow_inputs(),
+        )
+
+        outputs: Optional[BaseOutputs] = None
+        fulfilled_output_names: Set[str] = set()
+
+        for event in subworkflow_stream:
+            if event.name == "workflow.execution.streaming":
+                if event.output.is_fulfilled:
+                    fulfilled_output_names.add(event.output.name)
+                yield event.output
+            elif event.name == "workflow.execution.fulfilled":
+                outputs = event.outputs
+            elif event.name == "workflow.execution.rejected":
+                error = event.error
+                if error.code in VellumErrorCode._value2member_map_:
+                    raise NodeException(
+                        message=error.message,
+                        code=VellumErrorCode(error.code),
+                    )
+                else:
+                    raise NodeException(
+                        message=error.message,
+                        code=VellumErrorCode.INTERNAL_ERROR,
+                    )
+
+        if outputs is None:
+            raise NodeException(
+                message="Expected to receive outputs from Workflow Deployment",
+                code=VellumErrorCode.INTERNAL_ERROR,
+            )
+
+        # For any outputs somehow in our final fulfilled outputs array,
+        # but not fulfilled by the stream.
+        for output_descriptor, output_value in outputs:
+            if output_descriptor.name not in fulfilled_output_names:
+                yield BaseOutput(
+                    name=output_descriptor.name,
+                    value=output_value,
+                )
+
+    def _compile_subworkflow_inputs(self) -> WorkflowInputsType:
+        inputs_class = self.subworkflow.get_inputs_class()
+        return inputs_class(**self.subworkflow_inputs)

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

@@ -0,0 +1,5 @@
+from .node import MapNode
+
+__all__ = [
+    "MapNode",
+]