[APO-2122] Serialize workflow trigger id for scenarios (#3069)

* [APO-2122] Serialize workflow trigger id for scenarios

* udpate tests

---------

Co-authored-by: Harrison Ngo <[email protected]>

Author: NgoHarrison

ee/vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -1211,7 +1211,11 @@ def serialize_module(
                     for i, inputs_obj in enumerate(dataset_attr):
                         if isinstance(inputs_obj, DatasetRow):
                             serialized_inputs = json.loads(json.dumps(inputs_obj.inputs, cls=VellumJsonEncoder))
-                            dataset.append({"label": inputs_obj.label, "inputs": serialized_inputs})
+                            row_data = {"label": inputs_obj.label, "inputs": serialized_inputs}
+                            trigger_class = inputs_obj.workflow_trigger
+                            if trigger_class is not None:
+                                row_data["workflow_trigger_id"] = str(trigger_class.__id__)
+                            dataset.append(row_data)
                         elif isinstance(inputs_obj, BaseInputs):
                             serialized_inputs = json.loads(json.dumps(inputs_obj, cls=VellumJsonEncoder))
                             dataset.append({"label": f"Scenario {i + 1}", "inputs": serialized_inputs})

ee/vellum_ee/workflows/tests/test_serialize_module.py

@@ -3,6 +3,10 @@
 import shutil
 import sys
 import tempfile
+import uuid
+from uuid import UUID
+
+from pytest_mock import MockerFixture
 
 from vellum.workflows.exceptions import WorkflowInitializationException
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
@@ -20,6 +24,18 @@ def temp_module_path():
         shutil.rmtree(temp_dir)
 
 
+@pytest.fixture
+def metadata_trigger_factory(mocker: MockerFixture):
+    def _factory(metadata_trigger_id: UUID) -> UUID:
+        mocker.patch(
+            "vellum.workflows.triggers.base._get_trigger_id_from_metadata",
+            return_value=metadata_trigger_id,
+        )
+        return metadata_trigger_id
+
+    return _factory
+
+
 def test_serialize_module_with_dataset():
     """Test that serialize_module correctly serializes dataset from sandbox modules."""
     module_path = "tests.workflows.basic_inputs_and_outputs"
@@ -50,6 +66,25 @@ def test_serialize_module_with_actual_dataset():
     assert result.dataset[1]["inputs"]["message"] == "DatasetRow Test"
 
 
+def test_serialize_module_with_actual_dataset_with_trigger(metadata_trigger_factory):
+    """Test that serialize_module correctly serializes dataset with trigger"""
+    module_path = "tests.workflows.test_dataset_with_trigger_serialization"
+
+    metadata_trigger_id = uuid.uuid4()
+    metadata_trigger_factory(metadata_trigger_id)
+
+    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) == 1
+
+    assert result.dataset[0]["label"] == "Scenario 1"
+    assert result.dataset[0]["workflow_trigger_id"] == str(metadata_trigger_id)
+
+
 def test_serialize_module_happy_path():
     """Test that serialize_module works with a valid module path."""
     module_path = "tests.workflows.trivial"

src/vellum/workflows/inputs/dataset_row.py

@@ -1,9 +1,10 @@
-from typing import Any, Dict, Union
+from typing import Any, Dict, Optional, Type, Union
 
 from pydantic import Field, field_serializer
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.triggers import BaseTrigger
 
 
 class DatasetRow(UniversalBaseModel):
@@ -13,10 +14,12 @@ class DatasetRow(UniversalBaseModel):
     Attributes:
         label: String label for the dataset row
         inputs: BaseInputs instance or dict containing the input data
+        workflow_trigger_id: Optional Trigger identifying the workflow trigger class for this scenario
     """
 
     label: str
     inputs: Union[BaseInputs, Dict[str, Any]] = Field(default_factory=BaseInputs)
+    workflow_trigger: Optional[Type[BaseTrigger]] = None
 
     @field_serializer("inputs")
     def serialize_inputs(self, inputs: Union[BaseInputs, Dict[str, Any]]) -> Dict[str, Any]:

src/vellum/workflows/sandbox.py

@@ -1,11 +1,12 @@
-from typing import Any, Dict, Generic, Optional, Sequence, Union
+from typing import Any, Dict, Generic, Optional, Sequence, Type, Union
 
 import dotenv
 
 from vellum.workflows.events.workflow import WorkflowEventStream
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.inputs.dataset_row import DatasetRow
 from vellum.workflows.logging import load_logger
+from vellum.workflows.triggers.base import BaseTrigger
 from vellum.workflows.types.generics import WorkflowType
 from vellum.workflows.workflows.event_filters import root_workflow_event_filter
 
@@ -52,8 +53,10 @@ def run(self, index: int = 0):
         selected_inputs = self._inputs[index]
 
         raw_inputs: Union[BaseInputs, Dict[str, Any]]
+        trigger_class: Optional[Type[BaseTrigger]] = None
         if isinstance(selected_inputs, DatasetRow):
             raw_inputs = selected_inputs.inputs
+            trigger_class = selected_inputs.workflow_trigger
         else:
             raw_inputs = selected_inputs
 
@@ -64,9 +67,15 @@ def run(self, index: int = 0):
         else:
             inputs_for_stream = raw_inputs
 
+        trigger_instance: Optional[BaseTrigger] = None
+        if trigger_class is not None:
+            # Instantiate the trigger with the inputs
+            trigger_instance = trigger_class(**raw_inputs) if isinstance(raw_inputs, dict) else trigger_class()
+
         events = self._workflow.stream(
             inputs=inputs_for_stream,
             event_filter=root_workflow_event_filter,
+            trigger=trigger_instance,
         )
 
         self._process_events(events)

src/vellum/workflows/tests/test_sandbox.py

@@ -1,11 +1,13 @@
 import pytest
+from datetime import datetime
 from typing import List
 
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.inputs.dataset_row import DatasetRow
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.sandbox import WorkflowSandboxRunner
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.triggers import ScheduleTrigger
 from vellum.workflows.workflows.base import BaseWorkflow
 
 
@@ -100,3 +102,53 @@ class Outputs(BaseWorkflow.Outputs):
         "----------------------------------",
         "final_output: Hello from dict",
     ]
+
+
+def test_sandbox_runner_with_workflow_trigger(mock_logger):
+    """
+    Test that WorkflowSandboxRunner can run with DatasetRow containing workflow_trigger.
+    """
+
+    # GIVEN we capture the logs to stdout
+    logs = []
+    mock_logger.return_value.info.side_effect = lambda msg: logs.append(msg)
+
+    class MySchedule(ScheduleTrigger):
+        class Config(ScheduleTrigger.Config):
+            cron = "* * * * *"
+            timezone = "UTC"
+
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result = MySchedule.current_run_at
+
+    class Workflow(BaseWorkflow):
+        graph = MySchedule >> StartNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_output = StartNode.Outputs.result
+
+    # AND a dataset with workflow_trigger
+    dataset = [
+        DatasetRow(
+            label="test_row",
+            inputs={"current_run_at": datetime.min, "next_run_at": datetime.now()},
+            workflow_trigger=MySchedule,
+        ),
+    ]
+
+    # WHEN we run the sandbox with the DatasetRow containing workflow_trigger
+    runner = WorkflowSandboxRunner(workflow=Workflow(), dataset=dataset)
+    runner.run()
+
+    # THEN the workflow should run successfully
+    assert logs == [
+        "Just started Node: StartNode",
+        "Just finished Node: StartNode",
+        "Workflow fulfilled!",
+        "----------------------------------",
+        "final_output: 0001-01-01 00:00:00",
+    ]
+
+    # AND the dataset row should still have the trigger class
+    assert dataset[0].workflow_trigger == MySchedule

tests/workflows/test_dataset_with_trigger_serialization/__init__.py

@@ -0,0 +1,22 @@
+from datetime import datetime
+from typing import List
+
+from vellum.workflows.inputs.dataset_row import DatasetRow
+from vellum.workflows.sandbox import WorkflowSandboxRunner
+
+from .workflow import MySchedule, TestDatasetWithTriggerSerializationWorkflow
+
+if __name__ == "__main__":
+    raise Exception("This file is not meant to be imported")
+
+dataset: List[DatasetRow] = [
+    DatasetRow(
+        label="Scenario 1",
+        inputs={"current_run_at": datetime.min, "next_run_at": datetime.now()},
+        workflow_trigger=MySchedule,
+    ),
+]
+
+runner = WorkflowSandboxRunner(workflow=TestDatasetWithTriggerSerializationWorkflow(), dataset=dataset)
+
+runner.run()

tests/workflows/test_dataset_with_trigger_serialization/sandbox.py

@@ -0,0 +1,27 @@
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.triggers import ScheduleTrigger
+
+
+class MySchedule(ScheduleTrigger):
+    class Config(ScheduleTrigger.Config):
+        cron = "* * * * *"
+        timezone = "UTC"
+
+
+class SimpleNode(BaseNode):
+    message = MySchedule.current_run_at
+
+    class Outputs(BaseOutputs):
+        result: str
+
+    def run(self) -> BaseOutputs:
+        return self.Outputs(result=f"Current run at: {str(self.message)}")
+
+
+class TestDatasetWithTriggerSerializationWorkflow(BaseWorkflow):
+    graph = MySchedule >> SimpleNode
+
+    class Outputs(BaseOutputs):
+        final_result = SimpleNode.Outputs.result