vellum-ai 0.14.46__py3-none-any.whl → 0.14.48__py3-none-any.whl

vellum_cli/__init__.py

@@ -354,9 +354,10 @@ def images() -> None:
     help="Tags the provided image inside of Vellum's repo. "
     "This field does not push multiple local tags of the passed in image.",
 )
-def image_push(image: str, tag: Optional[List[str]] = None) -> None:
+@click.option("--workspace", type=str, help="The specific Workspace config to use when pushing")
+def image_push(image: str, tag: Optional[List[str]] = None, workspace: Optional[str] = None) -> None:
     """Push Docker image to Vellum"""
-    image_push_command(image, tag)
+    image_push_command(image, tag, workspace)
 
 
 @workflows.command(name="init")

vellum_cli/image_push.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 import re
 import subprocess
 from typing import List, Optional
@@ -9,15 +10,26 @@ from docker import DockerClient
 from dotenv import load_dotenv
 
 from vellum.workflows.vellum_client import create_vellum_client, create_vellum_environment
+from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, load_vellum_cli_config
 from vellum_cli.logger import load_cli_logger
 
 _SUPPORTED_ARCHITECTURE = "amd64"
 
 
-def image_push_command(image: str, tags: Optional[List[str]] = None) -> None:
-    load_dotenv()
+def image_push_command(image: str, tags: Optional[List[str]] = None, workspace: Optional[str] = None) -> None:
+    load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
     logger = load_cli_logger()
-    vellum_client = create_vellum_client()
+    config = load_vellum_cli_config()
+    workspace_config = next((w for w in config.workspaces if w.name == workspace), DEFAULT_WORKSPACE_CONFIG)
+
+    api_key = os.getenv(workspace_config.api_key, None)
+    if not api_key:
+        raise ValueError(f"API key {workspace_config.api_key} for workspace {workspace} not found")
+
+    vellum_client = create_vellum_client(
+        api_key=api_key,
+        api_url=workspace_config.api_url,
+    )
 
     # Check if we are self hosted by looking at our base url
     api_url = create_vellum_environment().default

vellum_cli/tests/test_image_push.py

@@ -1,16 +1,37 @@
+import pytest
+import json
+import os
+import shutil
 import subprocess
+import tempfile
 from unittest.mock import MagicMock, patch
+from uuid import uuid4
+from typing import Generator
 
 from click.testing import CliRunner
+from httpx import Response
 
 from vellum_cli import main as cli_main
 
 
+@pytest.fixture
+def mock_temp_dir() -> Generator[str, None, None]:
+    current_dir = os.getcwd()
+    temp_dir = tempfile.mkdtemp()
+    os.chdir(temp_dir)
+
+    yield temp_dir
+
+    os.chdir(current_dir)
+    shutil.rmtree(temp_dir)
+
+
 @patch("subprocess.run")
 @patch("docker.from_env")
 def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
     # GIVEN a self hosted vellum api URL env var
     monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
 
     # Mock Docker client
     mock_docker_client = MagicMock()
@@ -35,6 +56,94 @@ def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_run, vell
     assert "Image successfully pushed" in result.output
 
 
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push__self_hosted_happy_path__workspace_option(
+    mock_docker_from_env, mock_run, mock_httpx_transport, mock_temp_dir
+):
+    # GIVEN a workspace config with a new env for url
+    with open(os.path.join(mock_temp_dir, "vellum.lock.json"), "w") as f:
+        f.write(
+            json.dumps(
+                {
+                    "workspaces": [
+                        {
+                            "name": "my_workspace",
+                            "api_url": "MY_WORKSPACE_VELLUM_API_URL",
+                            "api_key": "MY_WORKSPACE_VELLUM_API_KEY",
+                        }
+                    ]
+                }
+            )
+        )
+
+    # AND a .env file with the workspace api key and url
+    with open(os.path.join(mock_temp_dir, ".env"), "w") as f:
+        f.write(
+            "VELLUM_API_KEY=123456abcdef\n"
+            "VELLUM_API_URL=https://api.vellum.ai\n"
+            "MY_WORKSPACE_VELLUM_API_KEY=789012ghijkl\n"
+            "MY_WORKSPACE_VELLUM_API_URL=https://api.vellum.mycompany.ai\n"
+        )
+
+    # AND the Docker client returns the correct response
+    mock_docker_client = MagicMock()
+    mock_docker_from_env.return_value = mock_docker_client
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"sha256:hellosha"),
+    ]
+
+    # AND the vellum client returns the correct response for
+    mock_httpx_transport.handle_request.side_effect = [
+        # First call to get the docker service token
+        Response(
+            status_code=200,
+            text=json.dumps(
+                {
+                    "access_token": "345678mnopqr",
+                    "organization_id": str(uuid4()),
+                    "repository": "myrepo.net",
+                }
+            ),
+        ),
+        # Second call to push the image
+        Response(
+            status_code=200,
+            text=json.dumps(
+                {
+                    "id": str(uuid4()),
+                    "name": "myrepo.net/myimage",
+                    "visibility": "PRIVATE",
+                    "created": "2021-01-01T00:00:00Z",
+                    "modified": "2021-01-01T00:00:00Z",
+                    "repository": "myrepo.net",
+                    "sha": "sha256:hellosha",
+                    "tags": [],
+                }
+            ),
+        ),
+    ]
+
+    # WHEN the user runs the image push command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myrepo.net/myimage:latest", "--workspace", "my_workspace"])
+
+    # THEN the command exits successfully
+    assert result.exit_code == 0, (result.output, str(result.exception))
+
+    # AND gives the success message
+    assert "Image successfully pushed" in result.output
+
+    # AND the vellum client was called with the correct api key and url
+    request = mock_httpx_transport.handle_request.call_args[0][0]
+    assert request.headers["X-API-KEY"] == "789012ghijkl", result.stdout
+    assert str(request.url) == "https://api.vellum.mycompany.ai/v1/container-images/push"
+
+
 @patch("subprocess.run")
 @patch("docker.from_env")
 def test_image_push__self_hosted_blocks_repo(mock_docker_from_env, mock_run, vellum_client, monkeypatch):

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -9,6 +9,7 @@ from typing import (
     ForwardRef,
     Generic,
     Optional,
+    Set,
     Tuple,
     Type,
     TypeVar,
@@ -24,6 +25,7 @@ from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.utils import get_unadorned_node, get_wrapped_node
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference
+from vellum.workflows.references.node import NodeReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.types.utils import get_original_base
@@ -33,7 +35,7 @@ from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_var
 from vellum_ee.workflows.display.editor.types import NodeDisplayComment, NodeDisplayData
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
-from vellum_ee.workflows.display.utils.expressions import serialize_condition, serialize_value
+from vellum_ee.workflows.display.utils.expressions import serialize_value
 from vellum_ee.workflows.display.utils.registry import register_node_display_class
 
 if TYPE_CHECKING:
@@ -98,11 +100,18 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
     # Default values set by the metaclass
     output_display: Dict[OutputReference, NodeOutputDisplay]
     port_displays: Dict[Port, PortDisplayOverrides] = {}
-    node_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
+    attribute_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
+    # START: Attributes for backwards compatible serialization
     # Used to explicitly set the target handle id for a node
     # Once all nodes are Generic Nodes, we may replace this with a trigger_id or trigger attribute
     target_handle_id: ClassVar[Optional[UUID]] = None
+    # Used to explicitly set the input ids for each node input
+    node_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
+    # Used by each class extending BaseNodeDisplay to specify which attributes are meant to be serialized
+    # as the former `"inputs"` field
+    __serializable_inputs__: Set[NodeReference] = set()
+    # END: Attributes for backwards compatible serialization
 
     def serialize(self, display_context: "WorkflowDisplayContext", **kwargs: Any) -> JsonObject:
         node = self._node
@@ -114,7 +123,11 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 # We don't need to serialize generic node attributes containing a subworkflow
                 continue
 
-            id = str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
+            id = (
+                str(self.attribute_ids_by_name[attribute.name])
+                if self.attribute_ids_by_name
+                else str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
+            )
             try:
                 attributes.append(
                     {
@@ -191,9 +204,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                         "id": id,
                         "name": port.name,
                         "type": port._condition_type.value,
-                        "expression": (
-                            serialize_condition(display_context, port._condition) if port._condition else None
-                        ),
+                        "expression": (serialize_value(display_context, port._condition) if port._condition else None),
                     }
                 )
             else:
@@ -249,10 +260,15 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
     def get_source_handle_id(self, port_displays: Dict[Port, PortDisplay]) -> UUID:
         unadorned_node = get_unadorned_node(self._node)
-        default_port = unadorned_node.Ports.default
+        default_port = next((port for port in unadorned_node.Ports if port.default), None)
+        if default_port in port_displays:
+            return port_displays[default_port].id
+
+        first_port = next((port for port in unadorned_node.Ports), None)
+        if not first_port:
+            raise ValueError(f"Node {self._node.__name__} must have at least one port.")
 
-        default_port_display = port_displays[default_port]
-        return default_port_display.id
+        return port_displays[first_port].id
 
     def get_trigger_id(self) -> UUID:
         return self.get_target_handle_id()

vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -2,7 +2,6 @@ import types
 from uuid import UUID
 from typing import TYPE_CHECKING, Any, Dict, Generic, Type, TypeVar
 
-from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.registry import get_from_node_display_registry
@@ -30,14 +29,14 @@ def get_node_display_class(node_class: Type[NodeType]) -> Type["BaseNodeDisplay"
                 node_input_ids_by_name.update(_get_node_input_ids_by_ref(f"{path}.{key}", value))
             return node_input_ids_by_name
 
-        if isinstance(inst, BaseDescriptor):
-            return {path: uuid4_from_hash(f"{node_class.__id__}|{path}")}
-
-        return {}
+        return {path: uuid4_from_hash(f"{node_class.__id__}|{path}")}
 
     def exec_body(ns: Dict):
         node_input_ids_by_name: Dict[str, UUID] = {}
         for ref in node_class:
+            if ref not in base_node_display_class.__serializable_inputs__:
+                continue
+
             node_input_ids_by_name.update(_get_node_input_ids_by_ref(ref.name, ref.instance))
 
         if node_input_ids_by_name:

vellum_ee/workflows/display/nodes/vellum/api_node.py

@@ -20,6 +20,17 @@ class BaseAPINodeDisplay(BaseNodeDisplay[_APINodeType], Generic[_APINodeType]):
     # A mapping between node input keys and their ids for inputs representing additional header values
     additional_header_value_input_ids: ClassVar[Optional[Dict[str, UUID]]] = None
 
+    __serializable_inputs__ = {
+        APINode.url,
+        APINode.method,
+        APINode.json,
+        APINode.headers,
+        APINode.api_key_header_key,
+        APINode.api_key_header_value,
+        APINode.bearer_token_value,
+        APINode.authorization_type,
+    }
+
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
     ) -> JsonObject:

vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -18,6 +18,11 @@ class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Gene
     output_id: ClassVar[Optional[UUID]] = None
     log_output_id: ClassVar[Optional[UUID]] = None
 
+    __serializable_inputs__ = {
+        CodeExecutionNode.code,
+        CodeExecutionNode.code_inputs,
+    }
+
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

vellum_ee/workflows/display/nodes/vellum/error_node.py

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import ClassVar, Generic, Optional, TypeVar
+from typing import Any, ClassVar, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes import ErrorNode
 from vellum.workflows.types.core import JsonObject
@@ -9,47 +9,53 @@ from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _ErrorNodeType = TypeVar("_ErrorNodeType", bound=ErrorNode)
+LEGACY_INPUT_NAME = "error_source_input_id"
 
 
 class BaseErrorNodeDisplay(BaseNodeDisplay[_ErrorNodeType], Generic[_ErrorNodeType]):
+    # DEPRECATED: Remove in 0.15.0 once removed from the vellum-side
     error_output_id: ClassVar[Optional[UUID]] = None
+    # DEPRECATED: Remove in 0.15.0 once removed from the vellum-side
+    name: ClassVar[Optional[str]] = None
 
-    name: ClassVar[str] = "error-node"
+    __serializable_inputs__ = {ErrorNode.error}
 
-    def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
-    ) -> JsonObject:
+    def serialize(self, display_context: WorkflowDisplayContext, **kwargs) -> JsonObject:
         node_id = self.node_id
-        error_source_input_id = self.node_input_ids_by_name.get("error_source_input_id")
+        error_source_input_id = self.node_input_ids_by_name.get(
+            ErrorNode.error.name,
+        ) or self.node_input_ids_by_name.get(LEGACY_INPUT_NAME)
 
         error_attribute = raise_if_descriptor(self._node.error)
-        input_values_by_name = {
-            "error_source_input_id": error_attribute,
-        }
 
         node_inputs = [
             create_node_input(
                 node_id=node_id,
-                input_name=variable_name,
-                value=variable_value,
+                input_name=LEGACY_INPUT_NAME,
+                value=error_attribute,
                 display_context=display_context,
-                input_id=self.node_input_ids_by_name.get(variable_name),
+                input_id=error_source_input_id,
             )
-            for variable_name, variable_value in input_values_by_name.items()
         ]
 
-        return {
+        node_data: dict[str, Any] = {
             "id": str(node_id),
             "type": "ERROR",
             "inputs": [node_input.dict() for node_input in node_inputs],
             "data": {
-                "name": self.name,
                 "label": self.label,
                 "target_handle_id": str(self.get_target_handle_id()),
                 "error_source_input_id": str(error_source_input_id),
-                "error_output_id": str(self.error_output_id),
             },
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
         }
+
+        if self.name:
+            node_data["data"]["name"] = self.name
+
+        if self.error_output_id:
+            node_data["data"]["error_output_id"] = str(self.error_output_id)
+
+        return node_data

vellum_ee/workflows/display/nodes/vellum/guardrail_node.py

@@ -12,6 +12,8 @@ _GuardrailNodeType = TypeVar("_GuardrailNodeType", bound=GuardrailNode)
 
 
 class BaseGuardrailNodeDisplay(BaseNodeDisplay[_GuardrailNodeType], Generic[_GuardrailNodeType]):
+    __serializable_inputs__ = {GuardrailNode.metric_inputs}
+
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py

@@ -10,6 +10,7 @@ from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.expressions import serialize_value
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.vellum import NodeInput
 
@@ -17,6 +18,8 @@ _InlinePromptNodeType = TypeVar("_InlinePromptNodeType", bound=InlinePromptNode)
 
 
 class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generic[_InlinePromptNodeType]):
+    __serializable_inputs__ = {InlinePromptNode.prompt_inputs}
+
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
@@ -26,12 +29,14 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
         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.global_node_output_displays[node.Outputs.text]
-        _, array_display = display_context.global_node_output_displays[node.Outputs.results]
-        _, json_display = display_context.global_node_output_displays[node.Outputs.json]
-        node_blocks = raise_if_descriptor(node.blocks)
+        output_display = self.output_display[node.Outputs.text]
+        array_display = self.output_display[node.Outputs.results]
+        json_display = self.output_display[node.Outputs.json]
+        node_blocks = raise_if_descriptor(node.blocks) or []
         function_definitions = raise_if_descriptor(node.functions)
 
+        ml_model = str(raise_if_descriptor(node.ml_model))
+
         blocks: list = [
             self._generate_prompt_block(block, input_variable_id_by_name, [i]) for i, block in enumerate(node_blocks)
         ]
@@ -42,7 +47,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
         )
         blocks.extend(functions)
 
-        return {
+        serialized_node: JsonObject = {
             "id": str(node_id),
             "type": "PROMPT",
             "inputs": [node_input.dict() for node_input in node_inputs],
@@ -62,7 +67,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
                         "blocks": blocks,
                     },
                 },
-                "ml_model_name": raise_if_descriptor(node.ml_model),
+                "ml_model_name": ml_model,
             },
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
@@ -74,6 +79,10 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
             ],
             "ports": self.serialize_ports(display_context),
         }
+        attributes = self._serialize_attributes(display_context)
+        if attributes:
+            serialized_node["attributes"] = attributes
+        return serialized_node
 
     def _generate_node_and_prompt_inputs(
         self,
@@ -127,6 +136,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
         path: List[int],
     ) -> JsonObject:
         block: JsonObject
+        block_id = uuid4_from_hash(f"{self.node_id}-{prompt_block.block_type}-{'-'.join([str(i) for i in path])}")
         if prompt_block.block_type == "JINJA":
             block = {
                 "block_type": "JINJA",
@@ -162,10 +172,21 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
             }
 
         elif prompt_block.block_type == "VARIABLE":
-            block = {
-                "block_type": "VARIABLE",
-                "input_variable_id": input_variable_id_by_name[prompt_block.input_variable],
-            }
+            input_variable_id = input_variable_id_by_name.get(prompt_block.input_variable)
+            if input_variable_id:
+                block = {
+                    "block_type": "VARIABLE",
+                    "input_variable_id": input_variable_id,
+                }
+            else:
+                # Even though this will likely fail in runtime, we want to allow serialization to succeed
+                # in case the block is work in progress or the node is not yet part of the graph
+                block = {
+                    "block_type": "VARIABLE",
+                    "input_variable_id": str(
+                        uuid4_from_hash(f"{block_id}-input_variable-{prompt_block.input_variable}")
+                    ),
+                }
 
         elif prompt_block.block_type == "PLAIN_TEXT":
             block = {
@@ -184,9 +205,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
         else:
             raise NotImplementedError(f"Serialization for prompt block type {prompt_block.block_type} not implemented")
 
-        block["id"] = str(
-            uuid4_from_hash(f"{self.node_id}-{prompt_block.block_type}-{'-'.join([str(i) for i in path])}")
-        )
+        block["id"] = str(block_id)
         if prompt_block.cache_config:
             block["cache_config"] = prompt_block.cache_config.dict()
         else:
@@ -198,3 +217,18 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
             block["state"] = "ENABLED"
 
         return block
+
+    def _serialize_attributes(self, display_context: "WorkflowDisplayContext"):
+        attribute_instances_by_name = {}
+        for attribute in self._node:
+            if attribute.name in self.attribute_ids_by_name:
+                attribute_instances_by_name[attribute.name] = attribute.instance
+
+        return [
+            {
+                "id": str(attr_id),
+                "name": attr_name,
+                "value": serialize_value(display_context, attribute_instances_by_name[attr_name]),
+            }
+            for attr_name, attr_id in self.attribute_ids_by_name.items()
+        ]

vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py

@@ -21,6 +21,8 @@ class BaseInlineSubworkflowNodeDisplay(
 ):
     workflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
+    __serializable_inputs__ = {InlineSubworkflowNode.subworkflow_inputs}
+
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

vellum_ee/workflows/display/nodes/vellum/map_node.py

@@ -13,6 +13,8 @@ _MapNodeType = TypeVar("_MapNodeType", bound=MapNode)
 
 
 class BaseMapNodeDisplay(BaseAdornmentNodeDisplay[_MapNodeType], Generic[_MapNodeType]):
+    __serializable_inputs__ = {MapNode.items}  # type: ignore[misc]
+
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py

@@ -13,6 +13,8 @@ _PromptDeploymentNodeType = TypeVar("_PromptDeploymentNodeType", bound=PromptDep
 
 
 class BasePromptDeploymentNodeDisplay(BaseNodeDisplay[_PromptDeploymentNodeType], Generic[_PromptDeploymentNodeType]):
+    __serializable_inputs__ = {PromptDeploymentNode.prompt_inputs}
+
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

vellum_ee/workflows/display/nodes/vellum/search_node.py

@@ -32,6 +32,14 @@ class BaseSearchNodeDisplay(BaseNodeDisplay[_SearchNodeType], Generic[_SearchNod
     # A mapping between the id of the operand (e.g. "lhs_variable_id" or "rhs_variable_id") and the id of the node input
     # that the operand is pointing to.
     metadata_filter_input_id_by_operand_id: Dict[UUID, UUID] = {}
+    __serializable_inputs__ = {
+        SearchNode.query,
+        SearchNode.document_index,
+        SearchNode.weights,
+        SearchNode.chunk_separator,
+        SearchNode.limit,
+        SearchNode.result_merging,
+    }
 
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs

vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py

@@ -14,6 +14,7 @@ _SubworkflowDeploymentNodeType = TypeVar("_SubworkflowDeploymentNodeType", bound
 class BaseSubworkflowDeploymentNodeDisplay(
     BaseNodeDisplay[_SubworkflowDeploymentNodeType], Generic[_SubworkflowDeploymentNodeType]
 ):
+    __serializable_inputs__ = {SubworkflowDeploymentNode.subworkflow_inputs}
 
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs

vellum_ee/workflows/display/nodes/vellum/templating_node.py

@@ -15,6 +15,8 @@ TEMPLATE_INPUT_NAME = TemplatingNode.template.name
 
 
 class BaseTemplatingNodeDisplay(BaseNodeDisplay[_TemplatingNodeType], Generic[_TemplatingNodeType]):
+    __serializable_inputs__ = {TemplatingNode.inputs}
+
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

vellum_ee/workflows/display/nodes/vellum/tests/test_code_execution_node.py

@@ -29,7 +29,7 @@ def _display_class_with_node_input_ids_by_name_with_inputs_prefix(Node: Type[Cod
 @pytest.mark.parametrize(
     ["GetDisplayClass", "expected_input_id"],
     [
-        (_no_display_class, "e3cdb222-324e-4ad1-abb2-bdd7881b3a0e"),
+        (_no_display_class, "a5dbe403-0b00-4df6-b8f7-ed5f7794b003"),
         (_display_class_with_node_input_ids_by_name, "fba6a4d5-835a-4e99-afb7-f6a4aed15110"),
         (_display_class_with_node_input_ids_by_name_with_inputs_prefix, "fba6a4d5-835a-4e99-afb7-f6a4aed15110"),
     ],

vellum_ee/workflows/display/nodes/vellum/tests/test_error_node.py

@@ -41,3 +41,7 @@ def test_error_node_display__serialize_with_vellum_error() -> None:
             }
         ],
     }
+
+    # AND we serialize the DEPRECATED fields
+    assert "error_output_id" not in serialized_node["data"]
+    assert "name" not in serialized_node["data"]

vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_deployment_node.py

@@ -5,6 +5,7 @@ from typing import Type
 
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes import PromptDeploymentNode
+from vellum_ee.workflows.display.nodes.vellum.prompt_deployment_node import BasePromptDeploymentNodeDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
@@ -13,14 +14,14 @@ def _no_display_class(Node: Type[PromptDeploymentNode]):  # type: ignore
 
 
 def _display_class_with_node_input_ids_by_name(Node: Type[PromptDeploymentNode]):
-    class PromptDeploymentNodeDisplay(PromptDeploymentNode[Node]):  # type: ignore[valid-type]
+    class PromptDeploymentNodeDisplay(BasePromptDeploymentNodeDisplay[Node]):  # type: ignore[valid-type]
         node_input_ids_by_name = {"foo": UUID("6037747a-1d35-4094-b363-4369fc92c5d4")}
 
     return PromptDeploymentNodeDisplay
 
 
 def _display_class_with_node_input_ids_by_name_with_inputs_prefix(Node: Type[PromptDeploymentNode]):
-    class PromptDeploymentNodeDisplay(PromptDeploymentNode[Node]):  # type: ignore[valid-type]
+    class PromptDeploymentNodeDisplay(BasePromptDeploymentNodeDisplay[Node]):  # type: ignore[valid-type]
         node_input_ids_by_name = {"prompt_inputs.foo": UUID("6037747a-1d35-4094-b363-4369fc92c5d4")}
 
     return PromptDeploymentNodeDisplay
@@ -51,7 +52,7 @@ def mock_fetch_deployment(mocker):
 @pytest.mark.parametrize(
     ["GetDisplayClass", "expected_input_id"],
     [
-        (_no_display_class, "6037747a-1d35-4094-b363-4369fc92c5d4"),
+        (_no_display_class, "016187d6-2830-4256-a61d-e52f9bf6355e"),
         (_display_class_with_node_input_ids_by_name, "6037747a-1d35-4094-b363-4369fc92c5d4"),
         (_display_class_with_node_input_ids_by_name_with_inputs_prefix, "6037747a-1d35-4094-b363-4369fc92c5d4"),
     ],