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

vellum_cli/tests/test_push.py

@@ -2,18 +2,21 @@ import pytest
 import io
 import json
 import os
+import pathlib
 import tarfile
 from unittest import mock
 from uuid import uuid4
 
 from click.testing import CliRunner
 from httpx import Response
+from pydantic import ValidationError
 
 from vellum.client.core.api_error import ApiError
 from vellum.client.types.workflow_push_response import WorkflowPushResponse
 from vellum.utils.uuid import is_valid_uuid
 from vellum_cli import main as cli_main
 from vellum_ee.workflows.display.nodes.utils import to_kebab_case
+from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay, WorkflowSerializationResult
 
 
 def _extract_tar_gz(tar_gz_bytes: bytes) -> dict[str, str]:
@@ -45,9 +48,15 @@ def _ensure_file(temp_dir: str, module: str, file_name: str, content: str) -> st
 def _ensure_workflow_py(temp_dir: str, module: str) -> str:
     workflow_py_file_content = """\
 from vellum.workflows import BaseWorkflow
+from vellum.workflows.nodes.bases.base import BaseNode
 
-class ExampleWorkflow(BaseWorkflow):
+
+class ExampleNode(BaseNode):
     pass
+
+
+class ExampleWorkflow(BaseWorkflow):
+    graph = ExampleNode
 """
     return _ensure_file(temp_dir, module, "workflow.py", workflow_py_file_content)
 
@@ -173,6 +182,61 @@ def test_push__happy_path(mock_module, vellum_client, base_command):
     assert extracted_files["workflow.py"] == workflow_py_file_content
 
 
+def test_push__forwards_dataset_to_client(mock_module, vellum_client, mocker):
+    """
+    Tests that the push command forwards the serialized dataset (scenarios)
+    to the API client when present.
+    """
+
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND serialize_module returns an exec_config and a dataset with a workflow_trigger_id
+    fake_exec_config = {
+        "workflow_raw_data": {
+            "definition": {"name": "ExampleWorkflow"},
+        }
+    }
+    fake_dataset = [
+        {
+            "label": "Scenario 1",
+            "inputs": {},
+            "workflow_trigger_id": "trigger-123",
+        }
+    ]
+    mocker.patch.object(
+        BaseWorkflowDisplay,
+        "serialize_module",
+        return_value=WorkflowSerializationResult(
+            exec_config=fake_exec_config,
+            errors=[],
+            dataset=fake_dataset,
+        ),
+    )
+
+    # AND the push API call returns successfully
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=str(uuid4()),
+    )
+
+    # WHEN calling `vellum workflows push`
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "push", module])
+
+    # THEN it should succeed
+    assert result.exit_code == 0
+
+    # AND we should have called the push API with the dataset argument
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+    assert "dataset" in call_args
+    assert json.loads(call_args["dataset"]) == fake_dataset
+
+
 @pytest.mark.usefixtures("info_log_level")
 def test_push__verify_default_url_in_raw_httpx_transport(mock_module, mock_httpx_transport):
     # GIVEN a single workflow configured
@@ -659,7 +723,9 @@ def test_push__push_fails_due_to_404_error_with_id(mock_module, vellum_client):
     assert "Traceback" not in result.output
 
 
-def test_push__push_fails_due_to_404_error_without_id(tmp_path, vellum_client, monkeypatch):
+def test_push__push_fails_due_to_404_error_without_id(
+    tmp_path: pathlib.Path, vellum_client, monkeypatch: pytest.MonkeyPatch
+):
     """
     Tests that a 404 error without a workflow sandbox ID shows a helpful error message.
     """
@@ -680,6 +746,7 @@ def test_push__push_fails_due_to_404_error_without_id(tmp_path, vellum_client, m
     pyproject_file.write_text(f'[tool.vellum]\n[[tool.vellum.workflows]]\nmodule = "{module}"\n')
 
     monkeypatch.chdir(tmp_path)
+    monkeypatch.setenv("VELLUM_API_KEY", "abcdef123456")
 
     # AND the push API call returns a 404 response
     vellum_client.workflows.push.side_effect = ApiError(
@@ -1375,3 +1442,185 @@ def test_push__workspace_option__nonexistent_workspace_should_fail(mock_module):
     assert "Available workspaces:" in result.output
     assert "default" in result.output
     assert "my_custom_workspace" in result.output
+
+
+def test_push__validation_error_during_serialization(mock_module, mocker):
+    """
+    Tests that a ValidationError raised during workflow serialization is handled gracefully
+    with a user-friendly error message.
+    """
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND BaseWorkflowDisplay.serialize_module raises a ValidationError
+    mock_serialize_module = mocker.patch(
+        "vellum_ee.workflows.display.workflows.base_workflow_display.BaseWorkflowDisplay.serialize_module",
+        side_effect=ValidationError.from_exception_data(
+            "WorkflowConfig",
+            [
+                {
+                    "type": "missing",
+                    "loc": ("field",),
+                    "msg": "Field required",  # type: ignore[typeddict-item]
+                    "input": {"invalid": "data"},
+                }
+            ],
+        ),
+    )
+
+    # WHEN calling `vellum push`
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push", module])
+
+    # THEN it should exit with an error
+    assert result.exit_code == 1
+
+    # AND serialize_module should have been called
+    mock_serialize_module.assert_called_once()
+
+    # AND the error message should contain the validation error title
+    assert f"Validation error while trying to push {module}" in result.output
+
+    # AND the error message should contain the validation error details
+    assert "Field required" in result.output
+
+
+def test_push__includes_metadata_json_in_artifact(mock_module, vellum_client):
+    """
+    Tests that the push command includes metadata.json in the artifact.
+    """
+
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    workflow_py_file_content = _ensure_workflow_py(temp_dir, module)
+
+    # AND a metadata.json file exists in the module
+    metadata_json_content = _ensure_file(temp_dir, module, "metadata.json", '{"key": "value"}')
+
+    # AND the push API call returns successfully
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=str(uuid4()),
+    )
+
+    # WHEN calling `vellum workflows push`
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "push", module])
+
+    # THEN it should succeed
+    assert result.exit_code == 0
+
+    # AND we should have called the push API with the correct args
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+
+    # AND the artifact should contain both workflow.py and metadata.json
+    extracted_files = _extract_tar_gz(call_args["artifact"].read())
+    assert extracted_files["workflow.py"] == workflow_py_file_content
+    assert extracted_files["metadata.json"] == metadata_json_content
+
+    # AND metadata.json should NOT be in the exec_config's module_data additional_files
+    exec_config = json.loads(call_args["exec_config"])
+    module_data = exec_config.get("module_data") or {}
+    additional_files = module_data.get("additional_files") or {}
+    assert "metadata.json" not in additional_files
+
+
+def test_push__workspace_option__same_module_different_workspaces_in_lockfile_uses_correct_sandbox_id(
+    mock_module, vellum_client_class
+):
+    """
+    Tests that when pushing with --workspace flag, the correct workflow_sandbox_id
+    is used based on the workspace when configs are in the lockfile.
+    """
+
+    # GIVEN a module configured for two different workspaces with the same sandbox ID in the lockfile
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    workflow_sandbox_id = str(uuid4())
+
+    # AND the lockfile has the same module and sandbox ID configured for two workspaces
+    with open(os.path.join(temp_dir, "vellum.lock.json"), "w") as f:
+        json.dump(
+            {
+                "version": "1.0",
+                "workflows": [
+                    {
+                        "module": module,
+                        "workflow_sandbox_id": workflow_sandbox_id,
+                        "workspace": "default",
+                        "container_image_name": None,
+                        "container_image_tag": None,
+                        "deployments": [],
+                        "ignore": None,
+                        "target_directory": None,
+                    },
+                    {
+                        "module": module,
+                        "workflow_sandbox_id": workflow_sandbox_id,
+                        "workspace": "other",
+                        "container_image_name": None,
+                        "container_image_tag": None,
+                        "deployments": [],
+                        "ignore": None,
+                        "target_directory": None,
+                    },
+                ],
+                "workspaces": [],
+            },
+            f,
+            indent=2,
+        )
+
+    # AND the pyproject.toml has the first workflow config and workspace config
+    mock_module.set_pyproject_toml(
+        {
+            "workflows": [
+                {
+                    "module": module,
+                    "workflow_sandbox_id": workflow_sandbox_id,
+                }
+            ],
+            "workspaces": [
+                {
+                    "name": "other",
+                    "api_key": "OTHER_VELLUM_API_KEY",
+                }
+            ],
+        }
+    )
+
+    # AND the .env file has both api keys stored
+    with open(os.path.join(temp_dir, ".env"), "w") as f:
+        f.write(
+            """\
+VELLUM_API_KEY=default_api_key_123
+OTHER_VELLUM_API_KEY=other_api_key_456
+"""
+        )
+
+    # AND a workflow exists in the module
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND the push API call returns successfully
+    vellum_client_class.return_value.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=workflow_sandbox_id,
+    )
+
+    # WHEN calling `vellum push` with --workspace other
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push", module, "--workspace", "other"])
+
+    # THEN it should succeed
+    assert result.exit_code == 0, result.output
+
+    # AND we should have called the push API with the correct workflow_sandbox_id
+    vellum_client_class.return_value.workflows.push.assert_called_once()
+    call_args = vellum_client_class.return_value.workflows.push.call_args.kwargs
+    assert call_args["workflow_sandbox_id"] == workflow_sandbox_id

vellum_ee/assets/node-definitions.json

@@ -53,25 +53,40 @@
           "id": "fd55dc40-4e55-4327-8a89-d4123bb775df",
           "name": "json",
           "type": "JSON",
-          "value": null
+          "value": null,
+          "schema": {
+            "$ref": "#/$defs/vellum.workflows.types.core.Json"
+          }
         },
         {
           "id": "357a9805-da18-400c-9ce5-152be1dd2ff8",
           "name": "headers",
           "type": "JSON",
-          "value": null
+          "value": null,
+          "schema": {
+            "type": "object",
+            "additionalProperties": {
+              "type": "string"
+            }
+          }
         },
         {
           "id": "4e8c472c-5f32-4d41-9746-4085b3e8b6ef",
           "name": "status_code",
           "type": "NUMBER",
-          "value": null
+          "value": null,
+          "schema": {
+            "type": "integer"
+          }
         },
         {
           "id": "90d2ccc0-40da-40c5-895d-adf7d2de28d3",
           "name": "text",
           "type": "STRING",
-          "value": null
+          "value": null,
+          "schema": {
+            "type": "string"
+          }
         }
       ],
       "attributes": [
@@ -133,6 +148,26 @@
           "name": "default",
           "type": "DEFAULT"
         }
+      ],
+      "outputs": [
+        {
+          "id": "cd6cf124-b178-4852-b5d5-251f5835fc8f",
+          "name": "result",
+          "type": "STRING",
+          "value": null,
+          "schema": {
+            "type": "string"
+          }
+        },
+        {
+          "id": "c86c2620-9192-412c-9423-4f594411fd47",
+          "name": "log",
+          "type": "STRING",
+          "value": null,
+          "schema": {
+            "type": "string"
+          }
+        }
       ]
     },
     {
@@ -435,6 +470,13 @@
     },
     {
       "id": "62ba7250-8508-41f1-8f14-709902c0e734",
+      "attributes": [
+        {
+          "id": "eb7d528f-3dae-4eac-91ff-957766be58c7",
+          "name": "items",
+          "value": null
+        }
+      ],
       "display_data": {
         "position": {
           "x": 0.0,
@@ -617,7 +659,8 @@
           "name": "default",
           "type": "DEFAULT"
         }
-      ]
+      ],
+      "outputs": []
     },
     {
       "id": "e1c8d236-6131-4218-aed8-45f986723227",
@@ -798,7 +841,11 @@
           "id": "57687d69-eb14-48fe-b919-abb044c85c35",
           "name": "results",
           "type": "JSON",
-          "value": null
+          "value": null,
+          "schema": {
+            "type": "object",
+            "additionalProperties": {}
+          }
         }
       ]
     },
@@ -986,17 +1033,36 @@
         }
       ],
       "outputs": [
+        {
+          "id": "1e346d06-a0d9-45d4-8d74-7df1f2c6d1dc",
+          "name": "json",
+          "type": "JSON",
+          "value": null,
+          "schema": {
+            "type": "object",
+            "additionalProperties": {}
+          }
+        },
         {
           "id": "bd83c7f1-fc8d-4441-abcf-95641aa03fcb",
           "name": "text",
           "type": "STRING",
-          "value": null
+          "value": null,
+          "schema": {
+            "type": "string"
+          }
         },
         {
           "id": "bb08b1a9-9745-475a-b030-6bb5b2146984",
           "name": "chat_history",
           "type": "CHAT_HISTORY",
-          "value": null
+          "value": null,
+          "schema": {
+            "type": "array",
+            "items": {
+              "$ref": "#/$defs/vellum.client.types.chat_message.ChatMessage"
+            }
+          }
         }
       ]
     },
@@ -1065,19 +1131,35 @@
           "id": "71bb6f36-d0cc-4c01-8dc2-1a0b98429c81",
           "name": "text",
           "type": "STRING",
-          "value": null
+          "value": null,
+          "schema": {
+            "type": "string"
+          }
         },
         {
           "id": "9df3ddf0-675f-48d5-b4b2-e3a805bbf450",
           "name": "urls",
           "type": "JSON",
-          "value": null
+          "value": null,
+          "schema": {
+            "type": "array",
+            "items": {
+              "type": "string"
+            }
+          }
         },
         {
           "id": "87994bf1-d8d7-40d0-a715-90b1e399c5f2",
           "name": "results",
           "type": "JSON",
-          "value": null
+          "value": null,
+          "schema": {
+            "type": "array",
+            "items": {
+              "type": "object",
+              "additionalProperties": {}
+            }
+          }
         }
       ]
     },
@@ -1126,10 +1208,141 @@
           "id": "c65369bf-7dbc-4f66-87c5-381c965ffb5d",
           "name": "value",
           "type": "STRING",
-          "value": null
+          "value": null,
+          "schema": {
+            "type": "string"
+          }
         }
       ]
     }
   ],
+  "triggers": [
+    {
+      "type": "CHAT_MESSAGE",
+      "name": "ChatMessageTrigger",
+      "module": [
+        "vellum",
+        "workflows",
+        "triggers",
+        "chat_message"
+      ],
+      "attributes": [
+        {
+          "id": "88a07223-94d0-411c-acb4-37a2b2ab108d",
+          "key": "message",
+          "type": "ARRAY",
+          "required": true,
+          "default": {
+            "type": "ARRAY",
+            "value": null
+          },
+          "extensions": null,
+          "schema": {
+            "anyOf": [
+              {
+                "type": "string"
+              },
+              {
+                "type": "array",
+                "items": {
+                  "$ref": "#/$defs/vellum.client.types.array_chat_message_content_item.ArrayChatMessageContentItem"
+                }
+              }
+            ]
+          }
+        }
+      ],
+      "display_data": {
+        "label": "Chat Message",
+        "position": {
+          "x": 0.0,
+          "y": 0.0
+        },
+        "z_index": 0,
+        "icon": "vellum:icon:message-dots",
+        "color": "blue"
+      }
+    },
+    {
+      "type": "INTEGRATION",
+      "name": "IntegrationTrigger",
+      "module": [
+        "vellum",
+        "workflows",
+        "triggers",
+        "integration"
+      ],
+      "attributes": [],
+      "display_data": {
+        "label": "Trigger",
+        "position": {
+          "x": 0.0,
+          "y": 0.0
+        },
+        "z_index": 0
+      }
+    },
+    {
+      "type": "MANUAL",
+      "name": "ManualTrigger",
+      "module": [
+        "vellum",
+        "workflows",
+        "triggers",
+        "manual"
+      ],
+      "attributes": []
+    },
+    {
+      "type": "SCHEDULED",
+      "name": "ScheduleTrigger",
+      "module": [
+        "vellum",
+        "workflows",
+        "triggers",
+        "schedule"
+      ],
+      "attributes": [
+        {
+          "id": "77f5f59d-d9b8-4a9c-b223-e3f5d9f4e376",
+          "key": "current_run_at",
+          "type": "STRING",
+          "required": true,
+          "default": {
+            "type": "STRING",
+            "value": null
+          },
+          "extensions": null,
+          "schema": {
+            "type": "string",
+            "format": "date-time"
+          }
+        },
+        {
+          "id": "e75ec8ba-2635-4be2-9c1d-f04ddcb8a7a1",
+          "key": "next_run_at",
+          "type": "STRING",
+          "required": true,
+          "default": {
+            "type": "STRING",
+            "value": null
+          },
+          "extensions": null,
+          "schema": {
+            "type": "string",
+            "format": "date-time"
+          }
+        }
+      ],
+      "display_data": {
+        "label": "Trigger",
+        "position": {
+          "x": 0.0,
+          "y": 0.0
+        },
+        "z_index": 0
+      }
+    }
+  ],
   "errors": []
 }

vellum_ee/scripts/generate_node_definitions.py

@@ -8,6 +8,7 @@ import vellum.workflows.nodes.displayable as displayable_module
 from vellum.workflows.vellum_client import create_vellum_client
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.triggers import get_all_trigger_classes, serialize_trigger_definition
 
 logger = logging.getLogger(__name__)
 
@@ -49,7 +50,7 @@ def serialize_node_definition(
 
 
 def main() -> None:
-    """Main function to generate node definitions."""
+    """Main function to generate node and trigger definitions."""
     logger.info("Generating node definitions...")
 
     display_context = create_display_context_with_client()
@@ -72,7 +73,17 @@ def main() -> None:
             except Exception as e:
                 errors.append({"node": node_class.__name__, "error": f"{e.__class__.__name__}: {str(e)}"})
 
-    result = {"nodes": successful_nodes, "errors": errors}
+    # Generate trigger definitions
+    logger.info("Generating trigger definitions...")
+    trigger_classes = get_all_trigger_classes()
+    triggers = []
+
+    for trigger_class in trigger_classes:
+        logger.info(f"Serializing {trigger_class.__name__}...")
+        trigger_definition = serialize_trigger_definition(trigger_class)
+        triggers.append(trigger_definition)
+
+    result = {"nodes": successful_nodes, "triggers": triggers, "errors": errors}
 
     output_path = "ee/vellum_ee/assets/node-definitions.json"
     os.makedirs(os.path.dirname(output_path), exist_ok=True)
@@ -81,7 +92,8 @@ def main() -> None:
         json.dump(result, f, indent=2)
 
     logger.info(
-        f"Generated {len(successful_nodes)} successful node definitions and {len(errors)} errors in {output_path}"
+        f"Generated {len(successful_nodes)} node definitions, {len(triggers)} trigger definitions, "
+        f"and {len(errors)} errors in {output_path}"
     )
 
 

vellum_ee/workflows/display/base.py

@@ -18,6 +18,7 @@ class WorkflowTriggerType(Enum):
     MANUAL = "MANUAL"
     INTEGRATION = "INTEGRATION"
     SCHEDULED = "SCHEDULED"
+    CHAT_MESSAGE = "CHAT_MESSAGE"
 
 
 def get_trigger_type_mapping() -> Dict[Type["BaseTrigger"], WorkflowTriggerType]:
@@ -49,9 +50,9 @@ class WorkflowDisplayData(UniversalBaseModel):
 
 @dataclass
 class WorkflowMetaDisplay:
-    entrypoint_node_id: UUID
-    entrypoint_node_source_handle_id: UUID
-    entrypoint_node_display: NodeDisplayData = Field(default_factory=NodeDisplayData)
+    entrypoint_node_id: Optional[UUID] = None
+    entrypoint_node_source_handle_id: Optional[UUID] = None
+    entrypoint_node_display: Optional[NodeDisplayData] = None
     display_data: WorkflowDisplayData = field(default_factory=WorkflowDisplayData)
 
     @classmethod