vellum-ai 0.10.9__py3-none-any.whl → 0.11.0__py3-none-any.whl

vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -1,7 +1,7 @@
 from typing import Optional, Type
 
-from vellum_ee.workflows.display.types import NodeDisplayType
 from vellum.workflows.types.generics import NodeType
+from vellum_ee.workflows.display.types import NodeDisplayType
 
 
 def get_node_display_class(

vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -1,8 +1,8 @@
 import pytest
 from uuid import UUID
 
-from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum.workflows.nodes.bases import BaseNode
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 
 
 @pytest.fixture

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

@@ -32,5 +32,5 @@ __all__ = [
     "BaseSearchNodeDisplay",
     "BaseSubworkflowDeploymentNodeDisplay",
     "BaseTemplatingNodeDisplay",
-    "BaseTryNodeDisplay"
+    "BaseTryNodeDisplay",
 ]

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

@@ -22,10 +22,6 @@ class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeTy
     api_key_header_key_input_id: ClassVar[Optional[UUID]] = None
     api_key_header_value_input_id: ClassVar[Optional[UUID]] = None
 
-    text_output_id: ClassVar[Optional[UUID]] = None
-    json_output_id: ClassVar[Optional[UUID]] = None
-    status_code_output_id: ClassVar[Optional[UUID]] = None
-
     # A mapping between node input keys and their ids for inputs representing additional header keys
     additional_header_key_input_ids: ClassVar[Optional[Dict[str, UUID]]] = None
 
@@ -68,6 +64,7 @@ class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeTy
 
         headers = raise_if_descriptor(node.headers)
         api_key_header_key = raise_if_descriptor(node.api_key_header_key)
+        api_key_header_value = raise_if_descriptor(node.api_key_header_value)
         authorization_type = raise_if_descriptor(node.authorization_type)
         bearer_token_value = raise_if_descriptor(node.bearer_token_value)
 
@@ -112,7 +109,7 @@ class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeTy
                 display_context=display_context,
                 input_id=self.api_key_header_value_input_id,
             )
-            if headers and api_key_header_key
+            if headers and api_key_header_key and api_key_header_value
             else None
         )
 
@@ -123,7 +120,7 @@ class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeTy
                 {
                     "header_key_input_id": create_node_input(
                         node_id=node_id,
-                        input_name=f"additional_header_key_{key}",
+                        input_name="additional_header_key",
                         value=key,
                         display_context=display_context,
                         input_id=(
@@ -134,7 +131,7 @@ class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeTy
                     ).id,
                     "header_value_input_id": create_node_input(
                         node_id=node_id,
-                        input_name=f"additional_header_value_{key}",
+                        input_name="additional_header_value",
                         value=value,
                         display_context=display_context,
                         input_id=(

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

@@ -1,12 +1,7 @@
 from dataclasses import dataclass
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, List, Optional, TypeVar, Union, Tuple
+from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, TypeVar, Union
 
-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.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.vellum import NodeInput
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.expressions.and_ import AndExpression
 from vellum.workflows.expressions.begins_with import BeginsWithExpression
@@ -29,7 +24,12 @@ from vellum.workflows.expressions.not_between import NotBetweenExpression
 from vellum.workflows.expressions.not_in import NotInExpression
 from vellum.workflows.expressions.or_ import OrExpression
 from vellum.workflows.nodes.displayable import ConditionalNode
-from vellum.workflows.types.core import JsonObject, ConditionType
+from vellum.workflows.types.core import ConditionType, JsonObject
+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.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.vellum import NodeInput
 
 _ConditionalNodeType = TypeVar("_ConditionalNodeType", bound=ConditionalNode)
 
@@ -42,6 +42,7 @@ class RuleIdMap:
     field_node_input_id: Optional[UUID]
     value_node_input_id: Optional[UUID]
 
+
 @dataclass
 class ConditionId:
     id: UUID
@@ -65,7 +66,9 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
 
         if len(condition_ids) > ports_size:
             raise ValueError(
-                f"Too many defined condition ids. Ports are size {ports_size} but the defined conditions have length {len(condition_ids)}")
+                f"""Too many defined condition ids. Ports are size {ports_size} \
+but the defined conditions have length {len(condition_ids)}"""
+            )
 
         def serialize_rule(
             descriptor: BaseDescriptor, path: List[int], rule_id_map: Optional[RuleIdMap]
@@ -113,7 +116,11 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
                     node_id, f"{current_id}.field", descriptor._value, display_context, field_node_input_id
                 )
                 value_node_input = create_node_input(
-                    node_id, f"{current_id}.value", f"{descriptor._start},{descriptor._end}", display_context, value_node_input_id
+                    node_id,
+                    f"{current_id}.value",
+                    f"{descriptor._start},{descriptor._end}",
+                    display_context,
+                    value_node_input_id,
                 )
                 node_inputs.extend([field_node_input, value_node_input])
                 operator = self._convert_descriptor_to_operator(descriptor)
@@ -124,11 +131,15 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
                 lhs = descriptor._lhs  # type: ignore[attr-defined]
                 rhs = descriptor._rhs  # type: ignore[attr-defined]
 
-                lhs_node_input = create_node_input(node_id, f"{current_id}.field", lhs, display_context, field_node_input_id)
+                lhs_node_input = create_node_input(
+                    node_id, f"{current_id}.field", lhs, display_context, field_node_input_id
+                )
                 node_inputs.append(lhs_node_input)
 
                 if descriptor._rhs is not None:  # type: ignore[attr-defined]
-                    rhs_node_input = create_node_input(node_id, f"{current_id}.value", rhs, display_context, value_node_input_id)
+                    rhs_node_input = create_node_input(
+                        node_id, f"{current_id}.value", rhs, display_context, value_node_input_id
+                    )
                     node_inputs.append(rhs_node_input)
                     value_node_input_id = UUID(rhs_node_input.id)
 
@@ -148,18 +159,24 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
         conditions = []
         for idx, port in enumerate(node.Ports):
 
-            condition_id = str(condition_ids[idx].id if condition_ids else uuid4_from_hash(f"{node_id}|conditions|{idx}"))
-            rule_group_id = str(condition_ids[idx].rule_group_id if condition_ids else uuid4_from_hash(f"{condition_id}|rule_group"))
+            condition_id = str(
+                condition_ids[idx].id if condition_ids else uuid4_from_hash(f"{node_id}|conditions|{idx}")
+            )
+            rule_group_id = str(
+                condition_ids[idx].rule_group_id if condition_ids else uuid4_from_hash(f"{condition_id}|rule_group")
+            )
             source_handle_id = str(source_handle_ids.get(idx) or uuid4_from_hash(f"{node_id}|handles|{idx}"))
 
             if port._condition is None:
                 if port._condition_type == ConditionType.ELSE:
-                    conditions.append({
-                        "id": condition_id,
-                        "type": ConditionType.ELSE.value,
-                        "source_handle_id": source_handle_id,
-                        "data": None
-                    })
+                    conditions.append(
+                        {
+                            "id": condition_id,
+                            "type": ConditionType.ELSE.value,
+                            "source_handle_id": source_handle_id,
+                            "data": None,
+                        }
+                    )
                 else:
                     continue
 
@@ -241,7 +258,7 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
             raise ValueError(f"Unsupported descriptor type: {descriptor}")
 
     def get_nested_rule_details_by_path(
-            self, rule_ids: List[RuleIdMap], path: List[int]
+        self, rule_ids: List[RuleIdMap], path: List[int]
     ) -> Union[Tuple[UUID, Optional[UUID], Optional[UUID]], None]:
         current_rule = rule_ids[path[0]]
 
@@ -267,7 +284,7 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
         return (
             uuid4_from_hash(f"{node_id}|{rule_id}|current"),
             uuid4_from_hash(f"{node_id}|{rule_id}||field"),
-            uuid4_from_hash(f"{node_id}|{rule_id}||value")
+            uuid4_from_hash(f"{node_id}|{rule_id}||value"),
         )
 
     def _get_source_handle_ids(self) -> Dict[int, UUID]:
@@ -276,5 +293,5 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
     def _get_rule_ids(self) -> List[RuleIdMap]:
         return self._get_explicit_node_display_attr("rule_ids", List[RuleIdMap]) or []
 
-    def _get_condition_ids(self)-> List[ConditionId]:
+    def _get_condition_ids(self) -> List[ConditionId]:
         return self._get_explicit_node_display_attr("condition_ids", List[ConditionId]) or []

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

@@ -2,22 +2,22 @@ from uuid import UUID
 from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes import ErrorNode
-from vellum.workflows.types.core import EntityInputsInterface, Json, JsonObject
+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.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _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(
-            self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
-        node = self._node
         node_id = self.node_id
         error_source_input_id = self.node_input_ids_by_name.get("error_source_input_id")
 

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

@@ -1,9 +1,7 @@
 from uuid import UUID
 from typing import Any, ClassVar, Generic, Optional, TypeVar
 
-from vellum.workflows.nodes.core.map_node.node import MapNode
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
-from vellum.workflows.references.output import OutputReference
 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 to_kebab_case

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

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar
+from typing import ClassVar, Dict, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes import GuardrailNode
 from vellum.workflows.types.core import JsonObject

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

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast
+from typing import ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast
 
 from vellum import PromptBlock, RichTextChildBlock, VellumVariable
 from vellum.workflows.nodes import InlinePromptNode

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

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, cast
+from typing import ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, cast
 
 from vellum import VellumVariable
 from vellum.workflows.nodes import InlineSubworkflowNode
@@ -90,7 +90,9 @@ class BaseInlineSubworkflowNodeDisplay(
     ) -> List[VellumVariable]:
         workflow_outputs: List[VellumVariable] = []
         for output_descriptor in raise_if_descriptor(node.subworkflow).Outputs:  # type: ignore[union-attr]
-            workflow_output_display = cast(WorkflowOutputVellumDisplay, display_context.workflow_output_displays[output_descriptor])
+            workflow_output_display = cast(
+                WorkflowOutputVellumDisplay, display_context.workflow_output_displays[output_descriptor]
+            )
             output_type = infer_vellum_variable_type(output_descriptor)
             workflow_outputs.append(
                 VellumVariable(id=str(workflow_output_display.id), key=workflow_output_display.name, type=output_type)

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

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Collection, Dict, Generic, List, Optional, TypeVar, cast
+from typing import Dict, Generic, List, Optional, TypeVar, cast
 
 from vellum.workflows.nodes import MapNode
 from vellum.workflows.types.core import JsonObject
@@ -39,15 +39,21 @@ class BaseMapNodeDisplay(BaseNodeVellumDisplay[_MapNodeType], Generic[_MapNodeTy
         renamed_input_variables = []
         for input_variable in cast(List[Dict[str, str]], serialized_subworkflow["input_variables"]):
             if input_variable["key"] == "all_items":
-                renamed_item = { **input_variable, "key": "items" }
+                renamed_item = {**input_variable, "key": "items"}
                 renamed_input_variables.append(renamed_item)
             else:
                 renamed_input_variables.append(input_variable)
 
         # Note: This must match the items input ID for the map node's node input
-        items_workflow_input_id = next(input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "items")
-        item_workflow_input_id = next(input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "item")
-        index_workflow_input_id = next(input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "index")
+        items_workflow_input_id = next(
+            input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "items"
+        )
+        item_workflow_input_id = next(
+            input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "item"
+        )
+        index_workflow_input_id = next(
+            input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "index"
+        )
 
         return {
             "id": str(node_id),

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

@@ -1,12 +1,12 @@
 from uuid import UUID
 from typing import Any, ClassVar, Generic, List, Optional, TypeVar
 
+from vellum.workflows.nodes.displayable import MergeNode
+from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.vellum import EdgeVellumDisplay
-from vellum.workflows.nodes.displayable import MergeNode
-from vellum.workflows.types.core import JsonObject
 
 _MergeNodeType = TypeVar("_MergeNodeType", bound=MergeNode)
 

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

@@ -13,9 +13,7 @@ class BaseNoteNodeDisplay(BaseNodeVellumDisplay[_NoteNodeType], Generic[_NoteNod
     text: ClassVar[str] = ""
     style: ClassVar[Union[Dict[str, Any], None]] = None
 
-    def serialize(
-        self, display_context: WorkflowDisplayContext, **kwargs: Any
-    ) -> JsonObject:
+    def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
         node_id = self.node_id
 
         return {

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

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar, cast
+from typing import ClassVar, Dict, Generic, Optional, TypeVar, cast
 
 from vellum.workflows.nodes.displayable.prompt_deployment_node import PromptDeploymentNode
 from vellum.workflows.references import OutputReference

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

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from uuid import UUID
-from typing import Any, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast
+from typing import Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast
 
 from vellum import (
     MetadataFilterConfigRequest,

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

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar
+from typing import ClassVar, Dict, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes import SubworkflowDeploymentNode
 from vellum.workflows.types.core import JsonObject

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

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar
+from typing import ClassVar, Dict, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes.core.templating_node import TemplatingNode
 from vellum.workflows.types.core import JsonObject

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

@@ -2,6 +2,11 @@ import pytest
 from uuid import UUID, uuid4
 from typing import List, cast
 
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.references import OutputReference, WorkflowInputReference
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input_value_pointer_rules
@@ -19,11 +24,6 @@ from vellum_ee.workflows.display.vellum import (
     WorkflowMetaVellumDisplay,
 )
 from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
-from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.inputs import BaseInputs
-from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.outputs import BaseOutputs
-from vellum.workflows.references import OutputReference, WorkflowInputReference
 
 
 class Inputs(BaseInputs):

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

@@ -1,6 +1,10 @@
 from uuid import UUID
 from typing import Any, List, Optional, cast
 
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
+from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
+from vellum.workflows.references import NodeReference, OutputReference
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.vellum import create_node_input_value_pointer_rule, primitive_to_vellum_value
@@ -10,10 +14,6 @@ from vellum_ee.workflows.display.vellum import (
     NodeInputValuePointer,
     NodeInputValuePointerRule,
 )
-from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
-from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
-from vellum.workflows.references import NodeReference, OutputReference
 
 
 def create_node_input(

vellum_ee/workflows/display/tests/test_vellum_workflow_display.py

@@ -0,0 +1,45 @@
+from vellum.workflows.workflows.base import BaseWorkflow
+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_vellum_workflow_display__serialize_empty_workflow():
+    # GIVEN an empty workflow
+    class ExampleWorkflow(BaseWorkflow):
+        pass
+
+    display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay,
+        workflow_class=ExampleWorkflow,
+    )
+
+    # WHEN serializing the workflow
+    exec_config = display.serialize()
+
+    # THEN it should return the expected config
+    assert exec_config == {
+        "input_variables": [],
+        "output_variables": [],
+        "workflow_raw_data": {
+            "definition": {
+                "module": ["vellum_ee", "workflows", "display", "tests", "test_vellum_workflow_display"],
+                "name": "ExampleWorkflow",
+            },
+            "display_data": {"viewport": {"x": 0.0, "y": 0.0, "zoom": 1.0}},
+            "edges": [],
+            "nodes": [
+                {
+                    "data": {"label": "Entrypoint Node", "source_handle_id": "508b8b82-3517-4672-a155-18c9c7b9c545"},
+                    "definition": {
+                        "bases": [],
+                        "module": ["vellum", "workflows", "nodes", "bases", "base"],
+                        "name": "BaseNode",
+                    },
+                    "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                    "id": "9eef0c18-f322-4d56-aa89-f088d3e53f6a",
+                    "inputs": [],
+                    "type": "ENTRYPOINT",
+                }
+            ],
+        },
+    }

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

@@ -4,17 +4,18 @@ from uuid import uuid4
 from deepdiff import DeepDiff
 
 from vellum import WorkspaceSecretRead
-
-from tests.workflows.basic_vellum_api_node.workflow import SimpleAPIWorkflow
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
+from tests.workflows.basic_vellum_api_node.workflow import SimpleAPIWorkflow
+
 
 def test_serialize_workflow(vellum_client):
     # GIVEN a Workflow that uses a vellum API node
     # AND stubbed out API calls
+    workspace_secret_id = str(uuid4())
     workspace_secret = WorkspaceSecretRead(
-        id=str(uuid4()),
+        id=workspace_secret_id,
         modified=datetime.now(),
         name="MY_SECRET",
         label="My Secret",
@@ -133,7 +134,10 @@ def test_serialize_workflow(vellum_client):
                         "rules": [
                             {
                                 "type": "WORKSPACE_SECRET",
-                                "data": {"type": "STRING", "workspace_secret_id": workspace_secret.id},
+                                "data": {
+                                    "type": "STRING",
+                                    "workspace_secret_id": f"{workspace_secret_id}",
+                                },
                             }
                         ],
                         "combinator": "OR",
@@ -154,8 +158,8 @@ def test_serialize_workflow(vellum_client):
                 "api_key_header_value_input_id": "d794bb51-a419-4fd8-be63-dfaf4166e831",
                 "additional_headers": [
                     {
-                        "header_key_input_id": "cff297b8-8ef9-46ab-9ef4-1f1eb43a2298",
-                        "header_value_input_id": "0d41f210-f428-4849-9c21-0cde2cf54657",
+                        "header_key_input_id": "57c31247-998a-430d-bb62-bf50eca7df35",
+                        "header_value_input_id": "47b32274-f19b-4c15-b788-55c069c311c5",
                     }
                 ],
                 "text_output_id": "17342c21-12bb-49ab-88ce-f144e0376b32",
@@ -164,26 +168,11 @@ def test_serialize_workflow(vellum_client):
             },
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "definition": {
+                "name": "SimpleAPINode",
+                "module": ["tests", "workflows", "basic_vellum_api_node", "workflow"],
                 "bases": [
-                    {
-                        "module": [
-                            "vellum",
-                            "workflows",
-                            "nodes",
-                            "displayable",
-                            "api_node",
-                            "node",
-                        ],
-                        "name": "APINode",
-                    }
+                    {"name": "APINode", "module": ["vellum", "workflows", "nodes", "displayable", "api_node", "node"]}
                 ],
-                "module": [
-                    "tests",
-                    "workflows",
-                    "basic_vellum_api_node",
-                    "workflow",
-                ],
-                "name": "SimpleAPINode",
             },
         },
         api_node,