vellum-ai 0.14.63__py3-none-any.whl → 1.13.5__py3-none-any.whl

vellum/workflows/nodes/displayable/set_state_node/node.py

@@ -0,0 +1,71 @@
+from typing import Any, ClassVar, Dict, Generic
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.descriptors.utils import resolve_value
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.types.generics import StateType
+
+
+class SetStateNode(BaseNode[StateType], Generic[StateType]):
+    """
+    A node that sets multiple state values at once.
+
+    This node accepts a dictionary mapping state paths to values and sets them all.
+    Values can be actual values or descriptors.
+
+    Example:
+        operations = {
+            "chat_history": State.chat_history.concat(LazyReference("MyAgentNode.Outputs.chat_history")),
+            "counter": State.counter + 1,
+            "total_tokens": State.user_tokens.add(State.assistant_tokens)
+        }
+
+    Args:
+        operations: Dictionary mapping state attribute names to values (descriptors or actual values)
+    """
+
+    # Dictionary mapping state paths to values
+    operations: ClassVar[Dict[str, Any]] = {}
+
+    class Outputs(BaseOutputs):
+        """
+        The outputs of the SetStateNode.
+
+        result: Dict[str, Any] - Dictionary of all state updates
+        """
+
+        results: Dict[str, Any]
+
+    def run(self) -> Outputs:
+        """
+        Run the node and set all the state values.
+        Resolves descriptors to their actual values before setting state.
+        """
+        # First pass: validate and resolve all operations without mutating state
+        resolved_updates: Dict[str, Any] = {}
+        for path, value in self.operations.items():
+            # Validate the state attribute exists prior to any mutation
+            if not hasattr(self.state, path):
+                raise NodeException(
+                    f"State does not have attribute '{path}'. "
+                    f"Only existing state attributes can be set via SetStateNode.",
+                    code=WorkflowErrorCode.INVALID_STATE,
+                )
+
+            # Resolve the value if it's a descriptor against the current (unmodified) state
+            if isinstance(value, BaseDescriptor):
+                resolved_value = resolve_value(value, self.state)
+            else:
+                resolved_value = value
+
+            resolved_updates[path] = resolved_value
+
+        # Second pass: apply the resolved updates to the state atomically
+        with self.state.__atomic__():
+            for path, resolved_value in resolved_updates.items():
+                setattr(self.state, path, resolved_value)
+
+        return self.Outputs(results=resolved_updates)

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

@@ -0,0 +1,212 @@
+import pytest
+from typing import List
+
+from vellum import ChatMessage
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.displayable.set_state_node import SetStateNode
+from vellum.workflows.state.base import BaseState
+
+
+def test_set_state_node_with_single_value():
+    """Test that SetStateNode can set a single state value."""
+
+    # GIVEN a state with a string field
+    class TestState(BaseState):
+        message: str = ""
+
+    # AND a SetStateNode that sets a message
+    class SetMessage(SetStateNode[TestState]):
+        operations = {"message": "Hello, World!"}
+
+    # WHEN we run the node
+    state = TestState()
+    node = SetMessage(state=state)
+    outputs = node.run()
+
+    # THEN the results should be the set value
+    assert outputs.results == {"message": "Hello, World!"}
+
+    # AND the state should be updated
+    assert state.message == "Hello, World!"
+
+
+def test_set_state_node_with_multiple_values():
+    """Test that SetStateNode can set multiple state values at once with expressions."""
+
+    # GIVEN a state with multiple fields
+    class TestState(BaseState):
+        counter: int = 0
+        total_tokens: int = 0
+        user_tokens: int = 5
+        assistant_tokens: int = 10
+
+    # AND a SetStateNode that sets multiple values with expressions using add
+    class UpdateMultipleState(SetStateNode[TestState]):
+        operations = {
+            "counter": 1,
+            "total_tokens": TestState.user_tokens + TestState.assistant_tokens,
+        }
+
+    # WHEN we run the node
+    state = TestState()
+    node = UpdateMultipleState(state=state)
+    outputs = node.run()
+
+    # THEN all state values should be updated
+    assert state.counter == 1
+    assert state.total_tokens == 15
+
+    # AND the results should contain all updates
+    assert outputs.results == {"counter": 1, "total_tokens": 15}
+
+
+def test_set_state_node_with_chat_history_and_concat():
+    """Test that SetStateNode can concatenate chat history using concat."""
+
+    # GIVEN a state with chat history
+    class TestState(BaseState):
+        chat_history: List[ChatMessage] = []
+
+    # AND initial chat history in state
+    initial_messages = [ChatMessage(role="ASSISTANT", text="Hi there")]
+
+    # AND a SetStateNode that uses concat() method to add to chat history
+    class UpdateChatHistory(SetStateNode[TestState]):
+        operations = {"chat_history": TestState.chat_history.concat([ChatMessage(role="USER", text="Hello")])}
+
+    # WHEN we run the node with initial state
+    state = TestState(chat_history=initial_messages)
+
+    node = UpdateChatHistory(state=state)
+    outputs = node.run()
+
+    # THEN the state should have concatenated chat history
+    assert len(state.chat_history) == 2
+    assert state.chat_history[0].role == "ASSISTANT"
+    assert state.chat_history[1].role == "USER"
+    assert state.chat_history[1].text == "Hello"
+
+    # AND the results should contain the updated chat history
+    assert outputs.results == {
+        "chat_history": [ChatMessage(role="ASSISTANT", text="Hi there"), ChatMessage(role="USER", text="Hello")]
+    }
+
+
+def test_set_state_node_with_empty_dict():
+    """Test that SetStateNode handles empty operations gracefully."""
+
+    # GIVEN a state
+    class TestState(BaseState):
+        counter: int = 0
+
+    # AND a SetStateNode with empty operations
+    class EmptySetState(SetStateNode[TestState]):
+        operations = {}
+
+    # WHEN we run the node
+    state = TestState()
+    node = EmptySetState(state=state)
+    outputs = node.run()
+
+    # THEN the results should be an empty dict
+    assert outputs.results == {}
+
+    # AND state should be unchanged
+    assert state.counter == 0
+
+
+def test_set_state_node_modifies_existing_values():
+    """Test that SetStateNode can modify existing state values."""
+
+    # GIVEN a state with existing values
+    class TestState(BaseState):
+        counter: int = 10
+        message: str = "old"
+
+    # AND a SetStateNode that updates existing values
+    class UpdateState(SetStateNode[TestState]):
+        operations = {"counter": 20, "message": "new"}
+
+    # WHEN we run the node
+    state = TestState()
+    node = UpdateState(state=state)
+    outputs = node.run()
+
+    # THEN state should be updated
+    assert state.counter == 20
+    assert state.message == "new"
+
+    # AND results should contain the new values with the correct structure
+    assert outputs.results == {"counter": 20, "message": "new"}
+
+
+def test_set_state_not_existing_value():
+    """Test that SetStateNode raises an error when trying to set a value that doesn't exist in the state."""
+
+    # GIVEN a state
+    class TestState(BaseState):
+        counter: int = 0
+
+    # AND a SetStateNode that tries to set a value that doesn't exist in the state
+    class SetState(SetStateNode[TestState]):
+        operations = {"unknown_value": 10}
+
+    # WHEN we run the node
+    state = TestState()
+    node = SetState(state=state)
+
+    # THEN it should raise a NodeException
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # AND the error should be INVALID_STATE and the message should mention the non-existent attribute
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_STATE
+    assert (
+        "State does not have attribute 'unknown_value'. Only existing state attributes can be set via SetStateNode."
+        == str(exc_info.value)
+    )
+
+
+def test_set_state_node_atomic_order_independent_resolution():
+    """Both operations resolve against the original state before any mutation."""
+
+    class TestState(BaseState):
+        a: int = 1
+        b: int = 0
+
+    class UpdateAB(SetStateNode[TestState]):
+        # Both should see original a=1, so both resolve to 2
+        operations = {
+            "a": TestState.a + 1,
+            "b": TestState.a + 1,
+        }
+
+    state = TestState()
+    node = UpdateAB(state=state)
+    outputs = node.run()
+
+    assert state.a == 2
+    assert state.b == 2
+    assert outputs.results == {"a": 2, "b": 2}
+
+
+def test_set_state_node_no_partial_update_on_error():
+    """If a later operation is invalid, no earlier changes should be applied."""
+
+    class TestState(BaseState):
+        a: int = 5
+
+    class PartialFail(SetStateNode[TestState]):
+        operations = {
+            "a": 42,  # would change a
+            "missing": 1,  # invalid attribute triggers NodeException
+        }
+
+    state = TestState()
+    node = PartialFail(state=state)
+
+    with pytest.raises(NodeException):
+        node.run()
+
+    assert state.a == 5

vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py

@@ -1,6 +1,8 @@
 import json
-from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, Iterator, List, Optional, Set, Union, cast
+from uuid import UUID, uuid4
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Generic, Iterator, List, Optional, Set, Type, Union, cast
+
+import httpx
 
 from vellum import (
     ChatMessage,
@@ -12,19 +14,27 @@ from vellum import (
     WorkflowRequestNumberInputRequest,
     WorkflowRequestStringInputRequest,
 )
+from vellum.client.core import RequestOptions
 from vellum.client.core.api_error import ApiError
 from vellum.client.types.chat_message_request import ChatMessageRequest
-from vellum.core import RequestOptions
-from vellum.workflows.constants import LATEST_RELEASE_TAG, OMIT
-from vellum.workflows.context import get_execution_context
+from vellum.workflows.constants import LATEST_RELEASE_TAG, OMIT, undefined
+from vellum.workflows.context import execution_context, get_execution_context, get_parent_context
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.errors.types import workflow_event_error_to_workflow_error
-from vellum.workflows.events.types import default_serializer
-from vellum.workflows.exceptions import NodeException
+from vellum.workflows.events.exception_handling import stream_initialization_exception
+from vellum.workflows.events.types import WorkflowDeploymentParentContext, default_serializer
+from vellum.workflows.events.workflow import is_workflow_event
+from vellum.workflows.exceptions import NodeException, WorkflowInitializationException
+from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.outputs.base import BaseOutput
+from vellum.workflows.state.context import WorkflowContext, WorkflowDeploymentMetadata
 from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior
 from vellum.workflows.types.generics import StateType
+from vellum.workflows.workflows.event_filters import all_workflow_event_filter
+
+if TYPE_CHECKING:
+    from vellum.workflows.workflows.base import BaseWorkflow
 
 
 class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
@@ -41,9 +51,13 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
     request_options: Optional[RequestOptions] = None - The request options to use for the Workflow Execution
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:diagram-sankey"
+        color = "grass"
+
     # Either the Workflow Deployment's UUID or its name.
     deployment: ClassVar[Union[UUID, str]]
-    subworkflow_inputs: ClassVar[EntityInputsInterface] = {}
+    subworkflow_inputs: ClassVar[Union[EntityInputsInterface, BaseInputs]] = {}
 
     release_tag: str = LATEST_RELEASE_TAG
     external_id: Optional[str] = OMIT
@@ -62,64 +76,192 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
 
         compiled_inputs: List[WorkflowRequestInputRequest] = []
 
-        for input_name, input_value in self.subworkflow_inputs.items():
-            if isinstance(input_value, str):
-                compiled_inputs.append(
-                    WorkflowRequestStringInputRequest(
-                        name=input_name,
-                        value=input_value,
-                    )
+        if isinstance(self.subworkflow_inputs, BaseInputs):
+            for input_descriptor, input_value in self.subworkflow_inputs:
+                self._add_compiled_input(compiled_inputs, input_descriptor.name, input_value)
+        else:
+            for input_name, input_value in self.subworkflow_inputs.items():
+                self._add_compiled_input(compiled_inputs, input_name, input_value)
+
+        return compiled_inputs
+
+    def _add_compiled_input(
+        self, compiled_inputs: List[WorkflowRequestInputRequest], input_name: str, input_value: Any
+    ) -> None:
+        # Exclude inputs that resolved to be null. This ensure that we don't pass input values
+        # to optional subworkflow inputs whose values were unresolved.
+        if input_value is None:
+            return
+        if isinstance(input_value, str):
+            compiled_inputs.append(
+                WorkflowRequestStringInputRequest(
+                    name=input_name,
+                    value=input_value,
                 )
-            elif (
-                isinstance(input_value, list)
-                and len(input_value) > 0
-                and all(isinstance(message, (ChatMessage, ChatMessageRequest)) for message in input_value)
-            ):
-                chat_history = [
-                    (
-                        message
-                        if isinstance(message, ChatMessageRequest)
-                        else ChatMessageRequest.model_validate(message.model_dump())
-                    )
-                    for message in input_value
-                    if isinstance(message, (ChatMessage, ChatMessageRequest))
-                ]
-                compiled_inputs.append(
-                    WorkflowRequestChatHistoryInputRequest(
-                        name=input_name,
-                        value=chat_history,
-                    )
+            )
+        elif (
+            isinstance(input_value, list)
+            and len(input_value) > 0
+            and all(isinstance(message, (ChatMessage, ChatMessageRequest)) for message in input_value)
+        ):
+            chat_history = [
+                (
+                    message
+                    if isinstance(message, ChatMessageRequest)
+                    else ChatMessageRequest.model_validate(message.model_dump())
                 )
-            elif isinstance(input_value, dict):
-                compiled_inputs.append(
-                    WorkflowRequestJsonInputRequest(
-                        name=input_name,
-                        value=cast(Dict[str, Any], input_value),
-                    )
+                for message in input_value
+                if isinstance(message, (ChatMessage, ChatMessageRequest))
+            ]
+            compiled_inputs.append(
+                WorkflowRequestChatHistoryInputRequest(
+                    name=input_name,
+                    value=chat_history,
                 )
-            elif isinstance(input_value, (int, float)):
-                compiled_inputs.append(
-                    WorkflowRequestNumberInputRequest(
-                        name=input_name,
-                        value=input_value,
-                    )
+            )
+        elif isinstance(input_value, dict):
+            compiled_inputs.append(
+                WorkflowRequestJsonInputRequest(
+                    name=input_name,
+                    value=cast(Dict[str, Any], input_value),
                 )
-            else:
-                try:
-                    input_value = default_serializer(input_value)
-                except json.JSONDecodeError as e:
-                    raise NodeException(
-                        message=f"Failed to serialize input '{input_name}' of type '{input_value.__class__}': {e}",
-                        code=WorkflowErrorCode.INVALID_INPUTS,
-                    )
-                compiled_inputs.append(
-                    WorkflowRequestJsonInputRequest(
-                        name=input_name,
-                        value=input_value,
+            )
+        elif isinstance(input_value, (int, float)):
+            compiled_inputs.append(
+                WorkflowRequestNumberInputRequest(
+                    name=input_name,
+                    value=input_value,
+                )
+            )
+        else:
+            try:
+                input_value = default_serializer(input_value)
+            except json.JSONDecodeError as e:
+                raise NodeException(
+                    message=f"Failed to serialize input '{input_name}' of type '{input_value.__class__}': {e}",
+                    code=WorkflowErrorCode.INVALID_INPUTS,
+                )
+            compiled_inputs.append(
+                WorkflowRequestJsonInputRequest(
+                    name=input_name,
+                    value=input_value,
+                )
+            )
+
+    def _compile_subworkflow_inputs_for_direct_invocation(self, workflow: "BaseWorkflow") -> Any:
+        """Compile inputs for direct workflow invocation (similar to InlineSubworkflowNode)."""
+        inputs_class = workflow.get_inputs_class()
+
+        if isinstance(self.subworkflow_inputs, BaseInputs):
+            inputs_dict = {}
+            for input_descriptor, input_value in self.subworkflow_inputs:
+                if input_value is not None:
+                    inputs_dict[input_descriptor.name] = input_value
+            return inputs_class(**inputs_dict)
+        else:
+            # Filter out None values for direct invocation
+            filtered_inputs = {k: v for k, v in self.subworkflow_inputs.items() if v is not None}
+            return inputs_class(**filtered_inputs)
+
+    def _run_resolved_workflow(
+        self,
+        workflow_class: Type["BaseWorkflow"],
+        deployment_metadata: Optional[WorkflowDeploymentMetadata],
+    ) -> Iterator[BaseOutput]:
+        """Execute resolved workflow directly (similar to InlineSubworkflowNode)."""
+        # Construct the parent context hierarchy for the subworkflow
+        parent_context = get_parent_context()
+
+        # If we have deployment metadata, wrap the parent context with WorkflowDeploymentParentContext
+        if deployment_metadata:
+            parent_context = WorkflowDeploymentParentContext(
+                span_id=uuid4(),
+                deployment_id=deployment_metadata.deployment_id,
+                deployment_name=deployment_metadata.deployment_name,
+                deployment_history_item_id=deployment_metadata.deployment_history_item_id,
+                release_tag_id=deployment_metadata.release_tag_id,
+                release_tag_name=deployment_metadata.release_tag_name,
+                workflow_version_id=deployment_metadata.workflow_version_id,
+                external_id=self.external_id if self.external_id is not OMIT else None,
+                metadata=self.metadata if self.metadata is not OMIT else None,
+                parent=parent_context,
+            )
+
+        with execution_context(parent_context=parent_context):
+            # Instantiate the workflow inside the execution context so it captures the correct parent context
+            resolved_workflow = workflow_class(
+                context=WorkflowContext.create_from(self._context), parent_state=self.state
+            )
+
+            try:
+                # The stream creation and first event retrieval are wrapped in try/except because
+                # WorkflowInitializationException can be raised during stream creation (e.g., when
+                # inputs are invalid) or when getting the first event
+                subworkflow_inputs = self._compile_subworkflow_inputs_for_direct_invocation(resolved_workflow)
+                subworkflow_stream = resolved_workflow.stream(
+                    inputs=subworkflow_inputs,
+                    event_filter=all_workflow_event_filter,
+                    node_output_mocks=self._context._get_all_node_output_mocks(),
+                    event_max_size=self._context.event_max_size,
+                )
+                first_event = next(subworkflow_stream)
+                self._context._emit_subworkflow_event(first_event)
+            except WorkflowInitializationException as e:
+                # Emit initiated and rejected events for the subworkflow so that
+                # the parent workflow can see the subworkflow's lifecycle events
+                for init_failure_event in stream_initialization_exception(e):
+                    self._context._emit_subworkflow_event(init_failure_event)
+
+                hashed_module = e.definition.__module__
+                raise NodeException(
+                    message=e.message,
+                    code=e.code,
+                    raw_data={"hashed_module": hashed_module},
+                ) from e
+
+            outputs = None
+            exception = None
+            fulfilled_output_names: Set[str] = set()
+
+            for event in subworkflow_stream:
+                self._context._emit_subworkflow_event(event)
+                if exception:
+                    continue
+
+                if not is_workflow_event(event):
+                    continue
+                if event.workflow_definition != resolved_workflow.__class__:
+                    continue
+
+                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":
+                    exception = NodeException.of(event.error)
+                elif event.name == "workflow.execution.paused":
+                    exception = NodeException(
+                        code=WorkflowErrorCode.INVALID_OUTPUTS,
+                        message="Subworkflow unexpectedly paused",
                     )
+
+            if exception:
+                raise exception
+
+            if outputs is None:
+                raise NodeException(
+                    message="Expected to receive outputs from Workflow Deployment",
+                    code=WorkflowErrorCode.INVALID_OUTPUTS,
                 )
 
-        return compiled_inputs
+            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 run(self) -> Iterator[BaseOutput]:
         execution_context = get_execution_context()
@@ -129,6 +271,12 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
             **request_options.get("additional_body_parameters", {}),
         }
 
+        if self.deployment is undefined:
+            raise NodeException(
+                code=WorkflowErrorCode.NODE_EXECUTION,
+                message="Expected subworkflow deployment attribute to be either a UUID or STR, got `undefined` instead",
+            )
+
         try:
             deployment_id = str(self.deployment) if isinstance(self.deployment, UUID) else None
             deployment_name = self.deployment if isinstance(self.deployment, str) else None
@@ -138,6 +286,20 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
                 message="Expected subworkflow deployment attribute to be either a UUID or STR, got None instead",
             )
 
+        if not deployment_name:
+            raise NodeException(
+                code=WorkflowErrorCode.INVALID_INPUTS,
+                message="Expected deployment name to be provided for subworkflow execution.",
+            )
+
+        resolved_result = self._context.resolve_workflow_deployment(
+            deployment_name=deployment_name, release_tag=self.release_tag, state=self.state
+        )
+        if resolved_result:
+            workflow_class, deployment_metadata = resolved_result
+            yield from self._run_resolved_workflow(workflow_class, deployment_metadata)
+            return
+
         try:
             subworkflow_stream = self._context.vellum_client.execute_workflow_stream(
                 inputs=self._compile_subworkflow_inputs(),
@@ -149,6 +311,11 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
                 metadata=self.metadata,
                 request_options=request_options,
             )
+        except httpx.TransportError:
+            raise NodeException(
+                message="Failed to connect to Vellum server",
+                code=WorkflowErrorCode.INTERNAL_ERROR,
+            )
         except ApiError as e:
             self._handle_api_error(e)
 
@@ -156,6 +323,11 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
         # and use the exception handling to catch other api level errors
         try:
             next(subworkflow_stream)
+        except httpx.TransportError:
+            raise NodeException(
+                message="Failed to connect to Vellum server",
+                code=WorkflowErrorCode.INTERNAL_ERROR,
+            )
         except ApiError as e:
             self._handle_api_error(e)
 

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

@@ -265,7 +265,7 @@ def test_run_workflow__no_deployment():
 
     # 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(
+    assert "Expected subworkflow deployment attribute to be either a UUID or STR, got `undefined` instead" in str(
         exc_info.value
     )