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

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

@@ -0,0 +1,268 @@
+import pytest
+import sys
+from uuid import uuid4
+from typing import Callable
+
+from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
+from vellum_ee.workflows.server.namespaces import get_random_namespace
+from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
+
+
+def _generate_uuid_namespace() -> str:
+    """Generate a UUID namespace like the workflow server uses for some requests."""
+    return str(uuid4())
+
+
+@pytest.mark.parametrize(
+    "namespace_generator",
+    [
+        pytest.param(_generate_uuid_namespace, id="uuid_namespace"),
+        pytest.param(get_random_namespace, id="workflow_tmp_namespace"),
+    ],
+)
+def test_serialize_module__dataset_mocks_are_stable(namespace_generator: Callable[[], str]):
+    """
+    Tests that serialization produces stable, deterministic output for DatasetRow mocks
+    across multiple serializations, regardless of the namespace format used.
+    """
+
+    # GIVEN a workflow module with a sandbox.py that has DatasetRow with mocks
+    files = {
+        "__init__.py": "",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes import BaseNode
+from vellum.workflows.state import BaseState
+
+
+class Inputs(BaseInputs):
+    message: str
+
+
+class CodeExecution(BaseNode[BaseState]):
+    class Outputs(BaseNode.Outputs):
+        result: str
+        log: str
+
+    def run(self) -> Outputs:
+        return self.Outputs(result="executed", log="")
+
+
+class Workflow(BaseWorkflow[Inputs, BaseState]):
+    graph = CodeExecution
+
+    class Outputs(BaseWorkflow.Outputs):
+        final_result = CodeExecution.Outputs.result
+""",
+        "sandbox.py": """\
+from vellum.workflows import DatasetRow
+
+from .workflow import CodeExecution, Inputs
+
+dataset = [
+    DatasetRow(
+        label="With mocked code node",
+        inputs=Inputs(message="hello"),
+        mocks=[
+            CodeExecution.Outputs(result="MOCKED_RESULT", log=""),
+        ],
+    ),
+]
+""",
+    }
+
+    # AND two different namespaces (simulating how the workflow server creates a new namespace per request)
+    namespace_1 = namespace_generator()
+    namespace_2 = namespace_generator()
+
+    # AND the virtual file loaders are registered for both namespaces
+    sys.meta_path.append(VirtualFileFinder(files, namespace_1))
+    sys.meta_path.append(VirtualFileFinder(files, namespace_2))
+
+    try:
+        # WHEN we serialize the module with different namespaces (like the workflow server does)
+        result_1 = BaseWorkflowDisplay.serialize_module(namespace_1)
+        result_2 = BaseWorkflowDisplay.serialize_module(namespace_2)
+
+        # THEN both serializations should succeed without errors
+        assert len(result_1.errors) == 0, f"First serialization had errors: {result_1.errors}"
+        assert len(result_2.errors) == 0, f"Second serialization had errors: {result_2.errors}"
+
+        # AND both should have a dataset
+        assert result_1.dataset is not None, "First serialization should have a dataset"
+        assert result_2.dataset is not None, "Second serialization should have a dataset"
+
+        # AND the datasets should be identical (stable serialization)
+        assert result_1.dataset == result_2.dataset, (
+            f"Datasets should be identical across serializations.\n"
+            f"First: {result_1.dataset}\n"
+            f"Second: {result_2.dataset}"
+        )
+    finally:
+        # Clean up the virtual file finders
+        sys.meta_path = [finder for finder in sys.meta_path if not isinstance(finder, VirtualFileFinder)]
+
+
+def test_serialize_module__dataset_with_node_output_mocks_field_name():
+    """
+    Tests that DatasetRow with node_output_mocks field (legacy field name) properly
+    serializes the mocks. This reproduces an issue where using node_output_mocks
+    instead of mocks caused the dataset to return None for the mocks field.
+    """
+
+    # GIVEN a workflow module with parallel code execution nodes
+    files = {
+        "__init__.py": "",
+        "inputs.py": """\
+from typing import Any, Optional
+
+from vellum.workflows.inputs import BaseInputs
+
+
+class Inputs(BaseInputs):
+    json: Optional[Any] = None
+""",
+        "nodes/__init__.py": "",
+        "nodes/code_execution/__init__.py": """\
+from vellum.workflows.nodes.displayable import CodeExecutionNode
+from vellum.workflows.state import BaseState
+from vellum.workflows.types.core import MergeBehavior
+
+
+class CodeExecution(CodeExecutionNode[BaseState, str]):
+    filepath = "./script.py"
+    code_inputs = {}
+    runtime = "PYTHON_3_11_6"
+    packages = []
+
+    class Trigger(CodeExecutionNode.Trigger):
+        merge_behavior = MergeBehavior.AWAIT_ATTRIBUTES
+""",
+        "nodes/code_execution/script.py": '''\
+"""
+You must define a function called `main` whose arguments are named after the
+Input Variables.
+"""
+
+def main() -> str:
+    return "5".upper()
+''',
+        "nodes/code_execution_2/__init__.py": """\
+from vellum.workflows.nodes.displayable import CodeExecutionNode
+from vellum.workflows.state import BaseState
+from vellum.workflows.types.core import MergeBehavior
+
+
+class CodeExecution2(CodeExecutionNode[BaseState, str]):
+    filepath = "./script.py"
+    code_inputs = {}
+    runtime = "PYTHON_3_11_6"
+    packages = []
+
+    class Trigger(CodeExecutionNode.Trigger):
+        merge_behavior = MergeBehavior.AWAIT_ATTRIBUTES
+""",
+        "nodes/code_execution_2/script.py": '''\
+"""
+You must define a function called `main` whose arguments are named after the
+Input Variables.
+"""
+
+def main() -> str:
+    return "TODO"
+''',
+        "nodes/output.py": """\
+from vellum.workflows.nodes.displayable import FinalOutputNode
+from vellum.workflows.state import BaseState
+
+
+class Output(FinalOutputNode[BaseState, str]):
+    class Outputs(FinalOutputNode.Outputs):
+        value = None
+""",
+        "nodes/output_2.py": """\
+from vellum.workflows.nodes.displayable import FinalOutputNode
+from vellum.workflows.state import BaseState
+from vellum.workflows.types.core import MergeBehavior
+
+
+class Output2(FinalOutputNode[BaseState, str]):
+    class Outputs(FinalOutputNode.Outputs):
+        value = ""
+
+    class Trigger(FinalOutputNode.Trigger):
+        merge_behavior = MergeBehavior.AWAIT_ATTRIBUTES
+""",
+        "sandbox.py": """\
+from vellum.workflows.inputs import DatasetRow
+
+from .inputs import Inputs
+from .nodes.code_execution import CodeExecution
+from .nodes.code_execution_2 import CodeExecution2
+
+dataset = [
+    DatasetRow(
+        label="With mocked code nodes",
+        inputs=Inputs(json={}),
+        node_output_mocks=[
+            CodeExecution.Outputs(result="hello world", log=""),
+            CodeExecution2.Outputs(result="hello world", log=""),
+        ],
+    ),
+]
+""",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.state import BaseState
+
+from .inputs import Inputs
+from .nodes.code_execution import CodeExecution
+from .nodes.code_execution_2 import CodeExecution2
+from .nodes.output import Output
+from .nodes.output_2 import Output2
+
+
+class Workflow(BaseWorkflow[Inputs, BaseState]):
+    graph = {
+        CodeExecution >> Output,
+        CodeExecution2 >> Output2,
+    }
+
+    class Outputs(BaseWorkflow.Outputs):
+        output = Output.Outputs.value
+        output_2 = Output2.Outputs.value
+""",
+    }
+
+    # AND a namespace for the virtual file loader
+    namespace = str(uuid4())
+    sys.meta_path.append(VirtualFileFinder(files, namespace))
+
+    try:
+        # WHEN we serialize the module
+        result = BaseWorkflowDisplay.serialize_module(namespace)
+
+        # THEN the serialization should return exactly 1 deprecation error
+        assert len(result.errors) == 1, f"Expected 1 deprecation error, got {len(result.errors)}: {result.errors}"
+
+        # AND the error should prompt the user to use 'mocks' instead of 'node_output_mocks'
+        error_message = str(result.errors[0])
+        assert "node_output_mocks" in error_message, f"Error should mention 'node_output_mocks': {error_message}"
+        assert "mocks" in error_message, f"Error should mention 'mocks': {error_message}"
+        assert "deprecated" in error_message.lower(), f"Error should mention 'deprecated': {error_message}"
+
+        # AND the dataset should not be None (the deprecated field should still work)
+        assert result.dataset is not None, "Dataset should not be None"
+
+        # AND the dataset should have one row
+        assert len(result.dataset) == 1, f"Dataset should have 1 row, got {len(result.dataset)}"
+
+        # AND the row should have mocks serialized
+        row = result.dataset[0]
+        assert "mocks" in row, f"Row should have 'mocks' field, got keys: {row.keys()}"
+        assert row["mocks"] is not None, "Mocks should not be None"
+        assert len(row["mocks"]) == 2, f"Should have 2 mocks, got {len(row['mocks'])}"
+    finally:
+        # Clean up the virtual file finders
+        sys.meta_path = [finder for finder in sys.meta_path if not isinstance(finder, VirtualFileFinder)]

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

@@ -0,0 +1,49 @@
+from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
+
+
+def test_serialize_module__invalid_pdf_data_url():
+    """
+    Tests that serialization returns an error when sandbox.py has an invalid PDF data URL
+    without wrecking the whole serialize process, and that valid dataset rows are still returned.
+    """
+
+    # GIVEN a workflow module with a sandbox.py that has one invalid and one valid PDF data URL
+    module = "tests.workflows.test_sandbox_invalid_pdf_data_url"
+
+    # WHEN we serialize the module
+    result = BaseWorkflowDisplay.serialize_module(module)
+
+    # THEN the result should still contain a valid exec_config
+    assert result.exec_config is not None
+    assert isinstance(result.exec_config, dict)
+
+    # AND the result should contain exactly one error about the invalid PDF data URL
+    assert len(result.errors) == 1
+    error_message = result.errors[0].message
+
+    # AND the error message should include the dataset row label
+    assert 'Dataset row "Scenario with invalid PDF"' in error_message
+
+    # AND the error message should include the input name
+    assert 'input "document"' in error_message
+
+    # AND the error message should mention the invalid base64 encoding
+    assert "Invalid base64 encoding in PDF data URL" in error_message
+
+    # AND the result should contain both dataset rows
+    assert result.dataset is not None
+    assert len(result.dataset) == 2
+
+    # AND the first row should have the valid input (name) serialized but the invalid input (document) skipped
+    invalid_row = result.dataset[0]
+    assert invalid_row["label"] == "Scenario with invalid PDF"
+    assert "document" not in invalid_row["inputs"]
+    assert "name" in invalid_row["inputs"]
+    assert invalid_row["inputs"]["name"] == "Test User"
+
+    # AND the second row (valid) should have both inputs serialized
+    valid_row = result.dataset[1]
+    assert valid_row["label"] == "Scenario with valid PDF"
+    assert "document" in valid_row["inputs"]
+    assert "name" in valid_row["inputs"]
+    assert valid_row["inputs"]["name"] == "Another User"

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

@@ -0,0 +1,112 @@
+from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
+
+
+def test_serialize_module__invalid_mocks_reference():
+    """
+    Tests that serialization returns an error when sandbox.py has an invalid Mocks reference on a node.
+    """
+
+    # GIVEN a workflow module with a sandbox.py that references CodeExecution.Mocks.Outputs
+    # which doesn't exist - nodes don't have a Mocks nested class
+    module = "tests.workflows.test_sandbox_invalid_mocks_reference"
+
+    # WHEN we serialize the module
+    result = BaseWorkflowDisplay.serialize_module(module)
+
+    # THEN the result should contain an error about the invalid attribute access
+    assert len(result.errors) > 0
+
+    # AND the error should be an AttributeError from trying to access a non-existent attribute
+    error_messages = [error.message for error in result.errors]
+    assert any(
+        "__annotations__" in msg or "has no attribute" in msg for msg in error_messages
+    ), f"Expected attribute error in error messages, got: {error_messages}"
+
+
+def test_serialize_module__invalid_runner_kwarg():
+    """
+    Tests that serialization returns an error when sandbox.py uses an invalid kwarg on WorkflowSandboxRunner.
+    """
+
+    # GIVEN a workflow module with a sandbox.py that uses 'scenarios' kwarg instead of 'dataset'
+    module = "tests.workflows.test_sandbox_invalid_runner_kwarg"
+
+    # WHEN we serialize the module
+    result = BaseWorkflowDisplay.serialize_module(module)
+
+    # THEN the result should contain an error about the invalid keyword argument
+    assert len(result.errors) > 0
+
+    # AND the error message should mention the unexpected keyword argument 'scenarios'
+    error_messages = [error.message for error in result.errors]
+    assert any(
+        "scenarios" in msg for msg in error_messages
+    ), f"Expected 'scenarios' in error messages, got: {error_messages}"
+
+
+def test_serialize_module__state_mutable_default_validation():
+    """
+    Tests that serialization returns a validation error when a state variable uses a mutable default
+    (e.g., = []) instead of Field(default_factory=list).
+    """
+
+    # GIVEN a workflow module with a state class that uses a mutable default (= [])
+    module = "tests.workflows.test_state_mutable_default_validation"
+
+    # WHEN we serialize the module
+    result = BaseWorkflowDisplay.serialize_module(module)
+
+    # THEN the result should contain an error about mutable defaults
+    assert len(result.errors) > 0
+
+    # AND the error message should mention mutable default, Field(default_factory=list), and the attribute name
+    error_messages = [error.message for error in result.errors]
+    assert any(
+        "Mutable default value detected" in msg and "Field(default_factory=list)" in msg and "State.chat_history" in msg
+        for msg in error_messages
+    ), f"Expected mutable default error in error messages, got: {error_messages}"
+
+
+def test_serialize_module__orphan_node_in_workflow_file():
+    """
+    Tests that serialization returns an error when a node is defined in workflow.py
+    but not included in the workflow's graph or unused_graphs.
+    """
+
+    # GIVEN a workflow module with a node defined in workflow.py that is not in graph or unused_graphs
+    module = "tests.workflows.test_orphan_node_serialization_error"
+
+    # WHEN we serialize the module
+    result = BaseWorkflowDisplay.serialize_module(module)
+
+    # THEN the result should contain an error about the orphan node
+    assert len(result.errors) > 0
+
+    # AND the error message should mention the orphan node and that it's not in graph or unused_graphs
+    error_messages = [error.message for error in result.errors]
+    assert any(
+        "OrphanNode" in msg and "not included in" in msg and "graph or unused_graphs" in msg for msg in error_messages
+    ), f"Expected orphan node error in error messages, got: {error_messages}"
+
+
+def test_serialize_module__orphan_node_in_nodes_directory():
+    """
+    Tests that serialization returns an error when a node is defined in the nodes/ directory
+    but not included in the workflow's graph or unused_graphs.
+    """
+
+    # GIVEN a workflow module with a node defined in nodes/orphan_node.py that is not in graph or unused_graphs
+    module = "tests.workflows.test_orphan_node_in_nodes_dir"
+
+    # WHEN we serialize the module
+    result = BaseWorkflowDisplay.serialize_module(module)
+
+    # THEN the result should contain an error about the orphan node
+    assert len(result.errors) > 0
+
+    # AND the error message should mention the orphan node and that it's not in graph or unused_graphs
+    error_messages = [error.message for error in result.errors]
+    assert any(
+        "OrphanNodeInNodesDir" in msg and "not included in" in msg and "graph or unused_graphs" in msg
+        for msg in error_messages
+    ), f"Expected orphan node error in error messages, got: {error_messages}"

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

@@ -244,24 +244,33 @@ def test_scheduled_trigger_serialization_full():
     # THEN we get the expected trigger
     assert len(result["triggers"]) == 1
     trigger = result["triggers"][0]
+
+    # AND the trigger has the expected structure with attributes
+    assert trigger["id"] == "f3e5eddb-75da-42e6-9abf-d616f30c145c"
+    assert trigger["type"] == "SCHEDULED"
+    assert trigger["cron"] == "0 9 * * *"
+    assert trigger["timezone"] == "UTC"
+
+    # AND attributes are serialized (current_run_at and next_run_at from ScheduleTrigger)
+    assert "attributes" in trigger
+    attributes = trigger["attributes"]
+    assert isinstance(attributes, list)
+    assert len(attributes) == 2
+    attribute_keys = {attr["key"] for attr in attributes}
+    assert attribute_keys == {"current_run_at", "next_run_at"}
+
+    # AND display_data is serialized correctly
     assert not DeepDiff(
-        trigger,
+        trigger["display_data"],
         {
-            "id": "f3e5eddb-75da-42e6-9abf-d616f30c145c",
-            "type": "SCHEDULED",
-            "cron": "0 9 * * *",
-            "timezone": "UTC",
-            "attributes": [],
-            "display_data": {
-                "label": "Daily Schedule",
-                "position": {"x": 100.5, "y": 200.75},
-                "z_index": 3,
-                "icon": "vellum:icon:calendar",
-                "color": "#4A90E2",
-                "comment": {
-                    "value": "This is scheduled trigger",
-                    "expanded": True,
-                },
+            "label": "Daily Schedule",
+            "position": {"x": 100.5, "y": 200.75},
+            "z_index": 3,
+            "icon": "vellum:icon:calendar",
+            "color": "#4A90E2",
+            "comment": {
+                "value": "This is scheduled trigger",
+                "expanded": True,
             },
         },
     )

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

@@ -0,0 +1,53 @@
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
+from vellum.workflows.state.base import BaseState
+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):
+    input_value: str
+
+
+class UnusedTerminalNode(FinalOutputNode):
+    class Outputs(FinalOutputNode.Outputs):
+        value = Inputs.input_value
+
+
+class TerminalNodeInUnusedGraphsWorkflow(BaseWorkflow[Inputs, BaseState]):
+    unused_graphs = {UnusedTerminalNode}
+
+    class Outputs(BaseWorkflow.Outputs):
+        result = UnusedTerminalNode.Outputs.value
+
+
+def test_serialize_workflow__terminal_node_in_unused_graphs():
+    """
+    Tests that a workflow with a single terminal node in unused_graphs serializes correctly.
+    """
+
+    # GIVEN a Workflow with a single terminal node in unused_graphs
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=TerminalNodeInUnusedGraphsWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a serialized representation of the Workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "state_variables",
+        "output_variables",
+    }
+
+    # AND its output variables should contain the terminal node's output
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 1
+    assert output_variables[0]["key"] == "result"
+    assert output_variables[0]["type"] == "STRING"
+
+    # AND its output values should reference the terminal node's output
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    output_values = workflow_raw_data["output_values"]
+    assert len(output_values) == 1
+    assert output_values[0]["output_variable_id"] == output_variables[0]["id"]
+    assert output_values[0]["value"]["type"] == "NODE_OUTPUT"

vellum_ee/workflows/display/utils/exceptions.py

@@ -24,3 +24,37 @@ class InvalidInputReferenceError(UserFacingException):
         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}")
+
+
+class InvalidOutputReferenceError(UserFacingException):
+    """Exception raised when a node references a non-existent output."""
+
+    pass
+
+
+class WorkflowValidationError(UserFacingException):
+    """Exception raised when a workflow fails validation during display serialization."""
+
+    def __init__(self, message: str, workflow_class_name: str):
+        self.message = message
+        self.workflow_class_name = workflow_class_name
+        super().__init__(f"Workflow validation error in {workflow_class_name}: {message}")
+
+
+class TriggerValidationError(UserFacingException):
+    """Exception raised when a trigger fails validation during serialization."""
+
+    def __init__(self, message: str, trigger_class_name: str):
+        self.message = message
+        self.trigger_class_name = trigger_class_name
+        super().__init__(f"Trigger validation error in {trigger_class_name}: {message}")
+
+
+class StateValidationError(UserFacingException):
+    """Exception raised when a state class fails validation during serialization."""
+
+    def __init__(self, message: str, state_class_name: str, attribute_name: str):
+        self.message = message
+        self.state_class_name = state_class_name
+        self.attribute_name = attribute_name
+        super().__init__(f"State validation error in {state_class_name}.{attribute_name}: {message}")