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

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

@@ -20,11 +20,10 @@ def test_serialize_node__basic(serialize_node):
         pass
 
     serialized_node = serialize_node(BasicGenericNode)
-
     assert not DeepDiff(
         {
             "id": "8d7cbfe4-72ca-4367-a401-8d28723d2f00",
-            "label": "test_serialize_node__basic.<locals>.BasicGenericNode",
+            "label": "Basic Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -70,7 +69,7 @@ def test_serialize_node__if(serialize_node):
     assert not DeepDiff(
         {
             "id": "bba4b15a-dea0-48c9-a79b-4e12e99db00f",
-            "label": "test_serialize_node__if.<locals>.IfGenericNode",
+            "label": "If Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -132,7 +131,7 @@ def test_serialize_node__if_else(serialize_node):
     assert not DeepDiff(
         {
             "id": "25c9c3f1-4014-47ac-90cf-5216de10d05c",
-            "label": "test_serialize_node__if_else.<locals>.IfElseGenericNode",
+            "label": "If Else Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -201,7 +200,7 @@ def test_serialize_node__if_elif_else(serialize_node):
     assert not DeepDiff(
         {
             "id": "7b2b9cfc-12aa-432c-940d-cbe53e71de9c",
-            "label": "test_serialize_node__if_elif_else.<locals>.IfElifElseGenericNode",
+            "label": "If Elif Else Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -300,7 +299,7 @@ def test_serialize_node__node_output_reference(serialize_node):
     assert not DeepDiff(
         {
             "id": "ac067acc-6a6f-44b1-ae84-428e965ce691",
-            "label": "test_serialize_node__node_output_reference.<locals>.GenericNodeReferencingOutput",
+            "label": "Generic Node Referencing Output",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "definition": {
@@ -363,7 +362,7 @@ def test_serialize_node__vellum_secret_reference(serialize_node):
     assert not DeepDiff(
         {
             "id": "feb4b331-e25f-4a5c-9840-c5575b1efd5c",
-            "label": "test_serialize_node__vellum_secret_reference.<locals>.GenericNodeReferencingSecret",
+            "label": "Generic Node Referencing Secret",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "definition": {
@@ -429,7 +428,7 @@ def test_serialize_node__execution_count_reference(serialize_node):
     assert not DeepDiff(
         {
             "id": "0b4fe8a6-6d0c-464e-9372-10110e2b0e13",
-            "label": "test_serialize_node__execution_count_reference.<locals>.GenericNodeReferencingExecutions",
+            "label": "Generic Node Referencing Executions",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "definition": {
@@ -490,7 +489,7 @@ def test_serialize_node__null(serialize_node):
     assert not DeepDiff(
         {
             "id": "1838ce1f-9c07-4fd0-9fd4-2a3a841ea402",
-            "label": "test_serialize_node__null.<locals>.NullGenericNode",
+            "label": "Null Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -548,7 +547,7 @@ def test_serialize_node__between(serialize_node):
     assert not DeepDiff(
         {
             "id": "f2f5a1f2-a12d-4ce0-bfe9-42190ffe5328",
-            "label": "test_serialize_node__between.<locals>.BetweenGenericNode",
+            "label": "Between Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -617,7 +616,7 @@ def test_serialize_node__or(serialize_node):
     assert not DeepDiff(
         {
             "id": "5386abad-3378-4378-b3a8-831b4b77dc23",
-            "label": "test_serialize_node__or.<locals>.OrGenericNode",
+            "label": "Or Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -700,7 +699,7 @@ def test_serialize_node__and_then_or(serialize_node):
     assert not DeepDiff(
         {
             "id": "4d3995b1-437b-48d9-8878-9f57a8b725f1",
-            "label": "test_serialize_node__and_then_or.<locals>.AndThenOrGenericNode",
+            "label": "And Then Or Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -802,7 +801,7 @@ def test_serialize_node__parenthesized_and_then_or(serialize_node):
     assert not DeepDiff(
         {
             "id": "223864c9-0088-4c05-9b7d-e5b1c9ec936d",
-            "label": "test_serialize_node__parenthesized_and_then_or.<locals>.ParenthesizedAndThenOrGenericNode",
+            "label": "Parenthesized And Then Or Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -904,7 +903,7 @@ def test_serialize_node__or_then_and(serialize_node):
     assert not DeepDiff(
         {
             "id": "a946342e-4ede-4e96-8e3d-f396748d9f7c",
-            "label": "test_serialize_node__or_then_and.<locals>.OrThenAndGenericNode",
+            "label": "Or Then And Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -1005,7 +1004,7 @@ def test_serialize_node__parse_json(serialize_node):
     assert not DeepDiff(
         {
             "id": "bfc3f81b-242a-4f43-9e1c-648223d77768",
-            "label": "test_serialize_node__parse_json.<locals>.ParseJsonGenericNode",
+            "label": "Parse Json Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},

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

@@ -2,7 +2,11 @@ from deepdiff import DeepDiff
 
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
+from vellum.workflows.state.base import BaseState
 from vellum.workflows.types.core import MergeBehavior
+from vellum.workflows.workflows.base import BaseWorkflow
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
 class Inputs(BaseInputs):
@@ -17,7 +21,7 @@ def test_serialize_node__basic(serialize_node):
     assert not DeepDiff(
         {
             "id": "8d7cbfe4-72ca-4367-a401-8d28723d2f00",
-            "label": "test_serialize_node__basic.<locals>.BasicGenericNode",
+            "label": "Basic Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -59,7 +63,7 @@ def test_serialize_node__await_any(serialize_node):
     assert not DeepDiff(
         {
             "id": "42e17f0e-8496-415f-9c72-f85250ba6f0b",
-            "label": "test_serialize_node__await_any.<locals>.AwaitAnyGenericNode",
+            "label": "Await Any Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -101,7 +105,7 @@ def test_serialize_node__await_all(serialize_node):
     assert not DeepDiff(
         {
             "id": "b3e1145a-5f41-456b-9382-6d0a1e828c2f",
-            "label": "test_serialize_node__await_all.<locals>.AwaitAllGenericNode",
+            "label": "Await All Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -132,3 +136,54 @@ def test_serialize_node__await_all(serialize_node):
         serialized_node,
         ignore_order=True,
     )
+
+
+def test_serialize_node__inline_prompt_await_all():
+    """
+    Tests that InlinePromptNode with AWAIT_ALL merge behavior can be defined and serializes without errors.
+    """
+
+    # GIVEN an InlinePromptNode with AWAIT_ALL merge behavior
+    class AwaitAllInlinePromptNode(InlinePromptNode):
+        ml_model = "gpt-4o"
+        blocks = []
+
+        class Trigger(InlinePromptNode.Trigger):
+            merge_behavior = MergeBehavior.AWAIT_ALL
+
+    class TestWorkflow(BaseWorkflow[Inputs, BaseState]):
+        graph = AwaitAllInlinePromptNode
+
+    # WHEN we serialize the workflow containing the node
+    workflow_display = get_workflow_display(workflow_class=TestWorkflow)
+    serialized = workflow_display.serialize()
+
+    # THEN the workflow should serialize successfully
+    assert "workflow_raw_data" in serialized  # type: ignore
+    assert "nodes" in serialized["workflow_raw_data"]  # type: ignore
+
+    # AND the workflow should contain the InlinePromptNode
+    nodes = serialized["workflow_raw_data"]["nodes"]  # type: ignore
+    prompt_nodes = [node for node in nodes if node["type"] == "PROMPT"]  # type: ignore
+    assert len(prompt_nodes) == 1
+
+    prompt_node = prompt_nodes[0]
+
+    # AND the node should have the correct type and base
+    assert prompt_node["type"] == "PROMPT"  # type: ignore
+    assert prompt_node["base"]["name"] == "InlinePromptNode"  # type: ignore
+    assert prompt_node["base"]["module"] == [  # type: ignore
+        "vellum",
+        "workflows",
+        "nodes",
+        "displayable",
+        "inline_prompt_node",
+        "node",
+    ]
+
+    # AND the node should have the expected structure (InlinePromptNode doesn't serialize trigger info)
+    assert "data" in prompt_node  # type: ignore
+    assert "ml_model_name" in prompt_node["data"]  # type: ignore
+    assert prompt_node["data"]["ml_model_name"] == "gpt-4o"  # type: ignore
+
+    assert prompt_node["trigger"]["merge_behavior"] == "AWAIT_ALL"  # type: ignore

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

@@ -591,7 +591,7 @@ def test_serialize_workflow__try_wrapped():
         "adornments": [
             {
                 "id": "3344083c-a32c-4a32-920b-0fb5093448fa",
-                "label": "TryNode",
+                "label": "Try Node",
                 "base": {"name": "TryNode", "module": ["vellum", "workflows", "nodes", "core", "try_node", "node"]},
                 "attributes": [
                     {

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

@@ -165,6 +165,10 @@ def test_serialize_workflow():
                 "name": "ExampleBaseInlinePromptNodeWithFunctions",
                 "module": ["tests", "workflows", "basic_inline_prompt_node_with_functions", "workflow"],
             },
+            "trigger": {
+                "id": "c2dccecb-8a41-40a8-95af-325d3ab8bfe5",
+                "merge_behavior": "AWAIT_ANY",
+            },
             "outputs": [
                 {"id": "9557bd86-702d-4b45-b8c1-c3980bffe28f", "name": "json", "type": "JSON", "value": None},
                 {"id": "ead0ccb5-092f-4d9b-a9ec-5eb83d498188", "name": "text", "type": "STRING", "value": None},

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

@@ -129,7 +129,7 @@ def test_serialize_workflow():
                         },
                         {
                             "id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
-                            "label": "StartNode",
+                            "label": "Start Node",
                             "type": "GENERIC",
                             "display_data": {"position": {"x": 200.0, "y": -50.0}},
                             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},

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

@@ -40,7 +40,7 @@ def test_serialize_workflow():
     tool_calling_node = workflow_raw_data["nodes"][1]
     assert tool_calling_node == {
         "id": "21f29cac-da87-495f-bba1-093d423f4e46",
-        "label": "GetCurrentWeatherNode",
+        "label": "Get Current Weather Node",
         "type": "GENERIC",
         "display_data": {
             "position": {"x": 200.0, "y": -50.0},
@@ -153,7 +153,7 @@ def test_serialize_workflow():
                                             },
                                             {
                                                 "id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
-                                                "label": "StartNode",
+                                                "label": "Start Node",
                                                 "type": "GENERIC",
                                                 "display_data": {"position": {"x": 200.0, "y": -50.0}},
                                                 "base": {

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

@@ -40,7 +40,7 @@ def test_serialize_workflow():
     tool_calling_node = workflow_raw_data["nodes"][1]
     assert tool_calling_node == {
         "id": "21f29cac-da87-495f-bba1-093d423f4e46",
-        "label": "GetCurrentWeatherNode",
+        "label": "Get Current Weather Node",
         "type": "GENERIC",
         "display_data": {
             "position": {"x": 200.0, "y": -50.0},

vellum_ee/workflows/display/utils/expressions.py

@@ -11,6 +11,7 @@ from vellum.workflows.expressions.and_ import AndExpression
 from vellum.workflows.expressions.begins_with import BeginsWithExpression
 from vellum.workflows.expressions.between import BetweenExpression
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
+from vellum.workflows.expressions.concat import ConcatExpression
 from vellum.workflows.expressions.contains import ContainsExpression
 from vellum.workflows.expressions.does_not_begin_with import DoesNotBeginWithExpression
 from vellum.workflows.expressions.does_not_contain import DoesNotContainExpression
@@ -105,6 +106,8 @@ def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> LogicalOperato
         return "+"
     elif isinstance(descriptor, MinusExpression):
         return "-"
+    elif isinstance(descriptor, ConcatExpression):
+        return "concat"
     else:
         raise ValueError(f"Unsupported descriptor type: {descriptor}")
 
@@ -171,6 +174,7 @@ def _serialize_condition(display_context: "WorkflowDisplayContext", condition: B
             AndExpression,
             BeginsWithExpression,
             CoalesceExpression,
+            ConcatExpression,
             ContainsExpression,
             DoesNotBeginWithExpression,
             DoesNotContainExpression,
@@ -349,7 +353,11 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
     if is_workflow_class(value):
         from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
-        workflow_display = get_workflow_display(workflow_class=value)
+        # Pass the parent display context so the subworkflow can resolve parent workflow inputs
+        workflow_display = get_workflow_display(
+            workflow_class=value,
+            parent_display_context=display_context,
+        )
         serialized_value: dict = workflow_display.serialize()
         name = serialized_value["workflow_raw_data"]["definition"]["name"]
         description = value.__doc__ or ""

vellum_ee/workflows/display/utils/registry.py

@@ -1,10 +1,14 @@
+from uuid import UUID
 from typing import TYPE_CHECKING, Dict, Optional, Type
 
+from vellum.workflows.events.types import BaseEvent
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.workflows.base import BaseWorkflow
 
 if TYPE_CHECKING:
+    from vellum.workflows.events.types import ParentContext
     from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+    from vellum_ee.workflows.display.types import WorkflowDisplayContext
     from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
@@ -14,6 +18,9 @@ _workflow_display_registry: Dict[Type[BaseWorkflow], Type["BaseWorkflowDisplay"]
 # Used to store the mapping between node types and their display classes
 _node_display_registry: Dict[Type[BaseNode], Type["BaseNodeDisplay"]] = {}
 
+# Registry to store active workflow display contexts by span ID for nested workflow inheritance
+_active_workflow_display_contexts: Dict[UUID, "WorkflowDisplayContext"] = {}
+
 
 def get_from_workflow_display_registry(workflow_class: Type[BaseWorkflow]) -> Optional[Type["BaseWorkflowDisplay"]]:
     return _workflow_display_registry.get(workflow_class)
@@ -35,3 +42,42 @@ def get_from_node_display_registry(node_class: Type[BaseNode]) -> Optional[Type[
 
 def register_node_display_class(node_class: Type[BaseNode], node_display_class: Type["BaseNodeDisplay"]) -> None:
     _node_display_registry[node_class] = node_display_class
+
+
+def register_workflow_display_context(span_id: UUID, display_context: "WorkflowDisplayContext") -> None:
+    """Register a workflow display context by span ID for nested workflow inheritance."""
+    _active_workflow_display_contexts[span_id] = display_context
+
+
+def _get_parent_display_context_for_span(span_id: UUID) -> Optional["WorkflowDisplayContext"]:
+    """Get the parent display context for a given span ID."""
+    return _active_workflow_display_contexts.get(span_id)
+
+
+def get_parent_display_context_from_event(event: BaseEvent) -> Optional["WorkflowDisplayContext"]:
+    """Extract parent display context from an event by traversing the parent chain.
+
+    This function traverses up the parent chain starting from the event's parent,
+    looking for workflow parents and attempting to get their display context.
+
+    Args:
+        event: The event to extract parent display context from
+
+    Returns:
+        The parent workflow display context if found, None otherwise
+    """
+    if not event.parent:
+        return None
+
+    current_parent: Optional["ParentContext"] = event.parent
+    while current_parent:
+        if current_parent.type == "WORKFLOW":
+            # Found a parent workflow, try to get its display context
+            parent_span_id = current_parent.span_id
+            parent_display_context = _get_parent_display_context_for_span(parent_span_id)
+            if parent_display_context:
+                return parent_display_context
+        # Move up the parent chain
+        current_parent = current_parent.parent
+
+    return None

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -3,6 +3,7 @@ import fnmatch
 from functools import cached_property
 import importlib
 import inspect
+import json
 import logging
 import os
 from uuid import UUID
@@ -15,12 +16,14 @@ from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges import Edge
 from vellum.workflows.events.workflow import NodeEventDisplayContext, WorkflowEventDisplayContext
+from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.nodes.displayable.final_output_node.node import FinalOutputNode
 from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port, get_wrapped_node
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference, WorkflowInputReference
+from vellum.workflows.state.encoder import DefaultStateEncoder
 from vellum.workflows.types.core import Json, JsonArray, JsonObject
 from vellum.workflows.types.generics import WorkflowType
 from vellum.workflows.types.utils import get_original_base
@@ -72,6 +75,7 @@ IGNORE_PATTERNS = [
 class WorkflowSerializationResult(UniversalBaseModel):
     exec_config: Dict[str, Any]
     errors: List[str]
+    dataset: Optional[List[Dict[str, Any]]] = None
 
 
 class BaseWorkflowDisplay(Generic[WorkflowType]):
@@ -524,7 +528,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         workflow_input_displays: WorkflowInputsDisplays = {}
         # If we're dealing with a nested workflow, then it should have access to the inputs of its parents.
         global_workflow_input_displays = (
-            copy(self._parent_display_context.workflow_input_displays) if self._parent_display_context else {}
+            copy(self._parent_display_context.global_workflow_input_displays) if self._parent_display_context else {}
         )
         for workflow_input in self._workflow.get_inputs_class():
             workflow_input_display_overrides = self.inputs_display.get(workflow_input)
@@ -892,9 +896,25 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         if additional_files:
             exec_config["module_data"] = {"additional_files": cast(JsonObject, additional_files)}
 
+        dataset = None
+        try:
+            sandbox_module_path = f"{module}.sandbox"
+            sandbox_module = importlib.import_module(sandbox_module_path)
+            if hasattr(sandbox_module, "dataset"):
+                dataset_attr = getattr(sandbox_module, "dataset")
+                if dataset_attr and isinstance(dataset_attr, list):
+                    dataset = []
+                    for i, inputs_obj in enumerate(dataset_attr):
+                        if isinstance(inputs_obj, BaseInputs):
+                            serialized_inputs = json.loads(json.dumps(inputs_obj, cls=DefaultStateEncoder))
+                            dataset.append({"label": f"Scenario {i + 1}", "inputs": serialized_inputs})
+        except (ImportError, AttributeError):
+            pass
+
         return WorkflowSerializationResult(
             exec_config=exec_config,
             errors=[str(error) for error in workflow_display.display_context.errors],
+            dataset=dataset,
         )
 
     def _gather_additional_module_files(self, module_path: str) -> Dict[str, str]:

vellum_ee/workflows/tests/test_registry.py

@@ -0,0 +1,169 @@
+from datetime import datetime, timezone
+from uuid import uuid4
+
+from vellum.workflows.events.types import NodeParentContext, WorkflowParentContext
+from vellum.workflows.events.workflow import WorkflowExecutionInitiatedBody, WorkflowExecutionInitiatedEvent
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes import BaseNode
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.workflows.base import BaseWorkflow
+from vellum_ee.workflows.display.utils.registry import (
+    get_parent_display_context_from_event,
+    register_workflow_display_context,
+)
+
+
+class MockInputs(BaseInputs):
+    pass
+
+
+class MockState(BaseState):
+    pass
+
+
+class MockNode(BaseNode):
+    pass
+
+
+class MockWorkflow(BaseWorkflow[MockInputs, MockState]):
+    pass
+
+
+class MockWorkflowDisplayContext:
+    pass
+
+
+def test_get_parent_display_context_from_event__no_parent():
+    """Test event with no parent returns None"""
+    # GIVEN a workflow execution initiated event with no parent
+    event: WorkflowExecutionInitiatedEvent = WorkflowExecutionInitiatedEvent(
+        id=uuid4(),
+        timestamp=datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+        trace_id=uuid4(),
+        span_id=uuid4(),
+        body=WorkflowExecutionInitiatedBody(
+            workflow_definition=MockWorkflow,
+            inputs=MockInputs(),
+        ),
+        parent=None,  # No parent
+    )
+
+    # WHEN getting parent display context
+    result = get_parent_display_context_from_event(event)
+
+    # THEN it should return None
+    assert result is None
+
+
+def test_get_parent_display_context_from_event__non_workflow_parent():
+    """Test event with non-workflow parent continues traversal"""
+    # GIVEN an event with a non-workflow parent (NodeParentContext)
+    non_workflow_parent = NodeParentContext(node_definition=MockNode, span_id=uuid4(), parent=None)
+
+    event: WorkflowExecutionInitiatedEvent = WorkflowExecutionInitiatedEvent(
+        id=uuid4(),
+        timestamp=datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+        trace_id=uuid4(),
+        span_id=uuid4(),
+        body=WorkflowExecutionInitiatedBody(
+            workflow_definition=MockWorkflow,
+            inputs=MockInputs(),
+        ),
+        parent=non_workflow_parent,
+    )
+
+    # WHEN getting parent display context
+    result = get_parent_display_context_from_event(event)
+
+    # THEN it should return None (no workflow parent found)
+    assert result is None
+
+
+def test_get_parent_display_context_from_event__nested_workflow_parents():
+    """Test event with nested workflow parents traverses correctly"""
+    # GIVEN a chain of nested contexts:
+    # Event -> WorkflowParent -> NodeParent -> MiddleWorkflowParent -> NodeParent
+
+    # Top level workflow parent
+    top_workflow_span_id = uuid4()
+    top_context = MockWorkflowDisplayContext()
+    register_workflow_display_context(top_workflow_span_id, top_context)  # type: ignore[arg-type]
+
+    top_workflow_parent = WorkflowParentContext(
+        workflow_definition=MockWorkflow, span_id=top_workflow_span_id, parent=None
+    )
+
+    top_node_parent = NodeParentContext(node_definition=MockNode, span_id=uuid4(), parent=top_workflow_parent)
+
+    # AND middle workflow parent (no display context)
+    middle_workflow_span_id = uuid4()
+    middle_workflow_parent = WorkflowParentContext(
+        workflow_definition=MockWorkflow, span_id=middle_workflow_span_id, parent=top_node_parent
+    )
+
+    # AND node parent between middle workflow and event
+    node_parent = NodeParentContext(node_definition=MockNode, span_id=uuid4(), parent=middle_workflow_parent)
+
+    event: WorkflowExecutionInitiatedEvent = WorkflowExecutionInitiatedEvent(
+        id=uuid4(),
+        timestamp=datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+        trace_id=uuid4(),
+        span_id=uuid4(),
+        body=WorkflowExecutionInitiatedBody(
+            workflow_definition=MockWorkflow,
+            inputs=MockInputs(),
+        ),
+        parent=node_parent,
+    )
+
+    # WHEN getting parent display context
+    result = get_parent_display_context_from_event(event)
+
+    # THEN it should find the top-level workflow context
+    assert result == top_context
+
+
+def test_get_parent_display_context_from_event__middle_workflow_has_context():
+    """Test event returns middle workflow context when it's the first one with registered context"""
+    # GIVEN a chain of nested contexts:
+    # Event -> WorkflowParent -> NodeParent -> MiddleWorkflowParent -> NodeParent
+
+    top_workflow_span_id = uuid4()
+    top_context = MockWorkflowDisplayContext()
+    register_workflow_display_context(top_workflow_span_id, top_context)  # type: ignore[arg-type]
+
+    top_workflow_parent = WorkflowParentContext(
+        workflow_definition=MockWorkflow, span_id=top_workflow_span_id, parent=None
+    )
+
+    # AND node parent between top workflow and middle workflow
+    top_node_parent = NodeParentContext(node_definition=MockNode, span_id=uuid4(), parent=top_workflow_parent)
+
+    # AND middle workflow parent
+    middle_workflow_span_id = uuid4()
+    middle_context = MockWorkflowDisplayContext()
+    register_workflow_display_context(middle_workflow_span_id, middle_context)  # type: ignore[arg-type]
+
+    middle_workflow_parent = WorkflowParentContext(
+        workflow_definition=MockWorkflow, span_id=middle_workflow_span_id, parent=top_node_parent
+    )
+
+    node_parent = NodeParentContext(node_definition=MockNode, span_id=uuid4(), parent=middle_workflow_parent)
+
+    event: WorkflowExecutionInitiatedEvent = WorkflowExecutionInitiatedEvent(
+        id=uuid4(),
+        timestamp=datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+        trace_id=uuid4(),
+        span_id=uuid4(),
+        body=WorkflowExecutionInitiatedBody(
+            workflow_definition=MockWorkflow,
+            inputs=MockInputs(),
+        ),
+        parent=node_parent,
+    )
+
+    # WHEN getting parent display context
+    result = get_parent_display_context_from_event(event)
+
+    # THEN it should find the MIDDLE workflow context
+    assert result == middle_context

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"