vellum-ai 1.1.4__py3-none-any.whl → 1.2.0__py3-none-any.whl

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

vellum_ee/workflows/tests/test_serialize_module.py

@@ -1,6 +1,37 @@
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
+def test_serialize_module_with_dataset():
+    """Test that serialize_module correctly serializes dataset from sandbox modules."""
+    module_path = "tests.workflows.basic_inputs_and_outputs"
+
+    result = BaseWorkflowDisplay.serialize_module(module_path)
+
+    assert hasattr(result, "dataset")
+
+    assert result.dataset is None
+
+
+def test_serialize_module_with_actual_dataset():
+    """Test that serialize_module correctly serializes dataset when sandbox has dataset attribute."""
+    module_path = "tests.workflows.test_dataset_serialization"
+
+    result = BaseWorkflowDisplay.serialize_module(module_path)
+
+    assert hasattr(result, "dataset")
+
+    assert result.dataset is not None
+    assert isinstance(result.dataset, list)
+    assert len(result.dataset) == 2
+
+    for i, item in enumerate(result.dataset):
+        assert "label" in item
+        assert "inputs" in item
+        assert item["label"] == f"Scenario {i + 1}"
+        assert isinstance(item["inputs"], dict)
+        assert "message" in item["inputs"]
+
+
 def test_serialize_module_happy_path():
     """Test that serialize_module works with a valid module path."""
     module_path = "tests.workflows.trivial"