vellum-ai 0.14.35__py3-none-any.whl → 0.14.37__py3-none-any.whl

vellum_ee/workflows/display/nodes/vellum/tests/test_note_node.py

@@ -0,0 +1,33 @@
+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():
+    # GIVEN a note node
+    class MyNoteNode(NoteNode):
+        pass
+
+    # AND a display class for the note node
+    class MyNoteNodeDisplay(BaseNoteNodeDisplay[MyNoteNode]):
+        text = "This makes sense"
+        style = {
+            "fontSize": 24,
+        }
+
+    # AND a workflow with the code node
+    class Workflow(BaseWorkflow):
+        graph = MyNoteNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the inputs
+    my_note_node = next(node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "NOTE")
+
+    assert my_note_node["inputs"] == []
+    assert my_note_node["data"]["text"] == "This makes sense"
+    assert my_note_node["data"]["style"] == {"fontSize": 24}

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

@@ -1,11 +1,9 @@
 import pytest
 from uuid import uuid4
-from typing import Any, Dict, Type
+from typing import Any, Type
 
-from vellum.workflows.references.workflow_input import WorkflowInputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.generics import NodeType
-from vellum_ee.workflows.display.base import WorkflowInputsDisplayType
 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.get_node_display_class import get_node_display_class
@@ -14,6 +12,7 @@ from vellum_ee.workflows.display.types import (
     NodeOutputDisplays,
     StateValueDisplays,
     WorkflowDisplayContext,
+    WorkflowInputsDisplays,
 )
 from vellum_ee.workflows.display.vellum import WorkflowMetaVellumDisplay
 from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
@@ -24,7 +23,7 @@ def serialize_node():
     def _serialize_node(
         node_class: Type[NodeType],
         base_class: type[BaseNodeDisplay[Any]] = BaseNodeDisplay,
-        global_workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = {},
+        global_workflow_input_displays: WorkflowInputsDisplays = {},
         global_state_value_displays: StateValueDisplays = {},
         global_node_displays: NodeDisplays = {},
         global_node_output_displays: NodeOutputDisplays = {},

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

@@ -32,7 +32,7 @@ def test_serialize_workflow():
                 "key": "example",
                 "type": "STRING",
                 "default": {"type": "STRING", "value": "hello"},
-                "required": True,
+                "required": False,
                 "extensions": {"color": None},
             },
         ],

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

@@ -70,7 +70,6 @@ def test_serialize_workflow():
             "source_handle_id": "844d992e-60ab-4af2-a8ff-52cd858386f7",
         },
         "base": None,
-        "base": None,
         "definition": None,
         "display_data": {
             "position": {"x": 0.0, "y": 0.0},

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

@@ -144,6 +144,7 @@ def test_serialize_workflow(vellum_client):
             "variant": "DEPLOYMENT",
             "prompt_deployment_id": deployment.id,
             "release_tag": "LATEST",
+            "ml_model_fallbacks": ["gpt-4o", "gemini-1.5-pro"],
         },
         "display_data": {
             "position": {

vellum_ee/workflows/display/types.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Dict, Generic, Tuple, Type, TypeVar
+from typing import TYPE_CHECKING, Dict, Tuple, Type, TypeVar
 
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.events.workflow import WorkflowEventDisplayContext  # noqa: F401
@@ -10,7 +10,7 @@ from vellum_ee.workflows.display.base import (
     EdgeDisplay,
     EntrypointDisplay,
     StateValueDisplay,
-    WorkflowInputsDisplayType,
+    WorkflowInputsDisplay,
     WorkflowMetaDisplay,
     WorkflowOutputDisplay,
 )
@@ -22,6 +22,7 @@ if TYPE_CHECKING:
 
 WorkflowDisplayType = TypeVar("WorkflowDisplayType", bound="BaseWorkflowDisplay")
 
+WorkflowInputsDisplays = Dict[WorkflowInputReference, WorkflowInputsDisplay]
 StateValueDisplays = Dict[StateValueReference, StateValueDisplay]
 NodeDisplays = Dict[Type[BaseNode], BaseNodeDisplay]
 NodeOutputDisplays = Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]]
@@ -32,13 +33,11 @@ PortDisplays = Dict[Port, PortDisplay]
 
 
 @dataclass
-class WorkflowDisplayContext(Generic[WorkflowInputsDisplayType,]):
+class WorkflowDisplayContext:
     workflow_display_class: Type["BaseWorkflowDisplay"]
     workflow_display: WorkflowMetaDisplay
-    workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = field(default_factory=dict)
-    global_workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = field(
-        default_factory=dict
-    )
+    workflow_input_displays: WorkflowInputsDisplays = field(default_factory=dict)
+    global_workflow_input_displays: WorkflowInputsDisplays = field(default_factory=dict)
     state_value_displays: StateValueDisplays = field(default_factory=dict)
     global_state_value_displays: StateValueDisplays = field(default_factory=dict)
     node_displays: NodeDisplays = field(default_factory=dict)

vellum_ee/workflows/display/vellum.py

@@ -8,7 +8,6 @@ from vellum_ee.workflows.display.base import (
     EntrypointDisplayOverrides,
     StateValueDisplayOverrides,
     WorkflowInputsDisplay,
-    WorkflowInputsDisplayOverrides,
     WorkflowMetaDisplayOverrides,
     WorkflowOutputDisplayOverrides,
 )
@@ -44,10 +43,12 @@ class WorkflowMetaVellumDisplay(WorkflowMetaVellumDisplayOverrides):
 
 
 @dataclass
-class WorkflowInputsVellumDisplayOverrides(WorkflowInputsDisplay, WorkflowInputsDisplayOverrides):
-    name: Optional[str] = None
+class WorkflowInputsVellumDisplayOverrides(WorkflowInputsDisplay):
+    """
+    DEPRECATED: Use WorkflowInputsDisplay instead. Will be removed in 0.15.0
+    """
+
     required: Optional[bool] = None
-    color: Optional[str] = None
 
 
 @dataclass

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -21,8 +21,7 @@ from vellum_ee.workflows.display.base import (
     EdgeDisplay,
     EntrypointDisplay,
     StateValueDisplay,
-    WorkflowInputsDisplayOverridesType,
-    WorkflowInputsDisplayType,
+    WorkflowInputsDisplay,
     WorkflowMetaDisplay,
     WorkflowOutputDisplay,
 )
@@ -40,6 +39,7 @@ from vellum_ee.workflows.display.types import (
     PortDisplays,
     StateValueDisplays,
     WorkflowDisplayContext,
+    WorkflowInputsDisplays,
     WorkflowOutputDisplays,
 )
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
@@ -47,18 +47,12 @@ from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class imp
 logger = logging.getLogger(__name__)
 
 
-class BaseWorkflowDisplay(
-    Generic[
-        WorkflowType,
-        WorkflowInputsDisplayType,
-        WorkflowInputsDisplayOverridesType,
-    ]
-):
+class BaseWorkflowDisplay(Generic[WorkflowType]):
     # Used to specify the display data for a workflow.
     workflow_display: Optional[WorkflowMetaDisplay] = None
 
     # Used to explicitly specify display data for a workflow's inputs.
-    inputs_display: Dict[WorkflowInputReference, WorkflowInputsDisplayOverridesType] = {}
+    inputs_display: WorkflowInputsDisplays = {}
 
     # Used to explicitly specify display data for a workflow's state values.
     state_value_displays: StateValueDisplays = {}
@@ -86,7 +80,7 @@ class BaseWorkflowDisplay(
         self,
         workflow: Type[WorkflowType],
         *,
-        parent_display_context: Optional[WorkflowDisplayContext[WorkflowInputsDisplayType,]] = None,
+        parent_display_context: Optional[WorkflowDisplayContext] = None,
         dry_run: bool = False,
     ):
         self._workflow = workflow
@@ -174,9 +168,7 @@ class BaseWorkflowDisplay(
         return node_display
 
     @cached_property
-    def display_context(
-        self,
-    ) -> WorkflowDisplayContext[WorkflowInputsDisplayType]:
+    def display_context(self) -> WorkflowDisplayContext:
         workflow_meta_display = self._generate_workflow_meta_display()
 
         global_node_output_displays: NodeOutputDisplays = (
@@ -209,7 +201,7 @@ class BaseWorkflowDisplay(
                 port_displays=port_displays,
             )
 
-        workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = {}
+        workflow_input_displays: WorkflowInputsDisplays = {}
         # If we're dealing with a nested workflow, then it should have access to the inputs of its parents.
         global_workflow_input_displays = (
             copy(self._parent_display_context.workflow_input_displays) if self._parent_display_context else {}
@@ -311,11 +303,20 @@ class BaseWorkflowDisplay(
             entrypoint_node_display=NodeDisplayData(),
         )
 
-    @abstractmethod
     def _generate_workflow_input_display(
-        self, workflow_input: WorkflowInputReference, overrides: Optional[WorkflowInputsDisplayOverridesType] = None
-    ) -> WorkflowInputsDisplayType:
-        pass
+        self, workflow_input: WorkflowInputReference, overrides: Optional[WorkflowInputsDisplay] = None
+    ) -> WorkflowInputsDisplay:
+        workflow_input_id: UUID
+        name = None
+        color = None
+        if overrides:
+            workflow_input_id = overrides.id
+            name = overrides.name
+            color = overrides.color
+        else:
+            workflow_input_id = uuid4_from_hash(f"{self.workflow_id}|inputs|id|{workflow_input.name}")
+
+        return WorkflowInputsDisplay(id=workflow_input_id, name=name, color=color)
 
     def _generate_state_value_display(
         self, state_value: BaseDescriptor, overrides: Optional[StateValueDisplay] = None

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -2,10 +2,10 @@ import logging
 from uuid import UUID
 from typing import Optional, cast
 
+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 import WorkflowInputReference
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.generics import WorkflowType
@@ -15,38 +15,29 @@ 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.vellum import WorkflowInputsVellumDisplay, WorkflowInputsVellumDisplayOverrides
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 logger = logging.getLogger(__name__)
 
 
-class VellumWorkflowDisplay(
-    BaseWorkflowDisplay[
-        WorkflowType,
-        WorkflowInputsVellumDisplay,
-        WorkflowInputsVellumDisplayOverrides,
-    ]
-):
+class VellumWorkflowDisplay(BaseWorkflowDisplay[WorkflowType]):
     node_display_base_class = BaseNodeDisplay
 
     def serialize(self) -> JsonObject:
         input_variables: JsonArray = []
-        for workflow_input, workflow_input_display in self.display_context.workflow_input_displays.items():
-            default = primitive_to_vellum_value(workflow_input.instance) if workflow_input.instance else None
-            required = (
-                workflow_input_display.required
-                if workflow_input_display.required is not None
-                else type(None) not in workflow_input.types
+        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.name,
-                    "type": infer_vellum_variable_type(workflow_input),
+                    "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": required,
+                    "required": workflow_input_reference.instance is undefined,
                     "extensions": {"color": workflow_input_display.color},
                 }
             )
@@ -62,7 +53,7 @@ class VellumWorkflowDisplay(
                     "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 None,
+                    "required": state_value_reference.instance is undefined,
                     "extensions": {"color": state_value_display.color},
                 }
             )
@@ -269,20 +260,3 @@ class VellumWorkflowDisplay(
             "state_variables": state_variables,
             "output_variables": output_variables,
         }
-
-    def _generate_workflow_input_display(
-        self, workflow_input: WorkflowInputReference, overrides: Optional[WorkflowInputsVellumDisplayOverrides] = None
-    ) -> WorkflowInputsVellumDisplay:
-        workflow_input_id: UUID
-        name = None
-        required = None
-        color = None
-        if overrides:
-            workflow_input_id = overrides.id
-            name = overrides.name
-            required = overrides.required
-            color = overrides.color
-        else:
-            workflow_input_id = uuid4_from_hash(f"{self.workflow_id}|inputs|id|{workflow_input.name}")
-
-        return WorkflowInputsVellumDisplay(id=workflow_input_id, name=name, required=required, color=color)