vellum-ai 0.10.3__py3-none-any.whl → 0.10.6__py3-none-any.whl

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

@@ -2,11 +2,12 @@ from unittest import mock
 
 from deepdiff import DeepDiff
 
-from tests.workflows.basic_inline_subworkflow.workflow import BasicInlineSubworkflowWorkflow
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 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_inline_subworkflow.workflow import BasicInlineSubworkflowWorkflow
+
 
 def test_serialize_workflow():
     # GIVEN a Workflow
@@ -37,17 +38,17 @@ def test_serialize_workflow():
                 "id": "fa73da37-34c3-47a9-be58-69cc6cdbfca5",
                 "key": "city",
                 "type": "STRING",
-                "required": None,
+                "required": True,
                 "default": None,
-                "extensions": None,
+                "extensions": { "color": None },
             },
             {
                 "id": "aba1e6e0-dfa7-4c15-a4e6-aec6feebfaca",
                 "key": "date",
                 "type": "STRING",
-                "required": None,
+                "required": True,
                 "default": None,
-                "extensions": None,
+                "extensions": { "color": None },
             },
         ],
         input_variables,
@@ -306,8 +307,8 @@ def test_serialize_workflow():
                 },
                 "input_variables": [{"id": "704c4640-bfda-44f0-8da3-e9cfc4f21cf2", "key": "metro", "type": "STRING"}],
                 "output_variables": [
-                    {"id": "86dd0202-c141-48a3-8382-2da60372e77c", "key": "temperature", "type": "NUMBER"},
-                    {"id": "0a7192da-5576-4933-bba4-de8adf5d7996", "key": "reasoning", "type": "STRING"},
+                    {"id": "2fc57139-7420-49e5-96a6-dcbb3ff5d622", "key": "temperature", "type": "NUMBER"},
+                    {"id": "fad5dd9f-3328-4e70-ad55-65a5325a4a82", "key": "reasoning", "type": "STRING"},
                 ],
             },
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
@@ -338,64 +339,74 @@ def test_serialize_workflow():
         ignore_order=True,
     )
 
+    temperature_terminal_node = next(node for node in workflow_raw_data["nodes"][2:] if node['data']['name'] == 'temperature')
+    reasoning_terminal_node = next(node for node in workflow_raw_data["nodes"][2:] if node['data']['name'] == 'reasoning')
+    
     assert not DeepDiff(
-        [
-            {
-                "id": "31b74695-3f1c-47cf-8be8-a4d86cc589e8",
-                "type": "TERMINAL",
-                "definition": {
-                    "bases": [
-                        {
-                            "bases": [],
-                            "module": [
-                                "vellum",
-                                "workflows",
-                                "nodes",
-                                "bases",
-                                "base",
-                            ],
-                            "name": "BaseNode",
-                        },
-                    ],
-                    "module": [
-                        "vellum",
-                        "workflows",
-                        "nodes",
-                        "displayable",
-                        "final_output_node",
-                        "node",
-                    ],
-                    "name": "FinalOutputNode",
-                },
-                "data": {
-                    "label": "Final Output",
-                    "name": "reasoning",
-                    "target_handle_id": "8b525943-6c27-414b-a329-e29c0b217f72",
-                    "output_id": "7444a019-081a-4e10-a528-3249299159f7",
-                    "output_type": "STRING",
-                    "node_input_id": "736473c8-b0b4-4cdd-b743-6453dd5306fc",
-                },
-                "inputs": [
+        {
+            "id": "31b74695-3f1c-47cf-8be8-a4d86cc589e8",
+            "type": "TERMINAL",
+            "definition": {
+                "bases": [
                     {
-                        "id": "736473c8-b0b4-4cdd-b743-6453dd5306fc",
-                        "key": "node_input",
-                        "value": {
-                            "rules": [
-                                {
-                                    "type": "NODE_OUTPUT",
-                                    "data": {
-                                        "node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
-                                        "output_id": "0a7192da-5576-4933-bba4-de8adf5d7996",
-                                    },
-                                }
-                            ],
-                            "combinator": "OR",
-                        },
-                    }
+                        "bases": [],
+                        "module": [
+                            "vellum",
+                            "workflows",
+                            "nodes",
+                            "bases",
+                            "base",
+                        ],
+                        "name": "BaseNode",
+                    },
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "module": [
+                    "vellum",
+                    "workflows",
+                    "nodes",
+                    "displayable",
+                    "final_output_node",
+                    "node",
+                ],
+                "name": "FinalOutputNode",
             },
-            {
+            "data": {
+                "label": "Final Output",
+                "name": "reasoning",
+                "target_handle_id": "8b525943-6c27-414b-a329-e29c0b217f72",
+                "output_id": "7444a019-081a-4e10-a528-3249299159f7",
+                "output_type": "STRING",
+                "node_input_id": "736473c8-b0b4-4cdd-b743-6453dd5306fc",
+            },
+            "inputs": [
+                {
+                    "id": "736473c8-b0b4-4cdd-b743-6453dd5306fc",
+                    "key": "node_input",
+                    "value": {
+                        "rules": [
+                            {
+                                "type": "NODE_OUTPUT",
+                                "data": {
+                                    "node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
+                                    "output_id": "fad5dd9f-3328-4e70-ad55-65a5325a4a82",
+                                },
+                            }
+                        ],
+                        "combinator": "OR",
+                    },
+                }
+            ],
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        },
+        reasoning_terminal_node,
+        ignore_order=True,
+        # TODO: Make sure this output ID matches the workflow output ID of the subworkflow node's workflow
+        # https://app.shortcut.com/vellum/story/5660/fix-output-id-in-subworkflow-nodes
+        exclude_regex_paths=r"root\['inputs'\]\[0\]\['value'\]\['rules'\]\[0\]\['data'\]\['output_id'\]"
+    )
+
+    assert not DeepDiff(
+        {
                 "id": "0779b232-82ab-4dbe-a340-6a85e6ab3368",
                 "type": "TERMINAL",
                 "definition": {
@@ -440,7 +451,7 @@ def test_serialize_workflow():
                                     "type": "NODE_OUTPUT",
                                     "data": {
                                         "node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
-                                        "output_id": "86dd0202-c141-48a3-8382-2da60372e77c",
+                                        "output_id": "2fc57139-7420-49e5-96a6-dcbb3ff5d622",
                                     },
                                 }
                             ],
@@ -450,9 +461,11 @@ def test_serialize_workflow():
                 ],
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
             },
-        ],
-        workflow_raw_data["nodes"][2:],
+        temperature_terminal_node,
         ignore_order=True,
+        # TODO: Make sure this output ID matches the workflow output ID of the subworkflow node's workflow
+        # https://app.shortcut.com/vellum/story/5660/fix-output-id-in-subworkflow-nodes
+        exclude_regex_paths=r"root\['inputs'\]\[0\]\['value'\]\['rules'\]\[0\]\['data'\]\['output_id'\]"
     )
 
     # AND each edge should be serialized correctly

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

@@ -2,11 +2,12 @@ from unittest import mock
 
 from deepdiff import DeepDiff
 
-from tests.workflows.basic_map_node.workflow import SimpleMapExample
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 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_map_node.workflow import SimpleMapExample
+
 
 def test_serialize_workflow():
     # GIVEN a Workflow that uses a MapNode
@@ -35,9 +36,9 @@ def test_serialize_workflow():
                 "id": "db2eb237-38e4-417a-8bfc-5bda0f3165ca",
                 "key": "fruits",
                 "type": "JSON",
-                "required": None,
+                "required": True,
                 "default": None,
-                "extensions": None,
+                "extensions": { "color": None },
             },
         ],
         input_variables,

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

@@ -4,11 +4,11 @@ from uuid import uuid4
 from deepdiff import DeepDiff
 
 from vellum import DeploymentRead
-
-from tests.workflows.basic_text_prompt_deployment.workflow import BasicTextPromptDeployment
 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_text_prompt_deployment.workflow import BasicTextPromptDeployment
+
 
 def test_serialize_workflow(vellum_client):
     # GIVEN a Workflow with stubbed out API calls
@@ -46,17 +46,17 @@ def test_serialize_workflow(vellum_client):
                 "id": "52995b50-84c9-465f-8a4b-a4ee2a92e388",
                 "key": "city",
                 "type": "STRING",
-                "required": None,
+                "required": True,
                 "default": None,
-                "extensions": None,
+                "extensions": { "color": None },
             },
             {
                 "id": "aa3ca842-250c-4a3f-853f-23928c28d0f8",
                 "key": "date",
                 "type": "STRING",
-                "required": None,
+                "required": True,
                 "default": None,
-                "extensions": None,
+                "extensions": { "color": None },
             },
         ],
         input_variables,

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

@@ -4,11 +4,11 @@ from uuid import uuid4
 from deepdiff import DeepDiff
 
 from vellum import WorkflowDeploymentRead
-
-from tests.workflows.basic_subworkflow_deployment.workflow import BasicSubworkflowDeploymentWorkflow
 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_subworkflow_deployment.workflow import BasicSubworkflowDeploymentWorkflow
+
 
 def test_serialize_workflow(vellum_client):
     # GIVEN a Workflow with stubbed out API calls
@@ -46,17 +46,17 @@ def test_serialize_workflow(vellum_client):
                 "id": "693cc9a5-8d74-4a58-bdcf-2b4989cdf250",
                 "key": "city",
                 "type": "STRING",
-                "required": None,
+                "required": True,
                 "default": None,
-                "extensions": None,
+                "extensions": { "color": None },
             },
             {
                 "id": "19a78824-9a98-4ae8-a1fc-61f81a422a17",
                 "key": "date",
                 "type": "STRING",
-                "required": None,
+                "required": True,
                 "default": None,
-                "extensions": None,
+                "extensions": { "color": None },
             },
         ],
         input_variables,

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

@@ -1,7 +1,8 @@
-from tests.workflows.basic_final_output_node.workflow import BasicFinalOutputNodeWorkflow
 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_final_output_node.workflow import BasicFinalOutputNodeWorkflow
+
 
 def test_serialize_workflow():
     # GIVEN a Workflow that uses a Final Output Node
@@ -26,9 +27,9 @@ def test_serialize_workflow():
             "id": "e39a7b63-de15-490a-ae9b-8112c767aea0",
             "key": "input",
             "type": "STRING",
-            "required": None,
+            "required": True,
             "default": None,
-            "extensions": None,
+            "extensions": { "color": None },
         }
     ]
 

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

@@ -3,12 +3,13 @@ from unittest import mock
 
 from deepdiff import DeepDiff
 
-from tests.workflows.complex_final_output_node.missing_final_output_node import MissingFinalOutputNodeWorkflow
-from tests.workflows.complex_final_output_node.missing_workflow_output import MissingWorkflowOutputWorkflow
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 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.complex_final_output_node.missing_final_output_node import MissingFinalOutputNodeWorkflow
+from tests.workflows.complex_final_output_node.missing_workflow_output import MissingWorkflowOutputWorkflow
+
 
 def test_serialize_workflow__missing_final_output_node():
     # GIVEN a Workflow that is missing a Terminal Node
@@ -40,17 +41,17 @@ def test_serialize_workflow__missing_final_output_node():
                 "id": "da086239-d743-4246-b666-5c91e22fb88c",
                 "key": "alpha",
                 "type": "STRING",
+                "required": True,
                 "default": None,
-                "required": None,
-                "extensions": None,
+                "extensions": { "color": None },
             },
             {
                 "id": "a8b6c5d4-a0e9-4457-834b-46b633c466a6",
                 "key": "beta",
                 "type": "STRING",
+                "required": True,
                 "default": None,
-                "required": None,
-                "extensions": None,
+                "extensions": { "color": None },
             },
         ],
         input_variables,

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -5,6 +5,15 @@ import logging
 from uuid import UUID
 from typing import Any, Dict, Generic, Optional, Tuple, Type, get_args
 
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.edges import Edge
+from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
+from vellum.workflows.ports import Port
+from vellum.workflows.references import OutputReference, WorkflowInputReference
+from vellum.workflows.types.core import JsonObject
+from vellum.workflows.types.generics import WorkflowType
 from vellum_ee.workflows.display.base import (
     EdgeDisplayOverridesType,
     EdgeDisplayType,
@@ -21,15 +30,6 @@ from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_di
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
 from vellum_ee.workflows.display.types import NodeDisplayType, WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
-from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.edges import Edge
-from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
-from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
-from vellum.workflows.ports import Port
-from vellum.workflows.references import OutputReference, WorkflowInputReference
-from vellum.workflows.types.core import JsonObject
-from vellum.workflows.types.generics import WorkflowType
 
 logger = logging.getLogger(__name__)
 
@@ -128,6 +128,8 @@ class BaseWorkflowDisplay(
                     inner_node_display = self._get_node_display(inner_node)
                     self._enrich_node_output_displays(inner_node, inner_node_display, node_output_displays)
 
+            # TODO: Make sure this output ID matches the workflow output ID of the subworkflow node's workflow
+            # https://app.shortcut.com/vellum/story/5660/fix-output-id-in-subworkflow-nodes
             node_output_displays[node_output] = node, node_display.get_node_output_display(node_output)
 
     def _enrich_node_port_displays(
@@ -179,6 +181,9 @@ class BaseWorkflowDisplay(
 
         node_displays: Dict[Type[BaseNode], NodeDisplayType] = {}
         port_displays: Dict[Port, PortDisplay] = {}
+
+        # TODO: We should still serialize nodes that are in the workflow's directory but aren't used in the graph.
+        # https://app.shortcut.com/vellum/story/5394
         for node in self._workflow.get_nodes():
             node_display = self._get_node_display(node)
             node_displays[node] = node_display

vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py

@@ -1,7 +1,7 @@
 from typing import Optional, Type
 
-from vellum_ee.workflows.display.types import WorkflowDisplayContext, WorkflowDisplayType
 from vellum.workflows.types.generics import WorkflowType
+from vellum_ee.workflows.display.types import WorkflowDisplayContext, WorkflowDisplayType
 
 
 def get_workflow_display(
@@ -24,9 +24,4 @@ def get_workflow_display(
         except IndexError:
             return base_display_class(workflow_class)
 
-    if not issubclass(workflow_display_class, base_display_class):
-        raise TypeError(
-            f"Expected to find a subclass of '{base_display_class.__name__}' for workflow class '{workflow_class.__name__}'"  # noqa: E501
-        )
-
-    return workflow_display_class(workflow_class, parent_display_context=parent_display_context)
+    return workflow_display_class(workflow_class, parent_display_context=parent_display_context) # type: ignore[return-value]

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -2,11 +2,22 @@ import logging
 from uuid import UUID
 from typing import Dict, List, Optional, Type, cast
 
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.edges import Edge
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
+from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
+from vellum.workflows.ports import 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
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.types import PortDisplay
+from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
+from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type, primitive_to_vellum_value
 from vellum_ee.workflows.display.vellum import (
     EdgeVellumDisplay,
     EdgeVellumDisplayOverrides,
@@ -21,16 +32,6 @@ from vellum_ee.workflows.display.vellum import (
     WorkflowOutputVellumDisplayOverrides,
 )
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
-from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.edges import Edge
-from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
-from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
-from vellum.workflows.ports import 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
 
 logger = logging.getLogger(__name__)
 
@@ -56,16 +57,17 @@ class VellumWorkflowDisplay(
     def serialize(self, raise_errors: bool = True) -> JsonObject:
         input_variables: JsonArray = []
         for workflow_input, workflow_input_display in self.display_context.workflow_input_displays.items():
+            default = primitive_to_vellum_value(raise_if_descriptor(workflow_input.instance)) if workflow_input.instance else None
+            required = type(None) not in workflow_input.types
+
             input_variables.append(
                 {
                     "id": str(workflow_input_display.id),
                     "key": workflow_input.name,
                     "type": infer_vellum_variable_type(workflow_input),
-                    # TODO: Add support for serializing default, required, and extensions
-                    #   https://app.shortcut.com/vellum/story/5429
-                    "default": None,
-                    "required": None,
-                    "extensions": None,
+                    "default": default.dict() if default else None,
+                    "required": required,
+                    "extensions": { "color" : workflow_input_display.color },
                 }
             )