vellum-ai 0.14.17__py3-none-any.whl → 0.14.18__py3-none-any.whl

vellum/client/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.14.17",
+            "X-Fern-SDK-Version": "0.14.18",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/workflows/events/workflow.py

@@ -48,7 +48,7 @@ class NodeEventDisplayContext(UniversalBaseModel):
 
 
 class WorkflowEventDisplayContext(UniversalBaseModel):
-    node_displays: Dict[str, NodeEventDisplayContext]
+    node_displays: Dict[UUID, NodeEventDisplayContext]
     workflow_inputs: Dict[str, UUID]
     workflow_outputs: Dict[str, UUID]
 
@@ -194,6 +194,12 @@ WorkflowExecutionEvent = Union[
     WorkflowExecutionSnapshottedEvent,
 ]
 
+TerminalWorkflowExecutionEvent = Union[
+    WorkflowExecutionFulfilledEvent,
+    WorkflowExecutionRejectedEvent,
+    WorkflowExecutionPausedEvent,
+]
+
 
 def is_workflow_event(event: WorkflowEvent) -> TypeGuard[WorkflowExecutionEvent]:
     return (
@@ -205,3 +211,11 @@ def is_workflow_event(event: WorkflowEvent) -> TypeGuard[WorkflowExecutionEvent]
         or event.name == "workflow.execution.resumed"
         or event.name == "workflow.execution.rejected"
     )
+
+
+def is_terminal_workflow_execution_event(event: WorkflowEvent) -> TypeGuard[TerminalWorkflowExecutionEvent]:
+    return (
+        event.name == "workflow.execution.fulfilled"
+        or event.name == "workflow.execution.rejected"
+        or event.name == "workflow.execution.paused"
+    )

vellum/workflows/nodes/core/retry_node/node.py

@@ -1,9 +1,11 @@
 import time
 from typing import Callable, Generic, Optional, Type
 
+from vellum.workflows.context import execution_context, get_parent_context
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.descriptors.utils import resolve_value
 from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.events.workflow import is_terminal_workflow_execution_event
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
@@ -11,6 +13,7 @@ from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
 from vellum.workflows.nodes.utils import create_adornment
 from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.types.generics import StateType
+from vellum.workflows.workflows.event_filters import all_workflow_event_filter
 
 
 class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
@@ -38,47 +41,64 @@ class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
 
         for index in range(self.max_attempts):
             attempt_number = index + 1
-            context = WorkflowContext(vellum_client=self._context.vellum_client)
-            subworkflow = self.subworkflow(
-                parent_state=self.state,
-                context=context,
-            )
-            terminal_event = subworkflow.run(
-                inputs=self.SubworkflowInputs(attempt_number=attempt_number),
-                node_output_mocks=self._context._get_all_node_output_mocks(),
-            )
-            if terminal_event.name == "workflow.execution.fulfilled":
-                node_outputs = self.Outputs()
-                workflow_output_vars = vars(terminal_event.outputs)
-
-                for output_name in workflow_output_vars:
-                    setattr(node_outputs, output_name, workflow_output_vars[output_name])
-
-                return node_outputs
-            elif terminal_event.name == "workflow.execution.paused":
-                raise NodeException(
-                    code=WorkflowErrorCode.INVALID_OUTPUTS,
-                    message=f"Subworkflow unexpectedly paused on attempt {attempt_number}",
-                )
-            elif self.retry_on_error_code and self.retry_on_error_code != terminal_event.error.code:
-                raise NodeException(
-                    code=WorkflowErrorCode.INVALID_OUTPUTS,
-                    message=f"""Unexpected rejection on attempt {attempt_number}: {terminal_event.error.code.value}.
-Message: {terminal_event.error.message}""",
+            parent_context = get_parent_context()
+            with execution_context(parent_context=parent_context):
+                subworkflow = self.subworkflow(
+                    parent_state=self.state,
+                    context=WorkflowContext(vellum_client=self._context.vellum_client),
                 )
-            elif self.retry_on_condition and not resolve_value(self.retry_on_condition, self.state):
-                raise NodeException(
-                    code=WorkflowErrorCode.INVALID_OUTPUTS,
-                    message=f"""Rejection failed on attempt {attempt_number}: {terminal_event.error.code.value}.
-Message: {terminal_event.error.message}""",
+                subworkflow_stream = subworkflow.stream(
+                    inputs=self.SubworkflowInputs(attempt_number=attempt_number),
+                    event_filter=all_workflow_event_filter,
+                    node_output_mocks=self._context._get_all_node_output_mocks(),
                 )
-            else:
-                last_exception = NodeException(
-                    terminal_event.error.message,
-                    code=terminal_event.error.code,
-                )
-                if self.delay:
-                    time.sleep(self.delay)
+
+            node_outputs: Optional[BaseNode.Outputs] = None
+            exception: Optional[NodeException] = None
+            for event in subworkflow_stream:
+                self._context._emit_subworkflow_event(event)
+
+                if not is_terminal_workflow_execution_event(event):
+                    continue
+
+                if event.workflow_definition != self.subworkflow:
+                    continue
+
+                if event.name == "workflow.execution.fulfilled":
+                    node_outputs = self.Outputs()
+
+                    for output_descriptor, output_value in event.outputs:
+                        setattr(node_outputs, output_descriptor.name, output_value)
+                elif event.name == "workflow.execution.paused":
+                    exception = NodeException(
+                        code=WorkflowErrorCode.INVALID_OUTPUTS,
+                        message=f"Subworkflow unexpectedly paused on attempt {attempt_number}",
+                    )
+                elif self.retry_on_error_code and self.retry_on_error_code != event.error.code:
+                    exception = NodeException(
+                        code=WorkflowErrorCode.INVALID_OUTPUTS,
+                        message=f"""Unexpected rejection on attempt {attempt_number}: {event.error.code.value}.
+Message: {event.error.message}""",
+                    )
+                elif self.retry_on_condition and not resolve_value(self.retry_on_condition, self.state):
+                    exception = NodeException(
+                        code=WorkflowErrorCode.INVALID_OUTPUTS,
+                        message=f"""Rejection failed on attempt {attempt_number}: {event.error.code.value}.
+Message: {event.error.message}""",
+                    )
+                else:
+                    last_exception = NodeException(
+                        event.error.message,
+                        code=event.error.code,
+                    )
+                    if self.delay:
+                        time.sleep(self.delay)
+
+            if exception:
+                raise exception
+
+            if node_outputs:
+                return node_outputs
 
         raise last_exception
 

vellum/workflows/nodes/core/templating_node/node.py

@@ -48,10 +48,10 @@ class TemplatingNode(BaseNode[StateType], Generic[StateType, _OutputType], metac
     """
 
     # The Jinja template to render.
-    template: ClassVar[str]
+    template: ClassVar[str] = ""
 
     # The inputs to render the template with.
-    inputs: ClassVar[EntityInputsInterface]
+    inputs: ClassVar[EntityInputsInterface] = {}
 
     jinja_globals: Dict[str, Any] = DEFAULT_JINJA_GLOBALS
     jinja_custom_filters: Mapping[str, FilterFunc] = DEFAULT_JINJA_CUSTOM_FILTERS

vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/test_inline_prompt_node.py

@@ -1,4 +1,6 @@
 import pytest
+import json
+from unittest import mock
 from uuid import uuid4
 from typing import Any, Iterator, List
 
@@ -7,6 +9,7 @@ from vellum import (
     JinjaPromptBlock,
     PlainTextPromptBlock,
     PromptBlock,
+    PromptParameters,
     RichTextPromptBlock,
     VariablePromptBlock,
 )
@@ -18,7 +21,11 @@ from vellum.client.types.prompt_request_string_input import PromptRequestStringI
 from vellum.client.types.string_vellum_value import StringVellumValue
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes import InlinePromptNode
 from vellum.workflows.nodes.displayable.bases.inline_prompt_node import BaseInlinePromptNode
+from vellum.workflows.state import BaseState
+from vellum.workflows.state.base import StateMeta
 
 
 def test_validation_with_missing_variables():
@@ -180,3 +187,112 @@ def test_validation_with_extra_variables(vellum_adhoc_prompt_client):
         PromptRequestStringInput(key="required_var", type="STRING", value="value"),
         PromptRequestStringInput(key="extra_var", type="STRING", value="extra_value"),
     ]
+
+
+def test_inline_prompt_node__json_output(vellum_adhoc_prompt_client):
+    """Confirm that InlinePromptNodes output the expected JSON when run."""
+
+    # GIVEN a node that subclasses InlinePromptNode
+    class Inputs(BaseInputs):
+        input: str
+
+    class State(BaseState):
+        pass
+
+    class MyInlinePromptNode(InlinePromptNode):
+        ml_model = "gpt-4o"
+        blocks = []
+        parameters = PromptParameters(
+            stop=[],
+            temperature=0.0,
+            max_tokens=4096,
+            top_p=1.0,
+            top_k=0,
+            frequency_penalty=0.0,
+            presence_penalty=0.0,
+            logit_bias=None,
+            custom_parameters={
+                "json_mode": False,
+                "json_schema": {
+                    "name": "get_result",
+                    "schema": {
+                        "type": "object",
+                        "required": ["result"],
+                        "properties": {"result": {"type": "string", "description": ""}},
+                    },
+                },
+            },
+        )
+
+    # AND a known JSON response from invoking an inline prompt
+    expected_json = {"result": "Hello, world!"}
+    expected_outputs: List[PromptOutput] = [
+        StringVellumValue(value=json.dumps(expected_json)),
+    ]
+
+    def generate_prompt_events(*args: Any, **kwargs: Any) -> Iterator[ExecutePromptEvent]:
+        execution_id = str(uuid4())
+        events: List[ExecutePromptEvent] = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(
+                execution_id=execution_id,
+                outputs=expected_outputs,
+            ),
+        ]
+        yield from events
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.side_effect = generate_prompt_events
+
+    # WHEN the node is run
+    node = MyInlinePromptNode(
+        state=State(
+            meta=StateMeta(workflow_inputs=Inputs(input="Generate JSON.")),
+        )
+    )
+    outputs = [o for o in node.run()]
+
+    # THEN the node should have produced the outputs we expect
+    results_output = outputs[0]
+    assert results_output.name == "results"
+    assert results_output.value == expected_outputs
+
+    text_output = outputs[1]
+    assert text_output.name == "text"
+    assert text_output.value == '{"result": "Hello, world!"}'
+
+    json_output = outputs[2]
+    assert json_output.name == "json"
+    assert json_output.value == expected_json
+
+    # AND we should have made the expected call to Vellum search
+    vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.assert_called_once_with(
+        blocks=[],
+        expand_meta=Ellipsis,
+        functions=None,
+        input_values=[],
+        input_variables=[],
+        ml_model="gpt-4o",
+        parameters=PromptParameters(
+            stop=[],
+            temperature=0.0,
+            max_tokens=4096,
+            top_p=1.0,
+            top_k=0,
+            frequency_penalty=0.0,
+            presence_penalty=0.0,
+            logit_bias=None,
+            custom_parameters={
+                "json_mode": False,
+                "json_schema": {
+                    "name": "get_result",
+                    "schema": {
+                        "type": "object",
+                        "required": ["result"],
+                        "properties": {"result": {"type": "string", "description": ""}},
+                    },
+                },
+            },
+        ),
+        request_options=mock.ANY,
+        settings=None,
+    )

vellum/workflows/nodes/displayable/inline_prompt_node/node.py

@@ -46,14 +46,31 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
             )
 
         string_outputs = []
+        json_output = None
+
+        should_parse_json = False
+        if hasattr(self, "parameters"):
+            custom_params = self.parameters.custom_parameters
+            if custom_params and isinstance(custom_params, dict):
+                json_schema = custom_params.get("json_schema", {})
+                if (isinstance(json_schema, dict) and "schema" in json_schema) or custom_params.get("json_mode", {}):
+                    should_parse_json = True
+
         for output in outputs:
             if output.value is None:
                 continue
 
             if output.type == "STRING":
                 string_outputs.append(output.value)
+                if should_parse_json:
+                    try:
+                        parsed_json = json.loads(output.value)
+                        json_output = parsed_json
+                    except (json.JSONDecodeError, TypeError):
+                        pass
             elif output.type == "JSON":
                 string_outputs.append(json.dumps(output.value, indent=4))
+                json_output = output.value
             elif output.type == "FUNCTION_CALL":
                 string_outputs.append(output.value.model_dump_json(indent=4))
             else:
@@ -61,3 +78,6 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
 
         value = "\n".join(string_outputs)
         yield BaseOutput(name="text", value=value)
+
+        if json_output:
+            yield BaseOutput(name="json", value=json_output)

vellum/workflows/runner/runner.py

@@ -43,7 +43,7 @@ from vellum.workflows.events.workflow import (
     WorkflowExecutionSnapshottedEvent,
     WorkflowExecutionStreamingBody,
 )
-from vellum.workflows.exceptions import NodeException
+from vellum.workflows.exceptions import NodeException, WorkflowInitializationException
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import NodeRunResponse
 from vellum.workflows.nodes.mocks import MockNodeExecutionArg
@@ -332,6 +332,18 @@ class WorkflowRunner(Generic[StateType]):
                     parent=parent_context,
                 )
             )
+        except WorkflowInitializationException as e:
+            self._workflow_event_inner_queue.put(
+                NodeExecutionRejectedEvent(
+                    trace_id=node.state.meta.trace_id,
+                    span_id=span_id,
+                    body=NodeExecutionRejectedBody(
+                        node_definition=node.__class__,
+                        error=e.error,
+                    ),
+                    parent=parent_context,
+                )
+            )
         except Exception as e:
             logger.exception(f"An unexpected error occurred while running node {node.__class__.__name__}")
 
@@ -563,6 +575,9 @@ class WorkflowRunner(Generic[StateType]):
             except NodeException as e:
                 self._workflow_event_outer_queue.put(self._reject_workflow_event(e.error))
                 return
+            except WorkflowInitializationException as e:
+                self._workflow_event_outer_queue.put(self._reject_workflow_event(e.error))
+                return
             except Exception:
                 err_message = f"An unexpected error occurred while initializing node {node_cls.__name__}"
                 logger.exception(err_message)

{vellum_ai-0.14.17.dist-info → vellum_ai-0.14.18.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.14.17
+Version: 0.14.18
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

{vellum_ai-0.14.17.dist-info → vellum_ai-0.14.18.dist-info}/RECORD

@@ -39,7 +39,7 @@ vellum_ee/workflows/display/nodes/vellum/conditional_node.py,sha256=ybLIa4uclqVI
 vellum_ee/workflows/display/nodes/vellum/error_node.py,sha256=I1Jkp2htRINJATtv1e-zs9BrReFX842djpiVgBPHDYg,2186
 vellum_ee/workflows/display/nodes/vellum/final_output_node.py,sha256=p-PvlnxpBQ7IKskZi2A19jKAtKnSxJ8LPbGMA83VkFk,2805
 vellum_ee/workflows/display/nodes/vellum/guardrail_node.py,sha256=aYZSJTxknU4LMiQdWk9LcK6CkhdozeDEMiRxfAyUNEc,2202
-vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py,sha256=CPNW7HtKKJCGOvvE9DTn-56XVDLgAcBg2QHTfqU4N7k,8543
+vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py,sha256=grb2iDxZ8DiNlo1nfO3KPxs8saM_JZAskOFp-C4ZEq4,8908
 vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py,sha256=MU9I8CB1X1TgL1aa1eT6DHWwNJ-2v79t74xl0oy-fBo,5510
 vellum_ee/workflows/display/nodes/vellum/map_node.py,sha256=8CPnn06HIBxBOiECevUffeVmQmCpec6WtPQnNl9gj9Y,3748
 vellum_ee/workflows/display/nodes/vellum/merge_node.py,sha256=HkNMgdQELiON42jdO-xDLmqrEKdGx1RVqrz2DXNTLS8,3239
@@ -48,15 +48,16 @@ vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py,sha256=LFjLUr
 vellum_ee/workflows/display/nodes/vellum/retry_node.py,sha256=LgokATi7sSS38Fil-XjqoR4t7AMOJ-GzXRw6p606Svo,3397
 vellum_ee/workflows/display/nodes/vellum/search_node.py,sha256=TxcAGZDl_hvJ7Y1hUi9YVEVrj9Ie0hKkASdpfRL4_cs,9227
 vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py,sha256=62baAElKoRKIoba0lLhnrXGWWx96B73VxKGxh7BaIxc,2612
-vellum_ee/workflows/display/nodes/vellum/templating_node.py,sha256=wP-ZAZq6ZmkS1fj0WkFtRkYqi8c6MWKLh4CEw8JD9mg,3207
+vellum_ee/workflows/display/nodes/vellum/templating_node.py,sha256=5EWzdA3TSUPlbJhs4uo1KGH6eGtDXZucVhlUoc_xpl4,3242
 vellum_ee/workflows/display/nodes/vellum/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/nodes/vellum/tests/test_error_node.py,sha256=ulrpoYUW-5kIxfG4Lf5F2p0k_EoYKhmahEbF3P_eruM,1648
 vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_node.py,sha256=bg9INsXiWfyK047u8TD1oEOFYrqDq8GC7Hvgz69n7BE,1988
 vellum_ee/workflows/display/nodes/vellum/tests/test_retry_node.py,sha256=NuIw8Yb42KUdoGi3Ur8_7VPg50IC4hNrwAkCociwqNk,2091
+vellum_ee/workflows/display/nodes/vellum/tests/test_templating_node.py,sha256=Us32jf_FQnLuT4Bs2o5JyHxihCTAN8ozZghWIR0pl9k,3459
 vellum_ee/workflows/display/nodes/vellum/tests/test_try_node.py,sha256=mtzB8LJlFCHVFM4H5AanLp29gQfaVmnN4A4iaRGJHoI,2427
 vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py,sha256=4YUaTeD_OWF-UaPMyOTBTu9skGC1jgSHlAYrzbH7Z04,5039
 vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=0ipBZMYm521tuwefQFgHOvTyTgiWVTkzxpQtnnmOAI0,4203
-vellum_ee/workflows/display/nodes/vellum/utils.py,sha256=F_0BrlSszllK_BhryPbojIleLq2dGXOfQD1rVp3fNFg,4733
+vellum_ee/workflows/display/nodes/vellum/utils.py,sha256=lY3YsvzawDy15-8w3Brb4oSnVpVaMe5FQtgMi3BE3Mo,4737
 vellum_ee/workflows/display/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/test_vellum_workflow_display.py,sha256=VD-4USiRlCcdC3Qe9WfdkxwFdircai0vqvuZCbELR84,9556
 vellum_ee/workflows/display/tests/workflow_serialization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -81,7 +82,7 @@ vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_s
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_prompt_deployment_serialization.py,sha256=-bGp82C5_XgcVxeM80zWL9viogxvE4DGk7-5H__wwxA,8649
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_search_node_serialization.py,sha256=elc4ak2OyI4y5P6GPnK2RwuZLcx4FL6hHlmhY3ARk6k,13008
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_subworkflow_deployment_serialization.py,sha256=BzFNl9ECeGh0krm-CUjbBQQq0g7krANsp0Sh-j5dAkc,11322
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_templating_node_serialization.py,sha256=xXtW915v9yxxKlyu5ObzKHyJYMvobvev3ermX61SGx4,7818
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_templating_node_serialization.py,sha256=qjaXGw0JTnPnQmDAFm7jkB3bVkApf-WQ0Gd6kptqdPA,7818
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_node_serialization.py,sha256=NdhE3lm7RMQ8DqkraPSq24IbOxNla9unbs4tsMWRzm4,3781
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_try_node_serialization.py,sha256=eD5686C9nWC5s6t08vbAnm9qf9t53gYQM-E1FwAa75c,3035
 vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py,sha256=huKAOeMJ2MKmp6XtbvMJTUadqynoV40Ypoz9jsBEBEQ,7431
@@ -89,11 +90,11 @@ vellum_ee/workflows/display/types.py,sha256=ixfmcQn51Rhsm4_0hWfG0_WpzLE89ZrDZpeY
 vellum_ee/workflows/display/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/utils/expressions.py,sha256=9FpOslDI-RCR5m4TgAu9KCHh4aTVnh7CHR2ykyMUDw0,1151
 vellum_ee/workflows/display/utils/vellum.py,sha256=EVPQUSsZ3OIeLTEbV6LHPor37t9fnj9kJxDqP4PmTLQ,8234
-vellum_ee/workflows/display/vellum.py,sha256=7mqQaKZPPrLMcXSAQkPIxCy5x8HkKs5PbCu3GRaC2o8,8507
+vellum_ee/workflows/display/vellum.py,sha256=ROvXFoLmnEXmJJ_Nc0ZVz94LZf6opU4O1TWbGnAfHug,5028
 vellum_ee/workflows/display/workflows/__init__.py,sha256=kapXsC67VJcgSuiBMa86FdePG5A9kMB5Pi4Uy1O2ob4,207
-vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=QKMSmV--UoE8L7pYmKvZqtaatygrxmSm9PDSH71Yr0Y,19937
+vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=-14UlLZlmfWDqNtxhY9GxXlrlPcgzEpAIffd94QPEvU,19932
 vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=kp0u8LN_2IwshLrhMImhpZx1hRyAcD5gXY-kDuuaGMQ,1269
-vellum_ee/workflows/display/workflows/tests/test_workflow_display.py,sha256=eCBCssh2J0hhGAKiv5pVxYKxMEUQCtFX-FKn49xercA,8576
+vellum_ee/workflows/display/workflows/tests/test_workflow_display.py,sha256=Ny9VWjCxc1_w8Z5xiQEJmNmEwtjlDlxgFCrmCWSmFmA,8511
 vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=mbAzCpswOek34ITeTkesbVreCXpulj4NFjIg3RcdVZ8,18243
 vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/server/virtual_file_loader.py,sha256=X_DdNK7MfyOjKWekk6YQpOSCT6klKcdjT6nVJcBH1sM,1481
@@ -116,7 +117,7 @@ vellum_ee/workflows/tests/local_workflow/nodes/__init__.py,sha256=1F6jxUpSKfPXPj
 vellum_ee/workflows/tests/local_workflow/nodes/final_output.py,sha256=ZX7zBv87zirg0w9VKUW3QVDSdBLDqcqAMZjCL_oWbpU,297
 vellum_ee/workflows/tests/local_workflow/nodes/templating_node.py,sha256=NQwFN61QkHfI3Vssz-B0NKGfupK8PU0FDSAIAhYBLi0,325
 vellum_ee/workflows/tests/local_workflow/workflow.py,sha256=A4qOzOPNwePYxWbcAgIPLsmrVS_aVEZEc-wULSv787Q,393
-vellum_ee/workflows/tests/test_display_meta.py,sha256=xLQ7QtIXIiIm61pm2lyl588Ohzc3pjyq1nDp-qVu2Fs,2295
+vellum_ee/workflows/tests/test_display_meta.py,sha256=C25dErwghPNXio49pvSRxyOuc96srH6eYEwTAWdE2zY,2258
 vellum_ee/workflows/tests/test_server.py,sha256=M6vvQ2hjIpDWtQdDM9EPbMvUrZ93niAuYnxMNJWOjPA,511
 vellum_ee/workflows/tests/test_virtual_files.py,sha256=TJEcMR0v2S8CkloXNmCHA0QW0K6pYNGaIjraJz7sFvY,2762
 vellum/__init__.py,sha256=Vh9jxpBVXB2fTCyoUMlLdGL1Ujf0zNNNvqhNu5II8KI,36466
@@ -124,7 +125,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=tKtdM1_GqmGq1gpi9ydWD_T-MM7fPn8QdHh8ww19cNI,117564
 vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
 vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-vellum/client/core/client_wrapper.py,sha256=lp2PrDTuF-fVxozBobWF7i3EJRnKs7ggqOiSWj5wuL4,1869
+vellum/client/core/client_wrapper.py,sha256=uLm2G6JmYjhKSyX55XTjjoXumxodrRRDWg6-C5GmXLI,1869
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
 vellum/client/core/http_client.py,sha256=R0pQpCppnEtxccGvXl4uJ76s7ro_65Fo_erlNNLp_AI,19228
@@ -1319,7 +1320,7 @@ vellum/workflows/events/node.py,sha256=jbmNHjdp331Q1IRK-AWtAxwF6Lidb9R7__N5rQuil
 vellum/workflows/events/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/events/tests/test_event.py,sha256=sHcKhZPDPtzZfTmehL4NORA_StR4M6nZDcx9kz3Avo0,16866
 vellum/workflows/events/types.py,sha256=AeTJaQt_fNHDLI4nyBzo7XrW9QQybRC09AKzu3kEYEE,3575
-vellum/workflows/events/workflow.py,sha256=QoSHyIOpuVacbR7POf7h104miTOhCjtO2udnYximJGs,6851
+vellum/workflows/events/workflow.py,sha256=xdqU6WOexaAqzJbU2Zw42o2LJhK7SDPtTFO5REGv17I,7293
 vellum/workflows/exceptions.py,sha256=NiBiR3ggfmPxBVqD-H1SqmjI-7mIn0EStSN1BqApvCM,1213
 vellum/workflows/expressions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/expressions/accessor.py,sha256=ItZF7fMLzVTqsdAiaXb5SiDupXmX0X9xbIus1W6hRds,1870
@@ -1381,11 +1382,11 @@ vellum/workflows/nodes/core/map_node/node.py,sha256=dY27Xm11LHsqD7hnZnVYYDIazZ-X
 vellum/workflows/nodes/core/map_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/core/map_node/tests/test_node.py,sha256=uMR0AyIFn539LqTKHdwuBswnx1i-PHyqPpgtYrnmYMY,3496
 vellum/workflows/nodes/core/retry_node/__init__.py,sha256=lN2bIy5a3Uzhs_FYCrooADyYU6ZGShtvLKFWpelwPvo,60
-vellum/workflows/nodes/core/retry_node/node.py,sha256=eRfDj54idq4JXkzgFyHNh18qnEMI6vLOtNgn7MsFg1o,4359
+vellum/workflows/nodes/core/retry_node/node.py,sha256=SoDQaizjoj6kKLr454rkXu4bdzN7hXrBBn7T2_cp1CA,5261
 vellum/workflows/nodes/core/retry_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/core/retry_node/tests/test_node.py,sha256=fNgDufkIsrTC-6ftvogqSpWhqqBj9iNESdfK19B1Yx0,5159
 vellum/workflows/nodes/core/templating_node/__init__.py,sha256=GmyuYo81_A1_Bz6id69ozVFS6FKiuDsZTiA3I6MaL2U,70
-vellum/workflows/nodes/core/templating_node/node.py,sha256=-JIqLUv6Xpx_LTVZt7whQ2X2VatgHDdTxjMrz64luEs,3721
+vellum/workflows/nodes/core/templating_node/node.py,sha256=iqBmr2i-f-BqhisNQJiDfewjol0ur7-XpupLStyMJsg,3731
 vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py,sha256=MHofz-BwAgt7EXkab8VIyacYznDEIJ7Er7MJUaxNQQo,9614
 vellum/workflows/nodes/core/try_node/__init__.py,sha256=JVD4DrldTIqFQQFrubs9KtWCCc0YCAc7Fzol5ZWIWeM,56
 vellum/workflows/nodes/core/try_node/node.py,sha256=RbxL0NRXS0IxRP0MJAnLABolF6dkwVniiqsagzy-lwk,4445
@@ -1405,7 +1406,7 @@ vellum/workflows/nodes/displayable/bases/inline_prompt_node/__init__.py,sha256=H
 vellum/workflows/nodes/displayable/bases/inline_prompt_node/constants.py,sha256=fnjiRWLoRlC4Puo5oQcpZD5Hd-EesxsAo9l5tGAkpZQ,270
 vellum/workflows/nodes/displayable/bases/inline_prompt_node/node.py,sha256=fTQ7KS5u7xCO_nJsWT5_LWgwGnYZg4U06slvXxHAzTU,8449
 vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/test_inline_prompt_node.py,sha256=kodL6T6booHtdTLMJ9wSk_JHFgPNkTz3U8AJDB2cxtU,6696
+vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/test_inline_prompt_node.py,sha256=ZCXCZs-_OyPk4nqCpuWY-vw87lg92TDZ2tK_gckJ7mg,10450
 vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py,sha256=UIyNUUr8ii6rIx1gQL5vIPqTAww4ExxxJkFMoglc8LE,5794
 vellum/workflows/nodes/displayable/bases/search_node.py,sha256=3UtbqY3QO4kzfJHbmUNZGnEEfJmaoiF892u8H6TGjp8,5381
 vellum/workflows/nodes/displayable/bases/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1428,7 +1429,7 @@ vellum/workflows/nodes/displayable/final_output_node/tests/test_node.py,sha256=E
 vellum/workflows/nodes/displayable/guardrail_node/__init__.py,sha256=Ab5eXmOoBhyV4dMWdzh32HLUmnPIBEK_zFCT38C4Fng,68
 vellum/workflows/nodes/displayable/guardrail_node/node.py,sha256=h5nIBzQxbXTrdTq1wjDcekk1RV4-rKUNCshqdBAiJJY,4025
 vellum/workflows/nodes/displayable/inline_prompt_node/__init__.py,sha256=gSUOoEZLlrx35-tQhSAd3An8WDwBqyiQh-sIebLU9wU,74
-vellum/workflows/nodes/displayable/inline_prompt_node/node.py,sha256=BaBfIfQX10sfoJ03H6zAnheae8bcJgi6Ff3SkYvmCKo,2523
+vellum/workflows/nodes/displayable/inline_prompt_node/node.py,sha256=8RXZqWMzViUjFfbpmcy1gkSsKnEpci8BGwsuPYv4xMQ,3380
 vellum/workflows/nodes/displayable/inline_prompt_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/inline_prompt_node/tests/test_node.py,sha256=Qg1oGXecGH2Hp9oBAY42bTfKHBvNHHoJ6vUPPEj8Lq0,8539
 vellum/workflows/nodes/displayable/merge_node/__init__.py,sha256=J8IC08dSH7P76wKlNuxe1sn7toNGtSQdFirUbtPDEs0,60
@@ -1481,7 +1482,7 @@ vellum/workflows/references/workflow_input.py,sha256=86IuhlBz-9cGxeUzizyjdp482aj
 vellum/workflows/resolvers/__init__.py,sha256=eH6hTvZO4IciDaf_cf7aM2vs-DkBDyJPycOQevJxQnI,82
 vellum/workflows/resolvers/base.py,sha256=WHra9LRtlTuB1jmuNqkfVE2JUgB61Cyntn8f0b0WZg4,411
 vellum/workflows/runner/__init__.py,sha256=i1iG5sAhtpdsrlvwgH6B-m49JsINkiWyPWs8vyT-bqM,72
-vellum/workflows/runner/runner.py,sha256=fS21u0LIyhxT4T6YQPgonVuH8X8TgD5BrsbVxPUWfuM,31071
+vellum/workflows/runner/runner.py,sha256=pK_cRWS75x_sX7zrnltz3IykCTgoAWA7AW19Ih5BaB4,31759
 vellum/workflows/sandbox.py,sha256=GVJzVjMuYzOBnSrboB0_6MMRZWBluAyQ2o7syeaeBd0,2235
 vellum/workflows/state/__init__.py,sha256=yUUdR-_Vl7UiixNDYQZ-GEM_kJI9dnOia75TtuNEsnE,60
 vellum/workflows/state/base.py,sha256=Vkhneko3VlQrPsMLU1PYSzXU_W1u7_AraJsghiv5O-4,15512
@@ -1517,8 +1518,8 @@ vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnad
 vellum/workflows/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/workflows/tests/test_base_workflow.py,sha256=NRteiICyJvDM5zrtUfq2fZoXcGQVaWC9xmNlLLVW0cU,7979
 vellum/workflows/workflows/tests/test_context.py,sha256=VJBUcyWVtMa_lE5KxdhgMu0WYNYnUQUDvTF7qm89hJ0,2333
-vellum_ai-0.14.17.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.14.17.dist-info/METADATA,sha256=DPUjnu7ydH6sKPprPiI_A0qW2RVB6JwGh5LV0DUZwk4,5408
-vellum_ai-0.14.17.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.14.17.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.14.17.dist-info/RECORD,,
+vellum_ai-0.14.18.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.14.18.dist-info/METADATA,sha256=jfwatBq_nGnqnh-_g3EkOaoxQtBIU5O4Gn7n-v3F_do,5408
+vellum_ai-0.14.18.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.14.18.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.14.18.dist-info/RECORD,,

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

@@ -19,7 +19,6 @@ _InlinePromptNodeType = TypeVar("_InlinePromptNodeType", bound=InlinePromptNode)
 class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType], Generic[_InlinePromptNodeType]):
     output_id: ClassVar[Optional[UUID]] = None
     array_output_id: ClassVar[Optional[UUID]] = None
-    json_output_id: ClassVar[Optional[UUID]] = None
     prompt_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(
@@ -33,6 +32,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
 
         _, 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)
         function_definitions = raise_if_descriptor(node.functions)
 
@@ -71,6 +71,11 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "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},
+            ],
         }
 
     def _generate_node_and_prompt_inputs(

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

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import ClassVar, Generic, Optional, TypeVar
+from typing import Generic, Optional, TypeVar
 
 from vellum.workflows.nodes.core.templating_node import TemplatingNode
 from vellum.workflows.types.core import JsonObject
@@ -11,23 +11,21 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _TemplatingNodeType = TypeVar("_TemplatingNodeType", bound=TemplatingNode)
 
-TEMPLATE_INPUT_NAME = "template"
+TEMPLATE_INPUT_NAME = TemplatingNode.template.name
 
 
 class BaseTemplatingNodeDisplay(BaseNodeVellumDisplay[_TemplatingNodeType], Generic[_TemplatingNodeType]):
-    template_input_id: ClassVar[Optional[UUID]] = None
-
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
 
-        template_input_id = self.template_input_id or self.node_input_ids_by_name.get(TEMPLATE_INPUT_NAME)
+        template_input_id = self.node_input_ids_by_name.get(TEMPLATE_INPUT_NAME)
 
         template_node_input = create_node_input(
             node_id=node_id,
-            input_name="template",
+            input_name=TEMPLATE_INPUT_NAME,
             value=node.template,
             display_context=display_context,
             input_id=template_input_id,
@@ -39,7 +37,8 @@ class BaseTemplatingNodeDisplay(BaseNodeVellumDisplay[_TemplatingNodeType], Gene
                 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"{TemplatingNode.inputs.name}.{variable_name}")
+                or self.node_input_ids_by_name.get(variable_name),
             )
             for variable_name, variable_value in template_node_inputs.items()
             if variable_name != TEMPLATE_INPUT_NAME

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

@@ -0,0 +1,97 @@
+import pytest
+from uuid import UUID
+from typing import Type
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
+from vellum_ee.workflows.display.nodes.vellum.templating_node import BaseTemplatingNodeDisplay
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
+
+
+def _no_display_class(Node: Type[TemplatingNode]):
+    return None
+
+
+def _display_class_with_node_input_ids_by_name(Node: Type[TemplatingNode]):
+    class TemplatingNodeDisplay(BaseTemplatingNodeDisplay[Node]):  # type: ignore[valid-type]
+        node_input_ids_by_name = {"foo": UUID("fba6a4d5-835a-4e99-afb7-f6a4aed15110")}
+
+    return TemplatingNodeDisplay
+
+
+def _display_class_with_node_input_ids_by_name_with_inputs_prefix(Node: Type[TemplatingNode]):
+    class TemplatingNodeDisplay(BaseTemplatingNodeDisplay[Node]):  # type: ignore[valid-type]
+        node_input_ids_by_name = {"inputs.foo": UUID("fba6a4d5-835a-4e99-afb7-f6a4aed15110")}
+
+    return TemplatingNodeDisplay
+
+
+@pytest.mark.parametrize(
+    ["GetDisplayClass", "expected_input_id"],
+    [
+        (_no_display_class, "d3519cec-590c-416d-8eb1-96051aed5ddd"),
+        (_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"),
+    ],
+    ids=[
+        "no_display_class",
+        "display_class_with_node_input_ids_by_name",
+        "display_class_with_node_input_ids_by_name_with_inputs_prefix",
+    ],
+)
+def test_serialize_node__templating_node_inputs(GetDisplayClass, expected_input_id):
+    # GIVEN a templating node with inputs
+    class MyTemplatingNode(TemplatingNode):
+        inputs = {"foo": "bar"}
+
+    # AND a workflow with the templating node
+    class Workflow(BaseWorkflow):
+        graph = MyTemplatingNode
+
+    # AND a display class
+    GetDisplayClass(MyTemplatingNode)
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the inputs
+    my_templating_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "TEMPLATING"
+    )
+
+    assert my_templating_node["inputs"] == [
+        {
+            "id": "b7d971ee-9156-46b2-9bf0-9292875211f8",
+            "key": "template",
+            "value": {
+                "combinator": "OR",
+                "rules": [
+                    {
+                        "type": "CONSTANT_VALUE",
+                        "data": {
+                            "type": "STRING",
+                            "value": "",
+                        },
+                    }
+                ],
+            },
+        },
+        {
+            "id": expected_input_id,
+            "key": "foo",
+            "value": {
+                "combinator": "OR",
+                "rules": [
+                    {
+                        "type": "CONSTANT_VALUE",
+                        "data": {
+                            "type": "STRING",
+                            "value": "bar",
+                        },
+                    }
+                ],
+            },
+        },
+    ]

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

@@ -54,7 +54,7 @@ def create_node_input_value_pointer_rules(
 
     if isinstance(value, BaseDescriptor):
         if isinstance(value, NodeReference):
-            if not value.instance:
+            if value.instance is None:
                 raise ValueError(f"Expected NodeReference {value.name} to have an instance")
             value = cast(BaseDescriptor, value.instance)
 

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_templating_node_serialization.py

@@ -85,7 +85,7 @@ def test_serialize_workflow():
                     },
                 },
                 {
-                    "id": "bb712499-4265-4d71-bc14-0c3d8ca6a7de",
+                    "id": "e6c9d062-1a4f-4d4d-a258-59b75c6062fc",
                     "key": "info",
                     "value": {
                         "rules": [

vellum_ee/workflows/display/vellum.py

@@ -1,11 +1,10 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from uuid import UUID
 from typing import List, Literal, Optional, Union
 
 from pydantic import Field
 
-from vellum import PromptParameters, VellumVariable, VellumVariableType
+from vellum import VellumVariableType
 from vellum.client.types.array_vellum_value import ArrayVellumValue
 from vellum.client.types.vellum_value import VellumValue
 from vellum.core import UniversalBaseModel
@@ -138,24 +137,6 @@ class WorkflowOutputVellumDisplay(WorkflowOutputVellumDisplayOverrides):
     pass
 
 
-class WorkflowNodeType(str, Enum):
-    PROMPT = "PROMPT"
-    TEMPLATING = "TEMPLATING"
-    NOTE = "NOTE"
-    CODE_EXECUTION = "CODE_EXECUTION"
-    METRIC = "METRIC"
-    SEARCH = "SEARCH"
-    WEBHOOK = "WEBHOOK"
-    MERGE = "MERGE"
-    CONDITIONAL = "CONDITIONAL"
-    API = "API"
-    ENTRYPOINT = "ENTRYPOINT"
-    TERMINAL = "TERMINAL"
-    SUBWORKFLOW = "SUBWORKFLOW"
-    MAP = "MAP"
-    ERROR = "ERROR"
-
-
 class ConstantValuePointer(UniversalBaseModel):
     type: Literal["CONSTANT_VALUE"] = "CONSTANT_VALUE"
     data: VellumValue
@@ -220,131 +201,3 @@ class NodeInput(UniversalBaseModel):
     id: str
     key: str
     value: NodeInputValuePointer
-
-
-class BaseWorkflowNode(UniversalBaseModel):
-    id: str
-    inputs: List[NodeInput]
-    type: str
-    display_data: Optional[NodeDisplayData] = None
-    base: CodeResourceDefinition
-    definition: CodeResourceDefinition
-
-
-class EntrypointNodeData(UniversalBaseModel):
-    source_handle_id: str
-
-
-class EntrypointNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.ENTRYPOINT] = WorkflowNodeType.ENTRYPOINT
-    data: EntrypointNodeData
-
-
-class PromptTemplateBlockData(UniversalBaseModel):
-    version: Literal[1] = 1
-    # blocks: List[PromptBlockRequest]
-
-
-class PromptVersionExecConfig(UniversalBaseModel):
-    parameters: PromptParameters
-    input_variables: List[VellumVariable]
-    prompt_template_block_data: PromptTemplateBlockData
-
-
-class BasePromptNodeData(UniversalBaseModel):
-    label: str
-    output_id: str
-    error_output_id: Optional[str] = None
-    array_output_id: str
-    source_handle_id: str
-    target_handle_id: str
-
-
-class InlinePromptNodeData(BasePromptNodeData):
-    variant: Literal["INLINE"] = "INLINE"
-    exec_config: PromptVersionExecConfig
-    ml_model_name: str
-
-
-class DeploymentPromptNodeData(BasePromptNodeData):
-    variant: Literal["DEPLOYMENT"] = "DEPLOYMENT"
-    deployment_id: str
-    release_tag: str
-
-
-PromptNodeData = Union[
-    InlinePromptNodeData,
-    DeploymentPromptNodeData,
-]
-
-
-class PromptNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.PROMPT] = WorkflowNodeType.PROMPT
-    data: PromptNodeData
-
-
-class SearchNodeData(UniversalBaseModel):
-    label: str
-
-    results_output_id: str
-    text_output_id: str
-    error_output_id: Optional[str] = None
-
-    source_handle_id: str
-    target_handle_id: str
-
-    query_node_input_id: str
-    document_index_node_input_id: str
-    weights_node_input_id: str
-    limit_node_input_id: str
-    separator_node_input_id: str
-    result_merging_enabled_node_input_id: str
-    external_id_filters_node_input_id: str
-    metadata_filters_node_input_id: str
-
-
-class SearchNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.SEARCH] = WorkflowNodeType.SEARCH
-    data: SearchNodeData
-
-
-class FinalOutputNodeData(UniversalBaseModel):
-    label: str
-    name: str
-    target_handle_id: str
-    output_id: str
-    output_type: VellumVariableType
-    node_input_id: str
-
-
-class FinalOutputNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.TERMINAL] = WorkflowNodeType.TERMINAL
-    data: FinalOutputNodeData
-
-
-WorkflowNode = Union[
-    EntrypointNode,
-    PromptNode,
-    SearchNode,
-    FinalOutputNode,
-]
-
-
-class WorkflowEdge(UniversalBaseModel):
-    id: str
-    source_node_id: str
-    source_handle_id: str
-    target_node_id: str
-    target_handle_id: str
-
-
-class WorkflowRawData(UniversalBaseModel):
-    nodes: List[WorkflowNode]
-    edges: List[WorkflowEdge]
-    display_data: Optional[WorkflowDisplayData] = None
-
-
-class WorkflowVersionExecConfig(UniversalBaseModel):
-    workflow_raw_data: WorkflowRawData
-    input_variables: List[VellumVariable]
-    output_variables: List[VellumVariable]

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -429,7 +429,7 @@ class BaseWorkflowDisplay(
                     )
                     subworkflow_display_context = subworkflow_display.get_event_display_context()
 
-            node_event_displays[str(node_id)] = NodeEventDisplayContext(
+            node_event_displays[node_id] = NodeEventDisplayContext(
                 input_display=input_display,
                 output_display=output_display,
                 port_display=port_display_meta,

vellum_ee/workflows/display/workflows/tests/test_workflow_display.py

@@ -114,8 +114,8 @@ def test_get_event_display_context__node_display_filled_without_base_display():
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the node display should be included
-    assert str(StartNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(StartNode.__id__)]
+    assert StartNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[StartNode.__id__]
 
     # AND so should their output ids
     assert StartNode.__output_ids__ == node_event_display.output_display
@@ -137,8 +137,8 @@ def test_get_event_display_context__node_display_filled_without_output_display()
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the node display should be included
-    assert str(StartNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(StartNode.__id__)]
+    assert StartNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[StartNode.__id__]
 
     # AND so should their output ids
     assert node_event_display.output_display.keys() == {"foo"}
@@ -163,11 +163,11 @@ def test_get_event_display_context__node_display_to_include_subworkflow_display(
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
-    assert str(SubworkflowNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(SubworkflowNode.__id__)]
+    assert SubworkflowNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[SubworkflowNode.__id__]
 
     assert node_event_display.subworkflow_display is not None
-    assert str(InnerNode.__id__) in node_event_display.subworkflow_display.node_displays
+    assert InnerNode.__id__ in node_event_display.subworkflow_display.node_displays
 
 
 @pytest.mark.parametrize(
@@ -204,13 +204,13 @@ def test_get_event_display_context__node_display_for_adornment_nodes(
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
-    assert str(adornment_node_id) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(adornment_node_id)]
+    assert adornment_node_id in display_context.node_displays
+    node_event_display = display_context.node_displays[adornment_node_id]
     assert node_event_display.subworkflow_display is not None
-    assert str(inner_node_id) in node_event_display.subworkflow_display.node_displays
+    assert inner_node_id in node_event_display.subworkflow_display.node_displays
 
     # AND the inner node should have the correct outputs
-    inner_node_display = node_event_display.subworkflow_display.node_displays[str(inner_node_id)]
+    inner_node_display = node_event_display.subworkflow_display.node_displays[inner_node_id]
     assert inner_node_display.output_display.keys() == {"foo"}
     assert node_event_display.output_display.keys() == expected_adornment_output_names
 
@@ -231,7 +231,7 @@ def test_get_event_display_context__templating_node_input_display():
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
-    assert str(MyNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(MyNode.__id__)]
+    assert MyNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[MyNode.__id__]
 
     assert node_event_display.input_display.keys() == {"inputs.foo"}

vellum_ee/workflows/tests/test_display_meta.py

@@ -1,7 +1,7 @@
 import pytest
 import os
 import sys
-from uuid import UUID, uuid4
+from uuid import uuid4
 
 from vellum.workflows import BaseWorkflow
 from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
@@ -35,22 +35,22 @@ def test_base_class_dynamic_import(files):
         "node_displays": {
             "533c6bd8-6088-4abc-a168-8c1758abcd33": {
                 "input_display": {
-                    "example_var_1": UUID("a0d1d7cf-242a-4bd9-a437-d308a7ced9b3"),
-                    "template": UUID("f97d721a-e685-498e-90c3-9c3d9358fdad"),
+                    "example_var_1": "a0d1d7cf-242a-4bd9-a437-d308a7ced9b3",
+                    "template": "f97d721a-e685-498e-90c3-9c3d9358fdad",
                 },
-                "output_display": {"result": UUID("423bc529-1a1a-4f72-af4d-cbdb5f0a5929")},
-                "port_display": {"default": UUID("afda9a19-0618-42e1-9b63-5d0db2a88f62")},
+                "output_display": {"result": "423bc529-1a1a-4f72-af4d-cbdb5f0a5929"},
+                "port_display": {"default": "afda9a19-0618-42e1-9b63-5d0db2a88f62"},
                 "subworkflow_display": None,
             },
             "f3ef4b2b-fec9-4026-9cc6-e5eac295307f": {
-                "input_display": {"node_input": UUID("fe6cba85-2423-4b5e-8f85-06311a8be5fb")},
-                "output_display": {"value": UUID("5469b810-6ea6-4362-9e79-e360d44a1405")},
+                "input_display": {"node_input": "fe6cba85-2423-4b5e-8f85-06311a8be5fb"},
+                "output_display": {"value": "5469b810-6ea6-4362-9e79-e360d44a1405"},
                 "port_display": {},
                 "subworkflow_display": None,
             },
         },
-        "workflow_inputs": {"input_value": UUID("2268a996-bd17-4832-b3ff-f5662d54b306")},
-        "workflow_outputs": {"final_output": UUID("5469b810-6ea6-4362-9e79-e360d44a1405")},
+        "workflow_inputs": {"input_value": "2268a996-bd17-4832-b3ff-f5662d54b306"},
+        "workflow_outputs": {"final_output": "5469b810-6ea6-4362-9e79-e360d44a1405"},
     }
     assert display_meta
-    assert display_meta.dict() == expected_result
+    assert display_meta.model_dump(mode="json") == expected_result