vellum-ai 0.14.41__py3-none-any.whl → 1.11.2__py3-none-any.whl

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

@@ -1,10 +1,18 @@
 from uuid import UUID
-from typing import Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
+from typing import TYPE_CHECKING, Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
 
 from vellum import FunctionDefinition, PromptBlock, RichTextChildBlock, VellumVariable
+from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes import InlinePromptNode
 from vellum.workflows.types.core import JsonObject
-from vellum.workflows.utils.functions import compile_function_definition
+from vellum.workflows.types.definition import DeploymentDefinition, VellumIntegrationToolDefinition
+from vellum.workflows.types.generics import is_workflow_class
+from vellum.workflows.utils.functions import (
+    compile_function_definition,
+    compile_inline_workflow_function_definition,
+    compile_vellum_integration_tool_definition,
+    compile_workflow_deployment_function_definition,
+)
 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.utils import raise_if_descriptor
@@ -13,12 +21,38 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.vellum import NodeInput
 
+if TYPE_CHECKING:
+    from vellum.workflows.workflows.base import BaseWorkflow
+
+
+def _contains_descriptors(obj):
+    """Check if an object contains any descriptors or references that need special handling."""
+    if isinstance(obj, BaseDescriptor):
+        return True
+    elif isinstance(obj, dict):
+        return any(_contains_descriptors(v) for v in obj.values())
+    elif isinstance(obj, (list, tuple)):
+        return any(_contains_descriptors(item) for item in obj)
+    elif hasattr(obj, "__dict__"):
+        return any(_contains_descriptors(getattr(obj, field)) for field in obj.__dict__)
+    return False
+
+
 _InlinePromptNodeType = TypeVar("_InlinePromptNodeType", bound=InlinePromptNode)
 
 
 class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generic[_InlinePromptNodeType]):
+    __serializable_inputs__ = {
+        InlinePromptNode.prompt_inputs,
+    }
+    __unserializable_attributes__ = {
+        InlinePromptNode.settings,
+        InlinePromptNode.expand_meta,
+        InlinePromptNode.request_options,
+    }
+
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
@@ -26,23 +60,35 @@ 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)
 
-        blocks: list = [
-            self._generate_prompt_block(block, input_variable_id_by_name, [i]) for i, block in enumerate(node_blocks)
-        ]
+        ml_model = str(raise_if_descriptor(node.ml_model))
+
+        has_descriptors = _contains_descriptors(node_blocks)
+
+        blocks: list = []
+        if not has_descriptors:
+            blocks = [
+                self._generate_prompt_block(block, input_variable_id_by_name, [i])
+                for i, block in enumerate(node_blocks)
+                if not isinstance(block, BaseDescriptor)
+            ]
+
         functions = (
-            [self._generate_function_tools(function, i) for i, function in enumerate(function_definitions)]
-            if function_definitions
+            [
+                self._generate_function_tools(function, i, display_context)
+                for i, function in enumerate(function_definitions)
+            ]
+            if isinstance(function_definitions, list)
             else []
         )
         blocks.extend(functions)
 
-        return {
+        serialized_node: JsonObject = {
             "id": str(node_id),
             "type": "PROMPT",
             "inputs": [node_input.dict() for node_input in node_inputs],
@@ -55,24 +101,26 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
                 "target_handle_id": str(self.get_target_handle_id()),
                 "variant": "INLINE",
                 "exec_config": {
-                    "parameters": raise_if_descriptor(node.parameters).dict(),
+                    "parameters": self._serialize_parameters(node.parameters, display_context),
                     "input_variables": [prompt_input.dict() for prompt_input in prompt_inputs],
                     "prompt_template_block_data": {
                         "version": 1,
                         "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(),
-            "definition": self.get_definition().dict(),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
             "outputs": [
                 {"id": str(json_display.id), "name": "json", "type": "JSON", "value": None},
                 {"id": str(output_display.id), "name": "text", "type": "STRING", "value": None},
                 {"id": str(array_display.id), "name": "results", "type": "ARRAY", "value": None},
             ],
         }
+        attributes = self._serialize_attributes(display_context)
+        if attributes:
+            serialized_node["attributes"] = attributes
+        return serialized_node
 
     def _generate_node_and_prompt_inputs(
         self,
@@ -103,10 +151,26 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
 
         return node_inputs, prompt_inputs
 
-    def _generate_function_tools(self, function: Union[FunctionDefinition, Callable], index: int) -> JsonObject:
-        normalized_functions = (
-            function if isinstance(function, FunctionDefinition) else compile_function_definition(function)
-        )
+    def _generate_function_tools(
+        self,
+        function: Union[
+            FunctionDefinition, Callable, DeploymentDefinition, Type["BaseWorkflow"], VellumIntegrationToolDefinition
+        ],
+        index: int,
+        display_context: WorkflowDisplayContext,
+    ) -> JsonObject:
+        if isinstance(function, FunctionDefinition):
+            normalized_functions = function
+        elif is_workflow_class(function):
+            normalized_functions = compile_inline_workflow_function_definition(function)
+        elif callable(function):
+            normalized_functions = compile_function_definition(function)
+        elif isinstance(function, DeploymentDefinition):
+            normalized_functions = compile_workflow_deployment_function_definition(function, display_context.client)
+        elif isinstance(function, VellumIntegrationToolDefinition):
+            normalized_functions = compile_vellum_integration_tool_definition(function, display_context.client)
+        else:
+            raise ValueError(f"Unsupported function type: {type(function)}")
         return {
             "id": str(uuid4_from_hash(f"{self.node_id}-FUNCTION_DEFINITION-{index}")),
             "block_type": "FUNCTION_DEFINITION",
@@ -126,6 +190,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",
@@ -161,10 +226,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 = {
@@ -183,9 +259,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:
@@ -197,3 +271,14 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
             block["state"] = "ENABLED"
 
         return block
+
+    def _serialize_parameters(self, parameters, display_context: "WorkflowDisplayContext") -> JsonObject:
+        """Serialize parameters, returning empty object when nested descriptors are detected."""
+        params = raise_if_descriptor(parameters)
+        if not params:
+            return {}
+
+        if _contains_descriptors(params):
+            return {}
+
+        return params.dict()

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

@@ -2,13 +2,18 @@ from uuid import UUID
 from typing import ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar
 
 from vellum import VellumVariable
+from vellum.workflows.constants import undefined
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes import InlineSubworkflowNode
+from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.types.core import JsonObject
+from vellum.workflows.types.generics import is_workflow_class
+from vellum.workflows.workflows.base import BaseWorkflow
 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.exceptions import NodeValidationError
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.vellum import NodeInput
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
@@ -21,20 +26,33 @@ 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
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
 
-        node_inputs, workflow_inputs = self._generate_node_and_workflow_inputs(node_id, node, display_context)
+        subworkflow_class = raise_if_descriptor(node.subworkflow)
+        if subworkflow_class is undefined or not is_workflow_class(subworkflow_class):
+            display_context.add_error(
+                NodeValidationError(
+                    "InlineSubworkflowNode requires a subworkflow to be defined",
+                    node_class_name=node.__class__.__name__,
+                )
+            )
+            subworkflow_class = BaseWorkflow
 
+        node_inputs, workflow_inputs = self._generate_node_and_workflow_inputs(
+            node_id, node, display_context, subworkflow_class
+        )
         subworkflow_display = get_workflow_display(
             base_display_class=display_context.workflow_display_class,
-            workflow_class=raise_if_descriptor(node.subworkflow),
+            workflow_class=subworkflow_class,
             parent_display_context=display_context,
         )
-        workflow_outputs = self._generate_workflow_outputs(node, subworkflow_display.display_context)
+        workflow_outputs = self._generate_workflow_outputs(node, subworkflow_display.display_context, subworkflow_class)
         serialized_subworkflow = subworkflow_display.serialize()
 
         return {
@@ -51,9 +69,7 @@ class BaseInlineSubworkflowNodeDisplay(
                 "input_variables": [workflow_input.dict() for workflow_input in workflow_inputs],
                 "output_variables": [workflow_output.dict() for workflow_output in workflow_outputs],
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }
 
     def _generate_node_and_workflow_inputs(
@@ -61,8 +77,11 @@ class BaseInlineSubworkflowNodeDisplay(
         node_id: UUID,
         node: Type[InlineSubworkflowNode],
         display_context: WorkflowDisplayContext,
+        subworkflow: Type[BaseWorkflow],
     ) -> Tuple[List[NodeInput], List[VellumVariable]]:
-        subworkflow = raise_if_descriptor(node.subworkflow)
+        if subworkflow is undefined or not is_workflow_class(subworkflow):
+            subworkflow = BaseWorkflow
+
         subworkflow_inputs_class = subworkflow.get_inputs_class()
         subworkflow_inputs = raise_if_descriptor(node.subworkflow_inputs)
 
@@ -97,6 +116,12 @@ class BaseInlineSubworkflowNodeDisplay(
                 id=node_inputs_by_key[descriptor.name].id,
                 key=descriptor.name,
                 type=infer_vellum_variable_type(descriptor),
+                required=descriptor.instance is undefined,
+                default=(
+                    primitive_to_vellum_value(descriptor.instance).dict()
+                    if descriptor.instance is not undefined
+                    else None
+                ),
             )
             for descriptor in subworkflow_inputs_class
         ]
@@ -107,9 +132,10 @@ class BaseInlineSubworkflowNodeDisplay(
         self,
         node: Type[InlineSubworkflowNode],
         display_context: WorkflowDisplayContext,
+        subworkflow: Type[BaseWorkflow],
     ) -> List[VellumVariable]:
         workflow_outputs: List[VellumVariable] = []
-        for output_descriptor in raise_if_descriptor(node.subworkflow).Outputs:  # type: ignore[union-attr]
+        for output_descriptor in subworkflow.Outputs:  # type: ignore[union-attr]
             workflow_output_display = display_context.workflow_output_displays[output_descriptor]
             output_type = infer_vellum_variable_type(output_descriptor)
             workflow_outputs.append(

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

@@ -1,8 +1,11 @@
 from uuid import UUID
-from typing import Dict, Generic, List, Optional, TypeVar, cast
+from typing import Generic, Optional, Type, TypeVar, cast
 
 from vellum.workflows.nodes import MapNode
+from vellum.workflows.state.base import BaseState
 from vellum.workflows.types.core import JsonObject
+from vellum.workflows.types.generics import is_workflow_class
+from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -13,22 +16,22 @@ _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
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
 
-        subworkflow = raise_if_descriptor(node.subworkflow)
-
-        items_node_input = create_node_input(
-            node_id=node_id,
-            input_name="items",
-            value=node.items,
-            display_context=display_context,
-            input_id=self.node_input_ids_by_name.get("items"),
+        subworkflow_value = raise_if_descriptor(node.subworkflow)
+        subworkflow = (
+            cast(type[BaseWorkflow], subworkflow_value)
+            if subworkflow_value is not None and is_workflow_class(subworkflow_value)
+            else self._default_workflow_class()
         )
-        node_inputs = [items_node_input]
+
+        items = raise_if_descriptor(node.items)
 
         subworkflow_display = get_workflow_display(
             base_display_class=display_context.workflow_display_class,
@@ -36,24 +39,36 @@ class BaseMapNodeDisplay(BaseAdornmentNodeDisplay[_MapNodeType], Generic[_MapNod
             parent_display_context=display_context,
         )
         serialized_subworkflow = subworkflow_display.serialize()
+        if not isinstance(serialized_subworkflow["input_variables"], list):
+            raise ValueError("input_variables must be a list")
 
-        renamed_input_variables = []
-        for input_variable in cast(List[Dict[str, str]], serialized_subworkflow["input_variables"]):
-            if input_variable["key"] == "all_items":
-                renamed_item = {**input_variable, "key": "items"}
-                renamed_input_variables.append(renamed_item)
-            else:
-                renamed_input_variables.append(input_variable)
+        input_variables = serialized_subworkflow["input_variables"]
 
         # Note: This must match the items input ID for the map node's node input
         items_workflow_input_id = next(
-            input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "items"
+            input_variable["id"]
+            for input_variable in input_variables
+            if isinstance(input_variable, dict) and input_variable["key"] == "items"
         )
+
+        # Use the same ID for the map node's items input as the subworkflow's items input
+        items_node_input = create_node_input(
+            node_id=node_id,
+            input_name="items",
+            value=items or [],
+            display_context=display_context,
+            input_id=str(items_workflow_input_id),
+        )
+        node_inputs = [items_node_input]
         item_workflow_input_id = next(
-            input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "item"
+            input_variable["id"]
+            for input_variable in input_variables
+            if isinstance(input_variable, dict) and input_variable["key"] == "item"
         )
         index_workflow_input_id = next(
-            input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "index"
+            input_variable["id"]
+            for input_variable in input_variables
+            if isinstance(input_variable, dict) and input_variable["key"] == "index"
         )
 
         return {
@@ -67,14 +82,18 @@ class BaseMapNodeDisplay(BaseAdornmentNodeDisplay[_MapNodeType], Generic[_MapNod
                 "target_handle_id": str(self.get_target_handle_id()),
                 "variant": "INLINE",
                 "workflow_raw_data": serialized_subworkflow["workflow_raw_data"],
-                "input_variables": cast(JsonObject, renamed_input_variables),
+                "input_variables": cast(JsonObject, input_variables),
                 "output_variables": serialized_subworkflow["output_variables"],
                 "concurrency": raise_if_descriptor(node.max_concurrency),
                 "items_input_id": items_workflow_input_id,
                 "item_input_id": item_workflow_input_id,
                 "index_input_id": index_workflow_input_id,
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }
+
+    def _default_workflow_class(self) -> Type[BaseWorkflow]:
+        class MapNodeSubworkflow(BaseWorkflow[MapNode.SubworkflowInputs, BaseState]):
+            pass
+
+        return MapNodeSubworkflow

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

@@ -17,7 +17,7 @@ class BaseMergeNodeDisplay(BaseNodeDisplay[_MergeNodeType], Generic[_MergeNodeTy
         super().__init__()
         self._target_handle_iterator = 0
 
-    def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
+    def serialize(self, display_context: WorkflowDisplayContext, **_kwargs: Any) -> JsonObject:
         node = self._node
         node_id = self.node_id
 
@@ -47,9 +47,7 @@ class BaseMergeNodeDisplay(BaseNodeDisplay[_MergeNodeType], Generic[_MergeNodeTy
                 "target_handles": [{"id": str(target_handle_id)} for target_handle_id in target_handle_ids],
                 "source_handle_id": str(self.get_source_handle_id(display_context.port_displays)),
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }
 
     def get_target_handle_ids(self) -> Optional[List[UUID]]:

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

@@ -13,6 +13,7 @@ class BaseNoteNodeDisplay(BaseNodeDisplay[_NoteNodeType], Generic[_NoteNodeType]
     style: ClassVar[Union[Dict[str, Any], None]] = None
 
     def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
+        del kwargs  # Unused parameters
         node_id = self.node_id
 
         return {
@@ -24,7 +25,5 @@ class BaseNoteNodeDisplay(BaseNodeDisplay[_NoteNodeType], Generic[_NoteNodeType]
                 "text": self.text,
                 "style": self.style,
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }

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

@@ -1,10 +1,9 @@
 from uuid import UUID
-from typing import Generic, Optional, TypeVar, cast
+from typing import Generic, Optional, TypeVar
 
 from vellum.workflows.nodes.displayable.prompt_deployment_node import PromptDeploymentNode
-from vellum.workflows.references import OutputReference
 from vellum.workflows.types.core import JsonObject
-from vellum.workflows.vellum_client import create_vellum_client
+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.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -14,8 +13,10 @@ _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
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
@@ -28,7 +29,10 @@ class BasePromptDeploymentNodeDisplay(BaseNodeDisplay[_PromptDeploymentNodeType]
                     input_name=variable_name,
                     value=variable_value,
                     display_context=display_context,
-                    input_id=self.node_input_ids_by_name.get(variable_name),
+                    input_id=self.node_input_ids_by_name.get(
+                        f"{PromptDeploymentNode.prompt_inputs.name}.{variable_name}"
+                    )
+                    or self.node_input_ids_by_name.get(variable_name),
                 )
                 for variable_name, variable_value in prompt_inputs.items()
             ]
@@ -36,15 +40,19 @@ class BasePromptDeploymentNodeDisplay(BaseNodeDisplay[_PromptDeploymentNodeType]
             else []
         )
 
-        _, 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)]
+        output_display = self.output_display[node.Outputs.text]
+        array_display = self.output_display[node.Outputs.results]
+        json_display = self.output_display[node.Outputs.json]
 
-        # TODO: Pass through the name instead of retrieving the ID
-        # https://app.shortcut.com/vellum/story/4702
-        vellum_client = create_vellum_client()
-        deployment = vellum_client.deployments.retrieve(
-            id=str(raise_if_descriptor(node.deployment)),
-        )
+        deployment_descriptor_id = str(raise_if_descriptor(node.deployment))
+        try:
+            deployment = display_context.client.deployments.retrieve(
+                id=deployment_descriptor_id,
+            )
+            deployment_id = str(deployment.id)
+        except Exception as e:
+            display_context.add_error(e)
+            deployment_id = str(uuid4_from_hash(deployment_descriptor_id))
         ml_model_fallbacks = raise_if_descriptor(node.ml_model_fallbacks)
 
         return {
@@ -59,11 +67,14 @@ class BasePromptDeploymentNodeDisplay(BaseNodeDisplay[_PromptDeploymentNodeType]
                 "source_handle_id": str(self.get_source_handle_id(display_context.port_displays)),
                 "target_handle_id": str(self.get_target_handle_id()),
                 "variant": "DEPLOYMENT",
-                "prompt_deployment_id": str(deployment.id),
+                "prompt_deployment_id": deployment_id,
                 "release_tag": raise_if_descriptor(node.release_tag),
                 "ml_model_fallbacks": list(ml_model_fallbacks) if ml_model_fallbacks else None,
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
+            "outputs": [
+                {"id": str(json_display.id), "name": "json", "type": "JSON", "value": None},
+                {"id": str(output_display.id), "name": "text", "type": "STRING", "value": None},
+                {"id": str(array_display.id), "name": "results", "type": "ARRAY", "value": None},
+            ],
         }

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

@@ -1,8 +1,6 @@
 import inspect
-from typing import Any, Generic, Tuple, Type, TypeVar, cast
+from typing import Any, Generic, TypeVar
 
-from vellum.workflows.descriptors.base import BaseDescriptor
-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
@@ -13,12 +11,13 @@ from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_di
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.expressions import serialize_value
 
 _RetryNodeType = TypeVar("_RetryNodeType", bound=RetryNode)
 
 
 class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_RetryNodeType]):
-    def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
+    def serialize(self, display_context: WorkflowDisplayContext, **_kwargs: Any) -> JsonObject:
         node = self._node
         node_id = self.node_id
 
@@ -33,13 +32,13 @@ class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_Re
                 {
                     "id": id,
                     "name": attribute.name,
-                    "value": self.serialize_value(display_context, cast(BaseDescriptor, attribute.instance)),
+                    "value": serialize_value(node_id, display_context, attribute.instance),
                 }
             )
 
         adornment: JsonObject = {
             "id": str(node_id),
-            "label": node.__qualname__,
+            "label": self.label,
             "base": self.get_base().dict(),
             "attributes": attributes,
         }
@@ -66,7 +65,7 @@ class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_Re
 
         return serialized_node
 
-    def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
+    def get_node_output_display(self, output: OutputReference) -> NodeOutputDisplay:
         inner_node = self._node.__wrapped_node__
         if not inner_node:
             return super().get_node_output_display(output)