vellum-ai 1.1.3__py3-none-any.whl → 1.1.5__py3-none-any.whl

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -417,8 +417,8 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         if explicit_value and explicit_value.comment and docstring:
             comment = (
                 NodeDisplayComment(value=docstring, expanded=explicit_value.comment.expanded)
-                if explicit_value.comment.expanded
-                else NodeDisplayComment(value=docstring)
+                if explicit_value.comment.expanded is not None
+                else NodeDisplayComment(value=docstring, expanded=True)
             )
             return NodeDisplayData(
                 position=explicit_value.position,
@@ -432,7 +432,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
         if docstring:
             return NodeDisplayData(
-                comment=NodeDisplayComment(value=docstring),
+                comment=NodeDisplayComment(value=docstring, expanded=True),
             )
 
         return NodeDisplayData()

vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -98,5 +98,5 @@ def test_serialize_display_data():
     # THEN the condition should be serialized correctly
     assert data["display_data"] == {
         "position": {"x": 0.0, "y": 0.0},
-        "comment": {"value": "I hope this works"},
+        "comment": {"expanded": True, "value": "I hope this works"},
     }

vellum_ee/workflows/display/nodes/vellum/tests/test_tool_calling_node.py

@@ -1,9 +1,13 @@
+from vellum.client.types.prompt_parameters import PromptParameters
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
 from vellum.workflows.nodes.displayable.tool_calling_node.node import ToolCallingNode
+from vellum.workflows.nodes.displayable.tool_calling_node.state import ToolCallingState
+from vellum.workflows.nodes.displayable.tool_calling_node.utils import create_router_node, create_tool_prompt_node
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.types.definition import AuthorizationType, EnvironmentVariableReference, MCPServer
+from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
@@ -178,3 +182,153 @@ def test_serialize_node__tool_calling_node__mcp_server_api_key():
             },
         },
     }
+
+
+def test_serialize_tool_router_node():
+    """
+    Test that the tool router node created by create_router_node serializes successfully.
+    """
+
+    # GIVEN a simple function for tool calling
+    def my_function(arg1: str) -> str:
+        return f"Result: {arg1}"
+
+    tool_prompt_node = create_tool_prompt_node(
+        ml_model="gpt-4o-mini",
+        blocks=[],
+        functions=[my_function],
+        prompt_inputs=None,
+        parameters=PromptParameters(),
+    )
+
+    # WHEN we create a router node using create_router_node
+    router_node = create_router_node(
+        functions=[my_function],
+        tool_prompt_node=tool_prompt_node,
+    )
+
+    router_node_display_class = get_node_display_class(router_node)
+    router_node_display = router_node_display_class()
+
+    class Workflow(BaseWorkflow[BaseInputs, ToolCallingState]):
+        graph = tool_prompt_node >> router_node
+
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    display_context = workflow_display.display_context
+
+    # WHEN we serialize the router node
+    serialized_router_node = router_node_display.serialize(display_context)
+
+    # THEN the router node should serialize to the exact expected structure
+    assert serialized_router_node == {
+        "adornments": None,
+        "attributes": [
+            {
+                "id": "cd208919-c66b-451b-a739-bcf7d3451dea",
+                "name": "prompt_outputs",
+                "value": {
+                    "node_id": "19e664fc-3b57-48d2-b47a-b143b475406a",
+                    "node_output_id": "c2a5a7f7-a234-45dc-adee-bc6fc0bd28dd",
+                    "type": "NODE_OUTPUT",
+                },
+            }
+        ],
+        "base": {
+            "module": ["vellum", "workflows", "nodes", "displayable", "tool_calling_node", "utils"],
+            "name": "RouterNode",
+        },
+        "definition": {
+            "module": ["vellum", "workflows", "nodes", "displayable", "tool_calling_node", "utils"],
+            "name": "RouterNode",
+        },
+        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "id": "690c66e1-1e18-4984-b695-84beb0157541",
+        "label": "RouterNode",
+        "outputs": [],
+        "ports": [
+            {
+                "expression": {
+                    "lhs": {
+                        "lhs": {
+                            "lhs": {
+                                "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
+                                "type": "WORKFLOW_STATE",
+                            },
+                            "operator": "<",
+                            "rhs": {
+                                "lhs": {
+                                    "node_id": "19e664fc-3b57-48d2-b47a-b143b475406a",
+                                    "node_output_id": "c2a5a7f7-a234-45dc-adee-bc6fc0bd28dd",
+                                    "type": "NODE_OUTPUT",
+                                },
+                                "operator": "length",
+                                "type": "UNARY_EXPRESSION",
+                            },
+                            "type": "BINARY_EXPRESSION",
+                        },
+                        "operator": "and",
+                        "rhs": {
+                            "lhs": {
+                                "lhs": {
+                                    "lhs": {
+                                        "node_id": "19e664fc-3b57-48d2-b47a-b143b475406a",
+                                        "node_output_id": "c2a5a7f7-a234-45dc-adee-bc6fc0bd28dd",
+                                        "type": "NODE_OUTPUT",
+                                    },
+                                    "operator": "accessField",
+                                    "rhs": {
+                                        "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
+                                        "type": "WORKFLOW_STATE",
+                                    },
+                                    "type": "BINARY_EXPRESSION",
+                                },
+                                "operator": "accessField",
+                                "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "type"}},
+                                "type": "BINARY_EXPRESSION",
+                            },
+                            "operator": "=",
+                            "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "FUNCTION_CALL"}},
+                            "type": "BINARY_EXPRESSION",
+                        },
+                        "type": "BINARY_EXPRESSION",
+                    },
+                    "operator": "and",
+                    "rhs": {
+                        "lhs": {
+                            "lhs": {
+                                "lhs": {
+                                    "lhs": {
+                                        "node_id": "19e664fc-3b57-48d2-b47a-b143b475406a",
+                                        "node_output_id": "c2a5a7f7-a234-45dc-adee-bc6fc0bd28dd",
+                                        "type": "NODE_OUTPUT",
+                                    },
+                                    "operator": "accessField",
+                                    "rhs": {
+                                        "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
+                                        "type": "WORKFLOW_STATE",
+                                    },
+                                    "type": "BINARY_EXPRESSION",
+                                },
+                                "operator": "accessField",
+                                "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "value"}},
+                                "type": "BINARY_EXPRESSION",
+                            },
+                            "operator": "accessField",
+                            "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "name"}},
+                            "type": "BINARY_EXPRESSION",
+                        },
+                        "operator": "=",
+                        "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "my_function"}},
+                        "type": "BINARY_EXPRESSION",
+                    },
+                    "type": "BINARY_EXPRESSION",
+                },
+                "id": "afb4b09d-659b-459e-9a28-cf73ba6e0574",
+                "name": "my_function",
+                "type": "IF",
+            },
+            {"expression": None, "id": "4ecd916e-b5d0-407e-aab4-35551c76d02c", "name": "default", "type": "ELSE"},
+        ],
+        "trigger": {"id": "73a96f44-c2dd-40cc-96f6-49b9f914b166", "merge_behavior": "AWAIT_ATTRIBUTES"},
+        "type": "GENERIC",
+    }

vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py

@@ -670,3 +670,53 @@ def test_serialize_node__pydantic_with_node_output_reference(serialize_node):
     assert any(
         entry["key"] == "node_ref" and entry["value"]["type"] == "NODE_OUTPUT" for entry in attr_value["entries"]
     )
+
+
+def test_serialize_node__comment_expanded_true_when_content_exists(serialize_node):
+    """
+    Tests that node comment serialization sets expanded=True when comment has content.
+    """
+
+    class NodeWithComment(BaseNode):
+        """This is a test comment for the node."""
+
+        pass
+
+    serialized_node = serialize_node(NodeWithComment)
+
+    # WHEN the node is serialized
+    display_data = serialized_node["display_data"]
+
+    # THEN the comment should have expanded=True
+    assert "comment" in display_data
+    assert display_data["comment"]["value"] == "This is a test comment for the node."
+    assert display_data["comment"]["expanded"] is True
+
+
+def test_serialize_node__comment_expanded_preserved_when_explicitly_set(serialize_node):
+    """
+    Tests that explicitly set expanded value is preserved during serialization.
+    """
+    from vellum_ee.workflows.display.editor.types import NodeDisplayComment, NodeDisplayData
+    from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+
+    class NodeWithExplicitComment(BaseNode):
+        """This is a test comment."""
+
+        pass
+
+    class NodeWithExplicitCommentDisplay(BaseNodeDisplay[NodeWithExplicitComment]):
+        display_data = NodeDisplayData(comment=NodeDisplayComment(value="Custom comment", expanded=False))
+
+    serialized_node = serialize_node(
+        NodeWithExplicitComment,
+        global_node_displays={NodeWithExplicitComment: NodeWithExplicitCommentDisplay()},
+    )
+
+    # WHEN the node is serialized
+    display_data = serialized_node["display_data"]
+
+    # THEN the comment should preserve expanded=False
+    assert "comment" in display_data
+    assert display_data["comment"]["value"] == "This is a test comment."
+    assert display_data["comment"]["expanded"] is False

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

@@ -44,7 +44,10 @@ def test_serialize_workflow():
         "type": "GENERIC",
         "display_data": {
             "position": {"x": 200.0, "y": -50.0},
-            "comment": {"value": "\n    A tool calling node that calls the get_current_weather function.\n    "},
+            "comment": {
+                "expanded": True,
+                "value": "\n    A tool calling node that calls the get_current_weather function.\n    ",
+            },
         },
         "base": {
             "name": "ToolCallingNode",

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

@@ -44,7 +44,10 @@ def test_serialize_workflow():
         "type": "GENERIC",
         "display_data": {
             "position": {"x": 200.0, "y": -50.0},
-            "comment": {"value": "\n    A tool calling node that calls the get_current_weather function.\n    "},
+            "comment": {
+                "expanded": True,
+                "value": "\n    A tool calling node that calls the get_current_weather function.\n    ",
+            },
         },
         "base": {
             "name": "ToolCallingNode",

vellum_ee/workflows/display/types.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Dict, Iterator, List, 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
@@ -52,15 +52,26 @@ class WorkflowDisplayContext:
     edge_displays: EdgeDisplays = field(default_factory=dict)
     port_displays: PortDisplays = field(default_factory=dict)
     _errors: List[Exception] = field(default_factory=list)
+    _invalid_nodes: List[Type[BaseNode]] = field(default_factory=list)
     _dry_run: bool = False
 
-    def add_error(self, error: Exception) -> None:
+    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_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/workflows/base_workflow_display.py

@@ -194,6 +194,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 serialized_node = node_display.serialize(self.display_context)
             except (NotImplementedError, NodeValidationError) as e:
                 self.display_context.add_error(e)
+                self.display_context.add_invalid_node(node)
                 continue
 
             serialized_nodes[node_display.node_id] = serialized_node
@@ -319,6 +320,10 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         # Add an edge for each edge in the workflow
         for target_node, entrypoint_display in self.display_context.entrypoint_displays.items():
             unadorned_target_node = get_unadorned_node(target_node)
+            # Skip edges to invalid nodes
+            if self._is_node_invalid(unadorned_target_node):
+                continue
+
             target_node_display = self.display_context.node_displays[unadorned_target_node]
             edges.append(
                 {
@@ -335,6 +340,12 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             unadorned_source_node_port = get_unadorned_port(source_node_port)
             unadorned_target_node = get_unadorned_node(target_node)
 
+            # Skip edges that reference invalid nodes
+            if self._is_node_invalid(unadorned_target_node) or self._is_node_invalid(
+                unadorned_source_node_port.node_class
+            ):
+                continue
+
             source_node_port_display = self.display_context.port_displays[unadorned_source_node_port]
             target_node_display = self.display_context.node_displays[unadorned_target_node]
 
@@ -939,5 +950,9 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         is_required = not has_default and not is_optional
         return is_required
 
+    def _is_node_invalid(self, node: Type[BaseNode]) -> bool:
+        """Check if a node failed to serialize and should be considered invalid."""
+        return node in self.display_context.invalid_nodes
+
 
 register_workflow_display_class(workflow_class=BaseWorkflow, workflow_display_class=BaseWorkflowDisplay)

vellum_ee/workflows/display/workflows/tests/test_workflow_display.py

@@ -1,6 +1,6 @@
 import pytest
-from uuid import uuid4
-from typing import Optional
+from uuid import UUID, uuid4
+from typing import Any, Optional
 
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
@@ -11,11 +11,13 @@ from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.nodes.displayable.final_output_node.node import FinalOutputNode
 from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.core import JsonObject
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.nodes import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.vellum.retry_node import BaseRetryNodeDisplay
 from vellum_ee.workflows.display.nodes.vellum.try_node import BaseTryNodeDisplay
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
@@ -955,3 +957,43 @@ def test_serialize_workflow__state_variables():
         "required": False,
         "extensions": {"color": None},
     }
+
+
+def test_serialize_workflow__with_complete_node_failure_prunes_edges():
+    """Test that edges are pruned when a node completely fails to serialize (serialized_node is null)."""
+
+    # GIVEN a node that completely fails to serialize
+    class FailingNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class FailingNodeDisplay(BaseNodeDisplay[FailingNode]):
+        def serialize(
+            self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        ) -> JsonObject:
+            raise NotImplementedError("Complete node serialization failure")
+
+    # AND a workflow with the failing node connected to another node
+    class WorkingNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            value: str
+
+    class Workflow(BaseWorkflow):
+        graph = FailingNode >> WorkingNode
+
+    # WHEN we serialize the workflow with dry_run=True
+    workflow_display = get_workflow_display(workflow_class=Workflow, dry_run=True)
+
+    # AND we register the failing display class
+    workflow_display.display_context.node_displays[FailingNode] = FailingNodeDisplay()
+
+    data: dict = workflow_display.serialize()
+
+    # THEN the workflow should serialize but with no edges (pruned due to invalid node)
+    assert data["workflow_raw_data"]["edges"] == []
+
+    # AND only the working node and entrypoint should be in the serialized nodes
+    assert len(data["workflow_raw_data"]["nodes"]) == 2
+    node_types = [node["type"] for node in data["workflow_raw_data"]["nodes"]]
+    assert "ENTRYPOINT" in node_types
+    assert "GENERIC" in node_types  # This is the WorkingNode that should still be serialized