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

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

@@ -11,7 +11,6 @@ from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDispla
 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.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
@@ -32,10 +31,7 @@ def test_serialize_workflow__node_referenced_in_workflow_outputs_not_in_graph():
             final = OutNode.Outputs.foo
 
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
-        workflow_class=Workflow,
-    )
+    workflow_display = get_workflow_display(workflow_class=Workflow)
 
     # THEN it should raise an error
     with pytest.raises(ValueError) as exc_info:
@@ -57,10 +53,7 @@ def test_serialize_workflow__workflow_outputs_reference_non_node_outputs():
             final = FirstWorkflow.Outputs.foo
 
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
-        workflow_class=Workflow,
-    )
+    workflow_display = get_workflow_display(workflow_class=Workflow)
 
     # THEN it should raise an error
     with pytest.raises(ValueError) as exc_info:
@@ -91,10 +84,7 @@ def test_serialize_workflow__node_display_class_not_registered():
             answer = StartNode.Outputs.result
 
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
-        workflow_class=MyWorkflow,
-    )
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
     data = workflow_display.serialize()
 
     # THEN it should should succeed
@@ -111,7 +101,7 @@ def test_get_event_display_context__node_display_filled_without_base_display():
         graph = StartNode
 
     # WHEN we gather the event display context
-    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+    display_context = get_workflow_display(workflow_class=MyWorkflow).get_event_display_context()
 
     # THEN the node display should be included
     assert StartNode.__id__ in display_context.node_displays
@@ -134,7 +124,7 @@ def test_get_event_display_context__node_display_filled_without_output_display()
         pass
 
     # WHEN we gather the event display context
-    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+    display_context = get_workflow_display(workflow_class=MyWorkflow).get_event_display_context()
 
     # THEN the node display should be included
     assert StartNode.__id__ in display_context.node_displays
@@ -160,7 +150,7 @@ def test_get_event_display_context__node_display_to_include_subworkflow_display(
         graph = SubworkflowNode
 
     # WHEN we gather the event display context
-    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+    display_context = get_workflow_display(workflow_class=MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
     assert SubworkflowNode.__id__ in display_context.node_displays
@@ -201,7 +191,7 @@ def test_get_event_display_context__node_display_for_adornment_nodes(
         node_id = inner_node_id
 
     # WHEN we gather the event display context
-    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+    display_context = get_workflow_display(workflow_class=MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
     assert adornment_node_id in display_context.node_displays
@@ -228,7 +218,7 @@ def test_get_event_display_context__templating_node_input_display():
         graph = DataNode >> MyNode
 
     # WHEN we gather the event display context
-    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+    display_context = get_workflow_display(workflow_class=MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
     assert MyNode.__id__ in display_context.node_displays
@@ -259,7 +249,7 @@ def test_get_event_display_context__node_display_for_mutiple_adornments():
         node_id = innermost_node_id
 
     # WHEN we gather the event display context
-    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+    display_context = get_workflow_display(workflow_class=MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
     assert node_id in display_context.node_displays
@@ -285,7 +275,55 @@ def test_get_event_display_context__workflow_output_display_with_none():
             bar = "baz"
 
     # WHEN we gather the event display context
-    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+    display_context = get_workflow_display(workflow_class=MyWorkflow).get_event_display_context()
 
     # THEN the workflow output display should be included
     assert display_context.workflow_outputs.keys() == {"foo", "bar"}
+
+
+def test_serialize_workflow__inherited_node_display_class_not_registered():
+    # GIVEN a node meant to be used as a base
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    # AND a node that inherits from it
+    class InheritedNode(StartNode):
+        foo: str
+
+    # AND a workflow that uses the inherited node
+    class MyWorkflow(BaseWorkflow):
+        graph = InheritedNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            answer = InheritedNode.Outputs.result
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    data = workflow_display.serialize()
+
+    # THEN it should should succeed
+    assert data is not None
+
+
+def test_serialize_workflow__inherited_workflow_display_class_not_registered():
+    # GIVEN a node
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    # AND a workflow that uses the node
+    class MyWorkflow(BaseWorkflow):
+        graph = StartNode
+
+    # AND a workflow that inherits from it
+    class InheritedWorkflow(MyWorkflow):
+        class Outputs(MyWorkflow.Outputs):
+            answer = StartNode.Outputs.result
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=InheritedWorkflow)
+    data = workflow_display.serialize()
+
+    # THEN it should should succeed
+    assert data is not None

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -1,262 +1,9 @@
-import logging
-from uuid import UUID
-from typing import Optional, cast
+from typing import Generic
 
-from vellum.workflows.constants import undefined
-from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
-from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
-from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port
-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.editor.types import NodeDisplayData
-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
-from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
-logger = logging.getLogger(__name__)
 
-
-class VellumWorkflowDisplay(BaseWorkflowDisplay[WorkflowType]):
-    node_display_base_class = BaseNodeDisplay
-
-    def serialize(self) -> JsonObject:
-        input_variables: JsonArray = []
-        for workflow_input_reference, workflow_input_display in self.display_context.workflow_input_displays.items():
-            default = (
-                primitive_to_vellum_value(workflow_input_reference.instance)
-                if workflow_input_reference.instance
-                else None
-            )
-            input_variables.append(
-                {
-                    "id": str(workflow_input_display.id),
-                    "key": workflow_input_display.name or workflow_input_reference.name,
-                    "type": infer_vellum_variable_type(workflow_input_reference),
-                    "default": default.dict() if default else None,
-                    "required": workflow_input_reference.instance is undefined,
-                    "extensions": {"color": workflow_input_display.color},
-                }
-            )
-
-        state_variables: JsonArray = []
-        for state_value_reference, state_value_display in self.display_context.state_value_displays.items():
-            default = (
-                primitive_to_vellum_value(state_value_reference.instance) if state_value_reference.instance else None
-            )
-            state_variables.append(
-                {
-                    "id": str(state_value_display.id),
-                    "key": state_value_display.name or state_value_reference.name,
-                    "type": infer_vellum_variable_type(state_value_reference),
-                    "default": default.dict() if default else None,
-                    "required": state_value_reference.instance is undefined,
-                    "extensions": {"color": state_value_display.color},
-                }
-            )
-
-        nodes: JsonArray = []
-        edges: JsonArray = []
-
-        # Add a single synthetic node for the workflow entrypoint
-        entrypoint_node_id = self.display_context.workflow_display.entrypoint_node_id
-        entrypoint_node_source_handle_id = self.display_context.workflow_display.entrypoint_node_source_handle_id
-        nodes.append(
-            {
-                "id": str(entrypoint_node_id),
-                "type": "ENTRYPOINT",
-                "inputs": [],
-                "data": {
-                    "label": "Entrypoint Node",
-                    "source_handle_id": str(entrypoint_node_source_handle_id),
-                },
-                "display_data": self.display_context.workflow_display.entrypoint_node_display.dict(),
-                "base": None,
-                "definition": None,
-            },
-        )
-
-        # Add all the nodes in the workflow
-        for node in self._workflow.get_nodes():
-            node_display = self.display_context.node_displays[node]
-
-            try:
-                serialized_node = node_display.serialize(self.display_context)
-            except NotImplementedError as e:
-                self.add_error(e)
-                continue
-
-            nodes.append(serialized_node)
-
-        # Add all unused nodes in the workflow
-        for node in self._workflow.get_unused_nodes():
-            node_display = self.display_context.node_displays[node]
-
-            try:
-                serialized_node = node_display.serialize(self.display_context)
-            except NotImplementedError as e:
-                self.add_error(e)
-                continue
-
-            nodes.append(serialized_node)
-
-        synthetic_output_edges: JsonArray = []
-        output_variables: JsonArray = []
-        final_output_nodes = [
-            node for node in self.display_context.node_displays.keys() if issubclass(node, FinalOutputNode)
-        ]
-        final_output_node_outputs = {node.Outputs.value for node in final_output_nodes}
-        unreferenced_final_output_node_outputs = final_output_node_outputs.copy()
-        final_output_node_base: JsonObject = {
-            "name": FinalOutputNode.__name__,
-            "module": cast(JsonArray, FinalOutputNode.__module__.split(".")),
-        }
-
-        # Add a synthetic Terminal Node and track the Workflow's output variables for each Workflow output
-        for workflow_output, workflow_output_display in self.display_context.workflow_output_displays.items():
-            final_output_node_id = uuid4_from_hash(f"{self.workflow_id}|node_id|{workflow_output.name}")
-            inferred_type = infer_vellum_variable_type(workflow_output)
-
-            # Remove the terminal node output from the unreferenced set
-            unreferenced_final_output_node_outputs.discard(cast(OutputReference, workflow_output.instance))
-
-            if workflow_output.instance not in final_output_node_outputs:
-                # Create a synthetic terminal node only if there is no terminal node for this output
-                try:
-                    node_input = create_node_input(
-                        final_output_node_id,
-                        "node_input",
-                        # This is currently the wrapper node's output, but we want the wrapped node
-                        workflow_output.instance,
-                        self.display_context,
-                    )
-                except ValueError as e:
-                    raise ValueError(f"Failed to serialize output '{workflow_output.name}': {str(e)}") from e
-
-                source_node_display: Optional[BaseNodeDisplay]
-                first_rule = node_input.value.rules[0]
-                if first_rule.type == "NODE_OUTPUT":
-                    source_node_id = UUID(first_rule.data.node_id)
-                    try:
-                        source_node_display = [
-                            node_display
-                            for node_display in self.display_context.node_displays.values()
-                            if node_display.node_id == source_node_id
-                        ][0]
-                    except IndexError:
-                        source_node_display = None
-
-                synthetic_target_handle_id = str(
-                    uuid4_from_hash(f"{self.workflow_id}|target_handle_id|{workflow_output_display.name}")
-                )
-                synthetic_display_data = NodeDisplayData().dict()
-                synthetic_node_label = "Final Output"
-                nodes.append(
-                    {
-                        "id": str(final_output_node_id),
-                        "type": "TERMINAL",
-                        "data": {
-                            "label": synthetic_node_label,
-                            "name": workflow_output_display.name,
-                            "target_handle_id": synthetic_target_handle_id,
-                            "output_id": str(workflow_output_display.id),
-                            "output_type": inferred_type,
-                            "node_input_id": str(node_input.id),
-                        },
-                        "inputs": [node_input.dict()],
-                        "display_data": synthetic_display_data,
-                        "base": final_output_node_base,
-                        "definition": None,
-                    }
-                )
-
-                if source_node_display:
-                    if isinstance(source_node_display, BaseNodeVellumDisplay):
-                        source_handle_id = source_node_display.get_source_handle_id(
-                            port_displays=self.display_context.port_displays
-                        )
-                    else:
-                        source_handle_id = source_node_display.get_node_port_display(
-                            source_node_display._node.Ports.default
-                        ).id
-
-                    synthetic_output_edges.append(
-                        {
-                            "id": str(uuid4_from_hash(f"{self.workflow_id}|edge_id|{workflow_output_display.name}")),
-                            "source_node_id": str(source_node_display.node_id),
-                            "source_handle_id": str(source_handle_id),
-                            "target_node_id": str(final_output_node_id),
-                            "target_handle_id": synthetic_target_handle_id,
-                            "type": "DEFAULT",
-                        }
-                    )
-
-            output_variables.append(
-                {
-                    "id": str(workflow_output_display.id),
-                    "key": workflow_output_display.name,
-                    "type": inferred_type,
-                }
-            )
-
-        # If there are terminal nodes with no workflow output reference,
-        # raise a serialization error
-        if len(unreferenced_final_output_node_outputs) > 0:
-            self.add_error(
-                ValueError("Unable to serialize terminal nodes that are not referenced by workflow outputs.")
-            )
-
-        # Add an edge for each edge in the workflow
-        for target_node, entrypoint_display in self.display_context.entrypoint_displays.items():
-            unadorned_target_node = get_unadorned_node(target_node)
-            target_node_display = self.display_context.node_displays[unadorned_target_node]
-            edges.append(
-                {
-                    "id": str(entrypoint_display.edge_display.id),
-                    "source_node_id": str(entrypoint_node_id),
-                    "source_handle_id": str(entrypoint_node_source_handle_id),
-                    "target_node_id": str(target_node_display.node_id),
-                    "target_handle_id": str(target_node_display.get_trigger_id()),
-                    "type": "DEFAULT",
-                }
-            )
-
-        for (source_node_port, target_node), edge_display in self.display_context.edge_displays.items():
-            unadorned_source_node_port = get_unadorned_port(source_node_port)
-            unadorned_target_node = get_unadorned_node(target_node)
-
-            source_node_port_display = self.display_context.port_displays[unadorned_source_node_port]
-            target_node_display = self.display_context.node_displays[unadorned_target_node]
-
-            edges.append(
-                {
-                    "id": str(edge_display.id),
-                    "source_node_id": str(source_node_port_display.node_id),
-                    "source_handle_id": str(source_node_port_display.id),
-                    "target_node_id": str(target_node_display.node_id),
-                    "target_handle_id": str(
-                        target_node_display.get_target_handle_id_by_source_node_id(source_node_port_display.node_id)
-                    ),
-                    "type": "DEFAULT",
-                }
-            )
-
-        edges.extend(synthetic_output_edges)
-
-        return {
-            "workflow_raw_data": {
-                "nodes": nodes,
-                "edges": edges,
-                "display_data": self.display_context.workflow_display.display_data.dict(),
-                "definition": {
-                    "name": self._workflow.__name__,
-                    "module": cast(JsonArray, self._workflow.__module__.split(".")),
-                },
-            },
-            "input_variables": input_variables,
-            "state_variables": state_variables,
-            "output_variables": output_variables,
-        }
+# DEPRECATED: Use BaseWorkflowDisplay instead - This file will be removed in 0.15.0
+class VellumWorkflowDisplay(BaseWorkflowDisplay[WorkflowType], Generic[WorkflowType]):
+    pass

vellum_ee/workflows/tests/local_workflow/display/workflow.py

@@ -11,7 +11,7 @@ from vellum_ee.workflows.display.vellum import (
     WorkflowMetaVellumDisplayOverrides,
     WorkflowOutputVellumDisplayOverrides,
 )
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
+from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
 
 from ..inputs import Inputs
 from ..nodes.final_output import FinalOutput
@@ -19,7 +19,7 @@ from ..nodes.templating_node import TemplatingNode
 from ..workflow import Workflow
 
 
-class WorkflowDisplay(VellumWorkflowDisplay[Workflow]):
+class WorkflowDisplay(BaseWorkflowDisplay[Workflow]):
     workflow_display = WorkflowMetaVellumDisplayOverrides(
         entrypoint_node_id=UUID("0bf86989-13f2-438c-ab9c-d172e5771d31"),
         entrypoint_node_source_handle_id=UUID("23448f09-81ad-4378-abbd-1cccff350627"),

vellum_ee/workflows/display/nodes/base_node_vellum_display.py

@@ -1,40 +0,0 @@
-from uuid import UUID
-from typing import ClassVar, Dict, Optional
-
-from vellum.workflows.nodes.utils import get_unadorned_node
-from vellum.workflows.ports import Port
-from vellum.workflows.types.generics import NodeType
-from vellum_ee.workflows.display.editor.types import NodeDisplayComment, NodeDisplayData
-from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
-from vellum_ee.workflows.display.nodes.types import PortDisplay
-
-
-class BaseNodeVellumDisplay(BaseNodeDisplay[NodeType]):
-    # Used to explicitly set display data for a node
-    display_data: ClassVar[Optional[NodeDisplayData]] = None
-
-    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_source_handle_id(self, port_displays: Dict[Port, PortDisplay]) -> UUID:
-        unadorned_node = get_unadorned_node(self._node)
-        default_port = unadorned_node.Ports.default
-
-        default_port_display = port_displays[default_port]
-        return default_port_display.id