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

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

@@ -0,0 +1,147 @@
+from collections import defaultdict
+from queue import Empty, Queue
+from threading import Thread
+from typing import TYPE_CHECKING, Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, overload
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.generics import NodeType, StateType
+
+if TYPE_CHECKING:
+    from vellum.workflows import BaseWorkflow
+    from vellum.workflows.events.workflow import WorkflowEvent
+
+MapNodeItemType = TypeVar("MapNodeItemType")
+
+
+class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
+    """
+    Used to map over a list of items and execute a Subworkflow on each iteration.
+
+    items: List[MapNodeItemType] - The items to map over
+    subworkflow: Type["BaseWorkflow[SubworkflowInputs, BaseState]"] - The Subworkflow to execute on each iteration
+    concurrency: Optional[int] = None - The maximum number of concurrent subworkflow executions
+    """
+
+    items: List[MapNodeItemType]
+    subworkflow: Type["BaseWorkflow"]
+    concurrency: Optional[int] = None
+
+    class Outputs(BaseOutputs):
+        mapped_items: list
+
+    class SubworkflowInputs(BaseInputs):
+        # TODO: Both type: ignore's below are believed to be incorrect and both have the following error:
+        # Type variable "workflows.nodes.map_node.map_node.MapNodeItemType" is unbound
+        # https://app.shortcut.com/vellum/story/4118
+
+        item: MapNodeItemType  # type: ignore[valid-type]
+        index: int
+        all_items: List[MapNodeItemType]  # type: ignore[valid-type]
+
+    def run(self) -> Outputs:
+        mapped_items: Dict[str, List] = defaultdict(list)
+        for output_descripter in self.subworkflow.Outputs:
+            mapped_items[output_descripter.name] = [None] * len(self.items)
+
+        self._event_queue: Queue[Tuple[int, WorkflowEvent]] = Queue()
+        fulfilled_iterations: List[bool] = []
+        for index, item in enumerate(self.items):
+            fulfilled_iterations.append(False)
+            thread = Thread(target=self._run_subworkflow, kwargs={"item": item, "index": index})
+            thread.start()
+
+        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]
+
+                if terminal_event.name == "workflow.execution.fulfilled":
+                    workflow_output_vars = vars(terminal_event.outputs)
+
+                    for output_name in workflow_output_vars:
+                        output_mapped_items = mapped_items[output_name]
+                        output_mapped_items[index] = workflow_output_vars[output_name]
+
+                    fulfilled_iterations[index] = True
+                    if all(fulfilled_iterations):
+                        break
+                elif terminal_event.name == "workflow.execution.paused":
+                    raise NodeException(
+                        code=VellumErrorCode.INVALID_OUTPUTS,
+                        message=f"Subworkflow unexpectedly paused on iteration {index}",
+                    )
+                elif terminal_event.name == "workflow.execution.rejected":
+                    raise NodeException(
+                        f"Subworkflow failed on iteration {index} with error: {terminal_event.error.message}",
+                        code=terminal_event.error.code,
+                    )
+        except Empty:
+            pass
+
+        return self.Outputs(**mapped_items)
+
+    def _run_subworkflow(self, *, item: MapNodeItemType, index: int) -> None:
+        subworkflow = self.subworkflow(parent_state=self.state)
+        events = subworkflow.stream(inputs=self.SubworkflowInputs(index=index, item=item, all_items=self.items))
+
+        for event in events:
+            self._event_queue.put((index, event))
+
+    @overload
+    @classmethod
+    def wrap(cls, items: List[MapNodeItemType]) -> Callable[..., Type["MapNode[StateType, MapNodeItemType]"]]: ...
+
+    # TODO: We should be able to do this overload automatically as we do with node attributes
+    # https://app.shortcut.com/vellum/story/5289
+    @overload
+    @classmethod
+    def wrap(
+        cls, items: BaseDescriptor[List[MapNodeItemType]]
+    ) -> Callable[..., Type["MapNode[StateType, MapNodeItemType]"]]: ...
+
+    @classmethod
+    def wrap(
+        cls, items: Union[List[MapNodeItemType], BaseDescriptor[List[MapNodeItemType]]]
+    ) -> Callable[..., Type["MapNode[StateType, MapNodeItemType]"]]:
+        _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
+
+        return decorator

vellum/workflows/nodes/core/map_node/tests/test_node.py

@@ -0,0 +1,65 @@
+import time
+
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.core.map_node.node import MapNode
+from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.state.base import BaseState, StateMeta
+
+
+def test_map_node__use_parent_inputs_and_state():
+    # GIVEN a parent workflow Inputs and State
+    class Inputs(BaseInputs):
+        foo: str
+
+    class State(BaseState):
+        bar: str
+
+    # AND a map node that is configured to use the parent's inputs and state
+    @MapNode.wrap(items=[1, 2, 3])
+    class TestNode(BaseNode):
+        item = MapNode.SubworkflowInputs.item
+        foo = Inputs.foo
+        bar = State.bar
+
+        class Outputs(BaseOutputs):
+            value: str
+
+        def run(self) -> Outputs:
+            return self.Outputs(value=f"{self.foo} {self.bar} {self.item}")
+
+    # WHEN the node is run
+    node = TestNode(
+        state=State(
+            bar="bar",
+            meta=StateMeta(workflow_inputs=Inputs(foo="foo")),
+        )
+    )
+    outputs = node.run()
+
+    # THEN the data is used successfully
+    assert outputs.value == ["foo bar 1", "foo bar 2", "foo bar 3"]
+
+
+def test_map_node__use_parallelism():
+    # GIVEN a map node that is configured to use the parent's inputs and state
+    @MapNode.wrap(items=list(range(10)))
+    class TestNode(BaseNode):
+        item = MapNode.SubworkflowInputs.item
+
+        class Outputs(BaseOutputs):
+            value: int
+
+        def run(self) -> Outputs:
+            time.sleep(0.03)
+            return self.Outputs(value=self.item + 1)
+
+    # WHEN the node is run
+    node = TestNode(state=BaseState())
+    start_ts = time.time_ns()
+    node.run()
+    end_ts = time.time_ns()
+
+    # THEN the node should have ran in parallel
+    run_time = (end_ts - start_ts) / 10**9
+    assert run_time < 0.1

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

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

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

@@ -0,0 +1,106 @@
+from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Optional, Type
+
+from vellum.workflows.errors.types import VellumErrorCode
+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.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):
+    """
+    Used to retry a Subworkflow a specified number of times.
+
+    max_attempts: int - The maximum number of attempts to retry the Subworkflow
+    retry_on_error_code: Optional[VellumErrorCode] = None - The error code to retry on
+    subworkflow: Type["BaseWorkflow[SubworkflowInputs, BaseState]"] - The Subworkflow to execute
+    """
+
+    max_attempts: int
+    retry_on_error_code: Optional[VellumErrorCode] = None
+    subworkflow: Type["BaseWorkflow[SubworkflowInputs, BaseState]"]
+
+    class SubworkflowInputs(BaseInputs):
+        attempt_number: int
+
+    def run(self) -> BaseNode.Outputs:
+        last_exception = Exception("max_attempts must be greater than 0")
+        for index in range(self.max_attempts):
+            attempt_number = index + 1
+            subworkflow = self.subworkflow(
+                parent_state=self.state,
+            )
+            terminal_event = subworkflow.run(
+                inputs=self.SubworkflowInputs(attempt_number=attempt_number),
+            )
+            if terminal_event.name == "workflow.execution.fulfilled":
+                node_outputs = self.Outputs()
+                workflow_output_vars = vars(terminal_event.outputs)
+
+                for output_name in workflow_output_vars:
+                    setattr(node_outputs, output_name, workflow_output_vars[output_name])
+
+                return node_outputs
+            elif terminal_event.name == "workflow.execution.paused":
+                last_exception = NodeException(
+                    code=VellumErrorCode.INVALID_OUTPUTS,
+                    message=f"Subworkflow unexpectedly paused on attempt {attempt_number}",
+                )
+                break
+            elif self.retry_on_error_code and self.retry_on_error_code != terminal_event.error.code:
+                last_exception = NodeException(
+                    code=VellumErrorCode.INVALID_OUTPUTS,
+                    message=f"""Unexpected rejection on attempt {attempt_number}: {terminal_event.error.code.value}.
+Message: {terminal_event.error.message}""",
+                )
+                break
+            else:
+                last_exception = Exception(terminal_event.error.message)
+
+        raise last_exception
+
+    @classmethod
+    def wrap(
+        cls, max_attempts: int, retry_on_error_code: Optional[VellumErrorCode] = 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

vellum/workflows/nodes/core/retry_node/tests/test_node.py

@@ -0,0 +1,93 @@
+import pytest
+
+from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.core.retry_node.node import RetryNode
+from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.state.base import BaseState, StateMeta
+
+
+def test_retry_node__retry_on_error_code__successfully_retried():
+    # GIVEN a retry node that is configured to retry on PROVIDER_ERROR
+    @RetryNode.wrap(max_attempts=3, retry_on_error_code=VellumErrorCode.PROVIDER_ERROR)
+    class TestNode(BaseNode):
+        attempt_number = RetryNode.SubworkflowInputs.attempt_number
+
+        class Outputs(BaseOutputs):
+            execution_count: int
+
+        def run(self) -> Outputs:
+            if self.attempt_number < 3:
+                raise NodeException(message="This will be retried", code=VellumErrorCode.PROVIDER_ERROR)
+
+            return self.Outputs(execution_count=self.attempt_number)
+
+    # WHEN the node is run and throws a PROVIDER_ERROR
+    node = TestNode(state=BaseState())
+    outputs = node.run()
+
+    # THEN the exception is retried
+    assert outputs.execution_count == 3
+
+
+def test_retry_node__retry_on_error_code__missed():
+    # GIVEN a retry node that is configured to retry on PROVIDER_ERROR
+    @RetryNode.wrap(max_attempts=3, retry_on_error_code=VellumErrorCode.PROVIDER_ERROR)
+    class TestNode(BaseNode):
+        attempt_number = RetryNode.SubworkflowInputs.attempt_number
+
+        class Outputs(BaseOutputs):
+            execution_count: int
+
+        def run(self) -> Outputs:
+            if self.attempt_number < 3:
+                raise Exception("This will not be retried")
+
+            return self.Outputs(execution_count=self.attempt_number)
+
+    # WHEN the node is run and throws a different exception
+    node = TestNode(state=BaseState())
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN the exception is not retried
+    assert (
+        exc_info.value.message
+        == "Unexpected rejection on attempt 1: INTERNAL_ERROR.\nMessage: This will not be retried"
+    )
+    assert exc_info.value.code == VellumErrorCode.INVALID_OUTPUTS
+
+
+def test_retry_node__use_parent_inputs_and_state():
+    # GIVEN a parent workflow Inputs and State
+    class Inputs(BaseInputs):
+        foo: str
+
+    class State(BaseState):
+        bar: str
+
+    # AND a retry node that uses the parent's inputs and state
+    @RetryNode.wrap(max_attempts=3, retry_on_error_code=VellumErrorCode.PROVIDER_ERROR)
+    class TestNode(BaseNode):
+        foo = Inputs.foo
+        bar = State.bar
+
+        class Outputs(BaseOutputs):
+            value: str
+
+        def run(self) -> Outputs:
+            return self.Outputs(value=f"{self.foo} {self.bar}")
+
+    # WHEN the node is run
+    node = TestNode(
+        state=State(
+            bar="bar",
+            meta=StateMeta(workflow_inputs=Inputs(foo="foo")),
+        )
+    )
+    outputs = node.run()
+
+    # THEN the data is used successfully
+    assert outputs.value == "foo bar"

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

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

vellum/workflows/nodes/core/templating_node/custom_filters.py

@@ -0,0 +1,12 @@
+import json
+from typing import Union
+
+
+def is_valid_json_string(value: Union[str, bytes]) -> bool:
+    """Determines whether the given value is a valid JSON string."""
+
+    try:
+        json.loads(value)
+    except ValueError:
+        return False
+    return True

vellum/workflows/nodes/core/templating_node/exceptions.py

@@ -0,0 +1,2 @@
+class JinjaTemplateError(Exception):
+    pass

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

@@ -0,0 +1,123 @@
+import datetime
+import itertools
+import json
+import random
+import re
+from typing import Any, Callable, ClassVar, Dict, Generic, Mapping, Tuple, Type, TypeVar, Union, get_args
+
+import dateutil.parser
+import pydash
+import pytz
+import yaml
+
+from vellum.workflows.errors import VellumErrorCode
+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.core.templating_node.custom_filters import is_valid_json_string
+from vellum.workflows.nodes.core.templating_node.exceptions import JinjaTemplateError
+from vellum.workflows.nodes.core.templating_node.render import render_sandboxed_jinja_template
+from vellum.workflows.types.core import EntityInputsInterface
+from vellum.workflows.types.generics import StateType
+from vellum.workflows.types.utils import get_original_base
+
+_DEFAULT_JINJA_GLOBALS: Dict[str, Any] = {
+    "datetime": datetime,
+    "dateutil": dateutil,
+    "itertools": itertools,
+    "json": json,
+    "pydash": pydash,
+    "pytz": pytz,
+    "random": random,
+    "re": re,
+    "yaml": yaml,
+}
+
+_DEFAULT_JINJA_CUSTOM_FILTERS: Dict[str, Callable[[Union[str, bytes]], bool]] = {
+    "is_valid_json_string": is_valid_json_string,
+}
+
+_OutputType = TypeVar("_OutputType")
+
+
+# TODO: Consolidate all dynamic output metaclasses
+# https://app.shortcut.com/vellum/story/5533
+class _TemplatingNodeMeta(BaseNodeMeta):
+    def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
+        parent = super().__new__(mcs, name, bases, dct)
+
+        if not isinstance(parent, _TemplatingNodeMeta):
+            raise ValueError("TemplatingNode must be created with the TemplatingNodeMeta metaclass")
+
+        parent.__dict__["Outputs"].__annotations__["result"] = parent.get_output_type()
+        return parent
+
+    def get_output_type(cls) -> Type:
+        original_base = get_original_base(cls)
+        all_args = get_args(original_base)
+
+        if len(all_args) < 2 or isinstance(all_args[1], TypeVar):
+            return str
+        else:
+            return all_args[1]
+
+
+class TemplatingNode(BaseNode[StateType], Generic[StateType, _OutputType], metaclass=_TemplatingNodeMeta):
+    """Used to render a Jinja template.
+
+    Useful for lightweight data transformations and complex string templating.
+    """
+
+    # The Jinja template to render.
+    template: ClassVar[str]
+
+    # The inputs to render the template with.
+    inputs: ClassVar[EntityInputsInterface]
+
+    jinja_globals: Dict[str, Any] = _DEFAULT_JINJA_GLOBALS
+    jinja_custom_filters: Mapping[str, Callable[[Union[str, bytes]], bool]] = _DEFAULT_JINJA_CUSTOM_FILTERS
+
+    class Outputs(BaseNode.Outputs):
+        # We use our mypy plugin to override the _OutputType with the actual output type
+        # for downstream references to this output.
+        result: _OutputType  # type: ignore[valid-type]
+
+    def _cast_rendered_template(self, rendered_template: str) -> Any:
+        original_base = get_original_base(self.__class__)
+        all_args = get_args(original_base)
+
+        if len(all_args) < 2 or isinstance(all_args[1], TypeVar):
+            output_type = str
+        else:
+            output_type = all_args[1]
+
+        if output_type is str:
+            return rendered_template
+
+        if output_type is float:
+            return float(rendered_template)
+
+        if output_type is int:
+            return int(rendered_template)
+
+        if output_type is bool:
+            return bool(rendered_template)
+
+        raise ValueError(f"Unsupported output type: {output_type}")
+
+    def run(self) -> Outputs:
+        rendered_template = self._render_template()
+        result = self._cast_rendered_template(rendered_template)
+
+        return self.Outputs(result=result)
+
+    def _render_template(self) -> str:
+        try:
+            return render_sandboxed_jinja_template(
+                template=self.template,
+                input_values=self.inputs,
+                jinja_custom_filters={**self.jinja_custom_filters},
+                jinja_globals=self.jinja_globals,
+            )
+        except JinjaTemplateError as e:
+            raise NodeException(message=str(e), code=VellumErrorCode.INVALID_TEMPLATE)

vellum/workflows/nodes/core/templating_node/render.py

@@ -0,0 +1,55 @@
+import json
+from typing import Any, Callable, Dict, Optional, Union
+
+from jinja2.sandbox import SandboxedEnvironment
+
+from vellum.workflows.nodes.core.templating_node.exceptions import JinjaTemplateError
+from vellum.workflows.state.encoder import DefaultStateEncoder
+
+
+def finalize(obj: Any) -> str:
+    if isinstance(obj, dict):
+        return json.dumps(obj, cls=DefaultStateEncoder)
+
+    return str(obj)
+
+
+def render_sandboxed_jinja_template(
+    *,
+    template: str,
+    input_values: Dict[str, Any],
+    jinja_custom_filters: Optional[Dict[str, Callable[[Union[str, bytes]], bool]]] = None,
+    jinja_globals: Optional[Dict[str, Any]] = None,
+) -> str:
+    """Render a Jinja template within a sandboxed environment."""
+
+    try:
+        environment = SandboxedEnvironment(
+            keep_trailing_newline=True,
+            finalize=finalize,
+        )
+
+        if jinja_custom_filters:
+            environment.filters.update(jinja_custom_filters)
+
+        jinja_template = environment.from_string(template)
+
+        if jinja_globals:
+            jinja_template.globals.update(jinja_globals)
+
+        rendered_template = jinja_template.render(input_values)
+    except json.JSONDecodeError as e:
+        if e.msg == "Invalid control character at":
+            raise JinjaTemplateError(
+                "Unable to render jinja template:\n"
+                "Cannot run json.loads() on JSON containing control characters. "
+                "Use json.loads(input, strict=False) instead.",
+            )
+
+        raise JinjaTemplateError(
+            f"Unable to render jinja template:\nCannot run json.loads() on invalid JSON\n{e.args[0]}"
+        )
+    except Exception as e:
+        raise JinjaTemplateError(f"Unable to render jinja template:\n{e.args[0]}")
+
+    return rendered_template

vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py

@@ -0,0 +1,21 @@
+import json
+
+from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
+
+
+def test_templating_node__dict_output():
+    # GIVEN a templating node with a dict input that just returns it
+    class TemplateNode(TemplatingNode):
+        template = "{{ data }}"
+        inputs = {
+            "data": {
+                "key": "value",
+            }
+        }
+
+    # WHEN the node is run
+    node = TemplateNode()
+    outputs = node.run()
+
+    # THEN the output is json serializable
+    assert json.loads(outputs.result) == {"key": "value"}

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

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