vellum-ai 0.12.15__tar.gz → 0.12.16__tar.gz

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.12.15
+Version: 0.12.16
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -40,9 +40,11 @@ Requires-Dist: typing_extensions (>=4.0.0)
 Description-Content-Type: text/markdown
 
 <p align="center">
-  <h1 align="center">
-  Vellum
-  </h1>
+   <a href="https://vellum.ai">
+    <picture>
+      <img alt="Vellum README banner" src="https://storage.googleapis.com/vellum-public/assets/github-readme-banner.png" height="128">
+    </picture>
+  </a>
   <p align="center">
     <a href="https://vellum.ai">Learn more</a>
     ·
@@ -64,13 +66,13 @@ Description-Content-Type: text/markdown
 
 # Introduction
 
-[Vellum](https://www.vellum.ai/) is the end-to-end development platform for building production-grade AI applications
+[Vellum](https://www.vellum.ai/) is the end-to-end development platform for building production-grade AI applications.
 
 ### Core Features
 
 - **Orchestration:** A powerful SDK and IDE for defining and debugging the control flow of your AI applications
 - **Prompting:** A best-in-class prompt playground for iterating on and refining prompts between models from any provider
-- **Evaluations**: An evaluations framework that makes it easy to measure the quality of your AI systems at scale.
+- **Evaluations**: An evaluations framework that makes it easy to measure the quality of your AI systems at scale
 - **Retrieval:** A ready-to-go service for turning unstructured content into intelligent, context-aware solutions
   optimized for AI systems
 - **Deployment:** Decouple updates to your AI systems from your application code with an easy integration +
@@ -87,7 +89,7 @@ Description-Content-Type: text/markdown
 
 ## Get Started
 
-Most functionality within the SDKs here requires a Vellum account and API key. To sign up, [talk to us](https://www.vellum.ai/landing-pages/request-demo)
+Most functionality within the SDK requires a Vellum account and API key. To sign up, [talk to us](https://www.vellum.ai/landing-pages/request-demo)
 or visit our [pricing page](https://www.vellum.ai/pricing).
 
 Even without a Vellum account, you can use the Workflows SDK to define the control flow of your AI systems. [Learn
@@ -101,7 +103,7 @@ Learn more and get started by visiting the [Vellum Client SDK README](/src/vellu
 ## Workflows SDK
 
 The Vellum Workflows SDK is a high-level framework for defining and debugging the control flow of AI systems. At
-it's core, it's a powerful workflow engine with syntactic sugar for intuitively defining graphs, the nodes within,
+it's core, it's a powerful workflow engine with syntactic sugar for declaratively defining graphs, the nodes within,
 and the relationships between them.
 
 The Workflows SDK can be used with or without a Vellum account, but a Vellum account is required to use certain

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/README.md

@@ -1,7 +1,9 @@
 <p align="center">
-  <h1 align="center">
-  Vellum
-  </h1>
+   <a href="https://vellum.ai">
+    <picture>
+      <img alt="Vellum README banner" src="https://storage.googleapis.com/vellum-public/assets/github-readme-banner.png" height="128">
+    </picture>
+  </a>
   <p align="center">
     <a href="https://vellum.ai">Learn more</a>
     ·
@@ -23,13 +25,13 @@
 
 # Introduction
 
-[Vellum](https://www.vellum.ai/) is the end-to-end development platform for building production-grade AI applications
+[Vellum](https://www.vellum.ai/) is the end-to-end development platform for building production-grade AI applications.
 
 ### Core Features
 
 - **Orchestration:** A powerful SDK and IDE for defining and debugging the control flow of your AI applications
 - **Prompting:** A best-in-class prompt playground for iterating on and refining prompts between models from any provider
-- **Evaluations**: An evaluations framework that makes it easy to measure the quality of your AI systems at scale.
+- **Evaluations**: An evaluations framework that makes it easy to measure the quality of your AI systems at scale
 - **Retrieval:** A ready-to-go service for turning unstructured content into intelligent, context-aware solutions
   optimized for AI systems
 - **Deployment:** Decouple updates to your AI systems from your application code with an easy integration +
@@ -46,7 +48,7 @@
 
 ## Get Started
 
-Most functionality within the SDKs here requires a Vellum account and API key. To sign up, [talk to us](https://www.vellum.ai/landing-pages/request-demo)
+Most functionality within the SDK requires a Vellum account and API key. To sign up, [talk to us](https://www.vellum.ai/landing-pages/request-demo)
 or visit our [pricing page](https://www.vellum.ai/pricing).
 
 Even without a Vellum account, you can use the Workflows SDK to define the control flow of your AI systems. [Learn
@@ -60,7 +62,7 @@ Learn more and get started by visiting the [Vellum Client SDK README](/src/vellu
 ## Workflows SDK
 
 The Vellum Workflows SDK is a high-level framework for defining and debugging the control flow of AI systems. At
-it's core, it's a powerful workflow engine with syntactic sugar for intuitively defining graphs, the nodes within,
+it's core, it's a powerful workflow engine with syntactic sugar for declaratively defining graphs, the nodes within,
 and the relationships between them.
 
 The Workflows SDK can be used with or without a Vellum account, but a Vellum account is required to use certain

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_cli/__init__.py

@@ -49,6 +49,11 @@ def workflows():
 @click.option("--deployment-name", type=str, help="Unique name for the Deployment")
 @click.option("--deployment-description", type=str, help="Description for the Deployment")
 @click.option("--release-tag", type=list, help="Release Tag for the Deployment", multiple=True)
+@click.option(
+    "--dry-run",
+    is_flag=True,
+    help="Check the Workflow for errors and expected changes, without updating its state in Vellum.",
+)
 def workflows_push(
     module: Optional[str],
     deploy: Optional[bool],
@@ -56,6 +61,7 @@ def workflows_push(
     deployment_name: Optional[str],
     deployment_description: Optional[str],
     release_tag: Optional[List[str]],
+    dry_run: Optional[bool],
 ) -> None:
     """
     Push Workflows to Vellum. If a module is provided, only the Workflow for that module will be pushed.
@@ -69,6 +75,7 @@ def workflows_push(
         deployment_name=deployment_name,
         deployment_description=deployment_description,
         release_tags=release_tag,
+        dry_run=dry_run,
     )
 
 
@@ -79,6 +86,11 @@ def workflows_push(
 @click.option("--deployment-name", type=str, help="Unique name for the Deployment")
 @click.option("--deployment-description", type=str, help="Description for the Deployment")
 @click.option("--release-tag", type=list, help="Release Tag for the Deployment", multiple=True)
+@click.option(
+    "--dry-run",
+    is_flag=True,
+    help="Check the Workflow for errors and expected changes, without updating its state in Vellum.",
+)
 def push_module(
     ctx: click.Context,
     deploy: Optional[bool],
@@ -86,6 +98,7 @@ def push_module(
     deployment_name: Optional[str],
     deployment_description: Optional[str],
     release_tag: Optional[List[str]],
+    dry_run: Optional[bool],
 ) -> None:
     """Push a specific module to Vellum"""
 
@@ -97,6 +110,7 @@ def push_module(
             deployment_name=deployment_name,
             deployment_description=deployment_description,
             release_tags=release_tag,
+            dry_run=dry_run,
         )
 
 

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_cli/push.py

@@ -27,6 +27,7 @@ def push_command(
     deployment_name: Optional[str] = None,
     deployment_description: Optional[str] = None,
     release_tags: Optional[List[str]] = None,
+    dry_run: Optional[bool] = None,
 ) -> None:
     load_dotenv()
     logger = load_cli_logger()
@@ -49,7 +50,11 @@ def push_command(
     # Remove this once we could serialize using the artifact in Vembda
     # https://app.shortcut.com/vellum/story/5585
     workflow = BaseWorkflow.load_from_module(workflow_config.module)
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=workflow)
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay,
+        workflow_class=workflow,
+        dry_run=dry_run or False,
+    )
     exec_config = workflow_display.serialize()
 
     container_tag = workflow_config.container_image_tag
@@ -114,11 +119,28 @@ def push_command(
         # We should check with fern if we could auto-serialize typed object fields for us
         # https://app.shortcut.com/vellum/story/5568
         deployment_config=deployment_config_serialized,  # type: ignore[arg-type]
+        dry_run=dry_run,
     )
-    logger.info(
-        f"""Successfully pushed {label} to Vellum!
+
+    if dry_run:
+        error_messages = [str(e) for e in workflow_display.errors]
+        error_message = "\n".join(error_messages) if error_messages else "No errors found."
+        logger.info(
+            f"""\
+# Workflow Push Report
+
+## Errors
+{error_message}
+
+## Proposed Diffs
+{json.dumps(response.proposed_diffs, indent=2)}
+"""
+        )  # type: ignore[attr-defined]
+    else:
+        logger.info(
+            f"""Successfully pushed {label} to Vellum!
 Visit at: https://app.vellum.ai/workflow-sandboxes/{response.workflow_sandbox_id}"""
-    )
+        )
 
     requires_save = False
     if not workflow_config.workflow_sandbox_id:

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_cli/tests/test_push.py

@@ -171,3 +171,66 @@ class ExampleWorkflow(BaseWorkflow):
 
     extracted_files = _extract_tar_gz(call_args["artifact"].read())
     assert extracted_files["workflow.py"] == workflow_py_file_content
+
+
+def test_push__dry_run_option_returns_report(mock_module, vellum_client):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    base_dir = os.path.join(temp_dir, *module.split("."))
+    os.makedirs(base_dir, exist_ok=True)
+    workflow_py_file_content = """\
+from typing import Dict
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.nodes import BaseNode
+from vellum_ee.workflows.display.nodes import BaseNodeDisplay
+
+class NotSupportedNode(BaseNode):
+    pass
+
+class NotSupportedNodeDisplay(BaseNodeDisplay[NotSupportedNode]):
+    def serialize(self, display_context, **kwargs) -> Dict:
+        raise NotImplementedError(f"Serialization is not supported.")
+
+class ExampleWorkflow(BaseWorkflow):
+    graph = NotSupportedNode
+"""
+    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
+        f.write(workflow_py_file_content)
+
+    # AND the push API call returns successfully
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=str(uuid4()),
+        proposed_diffs={
+            "iterable_item_added": {
+                "root['raw_data']['nodes'][0]": {
+                    "id": str(uuid4()),
+                    "type": "GENERIC",
+                    "definition": None,
+                    "display_data": None,
+                    "trigger": {"id": str(uuid4()), "merge_behavior": "AWAIT_ATTRIBUTES"},
+                    "ports": [{"id": str(uuid4()), "name": "default", "type": "DEFAULT"}],
+                },
+            },
+        },
+    )
+
+    # WHEN calling `vellum push`
+    runner = CliRunner(mix_stderr=True)
+    result = runner.invoke(cli_main, ["push", module, "--dry-run"])
+
+    # THEN it should succeed
+    assert result.exit_code == 0
+
+    # AND we should have called the push API with the dry-run option
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+    assert call_args["dry_run"] is True
+
+    # AND the report should be in the output
+    assert "## Errors" in result.output
+    # assert "Serialization is not supported." in result.output
+    assert "## Proposed Diffs" in result.output
+    assert "iterable_item_added" in result.output

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_ee/workflows/display/nodes/vellum/api_node.py

@@ -161,9 +161,9 @@ class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeTy
         ]
         inputs.extend(additional_header_inputs)
 
-        _, text_output_display = display_context.node_output_displays[cast(OutputReference, node.Outputs.text)]
-        _, json_output_display = display_context.node_output_displays[cast(OutputReference, node.Outputs.json)]
-        _, status_code_output_display = display_context.node_output_displays[
+        _, text_output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.text)]
+        _, json_output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.json)]
+        _, status_code_output_display = display_context.global_node_output_displays[
             cast(OutputReference, node.Outputs.status_code)
         ]
 

vellum_ai-0.12.16/ee/vellum_ee/workflows/display/nodes/vellum/base_node.py

@@ -0,0 +1,35 @@
+from typing import Any, Generic, TypeVar
+
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.types.core import JsonObject
+from vellum.workflows.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.vellum import GenericNodeDisplayData
+
+_BaseNodeType = TypeVar("_BaseNodeType", bound=BaseNode)
+
+
+class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeType]):
+    def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
+        node = self._node
+        node_id = self.node_id
+
+        return {
+            "id": str(node_id),
+            "label": node.__qualname__,
+            "type": "GENERIC",
+            "display_data": self.get_generic_node_display_data().dict(),
+            "definition": self.get_definition().dict(),
+            "trigger": {
+                "id": str(uuid4_from_hash(f"{node_id}|trigger")),
+                "merge_behavior": node.Trigger.merge_behavior.value,
+            },
+            "ports": [],
+            "adornments": None,
+            "attributes": [],
+        }
+
+    def get_generic_node_display_data(self) -> GenericNodeDisplayData:
+        explicit_value = self._get_explicit_node_display_attr("display_data", GenericNodeDisplayData)
+        return explicit_value if explicit_value else GenericNodeDisplayData()

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -70,8 +70,8 @@ class BaseCodeExecutionNodeDisplay(BaseNodeVellumDisplay[_CodeExecutionNodeType]
 
         packages = raise_if_descriptor(node.packages)
 
-        _, output_display = display_context.node_output_displays[node.Outputs.result]
-        _, log_output_display = display_context.node_output_displays[node.Outputs.log]
+        _, output_display = display_context.global_node_output_displays[node.Outputs.result]
+        _, log_output_display = display_context.global_node_output_displays[node.Outputs.log]
 
         output_type = primitive_type_to_vellum_variable_type(node.get_output_type())
 

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py

@@ -29,8 +29,8 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
         node_inputs, prompt_inputs = self._generate_node_and_prompt_inputs(node_id, node, display_context)
         input_variable_id_by_name = {prompt_input.key: prompt_input.id for prompt_input in prompt_inputs}
 
-        _, output_display = display_context.node_output_displays[node.Outputs.text]
-        _, array_display = display_context.node_output_displays[node.Outputs.results]
+        _, output_display = display_context.global_node_output_displays[node.Outputs.text]
+        _, array_display = display_context.global_node_output_displays[node.Outputs.results]
         node_blocks = raise_if_descriptor(node.blocks)
 
         return {

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py

@@ -2,6 +2,7 @@ from uuid import UUID
 from typing import ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, cast
 
 from vellum import VellumVariable
+from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes import InlineSubworkflowNode
 from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
@@ -60,12 +61,25 @@ class BaseInlineSubworkflowNodeDisplay(
         node: Type[InlineSubworkflowNode],
         display_context: WorkflowDisplayContext,
     ) -> Tuple[List[NodeInput], List[VellumVariable]]:
+        subworkflow = raise_if_descriptor(node.subworkflow)
+        subworkflow_inputs_class = subworkflow.get_inputs_class()
         subworkflow_inputs = raise_if_descriptor(node.subworkflow_inputs)
-        subworkflow_entries = (
-            [(variable_name, variable_value) for variable_name, variable_value in subworkflow_inputs.items()]
-            if isinstance(subworkflow_inputs, dict)
-            else [(variable_ref.name, variable_value) for variable_ref, variable_value in subworkflow_inputs]
-        )
+
+        if isinstance(subworkflow_inputs, BaseInputs):
+            subworkflow_entries = [
+                (variable_ref.name, variable_value) for variable_ref, variable_value in subworkflow_inputs
+            ]
+        elif isinstance(subworkflow_inputs, dict):
+            subworkflow_entries = [
+                (variable_name, variable_value) for variable_name, variable_value in subworkflow_inputs.items()
+            ]
+        else:
+            subworkflow_entries = [
+                (descriptor.name, getattr(subworkflow_inputs_class, descriptor.name))
+                for descriptor in subworkflow_inputs_class
+                if hasattr(subworkflow_inputs_class, descriptor.name)
+            ]
+
         node_inputs = [
             create_node_input(
                 node_id=node_id,
@@ -83,7 +97,7 @@ class BaseInlineSubworkflowNodeDisplay(
                 key=descriptor.name,
                 type=infer_vellum_variable_type(descriptor),
             )
-            for descriptor in raise_if_descriptor(node.subworkflow).get_inputs_class()
+            for descriptor in subworkflow_inputs_class
         ]
 
         return node_inputs, workflow_inputs

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_ee/workflows/display/nodes/vellum/map_node.py

@@ -33,6 +33,7 @@ class BaseMapNodeDisplay(BaseNodeVellumDisplay[_MapNodeType], Generic[_MapNodeTy
         subworkflow_display = get_workflow_display(
             base_display_class=display_context.workflow_display_class,
             workflow_class=subworkflow,
+            parent_display_context=display_context,
         )
         serialized_subworkflow = subworkflow_display.serialize()
 

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py

@@ -38,8 +38,8 @@ class BasePromptDeploymentNodeDisplay(
             for variable_name, variable_value in prompt_inputs.items()
         ]
 
-        _, output_display = display_context.node_output_displays[cast(OutputReference, node.Outputs.text)]
-        _, array_display = display_context.node_output_displays[cast(OutputReference, node.Outputs.results)]
+        _, output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.text)]
+        _, array_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.results)]
 
         # TODO: Pass through the name instead of retrieving the ID
         # https://app.shortcut.com/vellum/story/4702

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_ee/workflows/display/nodes/vellum/search_node.py

@@ -39,8 +39,10 @@ class BaseSearchNodeDisplay(BaseNodeVellumDisplay[_SearchNodeType], Generic[_Sea
         node_id = self.node_id
         node_inputs = self._generate_search_node_inputs(node_id, node, display_context)
 
-        _, results_output_display = display_context.node_output_displays[cast(OutputReference, node.Outputs.results)]
-        _, text_output_display = display_context.node_output_displays[cast(OutputReference, node.Outputs.text)]
+        _, results_output_display = display_context.global_node_output_displays[
+            cast(OutputReference, node.Outputs.results)
+        ]
+        _, text_output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.text)]
 
         return {
             "id": str(node_id),

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_ee/workflows/display/nodes/vellum/templating_node.py

@@ -56,7 +56,7 @@ class BaseTemplatingNodeDisplay(BaseNodeVellumDisplay[_TemplatingNodeType], Gene
         # Misc type ignore is due to `node.Outputs` being generic
         # https://app.shortcut.com/vellum/story/4784
         output_descriptor = node.Outputs.result  # type: ignore [misc]
-        _, output_display = display_context.node_output_displays[output_descriptor]
+        _, output_display = display_context.global_node_output_displays[output_descriptor]
         inferred_output_type = primitive_type_to_vellum_variable_type(output_descriptor)
 
         return {

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py

@@ -91,18 +91,18 @@ def test_create_node_input_value_pointer_rules(
                 entrypoint_node_source_handle_id=uuid4(),
                 entrypoint_node_display=NodeDisplayData(),
             ),
-            workflow_input_displays={
+            global_workflow_input_displays={
                 cast(WorkflowInputReference, Inputs.example_workflow_input): WorkflowInputsVellumDisplayOverrides(
                     id=UUID("a154c29d-fac0-4cd0-ba88-bc52034f5470"),
                 ),
             },
-            node_output_displays={
+            global_node_output_displays={
                 cast(OutputReference, MyNodeA.Outputs.output): (
                     MyNodeA,
                     NodeOutputDisplay(id=UUID("4b16a629-11a1-4b3f-a965-a57b872d13b8"), name="output"),
                 ),
             },
-            node_displays={
+            global_node_displays={
                 MyNodeA: MyNodeADisplay(),
             },
         ),

vellum_ai-0.12.16/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py

@@ -0,0 +1,28 @@
+import pytest
+from uuid import uuid4
+
+from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
+from vellum_ee.workflows.display.nodes.vellum.base_node import BaseNodeDisplay
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.vellum import NodeDisplayData, WorkflowMetaVellumDisplay
+from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
+
+
+@pytest.fixture()
+def serialize_node():
+    def _serialize_node(node_class) -> dict:
+        node_display_class = get_node_display_class(BaseNodeDisplay, node_class)
+        node_display = node_display_class()
+
+        context: WorkflowDisplayContext = WorkflowDisplayContext(
+            workflow_display_class=VellumWorkflowDisplay,
+            workflow_display=WorkflowMetaVellumDisplay(
+                entrypoint_node_id=uuid4(),
+                entrypoint_node_source_handle_id=uuid4(),
+                entrypoint_node_display=NodeDisplayData(),
+            ),
+            node_displays={node_class: node_display},
+        )
+        return node_display.serialize(context)
+
+    return _serialize_node

vellum_ai-0.12.16/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py

@@ -0,0 +1,123 @@
+from deepdiff import DeepDiff
+
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.types.core import MergeBehavior
+
+
+class Inputs(BaseInputs):
+    input: str
+
+
+class BasicGenericNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+
+class AwaitAnyGenericNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+    class Trigger(BaseNode.Trigger):
+        merge_behavior = MergeBehavior.AWAIT_ANY
+
+
+class AwaitAllGenericNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+    class Trigger(BaseNode.Trigger):
+        merge_behavior = MergeBehavior.AWAIT_ALL
+
+
+def test_serialize_node__basic(serialize_node):
+    serialized_node = serialize_node(BasicGenericNode)
+    assert not DeepDiff(
+        {
+            "id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
+            "label": "BasicGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "definition": {
+                "name": "BasicGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_trigger_serialization",
+                ],
+                "bases": [{"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]}],
+            },
+            "trigger": {"id": "9d3a1b3d-4a38-4f2e-bbf1-dd8be152bce8", "merge_behavior": "AWAIT_ANY"},
+            "ports": [],
+            "adornments": None,
+            "attributes": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+def test_serialize_node__await_any(serialize_node):
+    serialized_node = serialize_node(AwaitAnyGenericNode)
+    assert not DeepDiff(
+        {
+            "id": "0ba67f76-aaff-4bd4-a20f-73a32ef5810d",
+            "label": "AwaitAnyGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "definition": {
+                "name": "AwaitAnyGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_trigger_serialization",
+                ],
+                "bases": [{"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]}],
+            },
+            "trigger": {"id": "ffa72187-9a18-453f-ae55-b77aad332630", "merge_behavior": "AWAIT_ANY"},
+            "ports": [],
+            "adornments": None,
+            "attributes": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+def test_serialize_node__await_all(serialize_node):
+    serialized_node = serialize_node(AwaitAllGenericNode)
+    assert not DeepDiff(
+        {
+            "id": "09d06cd3-06ea-40cc-afd8-17ad88542271",
+            "label": "AwaitAllGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "definition": {
+                "name": "AwaitAllGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_trigger_serialization",
+                ],
+                "bases": [{"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]}],
+            },
+            "trigger": {"id": "62074276-c817-476d-b59d-da523ae3f218", "merge_behavior": "AWAIT_ALL"},
+            "ports": [],
+            "adornments": None,
+            "attributes": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )

{vellum_ai-0.12.15 → vellum_ai-0.12.16}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py

@@ -448,17 +448,8 @@ def test_serialize_workflow():
         ignore_order=True,
     )
 
-    assert not DeepDiff(
-        [
-            {"id": "148c61bd-e8b0-4d4b-8734-b043a72b90ed", "type": "GENERIC"},
-            {"id": "ed7caf01-9ae7-47a3-b15a-16697abaf486", "type": "GENERIC"},
-            {"id": "0d959311-c836-4641-a867-58f63df9dfea", "type": "GENERIC"},
-            {"id": "8df781b1-ff28-48a5-98a2-d7d796b932b0", "type": "GENERIC"},
-            {"id": "68c02b7c-5077-4087-803d-841474a8081f", "type": "GENERIC"},
-        ],
-        workflow_raw_data["nodes"][2:7],
-        ignore_order=True,
-    )
+    passthrough_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "GENERIC"]
+    assert len(passthrough_nodes) == 5
 
     assert not DeepDiff(
         [