vellum-ai 0.12.13__py3-none-any.whl → 0.12.15__py3-none-any.whl

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

@@ -0,0 +1,174 @@
+from deepdiff import DeepDiff
+
+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_generic_node.workflow import BasicGenericNodeWorkflow
+
+
+def test_serialize_workflow(vellum_client):
+    # GIVEN a Workflow that uses a generic node
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=BasicGenericNodeWorkflow
+    )
+
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a serialized representation of the Workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "output_variables",
+    }
+
+    # AND its input variables should be what we expect
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 1
+    assert not DeepDiff(
+        [
+            {
+                "id": "a07c2273-34a7-42b5-bcad-143b6127cc8a",
+                "key": "input",
+                "type": "STRING",
+                "default": None,
+                "required": True,
+                "extensions": {"color": None},
+            },
+        ],
+        input_variables,
+        ignore_order=True,
+    )
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 1
+    assert not DeepDiff(
+        [
+            {"id": "2b6389d0-266a-4be4-843e-4e543dd3d727", "key": "output", "type": "STRING"},
+        ],
+        output_variables,
+        ignore_order=True,
+    )
+
+    # AND its raw data should be what we expect
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    assert workflow_raw_data.keys() == {"edges", "nodes", "display_data", "definition"}
+    assert len(workflow_raw_data["edges"]) == 2
+    assert len(workflow_raw_data["nodes"]) == 3
+
+    # AND each node should be serialized correctly
+    entrypoint_node = workflow_raw_data["nodes"][0]
+    assert entrypoint_node == {
+        "id": "f1e4678f-c470-400b-a40e-c8922cc99a86",
+        "type": "ENTRYPOINT",
+        "inputs": [],
+        "data": {"label": "Entrypoint Node", "source_handle_id": "40201804-8beb-43ad-8873-a027759512f1"},
+        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "definition": {
+            "name": "BaseNode",
+            "module": [
+                "vellum",
+                "workflows",
+                "nodes",
+                "bases",
+                "base",
+            ],
+            "bases": [],
+        },
+    }
+
+    api_node = workflow_raw_data["nodes"][1]
+    assert not DeepDiff(
+        {
+            "id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
+            "type": "GENERIC",
+        },
+        api_node,
+    )
+
+    final_output_node = workflow_raw_data["nodes"][2]
+    assert not DeepDiff(
+        {
+            "id": "50e3b446-afcd-4a5d-8c6f-5f05eaf2200e",
+            "type": "TERMINAL",
+            "data": {
+                "label": "Final Output",
+                "name": "output",
+                "target_handle_id": "8bd9f4f3-9f66-4d95-8e84-529b0002c531",
+                "output_id": "2b6389d0-266a-4be4-843e-4e543dd3d727",
+                "output_type": "STRING",
+                "node_input_id": "7a9f2d3a-0b23-4bd4-b567-e9493135b727",
+            },
+            "inputs": [
+                {
+                    "id": "7a9f2d3a-0b23-4bd4-b567-e9493135b727",
+                    "key": "node_input",
+                    "value": {
+                        "rules": [
+                            {
+                                "type": "NODE_OUTPUT",
+                                "data": {
+                                    "node_id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
+                                    "output_id": "0a9c7a80-fc89-4a71-aac0-66489e4ddb85",
+                                },
+                            }
+                        ],
+                        "combinator": "OR",
+                    },
+                }
+            ],
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "definition": {
+                "name": "FinalOutputNode",
+                "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
+                "bases": [
+                    {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"], "bases": []}
+                ],
+            },
+        },
+        final_output_node,
+        ignore_order=True,
+    )
+
+    # AND each edge should be serialized correctly
+    serialized_edges = workflow_raw_data["edges"]
+    assert not DeepDiff(
+        [
+            {
+                "id": "445dd2de-82b2-482b-89f6-5f49d8eb21a9",
+                "source_node_id": "f1e4678f-c470-400b-a40e-c8922cc99a86",
+                "source_handle_id": "40201804-8beb-43ad-8873-a027759512f1",
+                "target_node_id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
+                "target_handle_id": "b7bfb298-959a-4d2b-8b85-bbd0d2522703",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "b741c861-cf67-4649-b9ef-b43a4add72b1",
+                "source_node_id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
+                "source_handle_id": "89dccfa5-cc1a-4612-bd87-86cb444f6dd4",
+                "target_node_id": "50e3b446-afcd-4a5d-8c6f-5f05eaf2200e",
+                "target_handle_id": "8bd9f4f3-9f66-4d95-8e84-529b0002c531",
+                "type": "DEFAULT",
+            },
+        ],
+        serialized_edges,
+        ignore_order=True,
+    )
+
+    # AND the display data should be what we expect
+    display_data = workflow_raw_data["display_data"]
+    assert display_data == {
+        "viewport": {
+            "x": 0.0,
+            "y": 0.0,
+            "zoom": 1.0,
+        }
+    }
+
+    # AND the definition should be what we expect
+    definition = workflow_raw_data["definition"]
+    assert definition == {
+        "name": "BasicGenericNodeWorkflow",
+        "module": ["tests", "workflows", "basic_generic_node", "workflow"],
+    }

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

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-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
 
@@ -15,12 +12,7 @@ def test_serialize_workflow():
     workflow_display = get_workflow_display(
         base_display_class=VellumWorkflowDisplay, workflow_class=BasicInlineSubworkflowWorkflow
     )
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -156,7 +148,7 @@ def test_serialize_workflow():
                                 "bases": [],
                             },
                         },
-                        {"type": "MOCKED"},
+                        {"id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", "type": "GENERIC"},
                         {
                             "id": "a773c3a5-78cb-4250-8d29-7282e8a579d3",
                             "type": "TERMINAL",

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

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-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
 
@@ -13,12 +10,7 @@ def test_serialize_workflow():
     # GIVEN a Workflow that uses a MapNode
     # WHEN we serialize it
     workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=SimpleMapExample)
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -141,7 +133,7 @@ def test_serialize_workflow():
                                 "bases": [],
                             },
                         },
-                        {"type": "MOCKED"},
+                        {"id": "baf6d316-dc75-41e8-96c0-015aede96309", "type": "GENERIC"},
                         {
                             "id": "6f4883b2-70b1-4e1c-ae15-7d0f5aec810b",
                             "type": "TERMINAL",

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

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-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
 
@@ -15,12 +12,7 @@ def test_serialize_workflow__await_all():
     workflow_display = get_workflow_display(
         base_display_class=VellumWorkflowDisplay, workflow_class=AwaitAllPassingWorkflow
     )
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -76,18 +68,12 @@ def test_serialize_workflow__await_all():
         },
     }
 
-    passthrough_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "MOCKED"]
+    passthrough_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "GENERIC"]
     assert not DeepDiff(
         [
-            {
-                "type": "MOCKED",
-            },
-            {
-                "type": "MOCKED",
-            },
-            {
-                "type": "MOCKED",
-            },
+            {"id": "127ef456-91bc-43c6-bd8b-1772db5e3cb5", "type": "GENERIC"},
+            {"id": "59243c65-053f-4ea6-9157-3f3edb1477bf", "type": "GENERIC"},
+            {"id": "634f0202-9ea9-4c62-b152-1a58c595cffb", "type": "GENERIC"},
         ],
         passthrough_nodes,
         ignore_order=True,

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

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-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
 
@@ -13,10 +10,7 @@ def test_serialize_workflow():
     # GIVEN a Workflow with a TryNode
     # WHEN we serialize it
     workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=SimpleTryExample)
-
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -82,4 +76,4 @@ def test_serialize_workflow():
     }
 
     try_node = workflow_raw_data["nodes"][1]
-    assert try_node == {"type": "MOCKED"}
+    assert try_node == {"id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", "type": "GENERIC"}

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

@@ -1,9 +1,7 @@
 import pytest
-from unittest import mock
 
 from deepdiff import DeepDiff
 
-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
 
@@ -13,17 +11,12 @@ from tests.workflows.complex_final_output_node.missing_workflow_output import Mi
 
 def test_serialize_workflow__missing_final_output_node():
     # GIVEN a Workflow that is missing a Terminal Node
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=MissingFinalOutputNodeWorkflow
+    )
 
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        workflow_display = get_workflow_display(
-            base_display_class=VellumWorkflowDisplay, workflow_class=MissingFinalOutputNodeWorkflow
-        )
-
-        # WHEN we serialize it
-        serialized_workflow: dict = workflow_display.serialize()
+    # WHEN we serialize it
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -97,9 +90,10 @@ def test_serialize_workflow__missing_final_output_node():
         },
     }
 
-    passthrough_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "MOCKED")
+    passthrough_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "GENERIC")
     assert passthrough_node == {
-        "type": "MOCKED",
+        "id": "32d88cab-e9fa-4a56-9bc2-fb6e1fd0897f",
+        "type": "GENERIC",
     }
 
     final_output_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "TERMINAL"]
@@ -218,17 +212,12 @@ def test_serialize_workflow__missing_final_output_node():
 
 def test_serialize_workflow__missing_workflow_output():
     # GIVEN a Workflow that contains a terminal node that is unreferenced by the Workflow's Outputs
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=MissingWorkflowOutputWorkflow
+    )
 
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        workflow_display = get_workflow_display(
-            base_display_class=VellumWorkflowDisplay, workflow_class=MissingWorkflowOutputWorkflow
-        )
-
-        # WHEN we serialize it, it should throw an error
-        with pytest.raises(ValueError) as exc_info:
-            workflow_display.serialize()
+    # WHEN we serialize it, it should throw an error
+    with pytest.raises(ValueError) as exc_info:
+        workflow_display.serialize()
 
     assert exc_info.value.args[0] == "Unable to serialize terminal nodes that are not referenced by workflow outputs."

vellum_ee/workflows/server/virtual_file_loader.py

@@ -0,0 +1,42 @@
+import importlib
+
+
+class VirtualFileLoader(importlib.abc.Loader):
+    def __init__(self, code: str, is_package: bool):
+        self.code = code
+        self.is_package = is_package
+
+    def create_module(self, spec):
+        return None  # use default module creation
+
+    def exec_module(self, module):
+        if not self.is_package:
+            exec(self.code, module.__dict__)
+
+
+class VirtualFileFinder(importlib.abc.MetaPathFinder, importlib.abc.Loader):
+    def __init__(self, files: dict[str, str], namespace: str):
+        self.files = files
+        self.namespace = namespace
+
+    def find_spec(self, fullname, path, target=None):
+        # Do the namespacing on the fly to avoid having to copy the file dict
+        prefixed_name = fullname if fullname.startswith(self.namespace) else f"{self.namespace}.{fullname}"
+
+        key_name = "__init__" if fullname == self.namespace else fullname.replace(f"{self.namespace}.", "")
+
+        files_key = f"{key_name.replace('.', '/')}.py"
+        if not self.files.get(files_key):
+            files_key = f"{key_name.replace('.', '/')}/__init__.py"
+
+        file = self.files.get(files_key)
+        is_package = "__init__" in files_key
+
+        if file:
+            return importlib.machinery.ModuleSpec(
+                prefixed_name,
+                VirtualFileLoader(file, is_package),
+                origin=prefixed_name,
+                is_package=is_package,
+            )
+        return None