vellum-ai 0.13.10__tar.gz → 0.13.12__tar.gz

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/PKG-INFO

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

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/ee/vellum_cli/__init__.py

@@ -37,13 +37,15 @@ class PushGroup(ClickAliasedGroup):
 
 @main.group(invoke_without_command=True, cls=PushGroup)
 @click.pass_context
+@click.option("--workspace", type=str, help="The specific Workspace config to use when pushing")
 def push(
     ctx: click.Context,
+    workspace: Optional[str],
 ) -> None:
     """Push Resources to Vellum"""
 
     if ctx.invoked_subcommand is None:
-        push_command()
+        push_command(workspace=workspace)
 
 
 @main.group()
@@ -69,6 +71,7 @@ def workflows():
     is_flag=True,
     help="Raises an error if we detect an unexpected discrepancy in the generated artifact.",
 )
+@click.option("--workspace", type=str, help="The specific Workspace config to use when pushing")
 def workflows_push(
     module: Optional[str],
     deploy: Optional[bool],
@@ -78,6 +81,7 @@ def workflows_push(
     release_tag: Optional[List[str]],
     dry_run: Optional[bool],
     strict: Optional[bool],
+    workspace: Optional[str],
 ) -> None:
     """
     Push Workflows to Vellum. If a module is provided, only the Workflow for that module will be pushed.
@@ -93,6 +97,7 @@ def workflows_push(
         release_tags=release_tag,
         dry_run=dry_run,
         strict=strict,
+        workspace=workspace,
     )
 
 
@@ -113,6 +118,7 @@ def workflows_push(
     is_flag=True,
     help="Raises an error if we detect an unexpected discrepancy in the generated artifact.",
 )
+@click.option("--workspace", type=str, help="The specific Workspace config to use when pushing")
 def push_module(
     ctx: click.Context,
     deploy: Optional[bool],
@@ -122,6 +128,7 @@ def push_module(
     release_tag: Optional[List[str]],
     dry_run: Optional[bool],
     strict: Optional[bool],
+    workspace: Optional[str],
 ) -> None:
     """Push a specific module to Vellum"""
 
@@ -135,6 +142,7 @@ def push_module(
             release_tags=release_tag,
             dry_run=dry_run,
             strict=strict,
+            workspace=workspace,
         )
 
 

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/ee/vellum_cli/config.py

@@ -13,6 +13,17 @@ LOCKFILE_PATH = "vellum.lock.json"
 PYPROJECT_TOML_PATH = "pyproject.toml"
 
 
+class WorkspaceConfig(UniversalBaseModel):
+    name: str
+    api_key: str
+
+    def merge(self, other: "WorkspaceConfig") -> "WorkspaceConfig":
+        return WorkspaceConfig(name=self.name or other.name, api_key=self.api_key or other.api_key)
+
+
+DEFAULT_WORKSPACE_CONFIG = WorkspaceConfig(name="default", api_key="VELLUM_API_KEY")
+
+
 class WorkflowDeploymentConfig(UniversalBaseModel):
     id: Optional[UUID] = None
     label: Optional[str] = None
@@ -37,6 +48,7 @@ class WorkflowConfig(UniversalBaseModel):
     deployments: List[WorkflowDeploymentConfig] = field(default_factory=list)
     container_image_name: Optional[str] = None
     container_image_tag: Optional[str] = None
+    workspace: str = DEFAULT_WORKSPACE_CONFIG.name
 
     def merge(self, other: "WorkflowConfig") -> "WorkflowConfig":
         self_deployment_by_id = {
@@ -72,12 +84,14 @@ class WorkflowConfig(UniversalBaseModel):
             deployments=merged_deployments,
             container_image_tag=self.container_image_tag or other.container_image_tag,
             container_image_name=self.container_image_name or other.container_image_name,
+            workspace=self.workspace or other.workspace,
         )
 
 
 class VellumCliConfig(UniversalBaseModel):
     version: Literal["1.0"] = "1.0"
     workflows: List[WorkflowConfig] = field(default_factory=list)
+    workspaces: List[WorkspaceConfig] = field(default_factory=list)
 
     def save(self) -> None:
         lockfile_path = os.path.join(os.getcwd(), LOCKFILE_PATH)
@@ -102,7 +116,21 @@ class VellumCliConfig(UniversalBaseModel):
             elif other_workflow:
                 merged_workflows.append(other_workflow)
 
-        return VellumCliConfig(workflows=merged_workflows, version=self.version)
+        self_workspace_by_name = {workspace.name: workspace for workspace in self.workspaces}
+        other_workspace_by_name = {workspace.name: workspace for workspace in other.workspaces}
+        all_names = sorted(set(self_workspace_by_name.keys()).union(set(other_workspace_by_name.keys())))
+        merged_workspaces = []
+        for name in all_names:
+            self_workspace = self_workspace_by_name.get(name)
+            other_workspace = other_workspace_by_name.get(name)
+            if self_workspace and other_workspace:
+                merged_workspaces.append(self_workspace.merge(other_workspace))
+            elif self_workspace:
+                merged_workspaces.append(self_workspace)
+            elif other_workspace:
+                merged_workspaces.append(other_workspace)
+
+        return VellumCliConfig(workflows=merged_workflows, workspaces=merged_workspaces, version=self.version)
 
 
 def load_vellum_cli_config(root_dir: Optional[str] = None) -> VellumCliConfig:

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/ee/vellum_cli/push.py

@@ -15,7 +15,7 @@ from vellum.types import WorkflowPushDeploymentConfigRequest
 from vellum.workflows.utils.names import snake_to_title_case
 from vellum.workflows.vellum_client import create_vellum_client
 from vellum.workflows.workflows.base import BaseWorkflow
-from vellum_cli.config import WorkflowDeploymentConfig, load_vellum_cli_config
+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
@@ -30,8 +30,9 @@ def push_command(
     release_tags: Optional[List[str]] = None,
     dry_run: Optional[bool] = None,
     strict: Optional[bool] = None,
+    workspace: Optional[str] = None,
 ) -> None:
-    load_dotenv()
+    load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
     logger = load_cli_logger()
     config = load_vellum_cli_config()
 
@@ -46,7 +47,27 @@ def push_command(
         raise ValueError(f"No workflow config for '{module}' found in project to push.")
 
     logger.info(f"Loading workflow from {workflow_config.module}")
-    client = create_vellum_client()
+    resolved_workspace = workspace or workflow_config.workspace or DEFAULT_WORKSPACE_CONFIG.name
+    workspace_config = (
+        next((w for w in config.workspaces if w.name == resolved_workspace), DEFAULT_WORKSPACE_CONFIG)
+        if workspace
+        else DEFAULT_WORKSPACE_CONFIG
+    )
+    api_key = os.getenv(workspace_config.api_key)
+    if not api_key:
+        raise ValueError(f"No API key value found in environment for workspace '{workspace_config.name}'.")
+
+    if workspace_config.name != workflow_config.workspace:
+        # We are pushing to a new workspace, so we need a new workflow config
+        workflow_config = WorkflowConfig(
+            module=workflow_config.module,
+            workspace=workspace_config.name,
+        )
+        config.workflows.append(workflow_config)
+
+    client = create_vellum_client(
+        api_key=api_key,
+    )
     sys.path.insert(0, os.getcwd())
 
     # Remove this once we could serialize using the artifact in Vembda

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/ee/vellum_cli/tests/conftest.py

@@ -46,6 +46,9 @@ def mock_module(request) -> Generator[MockModuleResult, None, None]:
         }
     )
 
+    with open(os.path.join(temp_dir, ".env"), "w") as f:
+        f.write("VELLUM_API_KEY=abcdef123456")
+
     yield MockModuleResult(
         temp_dir=temp_dir,
         module=module,

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/ee/vellum_cli/tests/test_pull.py

@@ -124,6 +124,7 @@ def test_pull__sandbox_id_with_no_config(vellum_client):
         lock_data = json.loads(f.read())
         assert lock_data == {
             "version": "1.0",
+            "workspaces": [],
             "workflows": [
                 {
                     "module": "workflow_87654321",
@@ -132,6 +133,7 @@ def test_pull__sandbox_id_with_no_config(vellum_client):
                     "deployments": [],
                     "container_image_tag": None,
                     "container_image_name": None,
+                    "workspace": "default",
                 }
             ],
         }
@@ -208,8 +210,10 @@ def test_pull__workflow_deployment_with_no_config(vellum_client):
                     "deployments": [],
                     "container_image_tag": None,
                     "container_image_name": None,
+                    "workspace": "default",
                 }
             ],
+            "workspaces": [],
         }
 
 
@@ -449,6 +453,7 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                 ],
                 "container_image_name": None,
                 "container_image_tag": None,
+                "workspace": "default",
             },
             {
                 "module": "workflow_87654321",
@@ -457,6 +462,7 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                 "deployments": [],
                 "container_image_name": "test",
                 "container_image_tag": "1.0",
+                "workspace": "default",
             },
         ]
 

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/ee/vellum_cli/tests/test_push.py

@@ -3,6 +3,7 @@ import io
 import json
 import os
 import tarfile
+from unittest import mock
 from uuid import uuid4
 
 from click.testing import CliRunner
@@ -194,7 +195,7 @@ def test_push__dry_run_option_returns_report(mock_module, vellum_client):
 from typing import Dict
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes import BaseNode
-from vellum_ee.workflows.display.nodes import BaseNodeDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 
 class NotSupportedNode(BaseNode):
     pass
@@ -318,3 +319,89 @@ Files that were different between the original project and the generated artifac
 \x1b[0m
 """
     )
+
+
+def test_push__workspace_option__uses_different_api_key(mock_module, vellum_client_class):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
+    set_pyproject_toml = mock_module.set_pyproject_toml
+
+    # AND a different workspace is set in the pyproject.toml
+    set_pyproject_toml(
+        {
+            "workflows": [
+                {
+                    "module": module,
+                    "workflow_sandbox_id": workflow_sandbox_id,
+                }
+            ],
+            "workspaces": [
+                {
+                    "name": "my_other_workspace",
+                    "api_key": "MY_OTHER_VELLUM_API_KEY",
+                }
+            ],
+        }
+    )
+
+    # AND the .env file has the other api key stored
+    with open(os.path.join(temp_dir, ".env"), "w") as f:
+        f.write(
+            """\
+VELLUM_API_KEY=abcdef123456
+MY_OTHER_VELLUM_API_KEY=aaabbbcccddd
+"""
+        )
+
+    # AND a workflow exists in the module successfully
+    base_dir = os.path.join(temp_dir, *module.split("."))
+    os.makedirs(base_dir, exist_ok=True)
+    workflow_py_file_content = """\
+from vellum.workflows import BaseWorkflow
+
+class ExampleWorkflow(BaseWorkflow):
+    pass
+"""
+    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
+        f.write(workflow_py_file_content)
+
+    # AND the push API call returns a new workflow sandbox id
+    new_workflow_sandbox_id = str(uuid4())
+    vellum_client_class.return_value.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=new_workflow_sandbox_id,
+    )
+
+    # WHEN calling `vellum push` on strict mode
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push", module, "--workspace", "my_other_workspace"])
+
+    # THEN it should succeed
+    assert result.exit_code == 0, result.output
+
+    # AND we should have called the push API once
+    vellum_client_class.return_value.workflows.push.assert_called_once()
+
+    # AND the workflow sandbox id arg passed in should be `None`
+    call_args = vellum_client_class.return_value.workflows.push.call_args.kwargs
+    assert call_args["workflow_sandbox_id"] is None
+
+    # AND with the correct api key
+    vellum_client_class.assert_called_once_with(
+        api_key="aaabbbcccddd",
+        environment=mock.ANY,
+    )
+
+    # AND the vellum lock file should have been updated with the correct workspace
+    with open(os.path.join(temp_dir, "vellum.lock.json")) as f:
+        lock_file_content = json.load(f)
+        assert lock_file_content["workflows"][1] == {
+            "module": module,
+            "workflow_sandbox_id": new_workflow_sandbox_id,
+            "workspace": "my_other_workspace",
+            "container_image_name": None,
+            "container_image_tag": None,
+            "deployments": [],
+            "ignore": None,
+        }

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/ee/vellum_ee/workflows/display/nodes/base_node_display.py

@@ -16,6 +16,8 @@ from typing import (
     get_origin,
 )
 
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.constants import UNDEF
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.expressions.between import BetweenExpression
 from vellum.workflows.expressions.is_nil import IsNilExpression
@@ -26,19 +28,24 @@ from vellum.workflows.expressions.is_null import IsNullExpression
 from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
 from vellum.workflows.expressions.not_between import NotBetweenExpression
 from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.nodes.utils import get_wrapped_node
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference
+from vellum.workflows.references.constant import ConstantValueReference
 from vellum.workflows.references.execution_count import ExecutionCountReference
+from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.references.vellum_secret import VellumSecretReference
 from vellum.workflows.references.workflow_input import WorkflowInputReference
-from vellum.workflows.types.core import JsonObject
+from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.names import pascal_to_title_case
 from vellum.workflows.utils.uuids import uuid4_from_hash
+from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
+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.vellum import convert_descriptor_to_operator, primitive_to_vellum_value
-from vellum_ee.workflows.display.vellum import CodeResourceDefinition
+from vellum_ee.workflows.display.vellum import CodeResourceDefinition, GenericNodeDisplayData
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -65,7 +72,97 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
     _node_display_registry: Dict[Type[NodeType], Type["BaseNodeDisplay"]] = {}
 
     def serialize(self, display_context: "WorkflowDisplayContext", **kwargs: Any) -> JsonObject:
-        raise NotImplementedError(f"Serialization for nodes of type {self._node.__name__} is not supported.")
+        node = self._node
+        node_id = self.node_id
+
+        attributes: JsonArray = []
+        for attribute in node:
+            if inspect.isclass(attribute.instance) and issubclass(attribute.instance, BaseWorkflow):
+                # We don't need to serialize generic node attributes containing a subworkflow
+                continue
+
+            id = str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
+            attributes.append(
+                {
+                    "id": id,
+                    "name": attribute.name,
+                    "value": self.serialize_value(display_context, cast(BaseDescriptor, attribute.instance)),
+                }
+            )
+
+        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)
+
+            adornment: JsonObject = {
+                "id": str(node_id),
+                "label": node.__qualname__,
+                "base": self.get_base().dict(),
+                "attributes": attributes,
+            }
+
+            existing_adornments = adornments if adornments is not None else []
+            return display_class().serialize(display_context, adornments=existing_adornments + [adornment])
+
+        ports: JsonArray = []
+        for port in node.Ports:
+            id = str(self.get_node_port_display(port).id)
+
+            if port._condition_type:
+                ports.append(
+                    {
+                        "id": id,
+                        "name": port.name,
+                        "type": port._condition_type.value,
+                        "expression": (
+                            self.serialize_condition(display_context, port._condition) if port._condition else None
+                        ),
+                    }
+                )
+            else:
+                ports.append(
+                    {
+                        "id": id,
+                        "name": port.name,
+                        "type": "DEFAULT",
+                    }
+                )
+
+        outputs: JsonArray = []
+        for output in node.Outputs:
+            type = primitive_type_to_vellum_variable_type(output)
+            value = (
+                self.serialize_value(display_context, output.instance)
+                if output.instance is not None and output.instance != UNDEF
+                else None
+            )
+
+            outputs.append(
+                {
+                    "id": str(uuid4_from_hash(f"{node_id}|{output.name}")),
+                    "name": output.name,
+                    "type": type,
+                    "value": value,
+                }
+            )
+
+        return {
+            "id": str(node_id),
+            "label": node.__qualname__,
+            "type": "GENERIC",
+            "display_data": self._get_generic_node_display_data().dict(),
+            "base": self.get_base().dict(),
+            "definition": self.get_definition().dict(),
+            "trigger": {
+                "id": str(self.get_trigger_id()),
+                "merge_behavior": node.Trigger.merge_behavior.value,
+            },
+            "ports": ports,
+            "adornments": adornments,
+            "attributes": attributes,
+            "outputs": outputs,
+        }
 
     def get_base(self) -> CodeResourceDefinition:
         node = self._node
@@ -89,9 +186,6 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         )
         return node_definition
 
-    def get_trigger_id(self) -> UUID:
-        return uuid4_from_hash(f"{self.node_id}|trigger")
-
     def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
         explicit_display = self.output_display.get(output)
         if explicit_display:
@@ -110,6 +204,9 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
         return PortDisplay(id=port_id, node_id=self.node_id)
 
+    def get_trigger_id(self) -> UUID:
+        return uuid4_from_hash(f"{self.node_id}|trigger")
+
     @classmethod
     def get_from_node_display_registry(cls, node_class: Type[NodeType]) -> Type["BaseNodeDisplay"]:
         return cls._node_display_registry[node_class]
@@ -187,10 +284,19 @@ 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
 
+    def _get_generic_node_display_data(self) -> GenericNodeDisplayData:
+        explicit_value = self._get_explicit_node_display_attr("display_data", GenericNodeDisplayData)
+        return explicit_value if explicit_value else GenericNodeDisplayData()
+
     def serialize_condition(self, display_context: "WorkflowDisplayContext", condition: BaseDescriptor) -> JsonObject:
         if isinstance(
             condition,
@@ -233,6 +339,12 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
             }
 
     def serialize_value(self, display_context: "WorkflowDisplayContext", value: BaseDescriptor) -> JsonObject:
+        if isinstance(value, ConstantValueReference):
+            return self.serialize_value(display_context, value._value)
+
+        if isinstance(value, LazyReference):
+            return self.serialize_value(display_context, value._get())
+
         if isinstance(value, WorkflowInputReference):
             workflow_input_display = display_context.global_workflow_input_displays[value]
             return {

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/ee/vellum_ee/workflows/display/nodes/base_node_vellum_display.py

@@ -7,7 +7,7 @@ from vellum.workflows.types.generics import NodeType
 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.types import PortDisplay
-from vellum_ee.workflows.display.vellum import NodeDisplayData
+from vellum_ee.workflows.display.vellum import NodeDisplayComment, NodeDisplayData
 
 
 class BaseNodeVellumDisplay(BaseNodeDisplay[NodeType]):
@@ -26,6 +26,21 @@ class BaseNodeVellumDisplay(BaseNodeDisplay[NodeType]):
 
     def get_display_data(self) -> NodeDisplayData:
         explicit_value = self._get_explicit_node_display_attr("display_data", NodeDisplayData)
+        docstring = self._node.__doc__
+
+        if explicit_value and explicit_value.comment and docstring:
+            comment = (
+                NodeDisplayComment(value=docstring, expanded=explicit_value.comment.expanded)
+                if explicit_value.comment.expanded
+                else NodeDisplayComment(value=docstring)
+            )
+            return NodeDisplayData(
+                position=explicit_value.position,
+                width=explicit_value.width,
+                height=explicit_value.height,
+                comment=comment,
+            )
+
         return explicit_value if explicit_value else NodeDisplayData()
 
     def get_target_handle_id(self) -> UUID:

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/ee/vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -1,13 +1,15 @@
 import types
-from typing import Optional, Type
+from typing import TYPE_CHECKING, Optional, Type
 
 from vellum.workflows.types.generics import NodeType
-from vellum_ee.workflows.display.types import NodeDisplayType
+
+if TYPE_CHECKING:
+    from vellum_ee.workflows.display.types import NodeDisplayType
 
 
 def get_node_display_class(
-    base_class: Type[NodeDisplayType], node_class: Type[NodeType], root_node_class: Optional[Type[NodeType]] = None
-) -> Type[NodeDisplayType]:
+    base_class: Type["NodeDisplayType"], node_class: Type[NodeType], root_node_class: Optional[Type[NodeType]] = None
+) -> Type["NodeDisplayType"]:
     try:
         node_display_class = base_class.get_from_node_display_registry(node_class)
     except KeyError:

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/ee/vellum_ee/workflows/display/nodes/vellum/__init__.py

@@ -1,5 +1,4 @@
 from .api_node import BaseAPINodeDisplay
-from .base_node import BaseNodeDisplay
 from .code_execution_node import BaseCodeExecutionNodeDisplay
 from .conditional_node import BaseConditionalNodeDisplay
 from .error_node import BaseErrorNodeDisplay
@@ -28,7 +27,6 @@ __all__ = [
     "BaseInlineSubworkflowNodeDisplay",
     "BaseMapNodeDisplay",
     "BaseMergeNodeDisplay",
-    "BaseNodeDisplay",
     "BaseNoteNodeDisplay",
     "BasePromptDeploymentNodeDisplay",
     "BaseSearchNodeDisplay",

{vellum_ai-0.13.10 → vellum_ai-0.13.12}/ee/vellum_ee/workflows/display/nodes/vellum/error_node.py

@@ -1,9 +1,10 @@
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar
+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.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
 
@@ -12,7 +13,7 @@ _ErrorNodeType = TypeVar("_ErrorNodeType", bound=ErrorNode)
 
 class BaseErrorNodeDisplay(BaseNodeVellumDisplay[_ErrorNodeType], Generic[_ErrorNodeType]):
     error_output_id: ClassVar[Optional[UUID]] = None
-    error_inputs_by_name: ClassVar[Dict[str, Any]] = {}
+
     name: ClassVar[str] = "error-node"
 
     def serialize(
@@ -21,6 +22,11 @@ class BaseErrorNodeDisplay(BaseNodeVellumDisplay[_ErrorNodeType], Generic[_Error
         node_id = self.node_id
         error_source_input_id = self.node_input_ids_by_name.get("error_source_input_id")
 
+        error_attribute = raise_if_descriptor(self._node.error)
+        input_values_by_name = {
+            "error_source_input_id": error_attribute,
+        }
+
         node_inputs = [
             create_node_input(
                 node_id=node_id,
@@ -29,7 +35,7 @@ class BaseErrorNodeDisplay(BaseNodeVellumDisplay[_ErrorNodeType], Generic[_Error
                 display_context=display_context,
                 input_id=self.node_input_ids_by_name.get(variable_name),
             )
-            for variable_name, variable_value in self.error_inputs_by_name.items()
+            for variable_name, variable_value in input_values_by_name.items()
         ]
 
         return {