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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.12.14
+Version: 0.12.15
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -95,7 +95,7 @@ more below](#workflows-sdk).
 
 ## Client SDK
 
-The Vellum Client SDK, found within `src/client` is a low-level client used to interact directly with the Vellum API.
+The Vellum Client SDK, found within `src/vellum/client` is a low-level client used to interact directly with the Vellum API.
 Learn more and get started by visiting the [Vellum Client SDK README](/src/vellum/client/README.md).
 
 ## Workflows SDK

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

@@ -54,7 +54,7 @@ more below](#workflows-sdk).
 
 ## Client SDK
 
-The Vellum Client SDK, found within `src/client` is a low-level client used to interact directly with the Vellum API.
+The Vellum Client SDK, found within `src/vellum/client` is a low-level client used to interact directly with the Vellum API.
 Learn more and get started by visiting the [Vellum Client SDK README](/src/vellum/client/README.md).
 
 ## Workflows SDK

{vellum_ai-0.12.14 → vellum_ai-0.12.15}/ee/vellum_cli/config.py

@@ -35,6 +35,8 @@ class WorkflowConfig(UniversalBaseModel):
     workflow_sandbox_id: Optional[str] = None
     ignore: Optional[Union[str, List[str]]] = None
     deployments: List[WorkflowDeploymentConfig] = field(default_factory=list)
+    container_image_name: Optional[str] = None
+    container_image_tag: Optional[str] = None
 
     def merge(self, other: "WorkflowConfig") -> "WorkflowConfig":
         self_deployment_by_id = {
@@ -68,6 +70,8 @@ class WorkflowConfig(UniversalBaseModel):
             workflow_sandbox_id=self.workflow_sandbox_id or other.workflow_sandbox_id,
             ignore=self.ignore or other.ignore,
             deployments=merged_deployments,
+            container_image_tag=self.container_image_tag or other.container_image_tag,
+            container_image_name=self.container_image_name or other.container_image_name,
         )
 
 

{vellum_ai-0.12.14 → vellum_ai-0.12.15}/ee/vellum_cli/pull.py

@@ -1,4 +1,5 @@
 import io
+import json
 import os
 from pathlib import Path
 from uuid import UUID
@@ -14,6 +15,7 @@ from vellum_cli.config import VellumCliConfig, WorkflowConfig, load_vellum_cli_c
 from vellum_cli.logger import load_cli_logger
 
 ERROR_LOG_FILE_NAME = "error.log"
+METADATA_FILE_NAME = "metadata.json"
 
 
 def _is_valid_uuid(val: Union[str, UUID, None]) -> bool:
@@ -109,6 +111,7 @@ def pull_command(
     logger.info(f"Pulling workflow into {workflow_config.module}")
     client = create_vellum_client()
     query_parameters = {}
+
     if include_json:
         query_parameters["include_json"] = include_json
     if exclude_code:
@@ -128,6 +131,8 @@ def pull_command(
 
     target_dir = os.path.join(os.getcwd(), *workflow_config.module.split("."))
     error_content = ""
+    metadata_json: Optional[dict] = None
+
     with zipfile.ZipFile(zip_buffer) as zip_file:
         # Delete files in target_dir that aren't in the zip file
         if os.path.exists(target_dir):
@@ -152,16 +157,26 @@ def pull_command(
                     os.remove(file_path)
 
         for file_name in zip_file.namelist():
-            target_file = os.path.join(target_dir, file_name)
-            os.makedirs(os.path.dirname(target_file), exist_ok=True)
-            with zip_file.open(file_name) as source, open(target_file, "w") as target:
+            with zip_file.open(file_name) as source:
                 content = source.read().decode("utf-8")
                 if file_name == ERROR_LOG_FILE_NAME:
                     error_content = content
                     continue
+                if file_name == METADATA_FILE_NAME:
+                    metadata_json = json.loads(content)
+                    continue
 
-                logger.info(f"Writing to {target_file}...")
-                target.write(content)
+                target_file = os.path.join(target_dir, file_name)
+                os.makedirs(os.path.dirname(target_file), exist_ok=True)
+                with open(target_file, "w") as target:
+                    logger.info(f"Writing to {target_file}...")
+                    target.write(content)
+
+    if metadata_json:
+        workflow_config.container_image_name = metadata_json.get("runner_config", {}).get("container_image_name")
+        workflow_config.container_image_tag = metadata_json.get("runner_config", {}).get("container_image_tag")
+        if workflow_config.container_image_name and not workflow_config.container_image_tag:
+            workflow_config.container_image_tag = "latest"
 
     if include_json:
         logger.warning(

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

@@ -51,8 +51,15 @@ def push_command(
     workflow = BaseWorkflow.load_from_module(workflow_config.module)
     workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=workflow)
     exec_config = workflow_display.serialize()
+
+    container_tag = workflow_config.container_image_tag
+    if workflow_config.container_image_name and not workflow_config.container_image_tag:
+        container_tag = "latest"
+
     exec_config["runner_config"] = {
         "sdk_version": metadata.version("vellum-ai"),
+        "container_image_tag": container_tag,
+        "container_image_name": workflow_config.container_image_name,
     }
 
     label = snake_to_title_case(workflow_config.module.split(".")[-1])

{vellum_ai-0.12.14 → vellum_ai-0.12.15}/ee/vellum_cli/tests/test_pull.py

@@ -130,6 +130,8 @@ def test_pull__sandbox_id_with_no_config(vellum_client):
                     "workflow_sandbox_id": "87654321-0000-0000-0000-000000000000",
                     "ignore": None,
                     "deployments": [],
+                    "container_image_tag": None,
+                    "container_image_name": None,
                 }
             ],
         }
@@ -204,6 +206,8 @@ def test_pull__workflow_deployment_with_no_config(vellum_client):
                     "workflow_sandbox_id": None,
                     "ignore": None,
                     "deployments": [],
+                    "container_image_tag": None,
+                    "container_image_name": None,
                 }
             ],
         }
@@ -407,7 +411,17 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
     new_workflow_sandbox_id = "87654321-0000-0000-0000-000000000000"
 
     # 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')",
+                    "metadata.json": '{"runner_config": { "container_image_name": "test", '
+                    '"container_image_tag": "1.0" } }',
+                }
+            )
+        ]
+    )
 
     # WHEN the user runs the pull command with the new workflow sandbox id
     runner = CliRunner()
@@ -433,12 +447,16 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                         "release_tags": None,
                     },
                 ],
+                "container_image_name": None,
+                "container_image_tag": None,
             },
             {
                 "module": "workflow_87654321",
                 "workflow_sandbox_id": new_workflow_sandbox_id,
                 "ignore": None,
                 "deployments": [],
+                "container_image_name": "test",
+                "container_image_tag": "1.0",
             },
         ]
 

{vellum_ai-0.12.14 → vellum_ai-0.12.15}/ee/vellum_ee/workflows/display/nodes/vellum/__init__.py

@@ -1,4 +1,5 @@
 from .api_node import BaseAPINodeDisplay
+from .base_node import BaseNodeDisplay
 from .code_execution_node import BaseCodeExecutionNodeDisplay
 from .conditional_node import BaseConditionalNodeDisplay
 from .error_node import BaseErrorNodeDisplay
@@ -27,6 +28,7 @@ __all__ = [
     "BaseInlineSubworkflowNodeDisplay",
     "BaseMapNodeDisplay",
     "BaseMergeNodeDisplay",
+    "BaseNodeDisplay",
     "BaseNoteNodeDisplay",
     "BasePromptDeploymentNodeDisplay",
     "BaseSearchNodeDisplay",

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

@@ -0,0 +1,18 @@
+from typing import Any, Generic, TypeVar
+
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.types.core import JsonObject
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+
+_BaseNodeType = TypeVar("_BaseNodeType", bound=BaseNode)
+
+
+class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeType]):
+    def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
+        node_id = self.node_id
+
+        return {
+            "id": str(node_id),
+            "type": "GENERIC",
+        }

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

@@ -1,5 +1,4 @@
 import pytest
-from unittest import mock
 
 from deepdiff import DeepDiff
 
@@ -21,7 +20,6 @@ from vellum.workflows.expressions.less_than import LessThanExpression
 from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualToExpression
 from vellum.workflows.expressions.not_between import NotBetweenExpression
 from vellum.workflows.expressions.not_in import NotInExpression
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -33,12 +31,7 @@ def test_serialize_workflow():
     # GIVEN a Workflow that uses a ConditionalNode
     # WHEN we serialize it
     workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=CategoryWorkflow)
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -457,23 +450,14 @@ def test_serialize_workflow():
 
     assert not DeepDiff(
         [
-            {
-                "type": "MOCKED",
-            },
-            {
-                "type": "MOCKED",
-            },
-            {
-                "type": "MOCKED",
-            },
-            {
-                "type": "MOCKED",
-            },
-            {
-                "type": "MOCKED",
-            },
+            {"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,
     )
 
     assert not DeepDiff(
@@ -917,12 +901,7 @@ def test_conditional_node_serialize_all_operators_with_lhs_and_rhs(descriptor, o
     workflow_cls = create_simple_workflow(descriptor)
 
     workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=workflow_cls)
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -1041,12 +1020,7 @@ def test_conditional_node_serialize_all_operators_with_expression(descriptor, op
     workflow_cls = create_simple_workflow(descriptor)
 
     workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=workflow_cls)
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -1152,12 +1126,7 @@ def test_conditional_node_serialize_all_operators_with_value_and_start_and_end(d
     workflow_cls = create_simple_workflow(descriptor)
 
     workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=workflow_cls)
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {

{vellum_ai-0.12.14 → vellum_ai-0.12.15}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -16,11 +13,7 @@ def test_serialize_workflow():
         base_display_class=VellumWorkflowDisplay, workflow_class=BasicErrorNodeWorkflow
     )
 
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -137,14 +130,11 @@ def test_serialize_workflow():
 
     assert not DeepDiff(
         [
-            {
-                "type": "MOCKED",
-            },
-            {
-                "type": "MOCKED",
-            },
+            {"id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", "type": "GENERIC"},
+            {"id": "1eee9b4e-531f-45f2-a4b9-42207fac2c33", "type": "GENERIC"},
         ],
         mocked_base_nodes,
+        ignore_order=True,
     )
 
     terminal_node = workflow_raw_data["nodes"][-1]

vellum_ai-0.12.15/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_generic_node_serialization.py

@@ -0,0 +1,174 @@
+from deepdiff import DeepDiff
+
+from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.basic_generic_node.workflow import BasicGenericNodeWorkflow
+
+
+def test_serialize_workflow(vellum_client):
+    # GIVEN a Workflow that uses a generic node
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=BasicGenericNodeWorkflow
+    )
+
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a serialized representation of the Workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "output_variables",
+    }
+
+    # AND its input variables should be what we expect
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 1
+    assert not DeepDiff(
+        [
+            {
+                "id": "a07c2273-34a7-42b5-bcad-143b6127cc8a",
+                "key": "input",
+                "type": "STRING",
+                "default": None,
+                "required": True,
+                "extensions": {"color": None},
+            },
+        ],
+        input_variables,
+        ignore_order=True,
+    )
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 1
+    assert not DeepDiff(
+        [
+            {"id": "2b6389d0-266a-4be4-843e-4e543dd3d727", "key": "output", "type": "STRING"},
+        ],
+        output_variables,
+        ignore_order=True,
+    )
+
+    # AND its raw data should be what we expect
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    assert workflow_raw_data.keys() == {"edges", "nodes", "display_data", "definition"}
+    assert len(workflow_raw_data["edges"]) == 2
+    assert len(workflow_raw_data["nodes"]) == 3
+
+    # AND each node should be serialized correctly
+    entrypoint_node = workflow_raw_data["nodes"][0]
+    assert entrypoint_node == {
+        "id": "f1e4678f-c470-400b-a40e-c8922cc99a86",
+        "type": "ENTRYPOINT",
+        "inputs": [],
+        "data": {"label": "Entrypoint Node", "source_handle_id": "40201804-8beb-43ad-8873-a027759512f1"},
+        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "definition": {
+            "name": "BaseNode",
+            "module": [
+                "vellum",
+                "workflows",
+                "nodes",
+                "bases",
+                "base",
+            ],
+            "bases": [],
+        },
+    }
+
+    api_node = workflow_raw_data["nodes"][1]
+    assert not DeepDiff(
+        {
+            "id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
+            "type": "GENERIC",
+        },
+        api_node,
+    )
+
+    final_output_node = workflow_raw_data["nodes"][2]
+    assert not DeepDiff(
+        {
+            "id": "50e3b446-afcd-4a5d-8c6f-5f05eaf2200e",
+            "type": "TERMINAL",
+            "data": {
+                "label": "Final Output",
+                "name": "output",
+                "target_handle_id": "8bd9f4f3-9f66-4d95-8e84-529b0002c531",
+                "output_id": "2b6389d0-266a-4be4-843e-4e543dd3d727",
+                "output_type": "STRING",
+                "node_input_id": "7a9f2d3a-0b23-4bd4-b567-e9493135b727",
+            },
+            "inputs": [
+                {
+                    "id": "7a9f2d3a-0b23-4bd4-b567-e9493135b727",
+                    "key": "node_input",
+                    "value": {
+                        "rules": [
+                            {
+                                "type": "NODE_OUTPUT",
+                                "data": {
+                                    "node_id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
+                                    "output_id": "0a9c7a80-fc89-4a71-aac0-66489e4ddb85",
+                                },
+                            }
+                        ],
+                        "combinator": "OR",
+                    },
+                }
+            ],
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "definition": {
+                "name": "FinalOutputNode",
+                "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
+                "bases": [
+                    {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"], "bases": []}
+                ],
+            },
+        },
+        final_output_node,
+        ignore_order=True,
+    )
+
+    # AND each edge should be serialized correctly
+    serialized_edges = workflow_raw_data["edges"]
+    assert not DeepDiff(
+        [
+            {
+                "id": "445dd2de-82b2-482b-89f6-5f49d8eb21a9",
+                "source_node_id": "f1e4678f-c470-400b-a40e-c8922cc99a86",
+                "source_handle_id": "40201804-8beb-43ad-8873-a027759512f1",
+                "target_node_id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
+                "target_handle_id": "b7bfb298-959a-4d2b-8b85-bbd0d2522703",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "b741c861-cf67-4649-b9ef-b43a4add72b1",
+                "source_node_id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
+                "source_handle_id": "89dccfa5-cc1a-4612-bd87-86cb444f6dd4",
+                "target_node_id": "50e3b446-afcd-4a5d-8c6f-5f05eaf2200e",
+                "target_handle_id": "8bd9f4f3-9f66-4d95-8e84-529b0002c531",
+                "type": "DEFAULT",
+            },
+        ],
+        serialized_edges,
+        ignore_order=True,
+    )
+
+    # AND the display data should be what we expect
+    display_data = workflow_raw_data["display_data"]
+    assert display_data == {
+        "viewport": {
+            "x": 0.0,
+            "y": 0.0,
+            "zoom": 1.0,
+        }
+    }
+
+    # AND the definition should be what we expect
+    definition = workflow_raw_data["definition"]
+    assert definition == {
+        "name": "BasicGenericNodeWorkflow",
+        "module": ["tests", "workflows", "basic_generic_node", "workflow"],
+    }

{vellum_ai-0.12.14 → vellum_ai-0.12.15}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -15,12 +12,7 @@ def test_serialize_workflow():
     workflow_display = get_workflow_display(
         base_display_class=VellumWorkflowDisplay, workflow_class=BasicInlineSubworkflowWorkflow
     )
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -156,7 +148,7 @@ def test_serialize_workflow():
                                 "bases": [],
                             },
                         },
-                        {"type": "MOCKED"},
+                        {"id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", "type": "GENERIC"},
                         {
                             "id": "a773c3a5-78cb-4250-8d29-7282e8a579d3",
                             "type": "TERMINAL",

{vellum_ai-0.12.14 → vellum_ai-0.12.15}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -13,12 +10,7 @@ def test_serialize_workflow():
     # GIVEN a Workflow that uses a MapNode
     # WHEN we serialize it
     workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=SimpleMapExample)
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -141,7 +133,7 @@ def test_serialize_workflow():
                                 "bases": [],
                             },
                         },
-                        {"type": "MOCKED"},
+                        {"id": "baf6d316-dc75-41e8-96c0-015aede96309", "type": "GENERIC"},
                         {
                             "id": "6f4883b2-70b1-4e1c-ae15-7d0f5aec810b",
                             "type": "TERMINAL",

{vellum_ai-0.12.14 → vellum_ai-0.12.15}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -15,12 +12,7 @@ def test_serialize_workflow__await_all():
     workflow_display = get_workflow_display(
         base_display_class=VellumWorkflowDisplay, workflow_class=AwaitAllPassingWorkflow
     )
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -76,18 +68,12 @@ def test_serialize_workflow__await_all():
         },
     }
 
-    passthrough_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "MOCKED"]
+    passthrough_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "GENERIC"]
     assert not DeepDiff(
         [
-            {
-                "type": "MOCKED",
-            },
-            {
-                "type": "MOCKED",
-            },
-            {
-                "type": "MOCKED",
-            },
+            {"id": "127ef456-91bc-43c6-bd8b-1772db5e3cb5", "type": "GENERIC"},
+            {"id": "59243c65-053f-4ea6-9157-3f3edb1477bf", "type": "GENERIC"},
+            {"id": "634f0202-9ea9-4c62-b152-1a58c595cffb", "type": "GENERIC"},
         ],
         passthrough_nodes,
         ignore_order=True,

{vellum_ai-0.12.14 → vellum_ai-0.12.15}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_try_node_serialization.py

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -13,10 +10,7 @@ def test_serialize_workflow():
     # GIVEN a Workflow with a TryNode
     # WHEN we serialize it
     workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=SimpleTryExample)
-
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -82,4 +76,4 @@ def test_serialize_workflow():
     }
 
     try_node = workflow_raw_data["nodes"][1]
-    assert try_node == {"type": "MOCKED"}
+    assert try_node == {"id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", "type": "GENERIC"}