vellum-ai 0.14.38__tar.gz → 0.14.40__tar.gz

{vellum_ai-0.14.38 → vellum_ai-0.14.40}/PKG-INFO

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

{vellum_ai-0.14.38 → vellum_ai-0.14.40}/ee/vellum_cli/push.py

@@ -17,7 +17,6 @@ from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, WorkflowConfig, WorkflowDeploymentConfig, load_vellum_cli_config
 from vellum_cli.logger import load_cli_logger
 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 push_command(
@@ -36,9 +35,6 @@ def push_command(
     logger = load_cli_logger()
     config = load_vellum_cli_config()
 
-    if not config.workflows:
-        raise ValueError("No Workflows found in project to push.")
-
     workflow_configs = (
         [
             w
@@ -61,6 +57,8 @@ def push_command(
             )
             config.workflows.append(new_config)
             workflow_configs = [new_config]
+        elif not module:
+            raise ValueError("No Workflows found in project to push.")
         else:
             raise ValueError(f"No workflow config for '{module}' found in project to push.")
 
@@ -106,7 +104,6 @@ def push_command(
     # https://app.shortcut.com/vellum/story/5585
     workflow = BaseWorkflow.load_from_module(workflow_config.module)
     workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
         workflow_class=workflow,
         dry_run=dry_run or False,
     )

{vellum_ai-0.14.38 → vellum_ai-0.14.40}/ee/vellum_cli/tests/test_push.py

@@ -133,6 +133,58 @@ def test_push__happy_path(mock_module, vellum_client, base_command):
     assert extracted_files["workflow.py"] == workflow_py_file_content
 
 
+def test_push__no_config__module_found(mock_module, vellum_client):
+    # GIVEN no config file set
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    mock_module.set_pyproject_toml({"workflows": []})
+
+    # AND a workflow exists in the module successfully
+    workflow_py_file_content = _ensure_workflow_py(temp_dir, module)
+
+    # AND the push API call returns successfully
+    new_workflow_sandbox_id = str(uuid4())
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=new_workflow_sandbox_id,
+    )
+
+    # WHEN calling `vellum push`
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push", module])
+
+    # THEN it should successfully push the workflow
+    assert result.exit_code == 0, result.stdout
+
+    # Get the last part of the module path and format it
+    expected_artifact_name = f"{mock_module.module.replace('.', '__')}.tar.gz"
+
+    # AND we should have called the push API with the correct args
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+    assert json.loads(call_args["exec_config"])["workflow_raw_data"]["definition"]["name"] == "ExampleWorkflow"
+    assert call_args["workflow_sandbox_id"] is None
+    assert "deplyment_config" not in call_args
+
+    # AND we should have pushed the correct artifact
+    assert call_args["artifact"].name == expected_artifact_name
+    extracted_files = _extract_tar_gz(call_args["artifact"].read())
+    assert extracted_files["workflow.py"] == workflow_py_file_content
+
+    # AND there should be a new entry in the lock file
+    with open(os.path.join(temp_dir, "vellum.lock.json")) as f:
+        lock_file_content = json.load(f)
+        assert lock_file_content["workflows"][0] == {
+            "container_image_name": None,
+            "container_image_tag": None,
+            "deployments": [],
+            "ignore": None,
+            "module": module,
+            "target_directory": None,
+            "workflow_sandbox_id": new_workflow_sandbox_id,
+            "workspace": "default",
+        }
+
+
 @pytest.mark.parametrize(
     "base_command",
     [

{vellum_ai-0.14.38 → vellum_ai-0.14.40}/ee/vellum_ee/workflows/display/base.py

@@ -1,10 +1,12 @@
 from dataclasses import dataclass, field
 from uuid import UUID
-from typing import Optional
+from typing import Optional, Type
 
 from pydantic import Field
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.workflows.utils.uuids import uuid4_from_hash
+from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.editor.types import NodeDisplayData
 
 
@@ -25,6 +27,17 @@ class WorkflowMetaDisplay:
     entrypoint_node_display: NodeDisplayData = Field(default_factory=NodeDisplayData)
     display_data: WorkflowDisplayData = field(default_factory=WorkflowDisplayData)
 
+    @classmethod
+    def get_default(cls, workflow: Type[BaseWorkflow]) -> "WorkflowMetaDisplay":
+        entrypoint_node_id = uuid4_from_hash(f"{workflow.__id__}|entrypoint_node_id")
+        entrypoint_node_source_handle_id = uuid4_from_hash(f"{workflow.__id__}|entrypoint_node_source_handle_id")
+
+        return WorkflowMetaDisplay(
+            entrypoint_node_id=entrypoint_node_id,
+            entrypoint_node_source_handle_id=entrypoint_node_source_handle_id,
+            entrypoint_node_display=NodeDisplayData(),
+        )
+
 
 @dataclass
 class WorkflowMetaDisplayOverrides(WorkflowMetaDisplay):

{vellum_ai-0.14.38 → vellum_ai-0.14.40}/ee/vellum_ee/workflows/display/nodes/base_node_display.py

@@ -21,14 +21,31 @@ from vellum.client.types.code_resource_definition import CodeResourceDefinition
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.expressions.accessor import AccessorExpression
+from vellum.workflows.expressions.and_ import AndExpression
+from vellum.workflows.expressions.begins_with import BeginsWithExpression
 from vellum.workflows.expressions.between import BetweenExpression
+from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
+from vellum.workflows.expressions.contains import ContainsExpression
+from vellum.workflows.expressions.does_not_begin_with import DoesNotBeginWithExpression
+from vellum.workflows.expressions.does_not_contain import DoesNotContainExpression
+from vellum.workflows.expressions.does_not_end_with import DoesNotEndWithExpression
+from vellum.workflows.expressions.ends_with import EndsWithExpression
+from vellum.workflows.expressions.equals import EqualsExpression
+from vellum.workflows.expressions.greater_than import GreaterThanExpression
+from vellum.workflows.expressions.greater_than_or_equal_to import GreaterThanOrEqualToExpression
+from vellum.workflows.expressions.in_ import InExpression
 from vellum.workflows.expressions.is_nil import IsNilExpression
 from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
 from vellum.workflows.expressions.is_not_null import IsNotNullExpression
 from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
 from vellum.workflows.expressions.is_null import IsNullExpression
 from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
+from vellum.workflows.expressions.less_than import LessThanExpression
+from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualToExpression
 from vellum.workflows.expressions.not_between import NotBetweenExpression
+from vellum.workflows.expressions.not_in import NotInExpression
+from vellum.workflows.expressions.or_ import OrExpression
 from vellum.workflows.expressions.parse_json import ParseJsonExpression
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
@@ -50,7 +67,9 @@ from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_var
 from vellum_ee.workflows.display.editor.types import NodeDisplayData
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
+from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
 from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
+from vellum_ee.workflows.display.utils.registry import register_node_display_class
 from vellum_ee.workflows.display.utils.vellum import convert_descriptor_to_operator
 
 if TYPE_CHECKING:
@@ -121,9 +140,6 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
     # Once all nodes are Generic Nodes, we may replace this with a trigger_id or trigger attribute
     target_handle_id: ClassVar[Optional[UUID]] = None
 
-    # Used to store the mapping between node types and their display classes
-    _node_display_registry: Dict[Type[NodeType], Type["BaseNodeDisplay"]] = {}
-
     def serialize(self, display_context: "WorkflowDisplayContext", **kwargs: Any) -> JsonObject:
         node = self._node
         node_id = self.node_id
@@ -149,7 +165,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         adornments = kwargs.get("adornments", None)
         wrapped_node = get_wrapped_node(node)
         if wrapped_node is not None:
-            display_class = get_node_display_class(BaseNodeDisplay, wrapped_node)
+            display_class = get_node_display_class(wrapped_node)
 
             adornment: JsonObject = {
                 "id": str(node_id),
@@ -280,10 +296,6 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
         return self.get_target_handle_id()
 
-    @classmethod
-    def get_from_node_display_registry(cls, node_class: Type[NodeType]) -> Optional[Type["BaseNodeDisplay"]]:
-        return cls._node_display_registry.get(node_class)
-
     @cached_property
     def node_id(self) -> UUID:
         """Can be overridden as a class attribute to specify a custom node id."""
@@ -344,14 +356,12 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
     def __init_subclass__(cls, **kwargs: Any) -> None:
         super().__init_subclass__(**kwargs)
-        if not cls._node_display_registry:
-            cls._node_display_registry[BaseNode] = BaseNodeDisplay
 
         node_class = cls.infer_node_class()
         if node_class is BaseNode:
             return
 
-        cls._node_display_registry[node_class] = cls
+        register_node_display_class(node_class=node_class, node_display_class=cls)
 
     def _get_generic_node_display_data(self) -> NodeDisplayData:
         explicit_value = self._get_explicit_node_display_attr("display_data", NodeDisplayData)
@@ -388,9 +398,29 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 "lhs": lhs,
                 "rhs": rhs,
             }
-        else:
-            lhs = self.serialize_value(display_context, condition._lhs)  # type: ignore[attr-defined]
-            rhs = self.serialize_value(display_context, condition._rhs)  # type: ignore[attr-defined]
+        elif isinstance(
+            condition,
+            (
+                AndExpression,
+                BeginsWithExpression,
+                CoalesceExpression,
+                ContainsExpression,
+                DoesNotBeginWithExpression,
+                DoesNotContainExpression,
+                DoesNotEndWithExpression,
+                EndsWithExpression,
+                EqualsExpression,
+                GreaterThanExpression,
+                GreaterThanOrEqualToExpression,
+                InExpression,
+                LessThanExpression,
+                LessThanOrEqualToExpression,
+                NotInExpression,
+                OrExpression,
+            ),
+        ):
+            lhs = self.serialize_value(display_context, condition._lhs)
+            rhs = self.serialize_value(display_context, condition._rhs)
 
             return {
                 "type": "BINARY_EXPRESSION",
@@ -398,6 +428,15 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 "operator": convert_descriptor_to_operator(condition),
                 "rhs": rhs,
             }
+        elif isinstance(condition, AccessorExpression):
+            return {
+                "type": "BINARY_EXPRESSION",
+                "lhs": self.serialize_value(display_context, condition._base),
+                "operator": "accessField",
+                "rhs": self.serialize_value(display_context, condition._field),
+            }
+
+        raise UnsupportedSerializationException(f"Unsupported condition type: {condition.__class__.__name__}")
 
     def serialize_value(self, display_context: "WorkflowDisplayContext", value: Any) -> JsonObject:
         if isinstance(value, ConstantValueReference):
@@ -458,3 +497,6 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         # If it's not any of the references we know about,
         # then try to serialize it as a nested value
         return self.serialize_condition(display_context, value)
+
+
+register_node_display_class(node_class=BaseNode, node_display_class=BaseNodeDisplay)

{vellum_ai-0.14.38 → vellum_ai-0.14.40}/ee/vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -1,33 +1,27 @@
 import types
 from uuid import UUID
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type
+from typing import TYPE_CHECKING, Any, Dict, Generic, Type, TypeVar
 
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.utils.registry import get_from_node_display_registry
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 
 
-def get_node_display_class(
-    base_class: Type["BaseNodeDisplay"], node_class: Type[NodeType], root_node_class: Optional[Type[NodeType]] = None
-) -> Type["BaseNodeDisplay"]:
-    node_display_class = base_class.get_from_node_display_registry(node_class)
+def get_node_display_class(node_class: Type[NodeType]) -> Type["BaseNodeDisplay"]:
+    node_display_class = get_from_node_display_registry(node_class)
     if node_display_class:
-        if not issubclass(node_display_class, base_class):
-            raise TypeError(
-                f"Expected to find a subclass of '{base_class.__name__}' for node class '{node_class.__name__}'"
-            )
-
         return node_display_class
 
-    base_node_display_class = get_node_display_class(
-        base_class, node_class.__bases__[0], node_class if root_node_class is None else root_node_class
-    )
+    base_node_display_class = get_node_display_class(node_class.__bases__[0])
 
+    # mypy gets upset at dynamic TypeVar's, but it's technically allowed by python
+    _NodeClassType = TypeVar(f"_{node_class.__name__}Type", bound=node_class)  # type: ignore[misc]
     # `base_node_display_class` is always a Generic class, so it's safe to index into it
-    NodeDisplayBaseClass = base_node_display_class[node_class]  # type: ignore[index]
+    NodeDisplayBaseClass = base_node_display_class[_NodeClassType]  # type: ignore[index]
 
     def _get_node_input_ids_by_ref(path: str, inst: Any):
         if isinstance(inst, dict):
@@ -51,7 +45,7 @@ def get_node_display_class(
 
     NodeDisplayClass = types.new_class(
         f"{node_class.__name__}Display",
-        bases=(NodeDisplayBaseClass,),
+        bases=(NodeDisplayBaseClass, Generic[_NodeClassType]),
         exec_body=exec_body,
     )
 

vellum_ai-0.14.40/ee/vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -0,0 +1,84 @@
+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
+def node_with_implicit_properties():
+    class MyNode(BaseNode):
+        pass
+
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        pass
+
+    expected_id = UUID("ace7f746-4fe6-45c7-8207-fc8a4d0c7f6f")
+
+    return MyNodeDisplay, expected_id
+
+
+@pytest.fixture
+def node_with_explicit_properties():
+    explicit_id = UUID("a422f67a-1d37-43f0-bdfc-1e4618c9496d")
+
+    class MyNode(BaseNode):
+        pass
+
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        node_id = explicit_id
+
+    return MyNodeDisplay, explicit_id
+
+
+@pytest.fixture(
+    params=[
+        "node_with_implicit_properties",
+        "node_with_explicit_properties",
+    ]
+)
+def node_info(request):
+    return request.getfixturevalue(request.param)
+
+
+def test_get_id(node_info):
+    node_display, expected_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"}},
+            },
+        }
+    ]

{vellum_ai-0.14.38 → vellum_ai-0.14.40}/ee/vellum_ee/workflows/display/nodes/vellum/base_adornment_node.py

@@ -14,6 +14,7 @@ from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeV
 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 +32,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
@@ -89,7 +90,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_ai-0.14.38 → vellum_ai-0.14.40}/ee/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_ai-0.14.38 → vellum_ai-0.14.40}/ee/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_ai-0.14.38 → vellum_ai-0.14.40}/ee/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_ai-0.14.38 → vellum_ai-0.14.40}/ee/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_ai-0.14.38 → vellum_ai-0.14.40}/ee/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_ai-0.14.38 → vellum_ai-0.14.40}/ee/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_ai-0.14.38 → vellum_ai-0.14.40}/ee/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_ai-0.14.38 → vellum_ai-0.14.40}/ee/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