vellum-ai 1.7.3__py3-none-any.whl → 1.7.5__py3-none-any.whl

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

@@ -0,0 +1,113 @@
+"""Tests for serialization of workflows with ManualTrigger."""
+
+import pytest
+from typing import cast
+
+from deepdiff import DeepDiff
+
+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.base import BaseTrigger
+from vellum.workflows.triggers.manual import ManualTrigger
+from vellum.workflows.types.core import JsonArray, JsonObject
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+class Inputs(BaseInputs):
+    input: str
+
+
+class SimpleNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+
+def create_workflow(trigger=None):
+    """Factory for creating test workflows."""
+
+    class TestWorkflow(BaseWorkflow[Inputs, BaseState]):
+        graph = trigger >> SimpleNode if trigger else SimpleNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            output = SimpleNode.Outputs.output
+
+    return TestWorkflow
+
+
+def serialize(workflow_class) -> JsonObject:
+    """Helper to serialize a workflow."""
+    return get_workflow_display(workflow_class=workflow_class).serialize()
+
+
+def test_manual_trigger_serialization():
+    """Workflow with ManualTrigger serializes with trigger field."""
+    result = serialize(create_workflow(ManualTrigger))
+    workflow_raw_data = cast(JsonObject, result["workflow_raw_data"])
+    trigger = cast(JsonObject, workflow_raw_data["trigger"])
+
+    assert trigger["type"] == "MANUAL"
+    assert not DeepDiff(
+        {
+            "type": "MANUAL",
+            "definition": {"name": "ManualTrigger", "module": ["vellum", "workflows", "triggers", "manual"]},
+        },
+        trigger,
+        ignore_order=True,
+    )
+
+
+def test_no_trigger_serialization():
+    """Workflow without trigger has no trigger field."""
+    result = serialize(create_workflow())
+    workflow_raw_data = cast(JsonObject, result["workflow_raw_data"])
+    assert "trigger" not in workflow_raw_data
+
+
+def test_manual_trigger_multiple_entrypoints():
+    """ManualTrigger with multiple entrypoints."""
+
+    class NodeA(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            output = Inputs.input
+
+    class NodeB(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            output = Inputs.input
+
+    class MultiWorkflow(BaseWorkflow[Inputs, BaseState]):
+        graph = ManualTrigger >> {NodeA, NodeB}
+
+        class Outputs(BaseWorkflow.Outputs):
+            output_a = NodeA.Outputs.output
+            output_b = NodeB.Outputs.output
+
+    result = serialize(MultiWorkflow)
+    workflow_data = cast(JsonObject, result["workflow_raw_data"])
+    trigger = cast(JsonObject, workflow_data["trigger"])
+    nodes = cast(JsonArray, workflow_data["nodes"])
+
+    assert trigger["type"] == "MANUAL"
+    assert len([n for n in nodes if cast(JsonObject, n)["type"] == "GENERIC"]) >= 2
+
+
+def test_serialized_workflow_structure():
+    """Verify complete structure of serialized workflow."""
+    result = serialize(create_workflow(ManualTrigger))
+    workflow_raw_data = cast(JsonObject, result["workflow_raw_data"])
+    definition = cast(JsonObject, workflow_raw_data["definition"])
+
+    assert result.keys() == {"workflow_raw_data", "input_variables", "state_variables", "output_variables"}
+    assert workflow_raw_data.keys() == {"nodes", "edges", "display_data", "definition", "output_values", "trigger"}
+    assert definition["name"] == "TestWorkflow"
+
+
+def test_unknown_trigger_type():
+    """Unknown trigger types raise ValueError."""
+
+    class UnknownTrigger(BaseTrigger):
+        pass
+
+    with pytest.raises(ValueError, match="Unknown trigger type: UnknownTrigger"):
+        serialize(create_workflow(UnknownTrigger))

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -37,6 +37,7 @@ from vellum_ee.workflows.display.base import (
     WorkflowInputsDisplay,
     WorkflowMetaDisplay,
     WorkflowOutputDisplay,
+    get_trigger_type_mapping,
 )
 from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
@@ -408,22 +409,74 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             except Exception as e:
                 self.display_context.add_error(e)
 
-        return {
-            "workflow_raw_data": {
-                "nodes": cast(JsonArray, nodes_dict_list),
-                "edges": edges,
-                "display_data": self.display_context.workflow_display.display_data.dict(),
-                "definition": {
-                    "name": self._workflow.__name__,
-                    "module": cast(JsonArray, self._workflow.__module__.split(".")),
-                },
-                "output_values": output_values,
+        # Serialize workflow-level trigger if present
+        trigger_data: Optional[JsonObject] = self._serialize_workflow_trigger()
+
+        workflow_raw_data: JsonObject = {
+            "nodes": cast(JsonArray, nodes_dict_list),
+            "edges": edges,
+            "display_data": self.display_context.workflow_display.display_data.dict(),
+            "definition": {
+                "name": self._workflow.__name__,
+                "module": cast(JsonArray, self._workflow.__module__.split(".")),
             },
+            "output_values": output_values,
+        }
+
+        if trigger_data is not None:
+            workflow_raw_data["trigger"] = trigger_data
+
+        return {
+            "workflow_raw_data": workflow_raw_data,
             "input_variables": input_variables,
             "state_variables": state_variables,
             "output_variables": output_variables,
         }
 
+    def _serialize_workflow_trigger(self) -> Optional[JsonObject]:
+        """
+        Serialize workflow-level trigger information.
+
+        Returns:
+            JsonObject with trigger data if a trigger is present, None otherwise
+        """
+        # Get all trigger edges from the workflow's subgraphs
+        trigger_edges = []
+        for subgraph in self._workflow.get_subgraphs():
+            trigger_edges.extend(list(subgraph.trigger_edges))
+
+        if not trigger_edges:
+            # No workflow-level trigger defined
+            return None
+
+        # Get the trigger class from the first edge
+        trigger_class = trigger_edges[0].trigger_class
+
+        # Validate that all trigger edges use the same trigger type
+        trigger_type_mapping = get_trigger_type_mapping()
+        for edge in trigger_edges:
+            if edge.trigger_class != trigger_class:
+                raise ValueError(
+                    f"Mixed trigger types not supported. Found {trigger_class.__name__} and "
+                    f"{edge.trigger_class.__name__} in the same workflow."
+                )
+
+        # Get the trigger type from the mapping
+        trigger_type = trigger_type_mapping.get(trigger_class)
+        if trigger_type is None:
+            raise ValueError(
+                f"Unknown trigger type: {trigger_class.__name__}. "
+                f"Please add it to the trigger type mapping in get_trigger_type_mapping()."
+            )
+
+        return {
+            "type": trigger_type.value,
+            "definition": {
+                "name": trigger_class.__name__,
+                "module": cast(JsonArray, trigger_class.__module__.split(".")),
+            },
+        }
+
     def _serialize_edge_display_data(self, edge_display: EdgeDisplay) -> Optional[JsonObject]:
         """Serialize edge display data, returning None if no display data is present."""
         if edge_display.z_index is not None: