vellum-ai 1.7.9__py3-none-any.whl → 1.7.11__py3-none-any.whl

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

@@ -1,148 +1,13 @@
-"""Tests for serialization of workflows with SlackTrigger."""
-
-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.slack import SlackTrigger
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
+from tests.workflows.basic_trigger_slack.workflow import SlackTriggerWorkflow
 
-class Inputs(BaseInputs):
-    input: str
-
-
-class SimpleNode(BaseNode):
-    class Outputs(BaseNode.Outputs):
-        output = Inputs.input
-
-
-def test_slack_trigger_serialization() -> None:
-    """Workflow with SlackTrigger serializes with triggers field."""
-
-    class TestWorkflow(BaseWorkflow[Inputs, BaseState]):
-        graph = SlackTrigger >> SimpleNode
-
-        class Outputs(BaseWorkflow.Outputs):
-            output = SimpleNode.Outputs.output
-
-    result = get_workflow_display(workflow_class=TestWorkflow).serialize()
-
-    # Validate triggers structure
-    assert "triggers" in result
-    triggers = result["triggers"]
-    assert isinstance(triggers, list)
-    assert len(triggers) == 1
-
-    trigger = triggers[0]
-    assert isinstance(trigger, dict)
-    assert trigger["type"] == "SLACK_MESSAGE"
-    assert "id" in trigger
-
-    # Validate attributes
-    assert "attributes" in trigger
-    attributes = trigger["attributes"]
-    assert isinstance(attributes, list)
-    assert len(attributes) == 6
-
-    attribute_names = set()
-    for attribute in attributes:
-        assert isinstance(attribute, dict)
-        assert "name" in attribute
-        assert isinstance(attribute["name"], str)
-        attribute_names.add(attribute["name"])
-    assert attribute_names == {
-        "message",
-        "channel",
-        "user",
-        "timestamp",
-        "thread_ts",
-        "event_type",
-    }
-
-    for attribute in attributes:
-        assert isinstance(attribute, dict)
-        assert attribute["value"] is None
-        assert isinstance(attribute["id"], str)
-        assert attribute["id"]
-
-
-def test_slack_trigger_multiple_entrypoints() -> None:
-    """SlackTrigger with multiple entrypoints."""
-
-    class NodeA(BaseNode):
-        class Outputs(BaseNode.Outputs):
-            output = Inputs.input
 
-    class NodeB(BaseNode):
-        class Outputs(BaseNode.Outputs):
-            output = Inputs.input
+def test_serialize_slack_trigger_workflow():
+    workflow_display = get_workflow_display(workflow_class=SlackTriggerWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
 
-    class MultiWorkflow(BaseWorkflow[Inputs, BaseState]):
-        graph = SlackTrigger >> {NodeA, NodeB}
-
-        class Outputs(BaseWorkflow.Outputs):
-            output_a = NodeA.Outputs.output
-            output_b = NodeB.Outputs.output
-
-    result = get_workflow_display(workflow_class=MultiWorkflow).serialize()
-
-    # Validate triggers
-    assert "triggers" in result
-    triggers = result["triggers"]
-    assert isinstance(triggers, list)
-    assert len(triggers) == 1
-
-    trigger = triggers[0]
-    assert isinstance(trigger, dict)
-    assert trigger["type"] == "SLACK_MESSAGE"
-
-    # Validate attributes
-    assert "attributes" in trigger
-    attributes = trigger["attributes"]
-    assert isinstance(attributes, list)
-    attribute_names = set()
-    for attribute in attributes:
-        assert isinstance(attribute, dict)
-        assert "name" in attribute
-        assert isinstance(attribute["name"], str)
-        attribute_names.add(attribute["name"])
-
-    assert attribute_names == {
-        "message",
-        "channel",
-        "user",
-        "timestamp",
-        "thread_ts",
-        "event_type",
-    }
-
-    # Validate nodes
-    assert "workflow_raw_data" in result
-    workflow_data = result["workflow_raw_data"]
-    assert isinstance(workflow_data, dict)
-    assert "nodes" in workflow_data
-    nodes = workflow_data["nodes"]
-    assert isinstance(nodes, list)
-
-    generic_nodes = [node for node in nodes if isinstance(node, dict) and node.get("type") == "GENERIC"]
-    assert len(generic_nodes) >= 2
-
-
-def test_serialized_slack_workflow_structure() -> None:
-    """Verify complete structure of serialized workflow with SlackTrigger."""
-
-    class TestWorkflow(BaseWorkflow[Inputs, BaseState]):
-        graph = SlackTrigger >> SimpleNode
-
-        class Outputs(BaseWorkflow.Outputs):
-            output = SimpleNode.Outputs.output
-
-    result = get_workflow_display(workflow_class=TestWorkflow).serialize()
-
-    # Validate top-level structure
-    assert isinstance(result, dict)
-    assert set(result.keys()) == {
+    assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
         "state_variables",
@@ -150,18 +15,38 @@ def test_serialized_slack_workflow_structure() -> None:
         "triggers",
     }
 
-    # Validate workflow_raw_data structure
-    workflow_raw_data = result["workflow_raw_data"]
-    assert isinstance(workflow_raw_data, dict)
-    assert set(workflow_raw_data.keys()) == {
-        "nodes",
-        "edges",
-        "display_data",
-        "definition",
-        "output_values",
-    }
-
-    # Validate definition
-    definition = workflow_raw_data["definition"]
-    assert isinstance(definition, dict)
-    assert definition["name"] == "TestWorkflow"
+    triggers = serialized_workflow["triggers"]
+    assert triggers == [
+        {
+            "id": "45855aa4-27a0-426b-b399-a8ff2932a684",
+            "type": "SLACK_MESSAGE",
+            "attributes": [
+                {"id": "9d4bd7d7-314d-48b8-a483-f964ac3ca28a", "name": "channel", "type": "STRING", "value": None},
+                {"id": "af4aac3c-74f2-4250-801b-f2dbd7745277", "name": "event_type", "type": "STRING", "value": None},
+                {"id": "bdf8965f-b2f1-4f83-9a5a-e1532d73c795", "name": "message", "type": "STRING", "value": None},
+                {"id": "5a910518-f875-497c-ab5f-680eecce2d1d", "name": "thread_ts", "type": "STRING", "value": None},
+                {"id": "4aadb9ec-aabf-4a58-a9bb-41e89e8a20cb", "name": "timestamp", "type": "STRING", "value": None},
+                {"id": "c16971a0-73a3-4b81-93dc-2bcaafa3585a", "name": "user", "type": "STRING", "value": None},
+            ],
+        }
+    ]
+
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    nodes = workflow_raw_data["nodes"]
+
+    process_node = next(node for node in nodes if node["type"] == "GENERIC" and node["label"] == "Process Message Node")
+    assert "outputs" in process_node
+    assert process_node["outputs"] == [
+        {
+            "id": "a1208db6-2daf-48a4-acee-71c8b1f42656",
+            "name": "processed_message",
+            "type": "STRING",
+            "value": {
+                "type": "TRIGGER_ATTRIBUTE",
+                "trigger_id": "45855aa4-27a0-426b-b399-a8ff2932a684",
+                "attribute_id": "bdf8965f-b2f1-4f83-9a5a-e1532d73c795",
+            },
+        }
+    ]
+
+    assert triggers[0]["id"] == process_node["outputs"][0]["value"]["trigger_id"]

vellum_ee/workflows/display/utils/events.py

@@ -1,7 +1,7 @@
 from typing import Optional
 
 from vellum import Vellum
-from vellum.workflows.events.workflow import WorkflowExecutionInitiatedEvent
+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,
@@ -29,8 +29,8 @@ def _should_mark_workflow_dynamic(event: WorkflowExecutionInitiatedEvent) -> boo
 
 
 def event_enricher(
-    event: WorkflowExecutionInitiatedEvent, client: Optional[Vellum] = None
-) -> WorkflowExecutionInitiatedEvent:
+    event: WorkflowEvent, client: Optional[Vellum] = None, metadata: Optional[dict] = None
+) -> WorkflowEvent:
     if event.name != "workflow.execution.initiated":
         return event
 
@@ -49,4 +49,7 @@ def event_enricher(
         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
+
     return event

vellum_ee/workflows/display/utils/tests/test_events.py

@@ -112,6 +112,7 @@ def test_event_enricher_marks_subworkflow_deployment_as_dynamic(vellum_client):
 
     enriched_event = event_enricher(event, vellum_client)
 
+    assert enriched_event.name == "workflow.execution.initiated"
     assert hasattr(enriched_event.body, "workflow_version_exec_config")
     assert enriched_event.body.workflow_version_exec_config is not None
 
@@ -119,3 +120,31 @@ def test_event_enricher_marks_subworkflow_deployment_as_dynamic(vellum_client):
     assert hasattr(enriched_event.body.display_context, "node_displays")
     assert hasattr(enriched_event.body.display_context, "workflow_inputs")
     assert hasattr(enriched_event.body.display_context, "workflow_outputs")
+
+
+def test_event_enricher_with_metadata(vellum_client):
+    """Test that event_enricher attaches metadata to server_metadata field."""
+
+    # GIVEN a workflow class
+    class TestWorkflow(BaseWorkflow):
+        is_dynamic = False
+
+    # AND an event
+    event: WorkflowExecutionInitiatedEvent = WorkflowExecutionInitiatedEvent(
+        trace_id=uuid4(),
+        span_id=uuid4(),
+        body=WorkflowExecutionInitiatedBody(
+            workflow_definition=TestWorkflow,
+            inputs=BaseInputs(),
+        ),
+    )
+
+    # AND some metadata
+    metadata = {"custom_key": "custom_value", "another_key": 123}
+
+    # WHEN the event_enricher is called with metadata
+    enriched_event = event_enricher(event, vellum_client, metadata=metadata)
+
+    # THEN the metadata should be attached to server_metadata
+    assert enriched_event.name == "workflow.execution.initiated"
+    assert enriched_event.body.server_metadata == metadata

vellum_ee/workflows/server/virtual_file_loader.py

@@ -4,11 +4,14 @@ import re
 import sys
 from typing import Optional
 
+from vellum.workflows.loaders.base import BaseWorkflowFinder
+
 
 class VirtualFileLoader(importlib.abc.Loader):
-    def __init__(self, files: dict[str, str], namespace: str):
+    def __init__(self, files: dict[str, str], namespace: str, source_module: Optional[str] = None):
         self.files = files
         self.namespace = namespace
+        self.source_module = source_module
 
     def create_module(self, spec: ModuleSpec):
         """
@@ -65,9 +68,17 @@ class VirtualFileLoader(importlib.abc.Loader):
         return self.files.get(file_key_name)
 
 
-class VirtualFileFinder(importlib.abc.MetaPathFinder, importlib.abc.Loader):
-    def __init__(self, files: dict[str, str], namespace: str):
-        self.loader = VirtualFileLoader(files, namespace)
+class VirtualFileFinder(BaseWorkflowFinder):
+    def __init__(self, files: dict[str, str], namespace: str, source_module: Optional[str] = None):
+        self.loader = VirtualFileLoader(files, namespace, source_module)
+        self.source_module = source_module
+        self.namespace = namespace
+
+    def format_error_message(self, error_message: str) -> str:
+        """Format error message by replacing namespace with source_module."""
+        if self.source_module and self.namespace in error_message:
+            return error_message.replace(self.namespace, self.source_module)
+        return error_message
 
     def find_spec(self, fullname: str, path, target=None):
         module_info = self.loader._resolve_module(fullname)

vellum_ee/workflows/tests/test_serialize_module.py

@@ -96,6 +96,54 @@ def test_serialize_module_includes_additional_files():
     assert "CONSTANT_VALUE" in additional_files["utils/constants.py"]
 
 
+def test_serialize_module_with_pydantic_array():
+    """
+    Test that serialize_module correctly serializes arrays of Pydantic models in workflow inputs.
+
+    This test verifies that when a workflow has inputs containing a List[PydanticModel],
+    the serialization properly converts the Pydantic models to JSON format.
+    """
+    module_path = "tests.workflows.pydantic_array_serialization"
+
+    # WHEN we serialize it
+    result = BaseWorkflowDisplay.serialize_module(module_path)
+
+    assert hasattr(result, "exec_config")
+    assert hasattr(result, "errors")
+    assert isinstance(result.exec_config, dict)
+    assert isinstance(result.errors, list)
+
+    input_variables = result.exec_config["input_variables"]
+    assert len(input_variables) == 1
+
+    items_input = input_variables[0]
+    assert items_input["key"] == "items"
+    assert items_input["type"] == "JSON"
+    # TODO: In the future, this should be a custom type based on an OpenAPI schema (important-comment)
+
+    assert result.dataset is not None
+    assert isinstance(result.dataset, list)
+    assert len(result.dataset) == 2
+
+    first_scenario = result.dataset[0]
+    assert first_scenario["label"] == "Scenario 1"
+    assert "items" in first_scenario["inputs"]
+    items = first_scenario["inputs"]["items"]
+    assert isinstance(items, list)
+    assert len(items) == 3
+    assert items[0]["name"] == "item1"
+    assert items[0]["value"] == 10
+    assert items[0]["is_active"] is True
+
+    second_scenario = result.dataset[1]
+    assert second_scenario["label"] == "Custom Test"
+    assert "items" in second_scenario["inputs"]
+    test_items = second_scenario["inputs"]["items"]
+    assert len(test_items) == 2
+    assert test_items[0]["name"] == "test1"
+    assert test_items[0]["value"] == 100
+
+
 def test_serialize_module__with_invalid_nested_set_graph(temp_module_path):
     """
     Tests that serialize_module raises a clear user-facing exception for workflows with nested sets in graph attribute.

vellum_ee/workflows/tests/test_server.py

@@ -145,6 +145,11 @@ class CodeExecutionNode(BaseCodeExecutionNode[BaseState, int]):
     # AND we get the code execution result
     assert event.body.outputs == {"final_output": 5.0}
 
+    # AND the workflow definition module is correctly serialized as a list
+    serialized_event = event.model_dump(mode="json")
+    workflow_definition = serialized_event["body"]["workflow_definition"]
+    assert workflow_definition["module"] == [namespace, "workflow"]
+
 
 def test_load_from_module__simple_code_execution_node_with_try(
     vellum_client,
@@ -579,6 +584,111 @@ class BrokenNode(BaseNode):
     assert "UndefinedClass" in error_message or "not defined" in error_message
 
 
+def test_load_from_module__module_not_found_error():
+    """
+    Tests that a ModuleNotFoundError raises WorkflowInitializationException with user-facing message.
+    """
+    # GIVEN a workflow module that imports a non-existent module
+    files = {
+        "__init__.py": "",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from .non_existent_module import SomeClass
+
+class Workflow(BaseWorkflow):
+    graph = None
+""",
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    sys.meta_path.append(VirtualFileFinder(files, namespace, source_module="test"))
+
+    # WHEN we attempt to load the workflow
+    # THEN it should raise WorkflowInitializationException
+    with pytest.raises(WorkflowInitializationException) as exc_info:
+        BaseWorkflow.load_from_module(namespace)
+
+    # AND the error message should be user-friendly and show source_module instead of namespace
+    error_message = str(exc_info.value)
+    assert error_message == "Workflow module not found: No module named 'test.non_existent_module'"
+
+
+def test_load_from_module__module_not_found_error_with_external_package():
+    """
+    Tests that when ModuleNotFoundError occurs for an external package (not containing the namespace),
+    the exception includes vellum_on_error_action set to CREATE_CUSTOM_IMAGE in raw_data.
+    """
+
+    # GIVEN a workflow module that imports a non-existent external package
+    files = {
+        "__init__.py": "",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+import some_external_package
+
+class Workflow(BaseWorkflow):
+    graph = None
+""",
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    finder = VirtualFileFinder(files, namespace, source_module="test")
+    sys.meta_path.append(finder)
+
+    # WHEN we attempt to load the workflow
+    # THEN it should raise WorkflowInitializationException
+    with pytest.raises(WorkflowInitializationException) as exc_info:
+        BaseWorkflow.load_from_module(namespace)
+
+    # AND the error message should be user-friendly
+    error_message = str(exc_info.value)
+    assert "Workflow module not found:" in error_message
+    assert "some_external_package" in error_message
+
+    assert exc_info.value.raw_data is not None
+    assert exc_info.value.raw_data["vellum_on_error_action"] == "CREATE_CUSTOM_IMAGE"
+
+
+def test_load_from_module__module_not_found_error_with_internal_package():
+    """
+    Tests that when ModuleNotFoundError occurs for an internal module (containing the namespace),
+    the exception does NOT include vellum_on_error_action in raw_data.
+    """
+
+    # GIVEN a workflow module that imports a non-existent internal module
+    files = {
+        "__init__.py": "",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from .non_existent_module import SomeClass
+
+class Workflow(BaseWorkflow):
+    graph = None
+""",
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    finder = VirtualFileFinder(files, namespace, source_module="test")
+    sys.meta_path.append(finder)
+
+    # WHEN we attempt to load the workflow
+    # THEN it should raise WorkflowInitializationException
+    with pytest.raises(WorkflowInitializationException) as exc_info:
+        BaseWorkflow.load_from_module(namespace)
+
+    # AND the error message should be user-friendly
+    error_message = str(exc_info.value)
+    assert "Workflow module not found:" in error_message
+
+    assert exc_info.value.raw_data is None
+
+
 def test_serialize_module__tool_calling_node_with_single_tool():
     """Test that serialize_module works with a tool calling node that has a single tool."""