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

vellum_ee/workflows/display/tests/workflow_serialization/test_trigger_display_from_display_class.py

@@ -0,0 +1,153 @@
+"""Tests for BaseTrigger.Display serialization."""
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.triggers.schedule import ScheduleTrigger
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+def test_serialize_basetrigger_with_display_class():
+    """Tests that BaseTrigger.Display class attributes serialize correctly."""
+
+    class MyTrigger(ScheduleTrigger):
+        class Config(ScheduleTrigger.Config):
+            cron = "0 9 * * *"
+            timezone = "UTC"
+
+        class Display:
+            icon = "vellum:icon:gear"
+            color = "purple"
+
+    class SimpleNode(BaseNode):
+        pass
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = MyTrigger >> SimpleNode
+
+    # WHEN we serialize the workflow
+    result: dict = get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+    assert "triggers" in result
+    triggers = result["triggers"]
+    assert isinstance(triggers, list)
+    assert len(triggers) == 1
+
+    trigger = triggers[0]
+    assert isinstance(trigger, dict)
+    assert "display_data" in trigger
+    display_data = trigger["display_data"]
+    assert isinstance(display_data, dict)
+    assert display_data["icon"] == "vellum:icon:gear"
+    assert display_data["color"] == "purple"
+
+
+def test_serialize_trigger_with_only_icon():
+    """Tests that triggers with only icon serialize correctly."""
+
+    class IconOnlyTrigger(ScheduleTrigger):
+        class Config(ScheduleTrigger.Config):
+            cron = "0 9 * * *"
+            timezone = "UTC"
+
+        class Display:
+            icon = "vellum:icon:star"
+
+    class SimpleNode(BaseNode):
+        pass
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = IconOnlyTrigger >> SimpleNode
+
+    # WHEN we serialize the workflow
+    result: dict = get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+    triggers = result["triggers"]
+    assert len(triggers) == 1
+    trigger = triggers[0]
+    display_data = trigger["display_data"]
+    assert display_data["icon"] == "vellum:icon:star"
+    assert "color" not in display_data
+
+
+def test_serialize_trigger_with_only_color():
+    """Tests that triggers with only color serialize correctly."""
+
+    class ColorOnlyTrigger(ScheduleTrigger):
+        class Config(ScheduleTrigger.Config):
+            cron = "0 9 * * *"
+            timezone = "UTC"
+
+        class Display:
+            color = "#FF5733"
+
+    class SimpleNode(BaseNode):
+        pass
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = ColorOnlyTrigger >> SimpleNode
+
+    # WHEN we serialize the workflow
+    result: dict = get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+    triggers = result["triggers"]
+    assert len(triggers) == 1
+    trigger = triggers[0]
+    display_data = trigger["display_data"]
+    assert display_data["color"] == "#FF5733"
+    assert "icon" not in display_data
+
+
+def test_serialize_trigger_without_display_class():
+    """Tests that triggers without Display class don't have display_data."""
+
+    class PlainTrigger(ScheduleTrigger):
+        class Config(ScheduleTrigger.Config):
+            cron = "0 9 * * *"
+            timezone = "UTC"
+
+    class SimpleNode(BaseNode):
+        pass
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = PlainTrigger >> SimpleNode
+
+    # WHEN we serialize the workflow
+    result: dict = get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+    triggers = result["triggers"]
+    assert len(triggers) == 1
+    trigger = triggers[0]
+    display_data = trigger.get("display_data")
+    if display_data is not None:
+        assert display_data == {"label": "Trigger", "position": {"x": 0.0, "y": 0.0}, "z_index": 0}
+
+
+def test_serialize_trigger_with_none_values():
+    """Tests that triggers with None icon/color don't serialize those fields."""
+
+    class NoneValuesTrigger(ScheduleTrigger):
+        class Config(ScheduleTrigger.Config):
+            cron = "0 9 * * *"
+            timezone = "UTC"
+
+        class Display:
+            icon = None
+            color = None
+
+    class SimpleNode(BaseNode):
+        pass
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = NoneValuesTrigger >> SimpleNode
+
+    # WHEN we serialize the workflow
+    result: dict = get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+    triggers = result["triggers"]
+    assert len(triggers) == 1
+    trigger = triggers[0]
+    display_data = trigger.get("display_data")
+    if display_data is not None:
+        assert display_data == {}

vellum_ee/workflows/display/tests/workflow_serialization/test_web_search_node_serialization.py

@@ -0,0 +1,72 @@
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.web_search.workflow import WebSearchWorkflow
+
+
+def test_serialize_web_search_workflow():
+    # GIVEN a WebSearchWorkflow with a node that has various input types
+    # - query from workflow input
+
+    # WHEN we serialize the workflow through the display system
+    workflow_display = get_workflow_display(workflow_class=WebSearchWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a properly structured serialized workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "state_variables",
+        "output_variables",
+    }
+
+    # AND it should have input variables
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 1
+    assert input_variables[0]["key"] == "search_query"
+    assert input_variables[0]["type"] == "STRING"
+
+    # AND it should have output variables
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 2
+    output_keys = {var["key"] for var in output_variables}
+    assert output_keys == {"search_results", "result_urls"}
+
+    # AND the web search node should be serialized as a GENERIC type node
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    nodes = workflow_raw_data["nodes"]
+
+    web_search_node = next(
+        (
+            node
+            for node in nodes
+            if node and node.get("type") == "GENERIC" and node.get("base", {}).get("name") == "WebSearchNode"
+        ),
+        None,
+    )
+
+    assert web_search_node is not None, "WebSearchNode should be serialized as GENERIC type"
+
+    # AND it should have the correct base module reference
+    expected_base = {
+        "name": "WebSearchNode",
+        "module": ["vellum", "workflows", "nodes", "displayable", "web_search_node", "node"],
+    }
+    assert (
+        web_search_node["base"] == expected_base
+    ), f"Base mismatch. Expected: {expected_base}, Found: {web_search_node['base']}"
+
+    # AND it should have one serializable attribute defined in our display class
+    web_search_attributes = web_search_node["attributes"]
+    attribute_names = {attr["name"] for attr in web_search_attributes}
+
+    # AND the attributes should match exactly what we defined in __serializable_inputs__
+    assert attribute_names == {"query"}
+
+    # AND it should have all three expected outputs from WebSearchNode
+    assert {output["name"] for output in web_search_node["outputs"]} == {"text", "urls", "results"}
+
+    # AND each attribute should have the correct value type based on its source
+    attr_types_by_name = {attr["name"]: attr["value"]["type"] for attr in web_search_attributes}
+
+    # AND query should reference the workflow input
+    assert attr_types_by_name["query"] == "WORKFLOW_INPUT"

vellum_ee/workflows/display/tests/workflow_serialization/test_workflow_input_parameterization_error.py

@@ -0,0 +1,37 @@
+import pytest
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+def test_workflow_serialization_error_when_node_references_unparameterized_inputs():
+    """Test that a helpful error is raised when a node references inputs not parameterized by the workflow."""
+
+    class CustomInputs(BaseInputs):
+        custom_input: str
+
+    class TestNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            output = CustomInputs.custom_input
+
+    class TestWorkflow(BaseWorkflow):
+        graph = TestNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            result = TestNode.Outputs.output
+
+    workflow_display = get_workflow_display(workflow_class=TestWorkflow)
+
+    with pytest.raises(UnsupportedSerializationException) as exc_info:
+        workflow_display.serialize()
+
+    error_message = str(exc_info.value)
+    expected_message = (
+        "Inputs class 'CustomInputs' referenced during serialization of 'TestWorkflow' "
+        "without parameterizing this Workflow with this Inputs definition. "
+        "Update your Workflow definition to 'TestWorkflow(BaseWorkflow[CustomInputs, BaseState])'."
+    )
+    assert error_message == expected_message

vellum_ee/workflows/display/types.py

@@ -1,11 +1,13 @@
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Dict, Tuple, Type
+from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Type
 
+from vellum.client import Vellum as VellumClient
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.events.workflow import WorkflowEventDisplayContext  # noqa: F401
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference, StateValueReference, WorkflowInputReference
+from vellum.workflows.vellum_client import create_vellum_client
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.base import (
     EdgeDisplay,
@@ -26,7 +28,7 @@ if TYPE_CHECKING:
 WorkflowInputsDisplays = Dict[WorkflowInputReference, WorkflowInputsDisplay]
 StateValueDisplays = Dict[StateValueReference, StateValueDisplay]
 NodeDisplays = Dict[Type[BaseNode], BaseNodeDisplay]
-NodeOutputDisplays = Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]]
+NodeOutputDisplays = Dict[OutputReference, NodeOutputDisplay]
 EntrypointDisplays = Dict[Type[BaseNode], EntrypointDisplay]
 WorkflowOutputDisplays = Dict[BaseDescriptor, WorkflowOutputDisplay]
 EdgeDisplays = Dict[Tuple[Port, Type[BaseNode]], EdgeDisplay]
@@ -35,6 +37,7 @@ PortDisplays = Dict[Port, PortDisplay]
 
 @dataclass
 class WorkflowDisplayContext:
+    client: VellumClient = field(default_factory=create_vellum_client)
     workflow_display_class: Type["BaseWorkflowDisplay"] = field(default_factory=get_default_workflow_display_class)
     workflow_display: WorkflowMetaDisplay = field(default_factory=lambda: WorkflowMetaDisplay.get_default(BaseWorkflow))
     workflow_input_displays: WorkflowInputsDisplays = field(default_factory=dict)
@@ -48,3 +51,30 @@ class WorkflowDisplayContext:
     workflow_output_displays: WorkflowOutputDisplays = field(default_factory=dict)
     edge_displays: EdgeDisplays = field(default_factory=dict)
     port_displays: PortDisplays = field(default_factory=dict)
+    dry_run: bool = False
+    _errors: List[Exception] = field(default_factory=list)
+    _invalid_nodes: List[Type[BaseNode]] = field(default_factory=list)
+
+    def add_error(self, error: Exception, node: Optional[Type[BaseNode]] = None) -> None:
+        if self.dry_run:
+            self._errors.append(error)
+            return
+
+        raise error
+
+    def add_validation_error(self, error: Exception) -> None:
+        self._errors.append(error)
+
+    def add_invalid_node(self, node: Type[BaseNode]) -> None:
+        """Track a node that failed to serialize."""
+        if node not in self._invalid_nodes:
+            self._invalid_nodes.append(node)
+
+    @property
+    def errors(self) -> Iterator[Exception]:
+        return iter(self._errors)
+
+    @property
+    def invalid_nodes(self) -> Iterator[Type[BaseNode]]:
+        """Get an iterator over nodes that failed to serialize."""
+        return iter(self._invalid_nodes)

vellum_ee/workflows/display/utils/auto_layout.py

@@ -0,0 +1,130 @@
+from collections import defaultdict
+from typing import Dict, List, Set, Tuple
+
+from vellum_ee.workflows.display.base import EdgeDisplay
+from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
+
+
+def auto_layout_nodes(
+    nodes: List[Tuple[str, NodeDisplayData]],
+    edges: List[Tuple[str, str, EdgeDisplay]],
+    node_spacing: float = 150.0,
+    layer_spacing: float = 200.0,
+) -> List[Tuple[str, NodeDisplayData]]:
+    """
+    Auto-layout nodes in a hierarchical left-to-right arrangement.
+
+    Args:
+        nodes: List of (node_id, NodeDisplayData) tuples
+        edges: List of (source_node_id, target_node_id, EdgeDisplay) tuples
+        node_spacing: Vertical spacing between nodes in the same layer
+        layer_spacing: Horizontal spacing between layers
+
+    Returns:
+        List of (node_id, NodeDisplayData) tuples with updated positions
+    """
+    if not nodes:
+        return []
+
+    node_dict = {node_id: data for node_id, data in nodes}
+
+    graph: Dict[str, List[str]] = defaultdict(list)
+    in_degree: Dict[str, int] = defaultdict(int)
+    all_nodes = set(node_dict.keys())
+
+    for source, target, _ in edges:
+        if source in all_nodes and target in all_nodes:
+            graph[source].append(target)
+            in_degree[target] += 1
+
+    for node_id in all_nodes:
+        if node_id not in in_degree:
+            in_degree[node_id] = 0
+
+    layers = _topological_sort_layers(graph, in_degree, all_nodes)
+
+    positioned_nodes = []
+    current_x = 0.0
+
+    for layer in layers:
+        if not layer:
+            continue
+
+        layer_nodes = [(node_id, node_dict[node_id]) for node_id in layer]
+        total_height = _calculate_layer_height(layer_nodes, node_spacing)
+
+        current_y = -total_height / 2.0
+
+        for node_id in layer:
+            node_data = node_dict[node_id]
+            node_height = node_data.height or 100.0
+
+            new_position = NodeDisplayPosition(x=current_x, y=current_y)
+
+            updated_data = NodeDisplayData(
+                position=new_position, width=node_data.width, height=node_data.height, comment=node_data.comment
+            )
+
+            positioned_nodes.append((node_id, updated_data))
+            current_y += node_height + node_spacing
+
+        current_x += layer_spacing
+
+    return positioned_nodes
+
+
+def _topological_sort_layers(
+    graph: Dict[str, List[str]], in_degree: Dict[str, int], all_nodes: Set[str]
+) -> List[List[str]]:
+    """
+    Perform topological sorting and group nodes into layers.
+
+    Returns:
+        List of layers, where each layer is a list of node IDs
+    """
+    layers = []
+    remaining_nodes = set(all_nodes)
+
+    while remaining_nodes:
+        current_layer = []
+        for node in remaining_nodes:
+            if in_degree[node] == 0:
+                current_layer.append(node)
+
+        if not current_layer:
+            current_layer = [next(iter(remaining_nodes))]
+
+        layers.append(current_layer)
+
+        for node in current_layer:
+            remaining_nodes.remove(node)
+            for neighbor in graph[node]:
+                if neighbor in remaining_nodes:
+                    in_degree[neighbor] -= 1
+
+    return [sorted(layer) for layer in layers]
+
+
+def _calculate_layer_height(layer_nodes: List[Tuple[str, NodeDisplayData]], node_spacing: float) -> float:
+    """
+    Calculate the total height needed for a layer of nodes.
+
+    Args:
+        layer_nodes: List of (node_id, NodeDisplayData) tuples in the layer
+        node_spacing: Spacing between nodes
+
+    Returns:
+        Total height needed for the layer
+    """
+    if not layer_nodes:
+        return 0.0
+
+    total_height = 0.0
+    for i, (_, node_data) in enumerate(layer_nodes):
+        node_height = node_data.height or 100.0
+        total_height += node_height
+
+        if i < len(layer_nodes) - 1:
+            total_height += node_spacing
+
+    return total_height

vellum_ee/workflows/display/utils/events.py

@@ -0,0 +1,57 @@
+from typing import Optional
+
+from vellum import Vellum
+from vellum.workflows.events.workflow import WorkflowEvent, WorkflowExecutionInitiatedEvent
+from vellum_ee.workflows.display.utils.registry import (
+    get_parent_display_context_from_event,
+    register_workflow_display_class,
+    register_workflow_display_context,
+)
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+def _should_mark_workflow_dynamic(event: WorkflowExecutionInitiatedEvent) -> bool:
+    """
+    Check if workflow should be marked as dynamic based on execution context.
+    Returns True if parent.type == WORKFLOW_RELEASE_TAG and parent.parent.type == WORKFLOW_NODE.
+    """
+    if not event.parent:
+        return False
+
+    parent = event.parent
+    if parent.type != "WORKFLOW_RELEASE_TAG":
+        return False
+
+    if not parent.parent or parent.parent.type != "WORKFLOW_NODE":
+        return False
+
+    return True
+
+
+def event_enricher(
+    event: WorkflowEvent, client: Optional[Vellum] = None, metadata: Optional[dict] = None
+) -> WorkflowEvent:
+    if event.name == "workflow.execution.initiated":
+        workflow_definition = event.body.workflow_definition
+        workflow_display = get_workflow_display(
+            workflow_class=workflow_definition,
+            parent_display_context=get_parent_display_context_from_event(event),
+            client=client,
+            dry_run=True,
+        )
+        register_workflow_display_context(event.span_id, workflow_display.display_context)
+        event.body.display_context = workflow_display.get_event_display_context()
+
+        if event.body.workflow_definition.is_dynamic or _should_mark_workflow_dynamic(event):
+            register_workflow_display_class(workflow_definition, workflow_display.__class__)
+            workflow_version_exec_config = workflow_display.serialize()
+            setattr(event.body, "workflow_version_exec_config", workflow_version_exec_config)
+
+        if metadata is not None:
+            event.body.server_metadata = metadata
+
+    elif event.name == "workflow.execution.fulfilled":
+        if metadata is not None:
+            event.body.server_metadata = metadata
+
+    return event

vellum_ee/workflows/display/utils/exceptions.py

@@ -5,3 +5,22 @@ class UserFacingException(Exception):
 
 class UnsupportedSerializationException(UserFacingException):
     pass
+
+
+class NodeValidationError(UserFacingException):
+    """Exception raised when a node fails validation during workflow display serialization."""
+
+    def __init__(self, message: str, node_class_name: str):
+        self.message = message
+        self.node_class_name = node_class_name
+        super().__init__(f"Node validation error in {node_class_name}: {message}")
+
+
+class InvalidInputReferenceError(UserFacingException):
+    """Exception raised when a node references a non-existent workflow input."""
+
+    def __init__(self, message: str, inputs_class_name: str, attribute_name: str):
+        self.message = message
+        self.inputs_class_name = inputs_class_name
+        self.attribute_name = attribute_name
+        super().__init__(f"Invalid input reference in {inputs_class_name}.{attribute_name}: {message}")