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

vellum/workflows/graph/tests/test_graph.py

@@ -1,7 +1,31 @@
+import pytest
+from typing import Type
+
 from vellum.workflows.edges.edge import Edge
-from vellum.workflows.graph.graph import Graph
+from vellum.workflows.graph.graph import Graph, NoPortsNode
 from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
 from vellum.workflows.ports.port import Port
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.triggers import ManualTrigger
+from vellum.workflows.triggers.schedule import ScheduleTrigger
+
+
+def test_graph__empty():
+    # WHEN we create an empty graph
+    graph = Graph.empty()
+
+    # THEN the graph has no entrypoints
+    assert len(list(graph.entrypoints)) == 0
+
+    # AND no nodes
+    assert len(list(graph.nodes)) == 0
+
+    # AND no edges
+    assert len(list(graph.edges)) == 0
+
+    # AND string representation indicates empty graph
+    assert str(graph) == "Graph(empty)"
 
 
 def test_graph__from_node():
@@ -484,3 +508,506 @@ def test_graph__set_to_graph():
 
     # AND two edges
     assert len(list(graph.edges)) == 2
+
+
+def test_graph__str_simple_linear():
+    # GIVEN a simple linear graph: A -> B -> C
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    class NodeC(BaseNode):
+        pass
+
+    graph = NodeA >> NodeB >> NodeC
+
+    # WHEN we convert the graph to string
+    result = str(graph)
+
+    # THEN it shows the linear flow structure
+    expected_lines = ["Graph:", "  └─ NodeA", "     └─ NodeB", "        └─ NodeC"]
+    assert result == "\n".join(expected_lines)
+
+
+def test_graph__str_with_branching():
+    # GIVEN a graph with branching: A -> {B, C}
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    class NodeC(BaseNode):
+        pass
+
+    graph = NodeA >> {NodeB, NodeC}
+
+    # WHEN we convert the graph to string
+    result = str(graph)
+
+    # THEN it shows the branching structure
+    # Note: The order might vary due to set ordering, so we check for the structure
+    lines = result.split("\n")
+    assert lines[0] == "Graph:"
+    assert lines[1] == "  └─ NodeA"
+
+    # Should have two branches (order may vary)
+    branch_lines = [line.strip() for line in lines[2:] if line.strip()]
+    assert len(branch_lines) == 2
+    assert any("NodeB" in line for line in branch_lines)
+    assert any("NodeC" in line for line in branch_lines)
+    assert all(line.startswith("├─ ") or line.startswith("└─ ") for line in branch_lines)
+
+
+def test_graph__str_with_loop():
+    # GIVEN a graph with a loop: A -> B -> A (loop back)
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    # Create the loop manually using edges
+    edge1 = Edge(NodeA.Ports.default, NodeB)
+    edge2 = Edge(NodeB.Ports.default, NodeA)
+
+    graph = Graph(entrypoints={NodeA.Ports.default}, edges=[edge1, edge2], terminals={NodeA.Ports.default})
+
+    # WHEN we convert the graph to string
+    result = str(graph)
+
+    # THEN it shows the loop with cycle detection
+    expected_lines = ["Graph:", "  └─ NodeA", "     └─ NodeB", "        └─ NodeA ⟲ (loops back)"]
+    assert result == "\n".join(expected_lines)
+
+
+def test_graph__str_empty_graph():
+    # GIVEN an empty graph
+    graph = Graph(entrypoints=set(), edges=[], terminals=set())
+
+    # WHEN we convert the graph to string
+    result = str(graph)
+
+    # THEN it shows empty graph message
+    assert result == "Graph(empty)"
+
+
+def test_graph__str_single_node():
+    # GIVEN a graph with just one node and no edges
+    class SingleNode(BaseNode):
+        pass
+
+    graph = Graph.from_node(SingleNode)
+
+    # WHEN we convert the graph to string
+    result = str(graph)
+
+    # THEN it shows the single node
+    assert "SingleNode.default" in result
+    assert "Graph:" in result
+
+
+def test_graph__from_node_with_empty_ports():
+    """
+    Tests that building a graph from a single node with empty Ports class generates 1 node.
+    """
+
+    # GIVEN a node with an empty Ports class
+    class NodeWithEmptyPorts(BaseNode):
+        class Ports(BaseNode.Ports):
+            pass
+
+    # WHEN we create a graph from the node
+    graph = Graph.from_node(NodeWithEmptyPorts)
+
+    # THEN the graph should have exactly 1 node
+    assert len(list(graph.nodes)) == 1
+
+
+def test_graph__manual_trigger_to_node():
+    # GIVEN a node
+    class MyNode(BaseNode):
+        pass
+
+    # WHEN we create graph with ManualTrigger >> Node (class-level, no instantiation)
+    graph = ManualTrigger >> MyNode
+
+    # THEN the graph has one trigger edge
+    trigger_edges = list(graph.trigger_edges)
+    assert len(trigger_edges) == 1
+    assert trigger_edges[0].trigger_class == ManualTrigger
+    assert trigger_edges[0].to_node == MyNode
+
+    # AND the graph has one trigger
+    triggers = list(graph.triggers)
+    assert len(triggers) == 1
+    assert triggers[0] == ManualTrigger
+
+    # AND the graph has one node
+    assert len(list(graph.nodes)) == 1
+    assert MyNode in list(graph.nodes)
+
+
+def test_graph__manual_trigger_to_set_of_nodes():
+    # GIVEN two nodes
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    # WHEN we create graph with ManualTrigger >> {NodeA, NodeB}
+    graph = ManualTrigger >> {NodeA, NodeB}
+
+    # THEN the graph has two trigger edges
+    trigger_edges = list(graph.trigger_edges)
+    assert len(trigger_edges) == 2
+
+    # AND both edges connect to the same ManualTrigger class
+    assert all(edge.trigger_class == ManualTrigger for edge in trigger_edges)
+
+    # AND edges connect to both nodes
+    target_nodes = {edge.to_node for edge in trigger_edges}
+    assert target_nodes == {NodeA, NodeB}
+
+    # AND the graph has one unique trigger
+    triggers = list(graph.triggers)
+    assert len(triggers) == 1
+
+    # AND the graph has two nodes
+    assert len(list(graph.nodes)) == 2
+
+
+def test_graph__manual_trigger_to_graph():
+    # GIVEN a graph of nodes
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    node_graph = NodeA >> NodeB
+
+    # WHEN we create graph with ManualTrigger >> Graph
+    graph = ManualTrigger >> node_graph
+
+    # THEN the graph has a trigger edge to the entrypoint
+    trigger_edges = list(graph.trigger_edges)
+    assert len(trigger_edges) == 1
+    assert trigger_edges[0].to_node == NodeA
+
+    # AND the graph preserves the original edges
+    edges = list(graph.edges)
+    assert len(edges) == 1
+    assert edges[0].to_node == NodeB
+
+    # AND the graph has both nodes
+    nodes = list(graph.nodes)
+    assert len(nodes) == 2
+    assert NodeA in nodes
+    assert NodeB in nodes
+
+
+def test_graph__manual_trigger_to_set_of_graphs_preserves_edges():
+    # GIVEN two graphs of nodes
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    class NodeC(BaseNode):
+        pass
+
+    class NodeD(BaseNode):
+        pass
+
+    graph_one = NodeA >> NodeB
+    graph_two = NodeC >> NodeD
+
+    # WHEN we create a graph with ManualTrigger >> {Graph1, Graph2}
+    combined_graph = ManualTrigger >> {graph_one, graph_two}
+
+    # THEN the combined graph has trigger edges to both entrypoints
+    trigger_edges = list(combined_graph.trigger_edges)
+    assert len(trigger_edges) == 2
+    assert {edge.to_node for edge in trigger_edges} == {NodeA, NodeC}
+
+    # AND the combined graph preserves all downstream edges
+    edges = list(combined_graph.edges)
+    assert len(edges) == 2
+    assert {(edge.from_port.node_class, edge.to_node) for edge in edges} == {
+        (NodeA, NodeB),
+        (NodeC, NodeD),
+    }
+
+    # AND the combined graph still exposes all nodes
+    nodes = list(combined_graph.nodes)
+    assert {NodeA, NodeB, NodeC, NodeD}.issubset(nodes)
+
+
+def test_graph__node_to_trigger_raises():
+    # GIVEN a node and trigger
+    class MyNode(BaseNode):
+        pass
+
+    # WHEN we try to create Node >> Trigger (class-level)
+    # THEN it raises TypeError
+    with pytest.raises(TypeError, match="Cannot create edge targeting trigger"):
+        MyNode >> ManualTrigger
+
+    # WHEN we try to create Node >> Trigger (instance-level)
+    # THEN it also raises TypeError
+    with pytest.raises(TypeError, match="Cannot create edge targeting trigger"):
+        MyNode >> ManualTrigger
+
+
+def test_graph__trigger_then_graph_then_node():
+    # GIVEN a trigger, a node, and another node
+    class StartNode(BaseNode):
+        pass
+
+    class EndNode(BaseNode):
+        pass
+
+    # WHEN we create Trigger >> Node >> Node
+    graph = ManualTrigger >> StartNode >> EndNode
+
+    # THEN the graph has one trigger edge
+    trigger_edges = list(graph.trigger_edges)
+    assert len(trigger_edges) == 1
+    assert trigger_edges[0].to_node == StartNode
+
+    # AND the graph has one regular edge
+    edges = list(graph.edges)
+    assert len(edges) == 1
+    assert edges[0].to_node == EndNode
+
+    # AND the graph has both nodes
+    nodes = list(graph.nodes)
+    assert len(nodes) == 2
+
+
+def test_graph__set_of_trigger_graphs_preserves_trigger_edges():
+    """Test that combining graphs with triggers via a set preserves trigger edges.
+
+    This tests the fix for Graph.from_set() not propagating _trigger_edges.
+    """
+
+    # GIVEN a custom scheduled trigger
+    class MyScheduledTrigger(ScheduleTrigger):
+        class Config:
+            cron = "0 9 * * *"
+
+    # AND two trigger-initiated graphs
+    class MyFirstNode(BaseNode):
+        pass
+
+    class MySecondNode(BaseNode):
+        pass
+
+    trigger_graph_a = MyScheduledTrigger >> MyFirstNode
+    trigger_graph_b = MyScheduledTrigger >> MySecondNode
+
+    # WHEN we combine them in a set
+    combined_graph = Graph.from_set({trigger_graph_a, trigger_graph_b})
+
+    # THEN the combined graph has both trigger edges
+    trigger_edges = list(combined_graph.trigger_edges)
+    assert len(trigger_edges) == 2
+
+    # AND both edges point to the correct nodes
+    target_nodes = {edge.to_node for edge in trigger_edges}
+    assert target_nodes == {MyFirstNode, MySecondNode}
+
+    # AND the graph exposes the trigger
+    triggers = list(combined_graph.triggers)
+    assert len(triggers) == 1
+    assert triggers[0] == MyScheduledTrigger
+
+
+def test_graph__set_of_trigger_graphs_to_node_preserves_trigger_edges():
+    # GIVEN a custom scheduled trigger
+    class MyScheduledTrigger(ScheduleTrigger):
+        class Config:
+            cron = "0 9 * * *"
+
+    # AND two trigger-initiated graphs
+    class MyFirstNode(BaseNode):
+        pass
+
+    class MySecondNode(BaseNode):
+        pass
+
+    class MyThirdNode(BaseNode):
+        pass
+
+    trigger_graph_a = MyScheduledTrigger >> MyFirstNode
+    trigger_graph_b = MyScheduledTrigger >> MySecondNode
+
+    # WHEN we combine them in a set and connect to another node
+    combined_graph = {trigger_graph_a, trigger_graph_b} >> MyThirdNode
+
+    # THEN the combined graph has both trigger edges
+    trigger_edges = list(combined_graph.trigger_edges)
+    assert len(trigger_edges) == 2
+
+    # AND both trigger edges point to the correct initial nodes
+    trigger_target_nodes = {edge.to_node for edge in trigger_edges}
+    assert trigger_target_nodes == {MyFirstNode, MySecondNode}
+
+    # AND the graph has regular edges to the third node
+    regular_edges = list(combined_graph.edges)
+    assert len(regular_edges) == 2
+    assert all(edge.to_node == MyThirdNode for edge in regular_edges)
+
+    # AND the graph has all three nodes
+    nodes = list(combined_graph.nodes)
+    assert len(nodes) == 3
+    assert set(nodes) == {MyFirstNode, MySecondNode, MyThirdNode}
+
+
+def test_graph__graph_rshift_graph_preserves_trigger_edges():
+    """Test that Graph >> Graph preserves trigger edges from the right-hand graph."""
+
+    # GIVEN a regular graph and a trigger-initiated graph
+    class StartNode(BaseNode):
+        pass
+
+    class TriggerNode(BaseNode):
+        pass
+
+    class EndNode(BaseNode):
+        pass
+
+    regular_graph = StartNode >> TriggerNode
+    trigger_graph = ManualTrigger >> EndNode
+
+    # WHEN we combine them with >>
+    combined = regular_graph >> trigger_graph
+
+    # THEN the combined graph has the trigger edge
+    trigger_edges = list(combined.trigger_edges)
+    assert len(trigger_edges) == 1
+    assert trigger_edges[0].to_node == EndNode
+
+    # AND the graph exposes the trigger
+    triggers = list(combined.triggers)
+    assert len(triggers) == 1
+
+
+def test_graph__graph_rshift_set_of_trigger_graphs_preserves_trigger_edges():
+    """Test that Graph >> {TriggerGraph1, TriggerGraph2} preserves trigger edges."""
+
+    # GIVEN a regular graph and two trigger-initiated graphs
+    class StartNode(BaseNode):
+        pass
+
+    class TriggerNodeA(BaseNode):
+        pass
+
+    class TriggerNodeB(BaseNode):
+        pass
+
+    regular_graph = Graph.from_node(StartNode)
+    trigger_graph_a = ManualTrigger >> TriggerNodeA
+    trigger_graph_b = ManualTrigger >> TriggerNodeB
+
+    # WHEN we combine them with >> and a set
+    combined = regular_graph >> {trigger_graph_a, trigger_graph_b}
+
+    # THEN the combined graph has both trigger edges
+    trigger_edges = list(combined.trigger_edges)
+    assert len(trigger_edges) == 2
+
+    # AND both trigger edges point to the correct nodes
+    trigger_target_nodes = {edge.to_node for edge in trigger_edges}
+    assert trigger_target_nodes == {TriggerNodeA, TriggerNodeB}
+
+
+class TestFinalOutputNoPortsNode:
+    """Test that we preserve final output nodes as NoPortsNode when constructing a Graph."""
+
+    def validate_graph(self, graph: Graph, final_output_node: Type[FinalOutputNode]):
+        assert len(list(graph._terminals)) == 1
+        node = list(graph._terminals)[0]
+        assert isinstance(node, NoPortsNode)
+        assert node.node_class == final_output_node
+
+    def test_from_edge(self):
+        # GIVEN
+        class StartNode(BaseNode[BaseState]):
+            pass
+
+        class MyFinalOutput(FinalOutputNode[BaseState, str]):
+            pass
+
+        # WHEN
+        edge = Edge(from_port=StartNode.Ports.default, to_node=MyFinalOutput)
+        graph = Graph.from_edge(edge)
+
+        # THEN
+        self.validate_graph(graph, MyFinalOutput)
+
+    def test_rshift__node(self):
+        # GIVEN
+        class StartNode(BaseNode[BaseState]):
+            pass
+
+        class MyFinalOutput(FinalOutputNode[BaseState, str]):
+            pass
+
+        # WHEN
+        initial_graph = Graph.from_node(StartNode)
+        graph = initial_graph >> MyFinalOutput
+
+        # THEN
+        self.validate_graph(graph, MyFinalOutput)
+
+    def test_rshift__graph(self):
+        # GIVEN
+        class StartNode(BaseNode[BaseState]):
+            pass
+
+        class MyFinalOutput(FinalOutputNode[BaseState, str]):
+            pass
+
+        # WHEN
+        subgraph = Graph.from_node(MyFinalOutput)
+        graph = Graph.from_node(StartNode) >> subgraph
+
+        # THEN
+        self.validate_graph(subgraph, MyFinalOutput)
+        self.validate_graph(graph, MyFinalOutput)
+
+    def test_rshift__set_with_node(self):
+        # GIVEN
+        class StartNode(BaseNode[BaseState]):
+            pass
+
+        class MyFinalOutput(FinalOutputNode[BaseState, str]):
+            pass
+
+        # WHEN
+        graph = Graph.from_node(StartNode) >> {MyFinalOutput}
+
+        # THEN
+        self.validate_graph(graph, MyFinalOutput)
+
+    def test_rshift__set_with_graph(self):
+        # GIVEN
+        class StartNode(BaseNode[BaseState]):
+            pass
+
+        class MyFinalOutput(FinalOutputNode[BaseState, str]):
+            pass
+
+        # WHEN
+        subgraph = Graph.from_node(MyFinalOutput)
+        graph = Graph.from_node(StartNode) >> {subgraph}
+
+        # THEN
+        self.validate_graph(subgraph, MyFinalOutput)
+        self.validate_graph(graph, MyFinalOutput)

vellum/workflows/inputs/__init__.py

@@ -1,5 +1,7 @@
 from .base import BaseInputs
+from .dataset_row import DatasetRow
 
 __all__ = [
     "BaseInputs",
+    "DatasetRow",
 ]

vellum/workflows/inputs/base.py

@@ -1,7 +1,8 @@
-from typing import Any, Dict, Iterator, Set, Tuple, Type, Union, get_args, get_origin
+from typing import Any, Callable, Dict, Iterator, Set, Tuple, Type, Union, cast, get_args, get_origin
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler
+from pydantic.fields import FieldInfo
 from pydantic_core import core_schema
 
 from vellum.workflows.constants import undefined
@@ -9,7 +10,7 @@ from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import WorkflowInitializationException
 from vellum.workflows.references import ExternalInputReference, WorkflowInputReference
 from vellum.workflows.references.input import InputReference
-from vellum.workflows.types.utils import get_class_attr_names, infer_types
+from vellum.workflows.types.utils import coerce_to_declared_type, get_class_attr_names, infer_types
 
 
 @dataclass_transform(kw_only_default=True)
@@ -41,7 +42,10 @@ class _BaseInputsMeta(type):
             else:
                 return WorkflowInputReference(name=name, types=types, instance=instance, inputs_class=cls)
 
-        return super().__getattribute__(name)
+        if getattr(cls, "__descriptor_class__", None) is ExternalInputReference:
+            return ExternalInputReference(name=name, types=(), instance=undefined, inputs_class=cls)
+        else:
+            return WorkflowInputReference(name=name, types=(), instance=undefined, inputs_class=cls)
 
     def __iter__(cls) -> Iterator[InputReference]:
         # We iterate through the inheritance hierarchy to find all the WorkflowInputReference attached to this
@@ -70,7 +74,7 @@ class BaseInputs(metaclass=_BaseInputsMeta):
         Initialize BaseInputs with provided keyword arguments.
 
         Validation logic:
-        1. Ensures all required fields (non-Optional types) either:
+        1. Ensures all required fields (values without defaults, regardless of Optional or not) either:
         - Have a value provided in kwargs, or
         - Have a default value defined in the class
         2. Validates that no None values are provided for required fields
@@ -84,21 +88,63 @@ class BaseInputs(metaclass=_BaseInputsMeta):
         """
         for name, field_type in self.__class__.__annotations__.items():
             # Get the value (either from kwargs or class default)
-            value = kwargs.get(name)
+            value = kwargs.get(name, undefined)
             has_default = name in vars(self.__class__)
 
-            if value is None and not has_default:
-                # Check if field_type allows None
-                origin = get_origin(field_type)
-                args = get_args(field_type)
-                if not (origin is Union and type(None) in args):
+            # If no value provided, check for default
+            if value is undefined and has_default:
+                default_value = vars(self.__class__)[name]
+                # Check if default is a FieldInfo with default_factory
+                if isinstance(default_value, FieldInfo):
+                    if default_value.default_factory is not None:
+                        value = cast(Callable[[], Any], default_value.default_factory)()
+                    elif hasattr(default_value, "default") and default_value.default is not ...:
+                        # Use the default value directly
+                        value = default_value.default
+                else:
+                    # Regular default value
+                    value = default_value
+
+            # Check if field_type allows None
+            origin = get_origin(field_type)
+            args = get_args(field_type)
+            is_optional = origin is Union and type(None) in args
+
+            if value is undefined and not has_default:
+                # All fields without defaults must be provided, even if Optional
+                raise WorkflowInitializationException(
+                    message=f"Required input variables '{name}' should have defined value",
+                    code=WorkflowErrorCode.INVALID_INPUTS,
+                    workflow_definition=self.__class__.__parent_class__,
+                )
+
+            # Validate that None is not provided for non-Optional fields
+            if value is None and not is_optional:
+                raise WorkflowInitializationException(
+                    message=f"Required input variables '{name}' should have defined value",
+                    code=WorkflowErrorCode.INVALID_INPUTS,
+                    workflow_definition=self.__class__.__parent_class__,
+                )
+
+            # Set the value on the instance (either from kwargs or default)
+            if value is not undefined:
+                # Coerce the value to the declared type if needed
+                # For Optional types, extract the non-None type for coercion
+                coercion_type = field_type
+                if is_optional:
+                    non_none_args = [arg for arg in args if arg is not type(None)]
+                    if len(non_none_args) == 1:
+                        coercion_type = non_none_args[0]
+
+                try:
+                    value = coerce_to_declared_type(value, coercion_type, name)
+                except ValueError as e:
                     raise WorkflowInitializationException(
-                        message=f"Required input variables {name} should have defined value",
+                        message=str(e),
                         code=WorkflowErrorCode.INVALID_INPUTS,
+                        workflow_definition=self.__class__.__parent_class__,
                     )
 
-            # If value provided in kwargs, set it on the instance
-            if name in kwargs:
                 setattr(self, name, value)
 
     def __iter__(self) -> Iterator[Tuple[InputReference, Any]]:
@@ -106,6 +152,10 @@ class BaseInputs(metaclass=_BaseInputsMeta):
             if hasattr(self, input_descriptor.name):
                 yield (input_descriptor, getattr(self, input_descriptor.name))
 
+    def __vellum_encode__(self) -> Dict[str, Any]:
+        """Return a JSON-serializable representation of this inputs object."""
+        return {descriptor.name: value for descriptor, value in self if value is not undefined}
+
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source_type: Type[Any], handler: GetCoreSchemaHandler