vellum-ai 1.2.4__py3-none-any.whl → 1.3.0__py3-none-any.whl

vellum_cli/tests/test_pull.py

@@ -1,34 +1,17 @@
 import pytest
-import io
 import json
 import os
 import tempfile
 from unittest import mock
 from uuid import uuid4
-import zipfile
 
 from click.testing import CliRunner
 
 from vellum.client.core.api_error import ApiError
+from vellum.workflows.utils.zip import zip_file_map
 from vellum_cli import main as cli_main
 
 
-def _zip_file_map(file_map: dict[str, str]) -> bytes:
-    # Create an in-memory bytes buffer to store the zip
-    zip_buffer = io.BytesIO()
-
-    # Create zip file and add files from file_map
-    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
-        for filename, content in file_map.items():
-            zip_file.writestr(filename, content)
-
-    # Get the bytes from the buffer
-    zip_bytes = zip_buffer.getvalue()
-    zip_buffer.close()
-
-    return zip_bytes
-
-
 @pytest.mark.parametrize(
     "base_command",
     [
@@ -44,7 +27,7 @@ def test_pull(vellum_client, mock_module, base_command):
     workflow_sandbox_id = mock_module.workflow_sandbox_id
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
 
     # WHEN the user runs the pull command
     runner = CliRunner()
@@ -89,7 +72,7 @@ def test_pull__second_module(vellum_client, mock_module):
     set_pyproject_toml = mock_module.set_pyproject_toml
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
 
     # AND the module we're about to pull is configured second
     set_pyproject_toml(
@@ -134,7 +117,7 @@ def test_pull__with_target_dir(vellum_client, mock_module, base_command):
     os.makedirs(target_dir, exist_ok=True)
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
 
     # WHEN the user runs the pull command with target-dir
     runner = CliRunner()
@@ -195,7 +178,7 @@ def test_pull__with_nested_target_dir(vellum_client, mock_module, base_command):
     nested_target_dir = os.path.join(temp_dir, "dir-1", "dir-2")
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
 
     # WHEN the user runs the pull command with nested target-dir
     runner = CliRunner()
@@ -250,11 +233,7 @@ def test_pull__sandbox_id_with_no_config(vellum_client):
 
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter(
-        [
-            _zip_file_map(
-                {"workflow.py": "print('hello')", "metadata.json": json.dumps({"label": "Super Cool Workflow"})}
-            )
-        ]
+        [zip_file_map({"workflow.py": "print('hello')", "metadata.json": json.dumps({"label": "Super Cool Workflow"})})]
     )
 
     # AND we are currently in a new directory
@@ -309,11 +288,7 @@ def test_pull__sandbox_id_with_other_workflow_configured(vellum_client, mock_mod
 
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter(
-        [
-            _zip_file_map(
-                {"workflow.py": "print('hello')", "metadata.json": json.dumps({"label": "Super Cool Workflow"})}
-            )
-        ]
+        [zip_file_map({"workflow.py": "print('hello')", "metadata.json": json.dumps({"label": "Super Cool Workflow"})})]
     )
 
     # WHEN the user runs the pull command with the new workflow sandbox id
@@ -343,7 +318,7 @@ def test_pull__workflow_deployment_with_no_config(vellum_client):
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter(
         [
-            _zip_file_map(
+            zip_file_map(
                 {
                     "workflow.py": "print('hello')",
                     "metadata.json": json.dumps(
@@ -414,7 +389,7 @@ def test_pull__both_workflow_sandbox_id_and_deployment(vellum_client):
     workflow_deployment = "my-deployment"
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
 
     # AND we are currently in a new directory
     current_dir = os.getcwd()
@@ -447,7 +422,7 @@ def test_pull__remove_missing_files(vellum_client, mock_module):
     module = mock_module.module
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
 
     # AND there is already a different file in the module directory
     other_file_path = os.path.join(temp_dir, *module.split("."), "other_file.py")
@@ -478,7 +453,7 @@ def test_pull__remove_missing_files__ignore_pattern(vellum_client, mock_module):
     set_pyproject_toml = mock_module.set_pyproject_toml
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
 
     # AND there is already a different file in the module directory
     other_file_path = os.path.join(temp_dir, *module.split("."), "other_file.py")
@@ -530,7 +505,7 @@ def test_pull__include_json(vellum_client, mock_module):
 
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter(
-        [_zip_file_map({"workflow.py": "print('hello')", "workflow.json": "{}"})]
+        [zip_file_map({"workflow.py": "print('hello')", "workflow.json": "{}"})]
     )
 
     # WHEN the user runs the pull command
@@ -552,7 +527,7 @@ def test_pull__exclude_code(vellum_client, mock_module):
 
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter(
-        [_zip_file_map({"workflow.py": "print('hello')", "workflow.json": "{}"})]
+        [zip_file_map({"workflow.py": "print('hello')", "workflow.json": "{}"})]
     )
 
     # WHEN the user runs the pull command
@@ -606,7 +581,7 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter(
         [
-            _zip_file_map(
+            zip_file_map(
                 {
                     "workflow.py": "print('hello')",
                     "metadata.json": json.dumps(
@@ -671,7 +646,7 @@ def test_pull__handle_error_log(vellum_client, mock_module):
 
     # AND the workflow pull API call returns a zip file with an error log
     vellum_client.workflows.pull.return_value = iter(
-        [_zip_file_map({"workflow.py": "print('hello')", "error.log": "test error"})]
+        [zip_file_map({"workflow.py": "print('hello')", "error.log": "test error"})]
     )
 
     # WHEN the user runs the pull command with the new workflow sandbox id
@@ -693,7 +668,7 @@ def test_pull__strict__with_error(vellum_client, mock_module):
     workflow_sandbox_id = mock_module.workflow_sandbox_id
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
 
     # WHEN the user runs the pull command with the new workflow sandbox id
     runner = CliRunner()
@@ -715,7 +690,7 @@ def test_pull__include_sandbox(vellum_client, mock_module):
 
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter(
-        [_zip_file_map({"workflow.py": "print('hello')", "sandbox.py": "print('hello')"})]
+        [zip_file_map({"workflow.py": "print('hello')", "sandbox.py": "print('hello')"})]
     )
 
     # WHEN the user runs the pull command
@@ -744,7 +719,7 @@ def test_pull__same_pull_twice__one_entry_in_lockfile(vellum_client, mock_module
     workflow_sandbox_id = mock_module.workflow_sandbox_id
 
     # AND the workflow pull API call returns a zip file both times
-    zip_contents = _zip_file_map({"workflow.py": "print('hello')"})
+    zip_contents = zip_file_map({"workflow.py": "print('hello')"})
     responses = iter([zip_contents, zip_contents])
 
     def workflows_pull_side_effect(*_args, **_kwargs):
@@ -780,7 +755,7 @@ def test_pull__module_not_in_config(vellum_client, mock_module):
     set_pyproject_toml({"workflows": []})
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
 
     # WHEN the user runs the pull command again with the workflow sandbox id and module
     runner = CliRunner()
@@ -833,7 +808,7 @@ def test_pull__multiple_instances_of_same_module__keep_when_pulling_another_modu
         json.dump(lock_data, f)
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
 
     # WHEN the user runs the pull command on the new module
     runner = CliRunner()
@@ -856,7 +831,7 @@ def test_pull__module_name_from_deployment_name(vellum_client):
     deployment_name = "Test Deployment"
     vellum_client.workflows.pull.return_value = iter(
         [
-            _zip_file_map(
+            zip_file_map(
                 {
                     "workflow.py": "print('hello')",
                     "metadata.json": json.dumps({"deployment_name": deployment_name, "label": "Some Label"}),
@@ -995,7 +970,7 @@ def test_pull__workflow_deployment_adds_deployment_to_config(vellum_client, work
     # AND the workflow pull API call returns a zip file with metadata
     vellum_client.workflows.pull.return_value = iter(
         [
-            _zip_file_map(
+            zip_file_map(
                 {
                     "workflow.py": "print('hello')",
                     "metadata.json": json.dumps(
@@ -1082,7 +1057,7 @@ def test_pull__workflow_deployment_name_is_uuid(vellum_client):
     updated_label = "Updated Label"
     vellum_client.workflows.pull.return_value = iter(
         [
-            _zip_file_map(
+            zip_file_map(
                 {
                     "workflow.py": "print('hello')",
                     "metadata.json": json.dumps(
@@ -1170,7 +1145,7 @@ def test_pull__workflow_deployment_updates_existing_deployment(vellum_client, ge
     updated_label = "Updated Label"
     vellum_client.workflows.pull.return_value = iter(
         [
-            _zip_file_map(
+            zip_file_map(
                 {
                     "workflow.py": "print('hello')",
                     "metadata.json": json.dumps(
@@ -1218,7 +1193,7 @@ def test_pull__workflow_deployment_with_name_and_id(vellum_client):
     # AND the workflow pull API call returns a zip file with metadata
     vellum_client.workflows.pull.return_value = iter(
         [
-            _zip_file_map(
+            zip_file_map(
                 {
                     "workflow.py": "print('hello')",
                     "metadata.json": json.dumps(
@@ -1264,7 +1239,7 @@ def test_pull__workflow_deployment_with_name_and_id(vellum_client):
     # AND pull with name will not add a new deployment to the config
     vellum_client.workflows.pull.return_value = iter(
         [
-            _zip_file_map(
+            zip_file_map(
                 {
                     "workflow.py": "print('hello')",
                     "metadata.json": json.dumps(
@@ -1323,7 +1298,7 @@ MY_OTHER_VELLUM_API_KEY=aaabbbcccddd
 
     # AND the workflow pull API call returns a zip file
     vellum_client_class.return_value.workflows.pull.return_value = iter(
-        [_zip_file_map({"workflow.py": "print('hello')"})]
+        [zip_file_map({"workflow.py": "print('hello')"})]
     )
 
     # WHEN calling `vellum pull` with --workspace

vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py

@@ -142,41 +142,15 @@ def test_serialize_node__lazy_reference(serialize_node):
         attr: str = LazyReference(lambda: ConstantValueReference("hello"))
 
     serialized_node = serialize_node(LazyReferenceGenericNode)
+    attributes = serialized_node["attributes"]
 
-    assert not DeepDiff(
+    assert attributes == [
         {
-            "id": "3d6bfe3b-263a-40a6-8a05-98288e9559a4",
-            "label": "Lazy Reference Generic Node",
-            "type": "GENERIC",
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
-            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
-            "definition": {
-                "name": "LazyReferenceGenericNode",
-                "module": [
-                    "vellum_ee",
-                    "workflows",
-                    "display",
-                    "tests",
-                    "workflow_serialization",
-                    "generic_nodes",
-                    "test_attributes_serialization",
-                ],
-            },
-            "trigger": {"id": "14ec4d19-13e5-4db3-94fa-4e15274bffc7", "merge_behavior": "AWAIT_ATTRIBUTES"},
-            "ports": [{"id": "2dba7224-a376-4780-8414-2b50601f9283", "name": "default", "type": "DEFAULT"}],
-            "adornments": None,
-            "attributes": [
-                {
-                    "id": "7ae37eb4-18c8-49e1-b5ac-6369ce7ed5dd",
-                    "name": "attr",
-                    "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "hello"}},
-                }
-            ],
-            "outputs": [],
-        },
-        serialized_node,
-        ignore_order=True,
-    )
+            "id": "7ae37eb4-18c8-49e1-b5ac-6369ce7ed5dd",
+            "name": "attr",
+            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "hello"}},
+        }
+    ]
 
 
 def test_serialize_node__lazy_reference_with_string():

vellum_ee/workflows/display/utils/events.py

@@ -7,6 +7,24 @@ from vellum_ee.workflows.display.utils.registry import (
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
+def _should_mark_workflow_dynamic(event: WorkflowExecutionInitiatedEvent) -> bool:
+    """
+    Check if workflow should be marked as dynamic based on execution context.
+    Returns True if parent.type == WORKFLOW_RELEASE_TAG and parent.parent.type == WORKFLOW_NODE.
+    """
+    if not event.parent:
+        return False
+
+    parent = event.parent
+    if parent.type != "WORKFLOW_RELEASE_TAG":
+        return False
+
+    if not parent.parent or parent.parent.type != "WORKFLOW_NODE":
+        return False
+
+    return True
+
+
 def event_enricher(event: WorkflowExecutionInitiatedEvent) -> WorkflowExecutionInitiatedEvent:
     workflow_definition = event.body.workflow_definition
     workflow_display = get_workflow_display(
@@ -16,7 +34,7 @@ def event_enricher(event: WorkflowExecutionInitiatedEvent) -> WorkflowExecutionI
     )
     register_workflow_display_context(event.span_id, workflow_display.display_context)
 
-    if event.body.workflow_definition.is_dynamic:
+    if event.body.workflow_definition.is_dynamic or _should_mark_workflow_dynamic(event):
         register_workflow_display_class(workflow_definition, workflow_display.__class__)
         workflow_version_exec_config = workflow_display.serialize()
         setattr(event.body, "workflow_version_exec_config", workflow_version_exec_config)

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

@@ -2,8 +2,10 @@ import pytest
 from uuid import uuid4
 from typing import Optional
 
+from vellum.workflows.events.types import NodeParentContext, WorkflowDeploymentParentContext
 from vellum.workflows.events.workflow import WorkflowExecutionInitiatedBody, WorkflowExecutionInitiatedEvent
 from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.utils.events import event_enricher
 
@@ -67,3 +69,43 @@ def test_event_enricher_static_workflow(is_dynamic: bool, expected_config: Optio
 
     # AND workflow_version_exec_config is set to the expected config
     assert event.body.workflow_version_exec_config == expected_config
+
+
+def test_event_enricher_marks_subworkflow_deployment_as_dynamic():
+    """Test that event_enricher treats subworkflow deployments as dynamic."""
+
+    class TestWorkflow(BaseWorkflow):
+        is_dynamic = False
+
+        class Outputs(BaseOutputs):
+            pass
+
+    event: WorkflowExecutionInitiatedEvent = WorkflowExecutionInitiatedEvent(
+        trace_id=uuid4(),
+        span_id=uuid4(),
+        parent=WorkflowDeploymentParentContext(
+            span_id=uuid4(),
+            deployment_id=uuid4(),
+            deployment_name="test-deployment",
+            deployment_history_item_id=uuid4(),
+            release_tag_id=uuid4(),
+            release_tag_name="test-tag",
+            workflow_version_id=uuid4(),
+            external_id=None,
+            metadata=None,
+            parent=NodeParentContext(
+                span_id=uuid4(),
+                node_definition=TestWorkflow,
+                parent=None,
+            ),
+        ),
+        body=WorkflowExecutionInitiatedBody(
+            workflow_definition=TestWorkflow,
+            inputs=BaseInputs(),
+        ),
+    )
+
+    enriched_event = event_enricher(event)
+
+    assert hasattr(enriched_event.body, "workflow_version_exec_config")
+    assert enriched_event.body.workflow_version_exec_config is not None

vellum_ee/workflows/tests/test_server.py

@@ -9,6 +9,7 @@ from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.utils.uuids import generate_workflow_deployment_prefix
+from vellum.workflows.utils.zip import zip_file_map
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
 
@@ -661,3 +662,117 @@ __all__ = ["TestNode"]
 
     # AND the method should return a workflow (not None) - this will pass once implemented
     assert event.name == "workflow.execution.fulfilled"
+
+
+def test_resolve_workflow_deployment__uses_pull_api_with_inputs_deployment_name(vellum_client):
+    """
+    Test that resolve_workflow_deployment uses the pull API to fetch subworkflow files
+    when the deployment name comes from Inputs.deployment_name.
+    """
+    # GIVEN a deployment name and release tag
+    deployment_name = "test_deployment"
+    release_tag = "LATEST"
+
+    test_node_code = """
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.outputs import BaseOutputs
+
+class TestNode(BaseNode):
+    template = "Hello"
+
+    class Outputs(BaseOutputs):
+        result: str
+
+    def run(self):
+        return self.Outputs(result="Hello, {template}")
+"""
+
+    mock_workflow_code = """
+from vellum.workflows import BaseWorkflow
+from .nodes.test_node import TestNode
+
+class ResolvedWorkflow(BaseWorkflow):
+    graph = TestNode
+"""
+
+    inputs_code = """
+from vellum.workflows.inputs import BaseInputs
+
+class Inputs(BaseInputs):
+    deployment_name: str
+"""
+
+    parent_workflow_code = """
+from vellum.workflows import BaseWorkflow
+from .inputs import Inputs
+from .nodes.subworkflow_deployment_node import TestSubworkflowDeploymentNode
+from vellum.workflows.state import BaseState
+
+class ParentWorkflow(BaseWorkflow[Inputs, BaseState]):
+    graph = TestSubworkflowDeploymentNode
+"""
+
+    parent_node_code = """
+from vellum.workflows.nodes import SubworkflowDeploymentNode
+from vellum.workflows.outputs import BaseOutputs
+from ..inputs import Inputs
+
+class TestSubworkflowDeploymentNode(SubworkflowDeploymentNode):
+    deployment = Inputs.deployment_name
+
+    class Outputs(BaseOutputs):
+        result: str
+
+    subworkflow_inputs = {"message": "test"}
+"""
+
+    subworkflow_files = {
+        "__init__.py": "",
+        "workflow.py": mock_workflow_code,
+        "nodes/__init__.py": """
+from .test_node import TestNode
+
+__all__ = ["TestNode"]
+""",
+        "nodes/test_node.py": test_node_code,
+    }
+
+    parent_files = {
+        "__init__.py": "",
+        "inputs.py": inputs_code,
+        "workflow.py": parent_workflow_code,
+        "nodes/__init__.py": """
+from .subworkflow_deployment_node import TestSubworkflowDeploymentNode
+
+__all__ = ["TestSubworkflowDeploymentNode"]
+""",
+        "nodes/subworkflow_deployment_node.py": parent_node_code,
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered for the parent workflow
+    finder = VirtualFileFinder(parent_files, namespace)
+    sys.meta_path.append(finder)
+
+    vellum_client.workflows.pull.return_value = iter([zip_file_map(subworkflow_files)])
+
+    # WHEN we execute the root workflow with the mocked client
+    Workflow = BaseWorkflow.load_from_module(namespace)
+    Inputs = Workflow.get_inputs_class()
+
+    workflow = Workflow(context=WorkflowContext(namespace=namespace, generated_files=parent_files))
+    final_event = workflow.run(inputs=Inputs(deployment_name=deployment_name))
+
+    # THEN the method should return a workflow (not None) - this will pass once implemented
+    assert final_event.name == "workflow.execution.fulfilled", final_event
+
+    # AND the pull API should have been called with the correct deployment name, release tag, and version
+    args, kwargs = vellum_client.workflows.pull.call_args
+    assert args[0] == deployment_name
+    assert kwargs["release_tag"] == release_tag
+    assert kwargs["version"].startswith(">=")
+    assert ".0.0,<=" in kwargs["version"]
+
+    # AND the X-Vellum-Always-Success header should be included for graceful error handling
+    assert kwargs["request_options"]["additional_headers"]["X-Vellum-Always-Success"] == "true"