vellum-ai 0.13.15__py3-none-any.whl → 0.13.18__py3-none-any.whl

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -208,8 +208,8 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         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]
+    def get_from_node_display_registry(cls, node_class: Type[NodeType]) -> Optional[Type["BaseNodeDisplay"]]:
+        return cls._node_display_registry.get(node_class)
 
     @classmethod
     def infer_node_class(cls) -> Type[NodeType]:

vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -10,27 +10,23 @@ if TYPE_CHECKING:
 def get_node_display_class(
     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:
-        try:
-            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
+    node_display_class = base_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__}'"
             )
 
-            # `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]
-            NodeDisplayClass = types.new_class(
-                f"{node_class.__name__}Display",
-                bases=(NodeDisplayBaseClass,),
-            )
-            return NodeDisplayClass
-        except IndexError:
-            return base_class
+        return 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__}'"
-        )
+    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
+    )
 
-    return node_display_class
+    # `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]
+    NodeDisplayClass = types.new_class(
+        f"{node_class.__name__}Display",
+        bases=(NodeDisplayBaseClass,),
+    )
+    return NodeDisplayClass

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

@@ -10,6 +10,7 @@ from .map_node import BaseMapNodeDisplay
 from .merge_node import BaseMergeNodeDisplay
 from .note_node import BaseNoteNodeDisplay
 from .prompt_deployment_node import BasePromptDeploymentNodeDisplay
+from .retry_node import BaseRetryNodeDisplay
 from .search_node import BaseSearchNodeDisplay
 from .subworkflow_deployment_node import BaseSubworkflowDeploymentNodeDisplay
 from .templating_node import BaseTemplatingNodeDisplay
@@ -17,6 +18,7 @@ from .try_node import BaseTryNodeDisplay
 
 # All node display classes must be imported here to be registered in BaseNodeDisplay's node display registry
 __all__ = [
+    "BaseRetryNodeDisplay",
     "BaseAPINodeDisplay",
     "BaseCodeExecutionNodeDisplay",
     "BaseConditionalNodeDisplay",

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

@@ -0,0 +1,10 @@
+from typing import Generic, TypeVar
+
+from vellum.workflows.nodes.core.retry_node.node import RetryNode
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+
+_RetryNodeType = TypeVar("_RetryNodeType", bound=RetryNode)
+
+
+class BaseRetryNodeDisplay(BaseNodeDisplay[_RetryNodeType], Generic[_RetryNodeType]):
+    pass

vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py

@@ -7,10 +7,13 @@ from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.retry_node.node import RetryNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.base import WorkflowInputsDisplay
 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.vellum.try_node import BaseTryNodeDisplay
+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
 
 
 class Inputs(BaseInputs):
@@ -110,6 +113,26 @@ def test_serialize_node__retry(serialize_node):
     )
 
 
+def test_serialize_node__retry__no_display():  # GIVEN an adornment node
+    @RetryNode.wrap(max_attempts=5)
+    class StartNode(BaseNode):
+        pass
+
+    # AND a workflow that uses the adornment node
+    class MyWorkflow(BaseWorkflow):
+        graph = StartNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay,
+        workflow_class=MyWorkflow,
+    )
+    exec_config = workflow_display.serialize()
+
+    # THEN the workflow display is created successfully
+    assert exec_config is not None
+
+
 @TryNode.wrap()
 class InnerTryGenericNode(BaseNode):
     input = Inputs.input

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

@@ -93,7 +93,7 @@ def test_serialize_workflow_with_filepath():
             "code_input_id": "f2e8a4fa-b54e-41e9-b314-0e5443519ac7",
             "runtime_input_id": "19d64948-f22b-4103-a7f5-3add184b31cc",
             "output_type": "NUMBER",
-            "packages": [],
+            "packages": [{"name": "openai", "version": "1.0.0"}],
             "output_id": "0fde9607-353f-42c2-85c4-20f720ebc1ec",
             "log_output_id": "7cac05e3-b7c3-475e-8df8-422b496c3398",
         },
@@ -566,7 +566,7 @@ def test_serialize_workflow__try_wrapped():
             "code_input_id": "f2e8a4fa-b54e-41e9-b314-0e5443519ac7",
             "runtime_input_id": "19d64948-f22b-4103-a7f5-3add184b31cc",
             "output_type": "NUMBER",
-            "packages": [],
+            "packages": [{"name": "openai", "version": "1.0.0"}],
             "output_id": "0fde9607-353f-42c2-85c4-20f720ebc1ec",
             "log_output_id": "7cac05e3-b7c3-475e-8df8-422b496c3398",
         },