vellum-ai 0.14.14__py3-none-any.whl → 0.14.16__py3-none-any.whl

vellum_cli/image_push.py

@@ -8,7 +8,7 @@ import docker
 from docker import DockerClient
 from dotenv import load_dotenv
 
-from vellum.workflows.vellum_client import create_vellum_client
+from vellum.workflows.vellum_client import create_vellum_client, create_vellum_environment
 from vellum_cli.logger import load_cli_logger
 
 _SUPPORTED_ARCHITECTURE = "amd64"
@@ -19,6 +19,21 @@ def image_push_command(image: str, tags: Optional[List[str]] = None) -> None:
     logger = load_cli_logger()
     vellum_client = create_vellum_client()
 
+    # Check if we are self hosted by looking at our base url
+    api_url = create_vellum_environment().default
+    is_self_hosted = not api_url.endswith("vellum.ai") and "localhost:" not in api_url and "127.0.0.1" not in api_url
+    if is_self_hosted:
+        logger.info("Self hosted install detected...")
+
+    if is_self_hosted and "/" not in image:
+        logger.error(
+            "For adding images to your self hosted install you must include the "
+            "repository address in the provided image name. You must also have "
+            "already pushed the image to the docker repository that your self "
+            "hosted install is using."
+        )
+        exit(1)
+
     # We're using docker python SDK here instead of subprocess since it connects to the docker host directly
     # instead of using the command line so it seemed like it would possibly be a little more robust since
     # it might avoid peoples' wonky paths, unfortunately it doesn't support the manifest command which we need for
@@ -29,59 +44,78 @@ def image_push_command(image: str, tags: Optional[List[str]] = None) -> None:
     docker_client = docker.from_env()
     check_architecture(docker_client, image, logger)
 
-    logger.info("Authenticating...")
-    auth = vellum_client.container_images.docker_service_token()
-
-    docker_client.login(
-        username="oauth2accesstoken",
-        password=auth.access_token,
-        registry=auth.repository,
-    )
-
     repo_split = image.split("/")
     tag_split = repo_split[-1].split(":")
     image_name = tag_split[0]
     main_tag = tag_split[1] if len(tag_split) > 1 else "latest"
-
     all_tags = [main_tag, *(tags or [])]
-    for tag in all_tags:
-        vellum_image_name = f"{auth.repository}/{image_name}:{tag}"
-
-        docker_client.api.tag(image, vellum_image_name)
-
-        push_result = docker_client.images.push(repository=vellum_image_name, stream=True)
-
-        # Here were trying to mime the output you would get from a normal docker push, which
-        # the python sdk makes as hard as possible.
-        for raw_line in push_result:
-            try:
-                for sub_line in raw_line.decode("utf-8").split("\r\n"):
-                    line = json.loads(sub_line)
-                    error_message = line.get("errorDetail", {}).get("message")
-                    status = line.get("status")
-                    id = line.get("id", "")
-
-                    if error_message:
-                        logger.error(error_message)
-                        exit(1)
-                    elif status == "Waiting":
-                        continue
-                    elif status:
-                        logger.info(f"{id}{': ' if id else ''}{status}")
-                    else:
-                        logger.info(line)
-            except Exception:
-                continue
 
-    result = subprocess.run(
+    if is_self_hosted:
+        # Include the repo for self hosted since we skip pushing it to our internal repo and expect
+        # the user the push it themselves and provide us with the repo name.
+        image_name = image.split(":")[0]
+
+        logger.info("Checking if image already exists in repository...")
+        exists_result = subprocess.run(
+            ["docker", "manifest", "inspect", image],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+
+        if exists_result.returncode != 0:
+            logger.error(
+                "Image does not exist in repository. Push the image to the "
+                "repository your self hosted install is using and try again."
+            )
+            exit(1)
+    else:
+        logger.info("Authenticating...")
+        auth = vellum_client.container_images.docker_service_token()
+
+        docker_client.login(
+            username="oauth2accesstoken",
+            password=auth.access_token,
+            registry=auth.repository,
+        )
+
+        for tag in all_tags:
+            vellum_image_name = f"{auth.repository}/{image_name}:{tag}"
+
+            docker_client.api.tag(image, vellum_image_name)
+
+            push_result = docker_client.images.push(repository=vellum_image_name, stream=True)
+
+            # Here were trying to mime the output you would get from a normal docker push, which
+            # the python sdk makes as hard as possible.
+            for raw_line in push_result:
+                try:
+                    for sub_line in raw_line.decode("utf-8").split("\r\n"):
+                        line = json.loads(sub_line)
+                        error_message = line.get("errorDetail", {}).get("message")
+                        status = line.get("status")
+                        id = line.get("id", "")
+
+                        if error_message:
+                            logger.error(error_message)
+                            exit(1)
+                        elif status == "Waiting":
+                            continue
+                        elif status:
+                            logger.info(f"{id}{': ' if id else ''}{status}")
+                        else:
+                            logger.info(line)
+                except Exception:
+                    continue
+
+    inspect_result = subprocess.run(
         ["docker", "inspect", "--format='{{index .RepoDigests 0}}'", image],
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
     )
 
     sha = ""
-    if result.returncode == 0:
-        match = re.search(r"sha256[^']*", result.stdout.decode("utf-8"))
+    if inspect_result.returncode == 0:
+        match = re.search(r"sha256[^']*", inspect_result.stdout.decode("utf-8"))
         if match and match.group(0):
             sha = match.group(0)
 

vellum_cli/tests/test_image_push.py

@@ -0,0 +1,56 @@
+import subprocess
+from unittest.mock import MagicMock, patch
+
+from click.testing import CliRunner
+
+from vellum_cli import main as cli_main
+
+
+@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")
+
+    # Mock Docker client
+    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"manifest"),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"sha256:hellosha"),
+    ]
+
+    # WHEN the user runs the image push command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myrepo.net/myimage:latest"])
+
+    # THEN the command exits successfully
+    assert result.exit_code == 0, result.output
+
+    # AND gives the success message
+    assert "Image successfully pushed" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push__self_hosted_blocks_repo(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")
+
+    # Mock Docker client
+    mock_docker_client = MagicMock()
+    mock_docker_from_env.return_value = mock_docker_client
+
+    # WHEN the user runs the image push command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage"])
+
+    # THEN the command exits unsuccessfully
+    assert result.exit_code == 1, result.output
+
+    # AND gives the error message for self hosted installs not including the repo
+    assert "For adding images to your self hosted install you must include" in result.output

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -4,6 +4,7 @@ from uuid import UUID
 from typing import (
     TYPE_CHECKING,
     Any,
+    ClassVar,
     Dict,
     ForwardRef,
     Generic,
@@ -59,17 +60,20 @@ _NodeDisplayAttrType = TypeVar("_NodeDisplayAttrType")
 class BaseNodeDisplayMeta(type):
     def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
         cls = super().__new__(mcs, name, bases, dct)
+        return cls.__annotate_node__()
+
+    def __annotate_node__(cls):
         base_node_display_class = cast(Type["BaseNodeDisplay"], cls)
         node_class = base_node_display_class.infer_node_class()
         if not issubclass(node_class, BaseNode):
             return cls
 
-        display_node_id = dct.get("node_id")
+        display_node_id = getattr(cls, "node_id", None)
         if isinstance(display_node_id, UUID):
             # Display classes are able to override the id of the node class it's parameterized by
             node_class.__id__ = display_node_id
 
-        output_display = dct.get("output_display")
+        output_display = getattr(cls, "output_display", None)
         if isinstance(output_display, dict):
             # And the node class' output ids
             for reference, node_output_display in output_display.items():
@@ -86,6 +90,7 @@ class BaseNodeDisplayMeta(type):
 class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
     output_display: Dict[OutputReference, NodeOutputDisplay] = {}
     port_displays: Dict[Port, PortDisplayOverrides] = {}
+    node_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     # Used to store the mapping between node types and their display classes
     _node_display_registry: Dict[Type[NodeType], Type["BaseNodeDisplay"]] = {}
@@ -326,6 +331,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 IsNotNilExpression,
                 IsUndefinedExpression,
                 IsNotUndefinedExpression,
+                ParseJsonExpression,
             ),
         ):
             lhs = self.serialize_value(display_context, condition._expression)
@@ -403,9 +409,6 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 "node_id": str(node_class_display.node_id),
             }
 
-        if isinstance(value, ParseJsonExpression):
-            raise ValueError("ParseJsonExpression is not supported in the UI")
-
         if not isinstance(value, BaseDescriptor):
             vellum_value = primitive_to_vellum_value(value)
             return {

vellum_ee/workflows/display/nodes/base_node_vellum_display.py

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import ClassVar, Dict, Optional, Union
+from typing import ClassVar, Dict, Optional
 
 from vellum.workflows.nodes.utils import get_unadorned_node
 from vellum.workflows.ports import Port
@@ -17,9 +17,6 @@ class BaseNodeVellumDisplay(BaseNodeDisplay[NodeType]):
     # Used to explicitly set the target handle id for a node
     target_handle_id: ClassVar[Optional[UUID]] = None
 
-    # Used to explicitly set the node input ids by name for a node
-    node_input_ids_by_name: ClassVar[Dict[str, Union[UUID, str]]] = {}
-
     def _get_node_display_uuid(self, attribute: str) -> UUID:
         explicit_value = self._get_explicit_node_display_attr(attribute, UUID)
         return explicit_value if explicit_value else uuid4_from_hash(f"{self.node_id}|{attribute}")

vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -1,7 +1,10 @@
 import types
-from typing import TYPE_CHECKING, Optional, Type
+from uuid import UUID
+from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 
+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.nodes.types import NodeOutputDisplay
 
 if TYPE_CHECKING:
@@ -26,16 +29,39 @@ def get_node_display_class(
 
     # `base_node_display_class` is always a Generic class, so it's safe to index into it
     NodeDisplayBaseClass = base_node_display_class[node_class]  # type: ignore[index]
+
+    def _get_node_input_ids_by_ref(path: str, inst: Any):
+        if isinstance(inst, dict):
+            node_input_ids_by_name: Dict[str, UUID] = {}
+            for key, value in inst.items():
+                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 {}
+
+    def exec_body(ns: Dict):
+        output_display = {
+            ref: NodeOutputDisplay(id=node_class.__output_ids__[ref.name], name=ref.name)
+            for ref in node_class.Outputs
+            if ref.name in node_class.__output_ids__
+        }
+        if output_display:
+            ns["output_display"] = output_display
+
+        node_input_ids_by_name: Dict[str, UUID] = {}
+        for ref in node_class:
+            node_input_ids_by_name.update(_get_node_input_ids_by_ref(ref.name, ref.instance))
+
+        if node_input_ids_by_name:
+            ns["node_input_ids_by_name"] = node_input_ids_by_name
+
     NodeDisplayClass = types.new_class(
         f"{node_class.__name__}Display",
         bases=(NodeDisplayBaseClass,),
+        exec_body=exec_body,
     )
-    output_display = {
-        ref: NodeOutputDisplay(id=node_class.__output_ids__[ref.name], name=ref.name)
-        for ref in node_class.Outputs
-        if ref.name in node_class.__output_ids__
-    }
-    if output_display:
-        setattr(NodeDisplayClass, "output_display", output_display)
 
     return NodeDisplayClass

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

@@ -1,9 +1,14 @@
+import re
+import types
 from uuid import UUID
-from typing import Any, Callable, Generic, Optional, TypeVar, cast
+from typing import Any, Callable, Dict, Generic, Optional, Type, TypeVar, cast
 
+from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
 from vellum.workflows.nodes.utils import get_wrapped_node
 from vellum.workflows.types.core import JsonArray, JsonObject
+from vellum.workflows.types.utils import get_original_base
+from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
@@ -37,3 +42,52 @@ class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Ge
         serialized_wrapped_node["adornments"] = adornments + [adornment] if adornment else adornments
 
         return serialized_wrapped_node
+
+    @classmethod
+    def wrap(cls, **kwargs: Any) -> Callable[..., Type[BaseNodeDisplay]]:
+        NodeDisplayType = TypeVar("NodeDisplayType", bound=BaseNodeDisplay)
+
+        def decorator(inner_cls: Type[NodeDisplayType]) -> Type[NodeDisplayType]:
+            node_class = inner_cls.infer_node_class()
+            wrapped_node_class = cast(Type[BaseNode], node_class.__wrapped_node__)
+
+            # `mypy` is wrong here, `cls` is indexable bc it's Generic
+            BaseAdornmentDisplay = cls[node_class]  # type: ignore[index]
+
+            def exec_body(ns: Dict):
+                for key, kwarg in kwargs.items():
+                    ns[key] = kwarg
+
+                if "node_id" not in kwargs:
+                    ns["node_id"] = uuid4_from_hash(node_class.__qualname__)
+
+            AdornmentDisplay = types.new_class(
+                re.sub(r"^Base", "", cls.__name__), bases=(BaseAdornmentDisplay,), exec_body=exec_body
+            )
+
+            setattr(inner_cls, "__adorned_by__", AdornmentDisplay)
+
+            # We must edit the node display class to use __wrapped_node__ everywhere it
+            # references the adorned node class, which is three places:
+
+            # 1. The node display class' parameterized type
+            original_base_node_display = get_original_base(inner_cls)
+            original_base_node_display.__args__ = (wrapped_node_class,)
+            inner_cls._node_display_registry[wrapped_node_class] = inner_cls
+            inner_cls.__annotate_node__()
+
+            # 2. The node display class' output displays
+            old_outputs = list(inner_cls.output_display.keys())
+            for old_output in old_outputs:
+                new_output = getattr(wrapped_node_class.Outputs, old_output.name)
+                inner_cls.output_display[new_output] = inner_cls.output_display.pop(old_output)
+
+            # 3. The node display class' port displays
+            old_ports = list(inner_cls.port_displays.keys())
+            for old_port in old_ports:
+                new_port = getattr(wrapped_node_class.Ports, old_port.name)
+                inner_cls.port_displays[new_port] = inner_cls.port_displays.pop(old_port)
+
+            return inner_cls
+
+        return decorator

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

@@ -1,14 +1,12 @@
 import inspect
-from typing import Any, Callable, Generic, Optional, Tuple, Type, TypeVar, cast
+from typing import Any, Generic, Tuple, Type, TypeVar, cast
 
 from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.retry_node.node import RetryNode
 from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
-from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
@@ -79,54 +77,3 @@ class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_Re
 
         inner_output = getattr(inner_node.Outputs, output.name)
         return node_display.get_node_output_display(inner_output)
-
-    @classmethod
-    def wrap(
-        cls,
-        max_attempts: int,
-        delay: Optional[float] = None,
-        retry_on_error_code: Optional[WorkflowErrorCode] = None,
-        retry_on_condition: Optional[BaseDescriptor] = None,
-    ) -> Callable[..., Type["BaseRetryNodeDisplay"]]:
-        _max_attempts = max_attempts
-        _delay = delay
-        _retry_on_error_code = retry_on_error_code
-        _retry_on_condition = retry_on_condition
-
-        NodeDisplayType = TypeVar("NodeDisplayType", bound=BaseNodeDisplay)
-
-        def decorator(inner_cls: Type[NodeDisplayType]) -> Type[NodeDisplayType]:
-            node_class = inner_cls.infer_node_class()
-            wrapped_node_class = cast(Type[BaseNode], node_class.__wrapped_node__)
-
-            class RetryNodeDisplay(BaseRetryNodeDisplay[node_class]):  # type: ignore[valid-type]
-                max_attempts = _max_attempts
-                delay = _delay
-                retry_on_error_code = _retry_on_error_code
-                retry_on_condition = _retry_on_condition
-
-            setattr(inner_cls, "__adorned_by__", RetryNodeDisplay)
-
-            # We must edit the node display class to use __wrapped_node__ everywhere it
-            # references the adorned node class, which is three places:
-
-            # 1. The node display class' parameterized type
-            original_base_node_display = get_original_base(inner_cls)
-            original_base_node_display.__args__ = (wrapped_node_class,)
-            inner_cls._node_display_registry[wrapped_node_class] = inner_cls
-
-            # 2. The node display class' output displays
-            old_outputs = list(inner_cls.output_display.keys())
-            for old_output in old_outputs:
-                new_output = getattr(wrapped_node_class.Outputs, old_output.name)
-                inner_cls.output_display[new_output] = inner_cls.output_display.pop(old_output)
-
-            # 3. The node display class' port displays
-            old_ports = list(inner_cls.port_displays.keys())
-            for old_port in old_ports:
-                new_port = getattr(wrapped_node_class.Ports, old_port.name)
-                inner_cls.port_displays[new_port] = inner_cls.port_displays.pop(old_port)
-
-            return inner_cls
-
-        return decorator

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

@@ -23,14 +23,7 @@ class BaseTemplatingNodeDisplay(BaseNodeVellumDisplay[_TemplatingNodeType], Gene
         node = self._node
         node_id = self.node_id
 
-        template_input_id = self.template_input_id or next(
-            (
-                UUID(input_id) if isinstance(input_id, str) else input_id
-                for input_name, input_id in self.node_input_ids_by_name.items()
-                if input_name == TEMPLATE_INPUT_NAME
-            ),
-            None,
-        )
+        template_input_id = self.template_input_id or self.node_input_ids_by_name.get(TEMPLATE_INPUT_NAME)
 
         template_node_input = create_node_input(
             node_id=node_id,

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

@@ -1,6 +1,6 @@
 import inspect
 from uuid import UUID
-from typing import Any, Callable, ClassVar, Generic, Optional, Tuple, Type, TypeVar, cast
+from typing import Any, ClassVar, Generic, Optional, Tuple, Type, TypeVar, cast
 
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes.bases.base import BaseNode
@@ -8,7 +8,6 @@ from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
-from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
@@ -93,43 +92,3 @@ class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNod
 
         inner_output = getattr(inner_node.Outputs, output.name)
         return node_display.get_node_output_display(inner_output)
-
-    @classmethod
-    def wrap(cls, error_output_id: Optional[UUID] = None) -> Callable[..., Type["BaseTryNodeDisplay"]]:
-        _error_output_id = error_output_id
-
-        NodeDisplayType = TypeVar("NodeDisplayType", bound=BaseNodeDisplay)
-
-        def decorator(inner_cls: Type[NodeDisplayType]) -> Type[NodeDisplayType]:
-            node_class = inner_cls.infer_node_class()
-            wrapped_node_class = cast(Type[BaseNode], node_class.__wrapped_node__)
-
-            # Mypy gets mad about dynamic parameter types like this, but it's fine
-            class TryNodeDisplay(BaseTryNodeDisplay[node_class]):  # type: ignore[valid-type]
-                error_output_id = _error_output_id
-
-            setattr(inner_cls, "__adorned_by__", TryNodeDisplay)
-
-            # We must edit the node display class to use __wrapped_node__ everywhere it
-            # references the adorned node class, which is three places:
-
-            # 1. The node display class' parameterized type
-            original_base_node_display = get_original_base(inner_cls)
-            original_base_node_display.__args__ = (wrapped_node_class,)
-            inner_cls._node_display_registry[wrapped_node_class] = inner_cls
-
-            # 2. The node display class' output displays
-            old_outputs = list(inner_cls.output_display.keys())
-            for old_output in old_outputs:
-                new_output = getattr(wrapped_node_class.Outputs, old_output.name)
-                inner_cls.output_display[new_output] = inner_cls.output_display.pop(old_output)
-
-            # 3. The node display class' port displays
-            old_ports = list(inner_cls.port_displays.keys())
-            for old_port in old_ports:
-                new_port = getattr(wrapped_node_class.Ports, old_port.name)
-                inner_cls.port_displays[new_port] = inner_cls.port_displays.pop(old_port)
-
-            return inner_cls
-
-        return decorator

vellum_ee/workflows/display/tests/test_vellum_workflow_display.py

@@ -1,4 +1,3 @@
-import pytest
 from uuid import UUID
 from typing import Dict
 
@@ -231,11 +230,7 @@ def test_vellum_workflow_display__serialize_with_unused_nodes_and_edges():
     assert edge_found, "Edge between unused nodes NodeB and NodeC not found in serialized output"
 
 
-def test_parse_json_not_supported_in_ui():
-    """
-    Test that verifies ParseJsonExpression is not yet supported in the UI.
-    This test should fail once UI support is added, at which point it should be updated.
-    """
+def test_vellum_workflow_display__serialize_with_parse_json_expression():
     # GIVEN a workflow that uses the parse_json function
     from vellum.workflows.references.constant import ConstantValueReference
 
@@ -249,13 +244,55 @@ def test_parse_json_not_supported_in_ui():
         class Outputs(BaseWorkflow.Outputs):
             final = JsonNode.Outputs.json_result
 
-    # WHEN we attempt to serialize it
+    # AND a display class for this workflow
     workflow_display = get_workflow_display(
         base_display_class=VellumWorkflowDisplay,
         workflow_class=Workflow,
     )
 
-    with pytest.raises(ValueError) as exc_info:
-        workflow_display.serialize()
+    # WHEN we serialize the workflow
+    exec_config = workflow_display.serialize()
+
+    # THEN the serialized workflow contains the parse_json expression
+    raw_data = exec_config["workflow_raw_data"]
+    assert isinstance(raw_data, dict)
 
-    assert "ParseJsonExpression is not supported in the UI" == str(exc_info.value)
+    nodes = raw_data["nodes"]
+    assert isinstance(nodes, list)
+
+    json_node = None
+    for node in nodes:
+        assert isinstance(node, dict)
+        definition = node.get("definition")
+        if node.get("type") == "GENERIC" and isinstance(definition, dict) and definition.get("name") == "JsonNode":
+            json_node = node
+            break
+
+    assert json_node is not None
+
+    outputs = json_node.get("outputs", [])
+    assert isinstance(outputs, list)
+
+    json_result = None
+    for output in outputs:
+        assert isinstance(output, dict)
+        if output.get("name") == "json_result":
+            json_result = output
+            break
+
+    assert json_result == {
+        "id": "44c7d94c-a76a-4151-9b95-85a31764f18f",
+        "name": "json_result",
+        "type": "JSON",
+        "value": {
+            "type": "UNARY_EXPRESSION",
+            "lhs": {
+                "type": "CONSTANT_VALUE",
+                "value": {
+                    "type": "STRING",
+                    "value": '{"key": "value"}',
+                },
+            },
+            "operator": "parseJson",
+        },
+    }