vellum-ai 0.9.16__py3-none-any.whl → 0.9.16rc4__py3-none-any.whl

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

@@ -0,0 +1,123 @@
+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
+
+
+def test_serialize_workflow():
+    # GIVEN a Workflow that uses a Final Output Node
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=BasicFinalOutputNodeWorkflow
+    )
+    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 input_variables == [
+        {
+            "id": "e39a7b63-de15-490a-ae9b-8112c767aea0",
+            "key": "input",
+            "type": "STRING",
+            "required": None,
+            "default": None,
+            "extensions": None,
+        }
+    ]
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 1
+    assert output_variables == [
+        {
+            "id": "a34cd21e-40e5-47f4-8fdb-910593f3e9e2",
+            "key": "value",
+            "type": "STRING",
+        }
+    ]
+
+    # 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"]) == 1
+    assert len(workflow_raw_data["nodes"]) == 2
+
+    # AND each node should be serialized correctly
+    entrypoint_node = workflow_raw_data["nodes"][0]
+    assert entrypoint_node == {
+        "id": "631e2789-d60d-4088-9e3a-0ea93517075b",
+        "type": "ENTRYPOINT",
+        "inputs": [],
+        "data": {"label": "Entrypoint Node", "source_handle_id": "8b8d52a2-844f-44fe-a6c4-142fa70d391b"},
+        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "definition": {
+            "name": "BaseNode",
+            "module": [
+                "vellum",
+                "workflows",
+                "nodes",
+                "bases",
+                "base",
+            ],
+            "bases": [],
+        },
+    }
+
+    final_output_node = workflow_raw_data["nodes"][1]
+    assert final_output_node == {
+        "id": "620ec17b-e330-4212-b619-3c39dc63fb22",
+        "type": "TERMINAL",
+        "data": {
+            "label": "Basic Final Output Node",
+            "name": "basic-final-output-node",
+            "target_handle_id": "0173d3c6-11d1-44b7-b070-ca9ff5119046",
+            "output_id": "97349956-d228-4b51-a64b-1331f788373f",
+            "output_type": "STRING",
+            "node_input_id": "5322567a-f40c-400a-96b3-c3b054db543e",
+        },
+        "inputs": [
+            {
+                "id": "5322567a-f40c-400a-96b3-c3b054db543e",
+                "key": "node_input",
+                "value": {
+                    "rules": [
+                        {
+                            "type": "INPUT_VARIABLE",
+                            "data": {"input_variable_id": "e39a7b63-de15-490a-ae9b-8112c767aea0"},
+                        }
+                    ],
+                    "combinator": "OR",
+                },
+            }
+        ],
+        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "definition": {
+            "name": "BasicFinalOutputNode",
+            "module": [
+                "tests",
+                "workflows",
+                "basic_final_output_node",
+                "workflow",
+            ],
+            "bases": [
+                {
+                    "name": "FinalOutputNode",
+                    "module": [
+                        "vellum",
+                        "workflows",
+                        "nodes",
+                        "displayable",
+                        "final_output_node",
+                        "node",
+                    ],
+                }
+            ],
+        },
+    }

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

@@ -0,0 +1,84 @@
+from unittest import mock
+
+from deepdiff import DeepDiff
+
+from tests.workflows.basic_try_node.workflow import SimpleTryExample
+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
+
+
+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()
+
+    # 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) == 0
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 2
+    assert not DeepDiff(
+        [
+            {
+                "id": "d4a3117e-8ec3-4579-8eeb-5e4247bb086d",
+                "key": "error",
+                "type": "JSON",
+            },
+            {
+                "id": "a8b99024-cd32-42f6-bb2f-827189bf3a3c",
+                "key": "final_value",
+                "type": "NUMBER",
+            },
+        ],
+        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"]) == 3
+    assert len(workflow_raw_data["nodes"]) == 4
+
+    # AND each node should be serialized correctly
+    entrypoint_node = workflow_raw_data["nodes"][0]
+    assert entrypoint_node == {
+        "id": "c238508d-85ab-4644-8cbb-88eae457fe12",
+        "type": "ENTRYPOINT",
+        "inputs": [],
+        "definition": {
+            "bases": [],
+            "module": [
+                "vellum",
+                "workflows",
+                "nodes",
+                "bases",
+                "base",
+            ],
+            "name": "BaseNode",
+        },
+        "data": {
+            "label": "Entrypoint Node",
+            "source_handle_id": "04da0bb6-5b42-4dd1-a4e4-08f3ab03e1a3",
+        },
+        "display_data": {
+            "position": {"x": 0.0, "y": 0.0},
+        },
+    }
+
+    try_node = workflow_raw_data["nodes"][1]
+    assert try_node == {"type": "MOCKED"}

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

@@ -0,0 +1,233 @@
+import pytest
+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
+
+
+def test_serialize_workflow__missing_final_output_node():
+    # GIVEN a Workflow that is missing a Terminal Node
+
+    # 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()
+
+    # 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) == 2
+    assert not DeepDiff(
+        [
+            {
+                "id": "da086239-d743-4246-b666-5c91e22fb88c",
+                "key": "alpha",
+                "type": "STRING",
+                "default": None,
+                "required": None,
+                "extensions": None,
+            },
+            {
+                "id": "a8b6c5d4-a0e9-4457-834b-46b633c466a6",
+                "key": "beta",
+                "type": "STRING",
+                "default": None,
+                "required": None,
+                "extensions": 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) == 2
+    assert not DeepDiff(
+        [
+            {"id": "a360aef6-3bb4-4c56-b407-478042ef224d", "key": "alpha", "type": "STRING"},
+            {"id": "5e6d3ea6-ef91-4937-8fff-f33e07446e6a", "key": "beta", "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"]) == 3
+    assert len(workflow_raw_data["nodes"]) == 4
+
+    # AND each node should be serialized correctly
+    entrypoint_node = workflow_raw_data["nodes"][0]
+    assert entrypoint_node == {
+        "id": "b109349f-ca1b-4a5a-a66e-a1321cf297f7",
+        "type": "ENTRYPOINT",
+        "inputs": [],
+        "data": {"label": "Entrypoint Node", "source_handle_id": "943ac183-d107-4604-aed1-619bd7fef09c"},
+        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "definition": {
+            "name": "BaseNode",
+            "module": [
+                "vellum",
+                "workflows",
+                "nodes",
+                "bases",
+                "base",
+            ],
+            "bases": [],
+        },
+    }
+
+    passthrough_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "MOCKED")
+    assert passthrough_node == {
+        "type": "MOCKED",
+    }
+
+    final_output_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "TERMINAL"]
+    assert not DeepDiff(
+        [
+            {
+                "id": "acc91103-f761-47a5-bdd4-0e5e7650bb30",
+                "type": "TERMINAL",
+                "data": {
+                    "label": "First Final Output Node",
+                    "name": "first-final-output-node",
+                    "target_handle_id": "a0c2eb7a-398e-4f28-b63d-f3bae9b563ee",
+                    "output_id": "5517e50d-7f40-4f7c-acb2-e329d79a25bf",
+                    "output_type": "STRING",
+                    "node_input_id": "16363762-c14a-4162-8fab-525079d3cffe",
+                },
+                "inputs": [
+                    {
+                        "id": "16363762-c14a-4162-8fab-525079d3cffe",
+                        "key": "node_input",
+                        "value": {
+                            "rules": [
+                                {
+                                    "type": "INPUT_VARIABLE",
+                                    "data": {"input_variable_id": "da086239-d743-4246-b666-5c91e22fb88c"},
+                                }
+                            ],
+                            "combinator": "OR",
+                        },
+                    }
+                ],
+                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "definition": {
+                    "name": "FirstFinalOutputNode",
+                    "module": [
+                        "tests",
+                        "workflows",
+                        "complex_final_output_node",
+                        "missing_final_output_node",
+                    ],
+                    "bases": [
+                        {
+                            "name": "FinalOutputNode",
+                            "module": [
+                                "vellum",
+                                "workflows",
+                                "nodes",
+                                "displayable",
+                                "final_output_node",
+                                "node",
+                            ],
+                        }
+                    ],
+                },
+            },
+            {
+                "id": "bb88768d-472e-4997-b7ea-de09163d1b4c",
+                "type": "TERMINAL",
+                "data": {
+                    "label": "Final Output",
+                    "name": "beta",
+                    "target_handle_id": "5e337b19-cef6-45af-802b-46da4ad7e794",
+                    "output_id": "5e6d3ea6-ef91-4937-8fff-f33e07446e6a",
+                    "output_type": "STRING",
+                    "node_input_id": "47ba0ee9-4725-4065-a178-c929ac556be9",
+                },
+                "inputs": [
+                    {
+                        "id": "47ba0ee9-4725-4065-a178-c929ac556be9",
+                        "key": "node_input",
+                        "value": {
+                            "rules": [
+                                {
+                                    "type": "NODE_OUTPUT",
+                                    "data": {
+                                        "node_id": "32d88cab-e9fa-4a56-9bc2-fb6e1fd0897f",
+                                        "output_id": "04df0e76-690a-4ae1-ab52-fe825a334dcc",
+                                    },
+                                }
+                            ],
+                            "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_nodes,
+        ignore_order=True,
+    )
+
+
+def test_serialize_workflow__missing_workflow_output():
+    # GIVEN a Workflow that contains a terminal node that is unreferenced by the Workflow's Outputs
+
+    # 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()
+
+    assert exc_info.value.args[0] == "Unable to serialize terminal nodes that are not referenced by workflow outputs."

vellum_ee/workflows/display/types.py

@@ -0,0 +1,46 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Dict, Generic, Tuple, Type, TypeVar
+
+from vellum_ee.workflows.display.base import (
+    EdgeDisplayType,
+    EntrypointDisplayType,
+    WorkflowInputsDisplayType,
+    WorkflowMetaDisplayType,
+    WorkflowOutputDisplayType,
+)
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.nodes import BaseNode
+from vellum.workflows.ports import Port
+from vellum.workflows.references import OutputReference, WorkflowInputReference
+
+if TYPE_CHECKING:
+    from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+    from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
+
+NodeDisplayType = TypeVar("NodeDisplayType", bound="BaseNodeDisplay")
+WorkflowDisplayType = TypeVar("WorkflowDisplayType", bound="BaseWorkflowDisplay")
+
+
+@dataclass
+class WorkflowDisplayContext(
+    Generic[
+        WorkflowMetaDisplayType,
+        WorkflowInputsDisplayType,
+        NodeDisplayType,
+        EntrypointDisplayType,
+        WorkflowOutputDisplayType,
+        EdgeDisplayType,
+    ]
+):
+    workflow_display_class: Type["BaseWorkflowDisplay"]
+    workflow_display: WorkflowMetaDisplayType
+    workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = field(default_factory=dict)
+    node_displays: Dict[Type[BaseNode], "NodeDisplayType"] = field(default_factory=dict)
+    node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], "NodeOutputDisplay"]] = field(
+        default_factory=dict
+    )
+    entrypoint_displays: Dict[Type[BaseNode], EntrypointDisplayType] = field(default_factory=dict)
+    workflow_output_displays: Dict[BaseDescriptor, WorkflowOutputDisplayType] = field(default_factory=dict)
+    edge_displays: Dict[Tuple[Port, Type[BaseNode]], EdgeDisplayType] = field(default_factory=dict)
+    port_displays: Dict[Port, "PortDisplay"] = field(default_factory=dict)

vellum_ee/workflows/display/utils/tests/test_uuids.py

@@ -0,0 +1,16 @@
+import pytest
+from uuid import UUID
+
+from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
+
+
+@pytest.mark.parametrize(
+    ["input_str", "expected"],
+    [
+        ("MyExampleString", UUID("b2dadec1-bff4-4e26-8d65-e99e62628cd2")),
+        ("My Example String", UUID("a1e68bde-3263-4526-88bd-70f4bf800224")),
+    ],
+)
+def test_uuid4_from_hash(input_str, expected):
+    actual = uuid4_from_hash(input_str)
+    assert actual == expected

vellum_ee/workflows/display/utils/uuids.py

@@ -0,0 +1,24 @@
+import hashlib
+from uuid import UUID
+
+
+def uuid4_from_hash(input_str: str) -> UUID:
+    # Create a SHA-256 hash of the input string
+    hash_bytes = hashlib.sha256(input_str.encode()).digest()
+
+    # Modify the hash to follow UUID4 structure
+    # UUID4 has 8-4-4-4-12 hexadecimal characters
+    # Version bits (4 bits) should be set to 4 (for UUID4),
+    # and variant bits (2 bits) should start with 0b10
+
+    # Convert first 16 bytes into UUID
+    hash_list = list(hash_bytes[:16])
+
+    # Set the version to 4 (UUID4)
+    hash_list[6] = (hash_list[6] & 0x0F) | 0x40
+
+    # Set the variant to 0b10xxxxxx
+    hash_list[8] = (hash_list[8] & 0x3F) | 0x80
+
+    # Create a UUID from the modified bytes
+    return UUID(bytes=bytes(hash_list))

vellum_ee/workflows/display/utils/vellum.py

@@ -0,0 +1,121 @@
+import enum
+import json
+import typing
+from typing import Any, List, Union, cast
+
+from vellum import ChatMessage, SearchResult, SearchResultRequest, VellumVariableType
+
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.vellum import (
+    ChatHistoryVellumValue,
+    ConstantValuePointer,
+    InputVariableData,
+    InputVariablePointer,
+    JsonVellumValue,
+    NodeInputValuePointerRule,
+    NodeOutputData,
+    NodeOutputPointer,
+    NumberVellumValue,
+    SearchResultsVellumValue,
+    StringVellumValue,
+    VellumValue,
+    WorkspaceSecretData,
+    WorkspaceSecretPointer,
+)
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.references import OutputReference, WorkflowInputReference
+from vellum.workflows.references.node import NodeReference
+from vellum.workflows.references.vellum_secret import VellumSecretReference
+from vellum.workflows.types.core import VellumValuePrimitive
+from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
+from vellum.workflows.vellum_client import create_vellum_client
+
+_T = typing.TypeVar("_T")
+
+
+def infer_vellum_variable_type(value: Any) -> VellumVariableType:
+    inferred_type: VellumVariableType
+
+    if isinstance(value, BaseDescriptor):
+        descriptor: BaseDescriptor = value
+        if isinstance(descriptor, NodeReference):
+            if not isinstance(descriptor.instance, BaseDescriptor):
+                raise ValueError(
+                    f"Expected NodeReference {descriptor.name} to have an instance pointing to a descriptor"
+                )
+
+            descriptor = descriptor.instance
+
+        inferred_type = primitive_type_to_vellum_variable_type(descriptor)
+    else:
+        vellum_variable_value = primitive_to_vellum_value(value)
+        inferred_type = vellum_variable_value.type
+
+    return inferred_type
+
+
+def create_node_input_value_pointer_rule(
+    value: Any, display_context: WorkflowDisplayContext
+) -> NodeInputValuePointerRule:
+    if isinstance(value, OutputReference):
+        upstream_node, output_display = display_context.node_output_displays[value]
+        upstream_node_display = display_context.node_displays[upstream_node]
+        return NodeOutputPointer(
+            type="NODE_OUTPUT",
+            data=NodeOutputData(node_id=str(upstream_node_display.node_id), output_id=str(output_display.id)),
+        )
+    if isinstance(value, WorkflowInputReference):
+        workflow_input_display = display_context.workflow_input_displays[value]
+        return InputVariablePointer(
+            type="INPUT_VARIABLE", data=InputVariableData(input_variable_id=str(workflow_input_display.id))
+        )
+    if isinstance(value, VellumSecretReference):
+        # TODO: Pass through the name instead of retrieving the ID
+        # https://app.shortcut.com/vellum/story/5072
+        vellum_client = create_vellum_client()
+        workspace_secret = vellum_client.workspace_secrets.retrieve(
+            id=value.name,
+        )
+        return WorkspaceSecretPointer(
+            type="WORKSPACE_SECRET",
+            data=WorkspaceSecretData(
+                type="STRING",
+                workspace_secret_id=str(workspace_secret.id),
+            ),
+        )
+
+    if not isinstance(value, BaseDescriptor):
+        vellum_value = primitive_to_vellum_value(value)
+        return ConstantValuePointer(type="CONSTANT_VALUE", data=vellum_value)
+
+    raise ValueError(f"Unsupported descriptor type: {value.__class__.__name__}")
+
+
+def primitive_to_vellum_value(value: VellumValuePrimitive) -> VellumValue:
+    """Converts a python primitive to a VellumVariableValue"""
+
+    if isinstance(value, str):
+        return StringVellumValue(value=value)
+    elif isinstance(value, enum.Enum):
+        return StringVellumValue(value=value.value)
+    elif isinstance(value, (int, float)):
+        return NumberVellumValue(value=value)
+    elif isinstance(value, list) and (
+        all(isinstance(message, ChatMessage) for message in value)
+        or all(isinstance(message, ChatMessage) for message in value)
+    ):
+        chat_messages = cast(Union[List[ChatMessage], List[ChatMessage]], value)
+        return ChatHistoryVellumValue(value=chat_messages)
+    elif isinstance(value, list) and (
+        all(isinstance(search_result, SearchResultRequest) for search_result in value)
+        or all(isinstance(search_result, SearchResult) for search_result in value)
+    ):
+        search_results = cast(Union[List[SearchResultRequest], List[SearchResult]], value)
+        return SearchResultsVellumValue(value=search_results)
+
+    try:
+        json_value = json.dumps(value)
+    except json.JSONDecodeError:
+        raise ValueError(f"Unsupported variable type: {value.__class__.__name__}")
+
+    return JsonVellumValue(value=json.loads(json_value))