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

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

@@ -1,14 +1,15 @@
-from typing import Any, Dict, Generic, Tuple, Type, TypeVar, get_args, get_origin
+from typing import Any, Dict, Generic, Tuple, Type, TypeVar, get_args
 
 from vellum.workflows.constants import undefined
+from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import BaseNodeMeta
 from vellum.workflows.nodes.utils import cast_to_output_type
 from vellum.workflows.ports import NodePorts
-from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 from vellum.workflows.types.utils import get_original_base
+from vellum.workflows.utils.validate import validate_target_type
 
 _OutputType = TypeVar("_OutputType")
 
@@ -40,73 +41,6 @@ class _FinalOutputNodeMeta(BaseNodeMeta):
         else:
             return all_args[1]
 
-    def __validate__(cls) -> None:
-        cls._validate_output_type_consistency(cls)
-
-    @classmethod
-    def _validate_output_type_consistency(mcs, cls: Type) -> None:
-        """
-        Validates that the declared output type of FinalOutputNode matches
-        the type of the descriptor assigned to the 'value' attribute in its Outputs class.
-
-        Raises ValueError if there's a type mismatch.
-        """
-        if not hasattr(cls, "Outputs"):
-            return
-
-        outputs_class = cls.Outputs
-        if not hasattr(outputs_class, "value"):
-            return
-
-        declared_output_type = cls.get_output_type()
-
-        if declared_output_type is Any:
-            return
-
-        value_descriptor = None
-
-        if "value" in outputs_class.__dict__:
-            value_descriptor = outputs_class.__dict__["value"]
-        else:
-            value_descriptor = getattr(outputs_class, "value")
-
-        if isinstance(value_descriptor, OutputReference):
-            descriptor_types = value_descriptor.types
-
-            type_mismatch = True
-            for descriptor_type in descriptor_types:
-                if descriptor_type == declared_output_type:
-                    type_mismatch = False
-                    break
-                if (
-                    get_origin(descriptor_type) == declared_output_type
-                    or get_origin(declared_output_type) == descriptor_type
-                ):
-                    type_mismatch = False
-                    break
-                try:
-                    if issubclass(descriptor_type, declared_output_type) or issubclass(
-                        declared_output_type, descriptor_type
-                    ):
-                        type_mismatch = False
-                        break
-                except TypeError:
-                    # Handle cases where types aren't classes (e.g., Union)
-                    if str(descriptor_type) == str(declared_output_type):
-                        type_mismatch = False
-                        break
-
-            if type_mismatch:
-                declared_type_name = getattr(declared_output_type, "__name__", str(declared_output_type))
-                descriptor_type_names = [getattr(t, "__name__", str(t)) for t in descriptor_types]
-
-                raise ValueError(
-                    f"Output type mismatch in {cls.__name__}: "
-                    f"FinalOutputNode is declared with output type '{declared_type_name}' "
-                    f"but the 'value' descriptor has type(s) {descriptor_type_names}. "
-                    f"The output descriptor type must match the declared FinalOutputNode output type."
-                )
-
 
 class FinalOutputNode(BaseNode[StateType], Generic[StateType, _OutputType], metaclass=_FinalOutputNodeMeta):
     """
@@ -140,3 +74,24 @@ class FinalOutputNode(BaseNode[StateType], Generic[StateType, _OutputType], meta
         )
 
     __simulates_workflow_output__ = True
+
+    @classmethod
+    def __validate__(cls) -> None:
+        cls._validate_output_type_consistency()
+
+    @classmethod
+    def _validate_output_type_consistency(cls) -> None:
+        """
+        Validates that the declared output type of FinalOutputNode matches
+        the type of the descriptor assigned to the 'value' attribute in its Outputs class.
+
+        Raises ValueError if there's a type mismatch.
+        """
+        declared_output_type = cls.get_output_type()
+        value_descriptor = cls.Outputs.value.instance
+
+        if isinstance(value_descriptor, BaseDescriptor):
+            try:
+                validate_target_type(declared_output_type, value_descriptor.normalized_type)
+            except ValueError as e:
+                raise ValueError(f"Failed to validate output type for node '{cls.__name__}': {e}")

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

@@ -2,12 +2,14 @@ import pytest
 from typing import Any, Dict
 
 from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.core.map_node import MapNode
 from vellum.workflows.nodes.core.templating_node import TemplatingNode
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
 from vellum.workflows.nodes.displayable.inline_prompt_node import InlinePromptNode
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.types.core import Json
+from vellum.workflows.workflows.base import BaseWorkflow
 
 
 def test_final_output_node__mismatched_output_type_should_raise_exception_when_ran():
@@ -41,9 +43,31 @@ def test_final_output_node__mismatched_output_type_should_raise_exception():
     # AND the error message should indicate the type mismatch
     assert (
         str(exc_info.value)
-        == "Output type mismatch in Output: FinalOutputNode is declared with output type 'list' but "
-        "the 'value' descriptor has type(s) ['str']. The output descriptor type must match the "
-        "declared FinalOutputNode output type."
+        == "Failed to validate output type for node 'Output': Output type mismatch: declared type 'list' but "
+        "the 'value' Output has type(s) 'str'. "
+    )
+
+
+def test_final_output_node__mismatched_output_type_in_state_should_raise_exception():
+    # GIVEN a state with a str type
+    class State(BaseState):
+        foo: str
+
+    # AND a FinalOutputNode declared with list output type
+    class Output(FinalOutputNode[BaseState, list]):
+        class Outputs(FinalOutputNode.Outputs):
+            value = State.foo
+
+    # WHEN attempting to validate the node class
+    # THEN a ValueError should be raised during validation
+    with pytest.raises(ValueError) as exc_info:
+        Output.__validate__()
+
+    # AND the error message should indicate the type mismatch
+    assert (
+        str(exc_info.value)
+        == "Failed to validate output type for node 'Output': Output type mismatch: declared type 'list' but "
+        "the 'value' Output has type(s) 'str'. "
     )
 
 
@@ -113,3 +137,29 @@ def test_final_output_node__any_output_type_should_accept_json():
         AnyOutputNode.__validate__()
     except ValueError as e:
         pytest.fail(f"Validation should not raise an exception when Any accepts Json: {e}")
+
+
+def test_final_output_node__list_str_output_type_should_pass_validation():
+    """
+    Tests that FinalOutputNode with list[str] output type accepts a descriptor with List[str] type.
+    """
+
+    # GIVEN value descriptor has List[str] type
+    class MySubworkflow(BaseWorkflow):
+        class Outputs(BaseWorkflow.Outputs):
+            result: str
+
+    class MyMap(MapNode):
+        subworkflow = MySubworkflow
+
+    # AND a FinalOutputNode declared with list[str] output type
+    class ListStrOutputNode(FinalOutputNode[BaseState, list[str]]):
+        class Outputs(FinalOutputNode.Outputs):
+            value = MyMap.Outputs.result
+
+    # WHEN attempting to validate the node class
+    try:
+        ListStrOutputNode.__validate__()
+    except Exception as e:
+        # THEN validation should pass without raising an exception
+        pytest.fail(f"Validation should not raise an exception for list[str]/List[str] compatibility: {e}")

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

@@ -1,4 +1,4 @@
-from typing import Generic
+from typing import Any, Dict, Generic, Optional
 
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.types import MergeBehavior
@@ -10,6 +10,9 @@ class NoteNode(BaseNode[StateType], Generic[StateType]):
     A no-op Node purely used to display a note in the Vellum UI.
     """
 
+    text: str = ""
+    style: Optional[Dict[str, Any]] = None
+
     class Display(BaseNode.Display):
         icon = "vellum:icon:note"
         color = "cyan"

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

@@ -21,6 +21,7 @@ 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.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
@@ -191,16 +192,26 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
             resolved_workflow = workflow_class(
                 context=WorkflowContext.create_from(self._context), parent_state=self.state
             )
-            subworkflow_stream = resolved_workflow.stream(
-                inputs=self._compile_subworkflow_inputs_for_direct_invocation(resolved_workflow),
-                event_filter=all_workflow_event_filter,
-                node_output_mocks=self._context._get_all_node_output_mocks(),
-            )
 
             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,

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

@@ -1,3 +1,4 @@
+import pytest
 from uuid import uuid4
 from typing import Any, Iterator, List
 
@@ -6,8 +7,12 @@ from vellum import (
     FulfilledExecutePromptEvent,
     InitiatedExecutePromptEvent,
     PromptOutput,
+    RejectedExecutePromptEvent,
     StringVellumValue,
+    VellumError,
 )
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes import PromptDeploymentNode
 from vellum.workflows.state import BaseState
@@ -73,3 +78,45 @@ def test_text_prompt_deployment_node__basic(vellum_client):
     expand_meta = call_kwargs.get("expand_meta")
     assert expand_meta is not None
     assert expand_meta.finish_reason is True
+
+
+def test_prompt_deployment_node__provider_error_from_first_event(vellum_client):
+    """
+    Tests that PromptDeploymentNode raises NodeException with the actual provider error message
+    when the first event is REJECTED, rather than a generic "Expected to receive outputs" error.
+    """
+
+    # GIVEN a PromptDeploymentNode with basic configuration
+    class MyPromptDeploymentNode(PromptDeploymentNode):
+        deployment = "my-deployment"
+
+    # AND the API returns a REJECTED event as the first event with a provider quota error
+    provider_error = VellumError(
+        code="PROVIDER_QUOTA_EXCEEDED",
+        message="Google: You exceeded your current quota, please check your plan and billing details.",
+    )
+
+    def generate_prompt_events(*args: Any, **kwargs: Any) -> Iterator[ExecutePromptEvent]:
+        execution_id = str(uuid4())
+        events: List[ExecutePromptEvent] = [
+            RejectedExecutePromptEvent(
+                execution_id=execution_id,
+                error=provider_error,
+            ),
+        ]
+        yield from events
+
+    vellum_client.execute_prompt_stream.side_effect = generate_prompt_events
+
+    # WHEN the node is run
+    node = MyPromptDeploymentNode()
+
+    # THEN it should raise a NodeException with the actual provider error
+    with pytest.raises(NodeException) as excinfo:
+        list(node.run())
+
+    # AND the exception should have the correct error code
+    assert excinfo.value.code == WorkflowErrorCode.PROVIDER_QUOTA_EXCEEDED
+
+    # AND the exception message should contain the actual provider error message
+    assert "Google: You exceeded your current quota" in excinfo.value.message

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

@@ -1,14 +1,18 @@
-from typing import Any, ClassVar, Dict, Generic, Iterator, List, Optional, Set, Union
+from typing import Any, ClassVar, Dict, Generic, Iterator, List, Optional, Set, Type, Union, cast
 
 from vellum import ChatMessage, PromptBlock, PromptOutput
 from vellum.client.types.prompt_parameters import PromptParameters
 from vellum.client.types.prompt_settings import PromptSettings
 from vellum.client.types.string_chat_message_content import StringChatMessageContent
 from vellum.prompts.constants import DEFAULT_PROMPT_PARAMETERS
+from vellum.workflows.constants import undefined
 from vellum.workflows.context import execution_context, get_parent_context
+from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.events.node import NodeExecutionStreamingEvent
 from vellum.workflows.events.workflow import is_workflow_event
 from vellum.workflows.exceptions import NodeException
+from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.graph.graph import Graph
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
@@ -16,20 +20,30 @@ from vellum.workflows.nodes.displayable.tool_calling_node.state import ToolCalli
 from vellum.workflows.nodes.displayable.tool_calling_node.utils import (
     create_else_node,
     create_function_node,
-    create_mcp_tool_node,
     create_router_node,
     create_tool_prompt_node,
     get_function_name,
 )
 from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.types.core import EntityInputsInterface
-from vellum.workflows.types.definition import MCPServer, Tool
+from vellum.workflows.types.definition import MCPServer, MCPToolDefinition, Tool, ToolBase
 from vellum.workflows.types.generics import StateType
-from vellum.workflows.utils.functions import compile_mcp_tool_definition, get_mcp_tool_name
+from vellum.workflows.utils.functions import compile_mcp_tool_definition
 from vellum.workflows.workflows.event_filters import all_workflow_event_filter
 
 
+def _contains_reference_to_output(reference: BaseDescriptor, target_reference: OutputReference) -> bool:
+    if reference == target_reference:
+        return True
+    if isinstance(reference, CoalesceExpression):
+        return _contains_reference_to_output(reference.lhs, target_reference) or _contains_reference_to_output(
+            reference.rhs, target_reference
+        )
+    return False
+
+
 class ToolCallingNode(BaseNode[StateType], Generic[StateType]):
     """
     A Node that dynamically invokes the provided functions to the underlying Prompt
@@ -60,10 +74,12 @@ class ToolCallingNode(BaseNode[StateType], Generic[StateType]):
         The outputs of the ToolCallingNode.
 
         text: The final text response after tool calling
+        json: The result of the Prompt Execution in JSON format (if parseable)
         chat_history: The complete chat history including tool calls
         """
 
         text: str
+        json: Union[Dict[Any, Any], Type[undefined]] = undefined
         chat_history: List[ChatMessage]
 
     def run(self) -> Iterator[BaseOutput]:
@@ -86,6 +102,7 @@ class ToolCallingNode(BaseNode[StateType], Generic[StateType]):
 
                 class Outputs(BaseWorkflow.Outputs):
                     text: str = self.tool_prompt_node.Outputs.text
+                    json: Any = self.tool_prompt_node.Outputs.json
                     chat_history: List[ChatMessage] = ToolCallingState.chat_history
                     results: List[PromptOutput] = self.tool_prompt_node.Outputs.results
 
@@ -97,6 +114,7 @@ class ToolCallingNode(BaseNode[StateType], Generic[StateType]):
             subworkflow_stream = subworkflow.stream(
                 event_filter=all_workflow_event_filter,
                 node_output_mocks=self._context._get_all_node_output_mocks(),
+                event_max_size=self._context.event_max_size,
             )
 
         outputs: Optional[BaseOutputs] = None
@@ -139,17 +157,6 @@ class ToolCallingNode(BaseNode[StateType], Generic[StateType]):
                 exception = NodeException.of(event.error)
 
         if exception:
-            # In the case of the Prompt Node receiving invalid inputs or schemas, that is a bug with the
-            # internals of the Agent Node. In the future, we should distinguish further by also checking
-            # that this exception came from the Tool Calling Node's prompt node instead of anything
-            # internal to a user defined tool.
-            if exception.code == WorkflowErrorCode.INVALID_INPUTS:
-                raise NodeException(
-                    message="Internal server error",
-                    code=WorkflowErrorCode.INTERNAL_ERROR,
-                    raw_data=exception.raw_data,
-                    stacktrace=exception.stacktrace,
-                ) from exception
             raise exception
 
         if outputs is None:
@@ -170,10 +177,20 @@ class ToolCallingNode(BaseNode[StateType], Generic[StateType]):
         process_parameters_method = getattr(self.__class__, "process_parameters", None)
         process_blocks_method = getattr(self.__class__, "process_blocks", None)
 
+        # Hydrate MCP servers upfront and replace them with tool definitions
+        hydrated_functions: List[Union[ToolBase, MCPToolDefinition]] = []
+        for function in self.functions:
+            if isinstance(function, MCPServer):
+                hydrated_functions.extend(compile_mcp_tool_definition(function))
+            else:
+                # After checking, function is ToolBase (not MCPServer)
+                # Mypy doesn't narrow Union[ToolBase, MCPServer] to ToolBase, so we cast
+                hydrated_functions.append(cast(ToolBase, function))
+
         self.tool_prompt_node = create_tool_prompt_node(
             ml_model=self.ml_model,
             blocks=self.blocks,
-            functions=self.functions,
+            functions=hydrated_functions,
             prompt_inputs=self.prompt_inputs,
             parameters=self.parameters,
             max_prompt_iterations=self.max_prompt_iterations,
@@ -184,28 +201,17 @@ class ToolCallingNode(BaseNode[StateType], Generic[StateType]):
 
         # Create the router node (handles routing logic only)
         self.router_node = create_router_node(
-            functions=self.functions,
+            functions=hydrated_functions,
             tool_prompt_node=self.tool_prompt_node,
         )
 
         self._function_nodes = {}
-        for function in self.functions:
-            if isinstance(function, MCPServer):
-                tool_definitions = compile_mcp_tool_definition(function)
-                for tool_definition in tool_definitions:
-                    function_name = get_mcp_tool_name(tool_definition)
-
-                    self._function_nodes[function_name] = create_mcp_tool_node(
-                        tool_def=tool_definition,
-                        tool_prompt_node=self.tool_prompt_node,
-                    )
-            else:
-                function_name = get_function_name(function)
-
-                self._function_nodes[function_name] = create_function_node(
-                    function=function,
-                    tool_prompt_node=self.tool_prompt_node,
-                )
+        for hydrated_function in hydrated_functions:
+            function_name = get_function_name(hydrated_function)
+            self._function_nodes[function_name] = create_function_node(
+                function=hydrated_function,
+                tool_prompt_node=self.tool_prompt_node,
+            )
 
         graph: Graph = self.tool_prompt_node >> self.router_node
 
@@ -223,3 +229,37 @@ class ToolCallingNode(BaseNode[StateType], Generic[StateType]):
         graph._extend_edges(default_port_graph.edges)
 
         self._graph = graph
+
+    def __directly_emit_workflow_output__(
+        self,
+        event: NodeExecutionStreamingEvent,
+        workflow_output_descriptor: OutputReference,
+    ) -> bool:
+        """
+        Check if this ToolCallingNode should directly emit workflow output events.
+        Similar to BasePromptNode, this allows streaming text and chat_history outputs through FinalOutputNode.
+        """
+        # Only handle text and chat_history outputs
+        if event.output.name != "text":
+            return False
+
+        if not isinstance(event.output.delta, str) and not event.output.is_initiated:
+            return False
+
+        target_nodes = [e.to_node for port in self.Ports for e in port.edges if e.to_node.__simulates_workflow_output__]
+        target_node_output = next(
+            (
+                o
+                for target_node in target_nodes
+                for o in target_node.Outputs
+                if o == workflow_output_descriptor.instance
+            ),
+            None,
+        )
+        if not target_node_output:
+            return False
+
+        if not isinstance(target_node_output.instance, BaseDescriptor):
+            return False
+
+        return _contains_reference_to_output(target_node_output.instance, event.node_definition.Outputs.text)