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

vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -2,7 +2,11 @@ import pytest
 from uuid import UUID
 
 from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.ports.port import Port
+from vellum.workflows.references.constant import ConstantValueReference
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 
 @pytest.fixture
@@ -46,3 +50,53 @@ def test_get_id(node_info):
 
     assert node_display().node_id == expected_id
     assert node_display.infer_node_class().__id__ == expected_id
+
+
+def test_serialize_condition__accessor_expression():
+    # GIVEN a node with an accessor expression in a Port
+    class MyNode(BaseNode):
+        class Ports(BaseNode.Ports):
+            foo = Port.on_if(ConstantValueReference({"hello": "world"})["hello"])
+
+    # WHEN we serialize the node
+    node_display_class = get_node_display_class(MyNode)
+    data = node_display_class().serialize(WorkflowDisplayContext())
+
+    # THEN the condition should be serialized correctly
+    assert data["ports"] == [
+        {
+            "id": "7de6ea94-7f6c-475e-8f38-ec8ac317fd19",
+            "name": "foo",
+            "type": "IF",
+            "expression": {
+                "type": "BINARY_EXPRESSION",
+                "lhs": {
+                    "type": "CONSTANT_VALUE",
+                    "value": {
+                        "type": "JSON",
+                        "value": {"hello": "world"},
+                    },
+                },
+                "operator": "accessField",
+                "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "hello"}},
+            },
+        }
+    ]
+
+
+def test_serialize_display_data():
+    # GIVEN a node with an accessor expression in a Port
+    class MyNode(BaseNode):
+        """I hope this works"""
+
+        pass
+
+    # WHEN we serialize the node
+    node_display_class = get_node_display_class(MyNode)
+    data = node_display_class().serialize(WorkflowDisplayContext())
+
+    # THEN the condition should be serialized correctly
+    assert data["display_data"] == {
+        "position": {"x": 0.0, "y": 0.0},
+        "comment": {"value": "I hope this works"},
+    }

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

@@ -4,7 +4,7 @@ from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar, cast
 from vellum.workflows.nodes.displayable import APINode
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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
@@ -13,7 +13,7 @@ from vellum_ee.workflows.display.utils.vellum import WorkspaceSecretPointer
 _APINodeType = TypeVar("_APINodeType", bound=APINode)
 
 
-class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeType]):
+class BaseAPINodeDisplay(BaseNodeDisplay[_APINodeType], Generic[_APINodeType]):
     # A mapping between node input keys and their ids for inputs representing additional header keys
     additional_header_key_input_ids: ClassVar[Optional[Dict[str, UUID]]] = None
 

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

@@ -10,10 +10,10 @@ 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
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.registry import register_node_display_class
 
 _BaseAdornmentNodeType = TypeVar("_BaseAdornmentNodeType", bound=BaseAdornmentNode)
 
@@ -31,7 +31,7 @@ def _recursively_replace_wrapped_node(node_class: Type[BaseNode], wrapped_node_d
     # 1. The node display class' parameterized type
     original_base_node_display = get_original_base(wrapped_node_display_class)
     original_base_node_display.__args__ = (wrapped_node_class,)
-    wrapped_node_display_class._node_display_registry[wrapped_node_class] = wrapped_node_display_class
+    register_node_display_class(node_class=wrapped_node_class, node_display_class=wrapped_node_display_class)
     wrapped_node_display_class.__annotate_node__()
 
     # 2. The node display class' output displays
@@ -71,7 +71,7 @@ def _recursively_replace_wrapped_node(node_class: Type[BaseNode], wrapped_node_d
         )
 
 
-class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Generic[_BaseAdornmentNodeType]):
+class BaseAdornmentNodeDisplay(BaseNodeDisplay[_BaseAdornmentNodeType], Generic[_BaseAdornmentNodeType]):
     __wrapped_node_display__: Optional[Type[BaseNodeDisplay]] = None
 
     def serialize(
@@ -89,7 +89,7 @@ class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Ge
                 "Unable to serialize standalone adornment nodes. Please use adornment nodes as a decorator."
             )
 
-        wrapped_node_display_class = get_node_display_class(BaseNodeDisplay, wrapped_node)
+        wrapped_node_display_class = get_node_display_class(wrapped_node)
         wrapped_node_display = wrapped_node_display_class()
         additional_kwargs = get_additional_kwargs(wrapped_node_display.node_id) if get_additional_kwargs else {}
         serialized_wrapped_node = wrapped_node_display.serialize(display_context, **kwargs, **additional_kwargs)

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

@@ -6,7 +6,7 @@ from vellum.workflows.nodes.displayable.code_execution_node import CodeExecution
 from vellum.workflows.nodes.displayable.code_execution_node.utils import read_file_from_path
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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
@@ -14,7 +14,7 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 _CodeExecutionNodeType = TypeVar("_CodeExecutionNodeType", bound=CodeExecutionNode)
 
 
-class BaseCodeExecutionNodeDisplay(BaseNodeVellumDisplay[_CodeExecutionNodeType], Generic[_CodeExecutionNodeType]):
+class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Generic[_CodeExecutionNodeType]):
     output_id: ClassVar[Optional[UUID]] = None
     log_output_id: ClassVar[Optional[UUID]] = None
 

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

@@ -16,7 +16,7 @@ from vellum.workflows.expressions.or_ import OrExpression
 from vellum.workflows.nodes.displayable import ConditionalNode
 from vellum.workflows.types.core import ConditionType, JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 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.vellum import convert_descriptor_to_operator
@@ -40,7 +40,7 @@ class ConditionId:
     rule_group_id: Optional[str]
 
 
-class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Generic[_ConditionalNodeType]):
+class BaseConditionalNodeDisplay(BaseNodeDisplay[_ConditionalNodeType], Generic[_ConditionalNodeType]):
     source_handle_ids: ClassVar[Dict[int, UUID]]
     rule_ids: ClassVar[List[RuleIdMap]]
     condition_ids: ClassVar[list[ConditionId]]

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

@@ -3,7 +3,7 @@ from typing import ClassVar, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes import ErrorNode
 from vellum.workflows.types.core import JsonObject
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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
@@ -11,7 +11,7 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 _ErrorNodeType = TypeVar("_ErrorNodeType", bound=ErrorNode)
 
 
-class BaseErrorNodeDisplay(BaseNodeVellumDisplay[_ErrorNodeType], Generic[_ErrorNodeType]):
+class BaseErrorNodeDisplay(BaseNodeDisplay[_ErrorNodeType], Generic[_ErrorNodeType]):
     error_output_id: ClassVar[Optional[UUID]] = None
 
     name: ClassVar[str] = "error-node"

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

@@ -4,7 +4,7 @@ from typing import Any, ClassVar, Generic, Optional, TypeVar
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import to_kebab_case
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -13,7 +13,7 @@ from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 _FinalOutputNodeType = TypeVar("_FinalOutputNodeType", bound=FinalOutputNode)
 
 
-class BaseFinalOutputNodeDisplay(BaseNodeVellumDisplay[_FinalOutputNodeType], Generic[_FinalOutputNodeType]):
+class BaseFinalOutputNodeDisplay(BaseNodeDisplay[_FinalOutputNodeType], Generic[_FinalOutputNodeType]):
     output_id: ClassVar[Optional[UUID]] = None
     output_name: ClassVar[Optional[str]] = None
     node_input_id: ClassVar[Optional[UUID]] = None

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

@@ -3,7 +3,7 @@ from typing import Generic, Optional, TypeVar
 
 from vellum.workflows.nodes import GuardrailNode
 from vellum.workflows.types.core import JsonObject
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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
@@ -11,7 +11,7 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 _GuardrailNodeType = TypeVar("_GuardrailNodeType", bound=GuardrailNode)
 
 
-class BaseGuardrailNodeDisplay(BaseNodeVellumDisplay[_GuardrailNodeType], Generic[_GuardrailNodeType]):
+class BaseGuardrailNodeDisplay(BaseNodeDisplay[_GuardrailNodeType], Generic[_GuardrailNodeType]):
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

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

@@ -6,7 +6,7 @@ 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.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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
@@ -16,7 +16,7 @@ from vellum_ee.workflows.display.vellum import NodeInput
 _InlinePromptNodeType = TypeVar("_InlinePromptNodeType", bound=InlinePromptNode)
 
 
-class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType], Generic[_InlinePromptNodeType]):
+class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generic[_InlinePromptNodeType]):
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

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

@@ -5,7 +5,7 @@ from vellum import VellumVariable
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes import InlineSubworkflowNode
 from vellum.workflows.types.core import JsonObject
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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
@@ -17,7 +17,7 @@ _InlineSubworkflowNodeType = TypeVar("_InlineSubworkflowNodeType", bound=InlineS
 
 
 class BaseInlineSubworkflowNodeDisplay(
-    BaseNodeVellumDisplay[_InlineSubworkflowNodeType], Generic[_InlineSubworkflowNodeType]
+    BaseNodeDisplay[_InlineSubworkflowNodeType], Generic[_InlineSubworkflowNodeType]
 ):
     workflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 

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

@@ -4,13 +4,13 @@ from typing import Any, ClassVar, Generic, List, Optional, TypeVar
 from vellum.workflows.nodes.displayable import MergeNode
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _MergeNodeType = TypeVar("_MergeNodeType", bound=MergeNode)
 
 
-class BaseMergeNodeDisplay(BaseNodeVellumDisplay[_MergeNodeType], Generic[_MergeNodeType]):
+class BaseMergeNodeDisplay(BaseNodeDisplay[_MergeNodeType], Generic[_MergeNodeType]):
     target_handle_ids: ClassVar[List[UUID]]
 
     def __init__(self):

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

@@ -2,13 +2,13 @@ from typing import Any, ClassVar, Dict, Generic, TypeVar, Union
 
 from vellum.workflows.nodes import NoteNode
 from vellum.workflows.types.core import JsonObject
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _NoteNodeType = TypeVar("_NoteNodeType", bound=NoteNode)
 
 
-class BaseNoteNodeDisplay(BaseNodeVellumDisplay[_NoteNodeType], Generic[_NoteNodeType]):
+class BaseNoteNodeDisplay(BaseNodeDisplay[_NoteNodeType], Generic[_NoteNodeType]):
     text: ClassVar[str] = ""
     style: ClassVar[Union[Dict[str, Any], None]] = None
 

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

@@ -5,7 +5,7 @@ from vellum.workflows.nodes.displayable.prompt_deployment_node import PromptDepl
 from vellum.workflows.references import OutputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.vellum_client import create_vellum_client
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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
@@ -13,9 +13,7 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 _PromptDeploymentNodeType = TypeVar("_PromptDeploymentNodeType", bound=PromptDeploymentNode)
 
 
-class BasePromptDeploymentNodeDisplay(
-    BaseNodeVellumDisplay[_PromptDeploymentNodeType], Generic[_PromptDeploymentNodeType]
-):
+class BasePromptDeploymentNodeDisplay(BaseNodeDisplay[_PromptDeploymentNodeType], Generic[_PromptDeploymentNodeType]):
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

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

@@ -9,7 +9,6 @@ from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 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
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
@@ -72,7 +71,7 @@ class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_Re
         if not inner_node:
             return super().get_node_output_display(output)
 
-        node_display_class = get_node_display_class(BaseNodeDisplay, inner_node)
+        node_display_class = get_node_display_class(inner_node)
         node_display = node_display_class()
 
         inner_output = getattr(inner_node.Outputs, output.name)

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

@@ -11,7 +11,7 @@ from vellum.workflows.nodes.displayable.search_node import SearchNode
 from vellum.workflows.references import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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
@@ -28,7 +28,7 @@ class VariableIdMap:
     rhs: Optional["VariableIdMap"]
 
 
-class BaseSearchNodeDisplay(BaseNodeVellumDisplay[_SearchNodeType], Generic[_SearchNodeType]):
+class BaseSearchNodeDisplay(BaseNodeDisplay[_SearchNodeType], Generic[_SearchNodeType]):
     # A mapping between the id of the operand (e.g. "lhs_variable_id" or "rhs_variable_id") and the id of the node input
     # that the operand is pointing to.
     metadata_filter_input_id_by_operand_id: Dict[UUID, UUID] = {}

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

@@ -4,7 +4,7 @@ from typing import Generic, Optional, TypeVar
 from vellum.workflows.nodes import SubworkflowDeploymentNode
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.vellum_client import create_vellum_client
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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
@@ -13,7 +13,7 @@ _SubworkflowDeploymentNodeType = TypeVar("_SubworkflowDeploymentNodeType", bound
 
 
 class BaseSubworkflowDeploymentNodeDisplay(
-    BaseNodeVellumDisplay[_SubworkflowDeploymentNodeType], Generic[_SubworkflowDeploymentNodeType]
+    BaseNodeDisplay[_SubworkflowDeploymentNodeType], Generic[_SubworkflowDeploymentNodeType]
 ):
 
     def serialize(

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

@@ -4,7 +4,7 @@ from typing import Generic, Optional, TypeVar
 from vellum.workflows.nodes.core.templating_node import TemplatingNode
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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
@@ -14,7 +14,7 @@ _TemplatingNodeType = TypeVar("_TemplatingNodeType", bound=TemplatingNode)
 TEMPLATE_INPUT_NAME = TemplatingNode.template.name
 
 
-class BaseTemplatingNodeDisplay(BaseNodeVellumDisplay[_TemplatingNodeType], Generic[_TemplatingNodeType]):
+class BaseTemplatingNodeDisplay(BaseNodeDisplay[_TemplatingNodeType], Generic[_TemplatingNodeType]):
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

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

@@ -6,7 +6,6 @@ from vellum.workflows.nodes.displayable.code_execution_node.node import CodeExec
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes.vellum.code_execution_node import BaseCodeExecutionNodeDisplay
 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[CodeExecutionNode]):
@@ -53,7 +52,7 @@ def test_serialize_node__code_node_inputs(GetDisplayClass, expected_input_id):
     GetDisplayClass(MyCodeExecutionNode)
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN the node should properly serialize the inputs

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

@@ -4,7 +4,6 @@ from vellum.client.types.vellum_error import VellumError
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes.core.error_node.node import ErrorNode
 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 test_error_node_display__serialize_with_vellum_error() -> None:
@@ -20,7 +19,7 @@ def test_error_node_display__serialize_with_vellum_error() -> None:
         graph = MyNode
 
     # WHEN we serialize the workflow
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=MyWorkflow)
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
     serialized_workflow = cast(Dict[str, Any], workflow_display.serialize())
 
     # THEN the correct inputs should be serialized on the node

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

@@ -2,7 +2,6 @@ from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes.displayable.note_node.node import NoteNode
 from vellum_ee.workflows.display.nodes.vellum.note_node import BaseNoteNodeDisplay
 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 test_serialize_node__note_node():
@@ -22,7 +21,7 @@ def test_serialize_node__note_node():
         graph = MyNoteNode
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN the node should properly serialize the inputs

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

@@ -6,9 +6,9 @@ from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
 from vellum.workflows.references.lazy import LazyReference
+from vellum.workflows.state.base import BaseState
 from vellum_ee.workflows.display.nodes.vellum.inline_prompt_node import BaseInlinePromptNodeDisplay
 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 test_serialize_node__lazy_reference_in_prompt_inputs():
@@ -27,7 +27,7 @@ def test_serialize_node__lazy_reference_in_prompt_inputs():
         graph = LazyReferencePromptNode >> OtherNode
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN the node should properly serialize the attribute reference
@@ -103,7 +103,7 @@ def test_serialize_node__prompt_inputs(GetDisplayClass, expected_input_id):
     GetDisplayClass(MyPromptNode)
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN the node should properly serialize the inputs
@@ -129,3 +129,55 @@ def test_serialize_node__prompt_inputs(GetDisplayClass, expected_input_id):
             },
         }
     ]
+
+
+def test_serialize_node__prompt_inputs__state_reference():
+    # GIVEN a state definition
+    class MyState(BaseState):
+        foo: str
+
+    # AND a prompt node with inputs
+    class MyPromptNode(InlinePromptNode):
+        prompt_inputs = {"foo": MyState.foo, "bar": "baz"}
+        blocks = []
+        ml_model = "gpt-4o"
+
+    # AND a workflow with the prompt node
+    class Workflow(BaseWorkflow):
+        graph = MyPromptNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should skip the state reference input rule
+    my_prompt_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MyPromptNode.__id__)
+    )
+
+    assert my_prompt_node["inputs"] == [
+        {
+            "id": "e47e0a80-afbb-4888-b06b-8dc78edd8572",
+            "key": "foo",
+            "value": {
+                "rules": [],
+                "combinator": "OR",
+            },
+        },
+        {
+            "id": "3750feb9-5d5c-4150-b62d-a9924f466888",
+            "key": "bar",
+            "value": {
+                "rules": [
+                    {
+                        "type": "CONSTANT_VALUE",
+                        "data": {
+                            "type": "STRING",
+                            "value": "baz",
+                        },
+                    }
+                ],
+                "combinator": "OR",
+            },
+        },
+    ]

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

@@ -5,7 +5,6 @@ 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_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 test_retry_node_parameters():
@@ -21,7 +20,7 @@ def test_retry_node_parameters():
         graph = MyRetryNode
 
     # WHEN we serialize the workflow
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=MyWorkflow)
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
     serialized_workflow = cast(Dict[str, Any], workflow_display.serialize())
 
     # THEN the correct inputs should be serialized on the node

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

@@ -6,7 +6,6 @@ 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]):
@@ -53,7 +52,7 @@ def test_serialize_node__templating_node_inputs(GetDisplayClass, expected_input_
     GetDisplayClass(MyTemplatingNode)
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN the node should properly serialize the inputs

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

@@ -5,7 +5,6 @@ from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
 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 test_try_node_display__serialize_with_error_output() -> None:
@@ -28,7 +27,7 @@ def test_try_node_display__serialize_with_error_output() -> None:
         graph = MyNode >> OtherNode
 
     # WHEN we serialize the workflow
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=MyWorkflow)
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
     serialized_workflow = cast(Dict[str, Any], workflow_display.serialize())
 
     # THEN the correct inputs should be serialized on the node

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

@@ -9,7 +9,7 @@ from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.references import LazyReference
 from vellum_ee.workflows.display.editor.types import NodeDisplayData
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input_value_pointer_rules
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -22,7 +22,7 @@ from vellum_ee.workflows.display.utils.vellum import (
     NodeOutputPointer,
 )
 from vellum_ee.workflows.display.vellum import WorkflowInputsVellumDisplayOverrides, WorkflowMetaVellumDisplay
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
+from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
 class Inputs(BaseInputs):
@@ -36,7 +36,7 @@ class MyNodeA(BaseNode):
         output: str
 
 
-class MyNodeADisplay(BaseNodeVellumDisplay[MyNodeA]):
+class MyNodeADisplay(BaseNodeDisplay[MyNodeA]):
     pass
 
 
@@ -109,7 +109,7 @@ def test_create_node_input_value_pointer_rules(
     rules = create_node_input_value_pointer_rules(
         descriptor,
         WorkflowDisplayContext(
-            workflow_display_class=VellumWorkflowDisplay,
+            workflow_display_class=BaseWorkflowDisplay,
             workflow_display=WorkflowMetaVellumDisplay(
                 entrypoint_node_id=uuid4(),
                 entrypoint_node_source_handle_id=uuid4(),

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

@@ -10,7 +10,6 @@ from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 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
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
@@ -82,7 +81,7 @@ class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNod
         if not inner_node:
             return super().get_node_output_display(output)
 
-        node_display_class = get_node_display_class(BaseNodeDisplay, inner_node)
+        node_display_class = get_node_display_class(inner_node)
         node_display = node_display_class()
         if output.name == "error":
             return inner_node, NodeOutputDisplay(