vellum-ai 0.14.5__py3-none-any.whl → 0.14.7__py3-none-any.whl

vellum/workflows/nodes/displayable/subworkflow_deployment_node/tests/test_node.py

@@ -3,6 +3,7 @@ from datetime import datetime
 from uuid import uuid4
 from typing import Any, Iterator, List
 
+from vellum.client.core.api_error import ApiError
 from vellum.client.types.chat_message import ChatMessage
 from vellum.client.types.chat_message_request import ChatMessageRequest
 from vellum.client.types.workflow_execution_workflow_result_event import WorkflowExecutionWorkflowResultEvent
@@ -11,6 +12,8 @@ from vellum.client.types.workflow_request_chat_history_input_request import Work
 from vellum.client.types.workflow_request_json_input_request import WorkflowRequestJsonInputRequest
 from vellum.client.types.workflow_result_event import WorkflowResultEvent
 from vellum.client.types.workflow_stream_event import WorkflowStreamEvent
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.subworkflow_deployment_node.node import SubworkflowDeploymentNode
 
 
@@ -129,3 +132,165 @@ def test_run_workflow__any_array(vellum_client):
     assert call_kwargs["inputs"] == [
         WorkflowRequestJsonInputRequest(name="fruits", value=["apple", "banana", "cherry"]),
     ]
+
+
+def test_run_workflow__no_deployment():
+    """Confirm that we raise error when running a subworkflow deployment node with no deployment attribute set"""
+
+    # GIVEN a Subworkflow Deployment Node
+    class ExampleSubworkflowDeploymentNode(SubworkflowDeploymentNode):
+        subworkflow_inputs = {
+            "fruits": ["apple", "banana", "cherry"],
+        }
+
+    # WHEN/THEN running the node should raise a NodeException
+    node = ExampleSubworkflowDeploymentNode()
+    with pytest.raises(NodeException) as exc_info:
+        list(node.run())
+
+    # AND the error message should be correct
+    assert exc_info.value.code == WorkflowErrorCode.NODE_EXECUTION
+    assert "Expected subworkflow deployment attribute to be either a UUID or STR, got None instead" in str(
+        exc_info.value
+    )
+
+
+def test_run_workflow__hyphenated_output(vellum_client):
+    """Confirm that we can successfully handle subworkflow outputs with hyphenated names"""
+
+    # GIVEN a Subworkflow Deployment Node
+    class ExampleSubworkflowDeploymentNode(SubworkflowDeploymentNode):
+        deployment = "example_subworkflow_deployment"
+        subworkflow_inputs = {
+            "test_input": "test_value",
+        }
+
+        class Outputs(SubworkflowDeploymentNode.Outputs):
+            final_output_copy: str
+
+    # AND we know what the Subworkflow Deployment will respond with
+    def generate_subworkflow_events(*args: Any, **kwargs: Any) -> Iterator[WorkflowStreamEvent]:
+        execution_id = str(uuid4())
+        expected_events: List[WorkflowStreamEvent] = [
+            WorkflowExecutionWorkflowResultEvent(
+                execution_id=execution_id,
+                data=WorkflowResultEvent(
+                    id=str(uuid4()),
+                    state="INITIATED",
+                    ts=datetime.now(),
+                ),
+            ),
+            WorkflowExecutionWorkflowResultEvent(
+                execution_id=execution_id,
+                data=WorkflowResultEvent(
+                    id=str(uuid4()),
+                    state="FULFILLED",
+                    ts=datetime.now(),
+                    outputs=[
+                        WorkflowOutputString(
+                            id=str(uuid4()),
+                            name="final-output_copy",  # Note the hyphen here
+                            value="test success",
+                        )
+                    ],
+                ),
+            ),
+        ]
+        yield from expected_events
+
+    vellum_client.execute_workflow_stream.side_effect = generate_subworkflow_events
+
+    # WHEN we run the node
+    node = ExampleSubworkflowDeploymentNode()
+    events = list(node.run())
+
+    # THEN the node should have completed successfully
+    assert events[-1].name == "final_output_copy"  # Note the underscore here
+    assert events[-1].value == "test success"
+
+
+@pytest.mark.parametrize(
+    ["exception", "expected_code", "expected_message"],
+    [
+        (
+            ApiError(status_code=400, body={"detail": "Missing required input variable: 'foo'"}),
+            WorkflowErrorCode.INVALID_INPUTS,
+            "Missing required input variable: 'foo'",
+        ),
+        (
+            ApiError(status_code=400, body={"message": "Missing required input variable: 'foo'"}),
+            WorkflowErrorCode.INVALID_INPUTS,
+            "Failed to execute Subworkflow Deployment",
+        ),
+        (
+            ApiError(status_code=400, body="Missing required input variable: 'foo'"),
+            WorkflowErrorCode.INTERNAL_ERROR,
+            "Failed to execute Subworkflow Deployment",
+        ),
+        (
+            ApiError(status_code=None, body={"detail": "Missing required input variable: 'foo'"}),
+            WorkflowErrorCode.INTERNAL_ERROR,
+            "Failed to execute Subworkflow Deployment",
+        ),
+        (
+            ApiError(status_code=500, body={"detail": "Missing required input variable: 'foo'"}),
+            WorkflowErrorCode.INTERNAL_ERROR,
+            "Failed to execute Subworkflow Deployment",
+        ),
+    ],
+    ids=["400", "invalid_dict", "invalid_body", "no_status_code", "500"],
+)
+def test_subworkflow_deployment_node__api_error__invalid_inputs_node_exception(
+    vellum_client, exception, expected_code, expected_message
+):
+    # GIVEN a prompt node with an invalid model name
+    class MyNode(SubworkflowDeploymentNode):
+        deployment = "example_subworkflow_deployment"
+        subworkflow_inputs = {
+            "not_foo": "bar",
+        }
+
+    # AND the Subworkflow Deployment API call fails
+    def _side_effect(*args: Any, **kwargs: Any) -> Iterator[WorkflowStreamEvent]:
+        if kwargs.get("_mock_condition_to_induce_an_error"):
+            yield WorkflowExecutionWorkflowResultEvent(
+                execution_id=str(uuid4()),
+                data=WorkflowResultEvent(
+                    id=str(uuid4()),
+                    state="INITIATED",
+                    ts=datetime.now(),
+                ),
+            )
+        else:
+            raise exception
+
+    # AND the vellum client execute workflow stream raises a 4xx error
+    vellum_client.execute_workflow_stream.side_effect = _side_effect
+
+    # WHEN the node is run
+    with pytest.raises(NodeException) as e:
+        list(MyNode().run())
+
+    # THEN the node raises the correct NodeException
+    assert e.value.code == expected_code
+    assert e.value.message == expected_message
+
+
+def test_subworkflow_deployment_node__immediate_api_error__node_exception(vellum_client):
+    # GIVEN a prompt node with an invalid model name
+    class MyNode(SubworkflowDeploymentNode):
+        deployment = "example_subworkflow_deployment"
+        subworkflow_inputs = {
+            "not_foo": "bar",
+        }
+
+    # AND the vellum client execute workflow stream raises a 4xx error
+    vellum_client.execute_workflow_stream.side_effect = ApiError(status_code=404, body={"detail": "Not found"})
+
+    # WHEN the node is run
+    with pytest.raises(NodeException) as e:
+        list(MyNode().run())
+
+    # THEN the node raises the correct NodeException
+    assert e.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert e.value.message == "Not found"

vellum/workflows/nodes/displayable/tests/test_text_prompt_deployment_node.py

@@ -74,5 +74,7 @@ def test_text_prompt_deployment_node__basic(vellum_client):
         prompt_deployment_name="my-deployment",
         raw_overrides=OMIT,
         release_tag="LATEST",
-        request_options={"additional_body_parameters": {"execution_context": {"parent_context": None}}},
+        request_options={
+            "additional_body_parameters": {"execution_context": {"parent_context": None, "trace_id": None}}
+        },
     )

vellum/workflows/outputs/base.py

@@ -32,7 +32,7 @@ class BaseOutput(Generic[_Delta, _Accumulated]):
         if value is not undefined and delta is not undefined:
             raise ValueError("Cannot set both value and delta")
 
-        self._name = name
+        self._name = name.replace("-", "_")  # Convert hyphens to underscores for valid python variable names
         self._value = value
         self._delta = delta
 

vellum/workflows/runner/runner.py

@@ -7,7 +7,7 @@ from uuid import UUID
 from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, Optional, Sequence, Set, Tuple, Type, Union
 
 from vellum.workflows.constants import undefined
-from vellum.workflows.context import execution_context, get_parent_context
+from vellum.workflows.context import ExecutionContext, execution_context, get_execution_context, get_parent_context
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges.edge import Edge
 from vellum.workflows.errors import WorkflowError, WorkflowErrorCode
@@ -29,7 +29,7 @@ from vellum.workflows.events.node import (
     NodeExecutionRejectedBody,
     NodeExecutionStreamingBody,
 )
-from vellum.workflows.events.types import BaseEvent, NodeParentContext, ParentContext, WorkflowParentContext
+from vellum.workflows.events.types import BaseEvent, NodeParentContext, WorkflowParentContext
 from vellum.workflows.events.workflow import (
     WorkflowExecutionFulfilledBody,
     WorkflowExecutionInitiatedBody,
@@ -75,8 +75,8 @@ class WorkflowRunner(Generic[StateType]):
         external_inputs: Optional[ExternalInputsArg] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
         node_output_mocks: Optional[MockNodeExecutionArg] = None,
-        parent_context: Optional[ParentContext] = None,
         max_concurrency: Optional[int] = None,
+        init_execution_context: Optional[ExecutionContext] = None,
     ):
         if state and external_inputs:
             raise ValueError("Can only run a Workflow providing one of state or external inputs, not both")
@@ -98,6 +98,11 @@ class WorkflowRunner(Generic[StateType]):
         elif external_inputs:
             self._initial_state = self.workflow.get_most_recent_state()
             for descriptor, value in external_inputs.items():
+                if not any(isinstance(value, type_) for type_ in descriptor.types):
+                    raise NodeException(
+                        f"Invalid external input type for {descriptor.name}",
+                        code=WorkflowErrorCode.INVALID_INPUTS,
+                    )
                 self._initial_state.meta.external_inputs[descriptor] = value
 
             self._entrypoints = [
@@ -133,7 +138,8 @@ class WorkflowRunner(Generic[StateType]):
 
         self._active_nodes_by_execution_id: Dict[UUID, BaseNode[StateType]] = {}
         self._cancel_signal = cancel_signal
-        self._parent_context = get_parent_context() or parent_context
+        self._execution_context = init_execution_context or get_execution_context()
+        self._parent_context = self._execution_context.parent_context
 
         setattr(
             self._initial_state,
@@ -196,7 +202,7 @@ class WorkflowRunner(Generic[StateType]):
                     break
 
             if not was_mocked:
-                with execution_context(parent_context=updated_parent_context):
+                with execution_context(parent_context=updated_parent_context, trace_id=node.state.meta.trace_id):
                     node_run_response = node.run()
 
             ports = node.Ports()
@@ -243,7 +249,7 @@ class WorkflowRunner(Generic[StateType]):
                         ),
                     )
 
-                with execution_context(parent_context=updated_parent_context):
+                with execution_context(parent_context=updated_parent_context, trace_id=node.state.meta.trace_id):
                     for output in node_run_response:
                         invoked_ports = output > ports
                         if output.is_initiated:
@@ -346,7 +352,7 @@ class WorkflowRunner(Generic[StateType]):
         if parent_context is None:
             parent_context = get_parent_context() or self._parent_context
 
-        with execution_context(parent_context=parent_context):
+        with execution_context(parent_context=parent_context, trace_id=node.state.meta.trace_id):
             self._run_work_item(node, span_id)
 
     def _handle_invoked_ports(self, state: StateType, ports: Optional[Iterable[Port]]) -> None:
@@ -524,7 +530,7 @@ class WorkflowRunner(Generic[StateType]):
         for node_cls in self._entrypoints:
             try:
                 if not self._max_concurrency or len(self._active_nodes_by_execution_id) < self._max_concurrency:
-                    with execution_context(parent_context=current_parent):
+                    with execution_context(parent_context=current_parent, trace_id=self._initial_state.meta.trace_id):
                         self._run_node_if_ready(self._initial_state, node_cls)
                 else:
                     self._concurrency_queue.put((self._initial_state, node_cls, None))
@@ -551,7 +557,7 @@ class WorkflowRunner(Generic[StateType]):
 
             self._workflow_event_outer_queue.put(event)
 
-            with execution_context(parent_context=current_parent):
+            with execution_context(parent_context=current_parent, trace_id=self._initial_state.meta.trace_id):
                 rejection_error = self._handle_work_item_event(event)
 
             if rejection_error:
@@ -562,7 +568,7 @@ class WorkflowRunner(Generic[StateType]):
             while event := self._workflow_event_inner_queue.get_nowait():
                 self._workflow_event_outer_queue.put(event)
 
-                with execution_context(parent_context=current_parent):
+                with execution_context(parent_context=current_parent, trace_id=self._initial_state.meta.trace_id):
                     rejection_error = self._handle_work_item_event(event)
 
                 if rejection_error:

vellum/workflows/state/context.py

@@ -3,7 +3,7 @@ from queue import Queue
 from typing import TYPE_CHECKING, Dict, List, Optional, Type
 
 from vellum import Vellum
-from vellum.workflows.events.types import ParentContext
+from vellum.workflows.context import ExecutionContext, get_execution_context
 from vellum.workflows.nodes.mocks import MockNodeExecution, MockNodeExecutionArg
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references.constant import ConstantValueReference
@@ -18,12 +18,14 @@ class WorkflowContext:
         self,
         *,
         vellum_client: Optional[Vellum] = None,
-        parent_context: Optional[ParentContext] = None,
+        execution_context: Optional[ExecutionContext] = None,
     ):
         self._vellum_client = vellum_client
-        self._parent_context = parent_context
         self._event_queue: Optional[Queue["WorkflowEvent"]] = None
         self._node_output_mocks_map: Dict[Type[BaseOutputs], List[MockNodeExecution]] = {}
+        self._execution_context = get_execution_context()
+        if not self._execution_context.parent_context and execution_context:
+            self._execution_context = execution_context
 
     @cached_property
     def vellum_client(self) -> Vellum:
@@ -33,10 +35,8 @@ class WorkflowContext:
         return create_vellum_client()
 
     @cached_property
-    def parent_context(self) -> Optional[ParentContext]:
-        if self._parent_context:
-            return self._parent_context
-        return None
+    def execution_context(self) -> ExecutionContext:
+        return self._execution_context
 
     @cached_property
     def node_output_mocks_map(self) -> Dict[Type[BaseOutputs], List[MockNodeExecution]]:

vellum/workflows/workflows/base.py

@@ -24,6 +24,7 @@ from typing import (
     get_args,
 )
 
+from vellum.workflows.context import get_execution_context
 from vellum.workflows.edges import Edge
 from vellum.workflows.emitters.base import BaseWorkflowEmitter
 from vellum.workflows.errors import WorkflowError, WorkflowErrorCode
@@ -171,6 +172,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         self.resolvers = resolvers or (self.resolvers if hasattr(self, "resolvers") else [])
         self._context = context or WorkflowContext()
         self._store = Store()
+        self._execution_context = self._context.execution_context
 
         self.validate()
 
@@ -178,48 +180,64 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
     def context(self) -> WorkflowContext:
         return self._context
 
-    @classmethod
-    def get_subgraphs(cls) -> List[Graph]:
-        original_graph = cls.graph
-        if isinstance(original_graph, Graph):
-            return [original_graph]
-        if isinstance(original_graph, set):
-            return [
-                subgraph if isinstance(subgraph, Graph) else Graph.from_node(subgraph) for subgraph in original_graph
-            ]
-        if issubclass(original_graph, BaseNode):
-            return [Graph.from_node(original_graph)]
-
-        raise ValueError(f"Unexpected graph type: {original_graph.__class__}")
-
-    @classmethod
-    def get_edges(cls) -> Iterator[Edge]:
+    @staticmethod
+    def _resolve_graph(graph: GraphAttribute) -> List[Graph]:
         """
-        Returns an iterator over the edges in the workflow. We use a set to
-        ensure uniqueness, and the iterator to preserve order.
+        Resolves a single graph source to a list of Graph objects.
         """
+        if isinstance(graph, Graph):
+            return [graph]
+        if isinstance(graph, set):
+            graphs = []
+            for item in graph:
+                if isinstance(item, Graph):
+                    graphs.append(item)
+                elif issubclass(item, BaseNode):
+                    graphs.append(Graph.from_node(item))
+                else:
+                    raise ValueError(f"Unexpected graph type: {type(item)}")
+            return graphs
+        if issubclass(graph, BaseNode):
+            return [Graph.from_node(graph)]
+        raise ValueError(f"Unexpected graph type: {type(graph)}")
 
+    @staticmethod
+    def _get_edges_from_subgraphs(subgraphs: Iterable[Graph]) -> Iterator[Edge]:
         edges = set()
-        subgraphs = cls.get_subgraphs()
         for subgraph in subgraphs:
             for edge in subgraph.edges:
                 if edge not in edges:
                     edges.add(edge)
                     yield edge
 
+    @staticmethod
+    def _get_nodes_from_subgraphs(subgraphs: Iterable[Graph]) -> Iterator[Type[BaseNode]]:
+        nodes = set()
+        for subgraph in subgraphs:
+            for node in subgraph.nodes:
+                if node not in nodes:
+                    nodes.add(node)
+                    yield node
+
+    @classmethod
+    def get_subgraphs(cls) -> List[Graph]:
+        return cls._resolve_graph(cls.graph)
+
+    @classmethod
+    def get_edges(cls) -> Iterator[Edge]:
+        """
+        Returns an iterator over the edges in the workflow. We use a set to
+        ensure uniqueness, and the iterator to preserve order.
+        """
+        return cls._get_edges_from_subgraphs(cls.get_subgraphs())
+
     @classmethod
     def get_nodes(cls) -> Iterator[Type[BaseNode]]:
         """
         Returns an iterator over the nodes in the workflow. We use a set to
         ensure uniqueness, and the iterator to preserve order.
         """
-
-        nodes = set()
-        for subgraph in cls.get_subgraphs():
-            for node in subgraph.nodes:
-                if node not in nodes:
-                    nodes.add(node)
-                    yield node
+        return cls._get_nodes_from_subgraphs(cls.get_subgraphs())
 
     @classmethod
     def get_unused_subgraphs(cls) -> List[Graph]:
@@ -228,19 +246,9 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         """
         if not hasattr(cls, "unused_graphs"):
             return []
-
         graphs = []
         for item in cls.unused_graphs:
-            if isinstance(item, Graph):
-                graphs.append(item)
-            elif isinstance(item, set):
-                for subitem in item:
-                    if isinstance(subitem, Graph):
-                        graphs.append(subitem)
-                    elif issubclass(subitem, BaseNode):
-                        graphs.append(Graph.from_node(subitem))
-            elif issubclass(item, BaseNode):
-                graphs.append(Graph.from_node(item))
+            graphs.extend(cls._resolve_graph(item))
         return graphs
 
     @classmethod
@@ -248,29 +256,14 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         """
         Returns an iterator over the nodes that are defined but not used in the graph.
         """
-        if not hasattr(cls, "unused_graphs"):
-            yield from ()
-        else:
-            nodes = set()
-            subgraphs = cls.get_unused_subgraphs()
-            for subgraph in subgraphs:
-                for node in subgraph.nodes:
-                    if node not in nodes:
-                        nodes.add(node)
-                        yield node
+        return cls._get_nodes_from_subgraphs(cls.get_unused_subgraphs())
 
     @classmethod
     def get_unused_edges(cls) -> Iterator[Edge]:
         """
         Returns an iterator over edges that are defined but not used in the graph.
         """
-        edges = set()
-        subgraphs = cls.get_unused_subgraphs()
-        for subgraph in subgraphs:
-            for edge in subgraph.edges:
-                if edge not in edges:
-                    edges.add(edge)
-                    yield edge
+        return cls._get_edges_from_subgraphs(cls.get_unused_subgraphs())
 
     @classmethod
     def get_entrypoints(cls) -> Iterable[Type[BaseNode]]:
@@ -329,8 +322,8 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
             external_inputs=external_inputs,
             cancel_signal=cancel_signal,
             node_output_mocks=node_output_mocks,
-            parent_context=self._context.parent_context,
             max_concurrency=max_concurrency,
+            init_execution_context=self._execution_context,
         ).stream()
         first_event: Optional[Union[WorkflowExecutionInitiatedEvent, WorkflowExecutionResumedEvent]] = None
         last_event = None
@@ -440,8 +433,8 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
             external_inputs=external_inputs,
             cancel_signal=cancel_signal,
             node_output_mocks=node_output_mocks,
-            parent_context=self.context.parent_context,
             max_concurrency=max_concurrency,
+            init_execution_context=self._execution_context,
         ).stream():
             if should_yield(self.__class__, event):
                 yield event
@@ -488,10 +481,19 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         return self.get_inputs_class()()
 
     def get_default_state(self, workflow_inputs: Optional[InputsType] = None) -> StateType:
+        execution_context = get_execution_context()
         return self.get_state_class()(
-            meta=StateMeta(
-                parent=self._parent_state,
-                workflow_inputs=workflow_inputs or self.get_default_inputs(),
+            meta=(
+                StateMeta(
+                    parent=self._parent_state,
+                    workflow_inputs=workflow_inputs or self.get_default_inputs(),
+                    trace_id=execution_context.trace_id,
+                )
+                if execution_context and execution_context.trace_id
+                else StateMeta(
+                    parent=self._parent_state,
+                    workflow_inputs=workflow_inputs or self.get_default_inputs(),
+                )
             )
         )