vellum-ai 0.14.22__py3-none-any.whl → 0.14.24__py3-none-any.whl

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -10,7 +10,6 @@ from vellum.workflows import BaseWorkflow
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges import Edge
 from vellum.workflows.events.workflow import NodeEventDisplayContext, WorkflowEventDisplayContext
-from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.utils import get_unadorned_node, get_wrapped_node
 from vellum.workflows.ports import Port
@@ -26,16 +25,22 @@ from vellum_ee.workflows.display.base import (
     StateValueDisplayType,
     WorkflowInputsDisplayOverridesType,
     WorkflowInputsDisplayType,
-    WorkflowMetaDisplayOverridesType,
-    WorkflowMetaDisplayType,
+    WorkflowMetaDisplay,
     WorkflowOutputDisplay,
 )
 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, PortDisplay, PortDisplayOverrides
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
-from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.types import (
+    EdgeDisplays,
+    NodeDisplays,
+    NodeOutputDisplays,
+    PortDisplays,
+    WorkflowDisplayContext,
+    WorkflowOutputDisplays,
+)
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 logger = logging.getLogger(__name__)
@@ -44,8 +49,6 @@ logger = logging.getLogger(__name__)
 class BaseWorkflowDisplay(
     Generic[
         WorkflowType,
-        WorkflowMetaDisplayType,
-        WorkflowMetaDisplayOverridesType,
         WorkflowInputsDisplayType,
         WorkflowInputsDisplayOverridesType,
         StateValueDisplayType,
@@ -55,7 +58,7 @@ class BaseWorkflowDisplay(
     ]
 ):
     # Used to specify the display data for a workflow.
-    workflow_display: Optional[WorkflowMetaDisplayOverridesType] = None
+    workflow_display: Optional[WorkflowMetaDisplay] = None
 
     # Used to explicitly specify display data for a workflow's inputs.
     inputs_display: Dict[WorkflowInputReference, WorkflowInputsDisplayOverridesType] = {}
@@ -67,13 +70,13 @@ class BaseWorkflowDisplay(
     entrypoint_displays: Dict[Type[BaseNode], EntrypointDisplayOverridesType] = {}
 
     # Used to explicitly specify display data for a workflow's outputs.
-    output_displays: Dict[BaseDescriptor, WorkflowOutputDisplay] = {}
+    output_displays: WorkflowOutputDisplays = {}
 
     # Used to explicitly specify display data for a workflow's edges.
-    edge_displays: Dict[Tuple[Port, Type[BaseNode]], EdgeDisplay] = {}
+    edge_displays: EdgeDisplays = {}
 
     # Used to explicitly specify display data for a workflow's ports.
-    port_displays: Dict[Port, PortDisplayOverrides] = {}
+    port_displays: PortDisplays = {}
 
     # Used to store the mapping between workflows and their display classes
     _workflow_display_registry: Dict[Type[WorkflowType], Type["BaseWorkflowDisplay"]] = {}
@@ -88,7 +91,6 @@ class BaseWorkflowDisplay(
         *,
         parent_display_context: Optional[
             WorkflowDisplayContext[
-                WorkflowMetaDisplayType,
                 WorkflowInputsDisplayType,
                 StateValueDisplayType,
                 EntrypointDisplayType,
@@ -184,50 +186,41 @@ class BaseWorkflowDisplay(
     def display_context(
         self,
     ) -> WorkflowDisplayContext[
-        WorkflowMetaDisplayType,
         WorkflowInputsDisplayType,
         StateValueDisplayType,
         EntrypointDisplayType,
     ]:
         workflow_display = self._generate_workflow_meta_display()
 
-        global_node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]] = (
+        global_node_output_displays: NodeOutputDisplays = (
             copy(self._parent_display_context.global_node_output_displays) if self._parent_display_context else {}
         )
 
-        node_displays: Dict[Type[BaseNode], BaseNodeDisplay] = {}
+        node_displays: NodeDisplays = {}
 
-        global_node_displays: Dict[Type[BaseNode], BaseNodeDisplay] = (
+        global_node_displays: NodeDisplays = (
             copy(self._parent_display_context.global_node_displays) if self._parent_display_context else {}
         )
 
-        port_displays: Dict[Port, PortDisplay] = {}
+        port_displays: PortDisplays = {}
 
         for node in self._workflow.get_nodes():
-            extracted_node_displays = self._extract_node_displays(node)
-
-            for extracted_node, extracted_node_display in extracted_node_displays.items():
-                if extracted_node not in node_displays:
-                    node_displays[extracted_node] = extracted_node_display
-
-                if extracted_node not in global_node_displays:
-                    global_node_displays[extracted_node] = extracted_node_display
-
-            self._enrich_global_node_output_displays(node, extracted_node_displays[node], global_node_output_displays)
-            self._enrich_node_port_displays(node, extracted_node_displays[node], port_displays)
+            self._enrich_node_displays(
+                node=node,
+                node_displays=node_displays,
+                global_node_displays=global_node_displays,
+                global_node_output_displays=global_node_output_displays,
+                port_displays=port_displays,
+            )
 
         for node in self._workflow.get_unused_nodes():
-            extracted_node_displays = self._extract_node_displays(node)
-
-            for extracted_node, extracted_node_display in extracted_node_displays.items():
-                if extracted_node not in node_displays:
-                    node_displays[extracted_node] = extracted_node_display
-
-                if extracted_node not in global_node_displays:
-                    global_node_displays[extracted_node] = extracted_node_display
-
-            self._enrich_global_node_output_displays(node, extracted_node_displays[node], global_node_output_displays)
-            self._enrich_node_port_displays(node, extracted_node_displays[node], port_displays)
+            self._enrich_node_displays(
+                node=node,
+                node_displays=node_displays,
+                global_node_displays=global_node_displays,
+                global_node_output_displays=global_node_output_displays,
+                port_displays=port_displays,
+            )
 
         workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = {}
         # If we're dealing with a nested workflow, then it should have access to the inputs of its parents.
@@ -291,11 +284,6 @@ class BaseWorkflowDisplay(
             if not isinstance(workflow_output, OutputReference):
                 raise ValueError(f"{workflow_output} must be an {OutputReference.__name__}")
 
-            if not workflow_output.instance or not isinstance(
-                workflow_output.instance, (OutputReference, CoalesceExpression)
-            ):
-                raise ValueError("Expected to find a descriptor instance on the workflow output")
-
             workflow_output_display = self.output_displays.get(workflow_output)
             workflow_output_displays[workflow_output] = (
                 workflow_output_display or self._generate_workflow_output_display(workflow_output)
@@ -318,7 +306,7 @@ class BaseWorkflowDisplay(
         )
 
     @abstractmethod
-    def _generate_workflow_meta_display(self) -> WorkflowMetaDisplayType:
+    def _generate_workflow_meta_display(self) -> WorkflowMetaDisplay:
         pass
 
     @abstractmethod
@@ -337,7 +325,7 @@ class BaseWorkflowDisplay(
     def _generate_entrypoint_display(
         self,
         entrypoint: Type[BaseNode],
-        workflow_display: WorkflowMetaDisplayType,
+        workflow_display: WorkflowMetaDisplay,
         node_displays: Dict[Type[BaseNode], BaseNodeDisplay],
         overrides: Optional[EntrypointDisplayOverridesType] = None,
     ) -> EntrypointDisplayType:
@@ -424,6 +412,26 @@ class BaseWorkflowDisplay(
         )
         return display_meta
 
+    def _enrich_node_displays(
+        self,
+        node: Type[BaseNode],
+        node_displays: NodeDisplays,
+        global_node_displays: NodeDisplays,
+        global_node_output_displays: NodeOutputDisplays,
+        port_displays: PortDisplays,
+    ) -> None:
+        extracted_node_displays = self._extract_node_displays(node)
+
+        for extracted_node, extracted_node_display in extracted_node_displays.items():
+            if extracted_node not in node_displays:
+                node_displays[extracted_node] = extracted_node_display
+
+            if extracted_node not in global_node_displays:
+                global_node_displays[extracted_node] = extracted_node_display
+
+        self._enrich_global_node_output_displays(node, extracted_node_displays[node], global_node_output_displays)
+        self._enrich_node_port_displays(node, extracted_node_displays[node], port_displays)
+
     def _extract_node_displays(self, node: Type[BaseNode]) -> Dict[Type[BaseNode], BaseNodeDisplay]:
         node_display = self._get_node_display(node)
         additional_node_displays: Dict[Type[BaseNode], BaseNodeDisplay] = {

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

@@ -7,10 +7,10 @@ from vellum.workflows.nodes.core.retry_node.node import RetryNode
 from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.workflows.base import BaseWorkflow
+from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.nodes import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.vellum.retry_node import BaseRetryNodeDisplay
 from vellum_ee.workflows.display.nodes.vellum.try_node import BaseTryNodeDisplay
-from vellum_ee.workflows.display.vellum import NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -275,3 +275,17 @@ def test_get_event_display_context__node_display_for_mutiple_adornments():
     assert innermost_node_id in inner_node_event_display.subworkflow_display.node_displays
     innermost_node_event_display = inner_node_event_display.subworkflow_display.node_displays[innermost_node_id]
     assert not innermost_node_event_display.subworkflow_display
+
+
+def test_get_event_display_context__workflow_output_display_with_none():
+    # GIVEN a workflow with a workflow output that is None
+    class MyWorkflow(BaseWorkflow):
+        class Outputs(BaseWorkflow.Outputs):
+            foo = None
+            bar = "baz"
+
+    # WHEN we gather the event display context
+    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+
+    # THEN the workflow output display should be included
+    assert display_context.workflow_outputs.keys() == {"foo", "bar"}

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -12,6 +12,7 @@ from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.generics import WorkflowType
 from vellum.workflows.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.base import WorkflowMetaDisplay
 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.utils import create_node_input
@@ -25,7 +26,6 @@ from vellum_ee.workflows.display.vellum import (
     WorkflowInputsVellumDisplay,
     WorkflowInputsVellumDisplayOverrides,
     WorkflowMetaVellumDisplay,
-    WorkflowMetaVellumDisplayOverrides,
 )
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
@@ -35,8 +35,6 @@ logger = logging.getLogger(__name__)
 class VellumWorkflowDisplay(
     BaseWorkflowDisplay[
         WorkflowType,
-        WorkflowMetaVellumDisplay,
-        WorkflowMetaVellumDisplayOverrides,
         WorkflowInputsVellumDisplay,
         WorkflowInputsVellumDisplayOverrides,
         StateValueVellumDisplay,
@@ -346,7 +344,7 @@ class VellumWorkflowDisplay(
     def _generate_entrypoint_display(
         self,
         entrypoint: Type[BaseNode],
-        workflow_display: WorkflowMetaVellumDisplay,
+        workflow_display: WorkflowMetaDisplay,
         node_displays: Dict[Type[BaseNode], BaseNodeDisplay],
         overrides: Optional[EntrypointVellumDisplayOverrides] = None,
     ) -> EntrypointVellumDisplay: