vellum-ai 1.11.2__py3-none-any.whl → 1.13.5__py3-none-any.whl

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

@@ -1,7 +1,10 @@
+from typing import Any
+
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
 from vellum.workflows.state.base import BaseState
+from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
@@ -78,3 +81,136 @@ def test_final_output_node_display__serialize_with_invalid_types_should_raise_er
 
     # AND the serialized workflow should still be created
     assert "workflow_raw_data" in serialized_workflow
+
+
+def test_final_output_node_display__serialize_with_list_str_type():
+    """
+    Tests that FinalOutputNode with list[str] as second type parameter serializes correctly.
+    """
+
+    # GIVEN a node that outputs a list of strings
+    class ListNode(BaseNode):
+
+        class Outputs:
+            result: list[str]
+
+    # AND a FinalOutputNode with list[str] as the second type parameter
+    class ListOutput(FinalOutputNode[BaseState, list[str]]):
+        class Outputs(FinalOutputNode.Outputs):
+            value = ListNode.Outputs.result
+
+    # AND a workflow referencing the node
+    class MyWorkflow(BaseWorkflow):
+        graph = ListNode >> ListOutput
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_result = ListOutput.Outputs.value
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN serialization should succeed without raising validation errors
+    assert "workflow_raw_data" in serialized_workflow
+    assert "nodes" in serialized_workflow["workflow_raw_data"]
+
+    # AND the terminal node should be properly serialized
+    terminal_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "TERMINAL"
+    )
+    assert terminal_node is not None
+    assert terminal_node["id"] == str(ListOutput.__id__)
+
+    # AND the output type should be correctly serialized as JSON
+    assert terminal_node["data"]["output_type"] == "JSON"
+
+    # AND the outputs should contain the correct type information
+    assert len(terminal_node["outputs"]) == 1
+    assert terminal_node["outputs"][0]["type"] == "JSON"
+
+    # AND the output should have an inline type reference with JSON Schema
+    schema = terminal_node["outputs"][0]["schema"]
+    assert schema["type"] == "array"
+    assert schema["items"]["type"] == "string"
+
+
+def test_final_output_node_display__serialize_with_nested_node_output_reference():
+    """
+    Tests that FinalOutputNode with a dict containing nested node output references serializes correctly.
+    """
+
+    # GIVEN a node that outputs multiple values
+    class DataNode(BaseNode):
+        class Outputs:
+            url: str
+            count: int
+
+    # AND a FinalOutputNode with a dict containing nested references to node outputs
+    class NestedOutput(FinalOutputNode[BaseState, dict[str, Any]]):
+        class Outputs(FinalOutputNode.Outputs):
+            value = {
+                "download_url": DataNode.Outputs.url,
+                "row_count": DataNode.Outputs.count,
+                "message": "Data processed successfully",
+            }
+
+    # AND a workflow referencing the node
+    class MyWorkflow(BaseWorkflow):
+        graph = DataNode >> NestedOutput
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_result = NestedOutput.Outputs.value
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN serialization should succeed
+    assert "workflow_raw_data" in serialized_workflow
+    assert "nodes" in serialized_workflow["workflow_raw_data"]
+
+    # AND the terminal node should be properly serialized
+    terminal_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "TERMINAL"
+    )
+    assert terminal_node is not None
+    assert terminal_node["id"] == str(NestedOutput.__id__)
+
+    # AND the node should have inputs
+    assert len(terminal_node["inputs"]) == 1
+    node_input = terminal_node["inputs"][0]
+    assert node_input["key"] == "node_input"
+
+    # AND the input value should have rules with a CONSTANT_VALUE containing the serialized dict
+    assert node_input["value"]["combinator"] == "OR"
+    assert len(node_input["value"]["rules"]) == 1
+    rule = node_input["value"]["rules"][0]
+    assert rule["type"] == "CONSTANT_VALUE"
+    assert rule["data"]["type"] == "JSON"
+
+    # AND the serialized value should be a DICTIONARY_REFERENCE with entries
+    serialized_value = rule["data"]["value"]
+    assert serialized_value["type"] == "DICTIONARY_REFERENCE"
+    assert "entries" in serialized_value
+    assert len(serialized_value["entries"]) == 3
+
+    # AND the entries should contain the nested references and constant value
+    entry_keys = {entry["key"] for entry in serialized_value["entries"]}
+    assert entry_keys == {"download_url", "row_count", "message"}
+
+    # AND the outputs should contain the DICTIONARY_REFERENCE value
+    assert len(terminal_node["outputs"]) == 1
+    output = terminal_node["outputs"][0]
+    assert output["name"] == "value"
+    assert output["type"] == "JSON"
+    assert output["value"]["type"] == "DICTIONARY_REFERENCE"
+    assert len(output["value"]["entries"]) == 3
+
+    # AND the attributes should be empty for this node
+    assert terminal_node.get("attributes", []) == []
+
+    # AND a validation error should be recorded for the nested references
+    errors = list(workflow_display.display_context.errors)
+    assert len(errors) == 1
+    assert isinstance(errors[0], UnsupportedSerializationException)
+    assert "nested references" in str(errors[0]).lower()

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

@@ -1,22 +1,21 @@
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes.displayable.note_node.node import NoteNode
-from vellum_ee.workflows.display.nodes.vellum.note_node import BaseNoteNodeDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
 def test_serialize_node__note_node():
-    # GIVEN a note node
-    class MyNoteNode(NoteNode):
-        pass
+    """
+    Tests that a single note node is properly serialized in a workflow.
+    """
 
-    # AND a display class for the note node
-    class MyNoteNodeDisplay(BaseNoteNodeDisplay[MyNoteNode]):
+    # GIVEN a note node with text and style defined on the node class
+    class MyNoteNode(NoteNode):
         text = "This makes sense"
         style = {
             "fontSize": 24,
         }
 
-    # AND a workflow with the code node
+    # AND a workflow with the note node
     class Workflow(BaseWorkflow):
         graph = MyNoteNode
 
@@ -30,3 +29,41 @@ def test_serialize_node__note_node():
     assert my_note_node["inputs"] == []
     assert my_note_node["data"]["text"] == "This makes sense"
     assert my_note_node["data"]["style"] == {"fontSize": 24}
+
+
+def test_serialize_workflow__graph_with_only_note_nodes():
+    """
+    Tests that a workflow with only note nodes properly serializes all note nodes in the final payload.
+    """
+
+    # GIVEN multiple note nodes with text and style defined on the node classes
+    class FirstNoteNode(NoteNode):
+        text = "First note"
+        style = {"fontSize": 16}
+
+    class SecondNoteNode(NoteNode):
+        text = "Second note"
+        style = {"fontSize": 20}
+
+    # AND a workflow with only note nodes in unused_graphs
+    class Workflow(BaseWorkflow):
+        unused_graphs = {FirstNoteNode, SecondNoteNode}  # type: ignore[assignment]
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the serialized workflow should contain both note nodes
+    note_nodes = [node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "NOTE"]
+    assert len(note_nodes) == 2
+
+    # AND the note nodes should have the correct data
+    note_texts = {node["data"]["text"] for node in note_nodes}
+    assert note_texts == {"First note", "Second note"}
+
+    # AND each note node should have the correct style
+    first_note = next(node for node in note_nodes if node["data"]["text"] == "First note")
+    second_note = next(node for node in note_nodes if node["data"]["text"] == "Second note")
+
+    assert first_note["data"]["style"] == {"fontSize": 16}
+    assert second_note["data"]["style"] == {"fontSize": 20}

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

@@ -154,17 +154,17 @@ def test_serialize_node__prompt_inputs__state_reference():
     workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
-    # THEN the node should skip the state reference input rule
     my_prompt_node = next(
         node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MyPromptNode.__id__)
     )
-
     assert my_prompt_node["inputs"] == [
         {
             "id": "7c5d23b3-c5ed-4ed6-a685-43fbe9a9baf8",
             "key": "foo",
             "value": {
-                "rules": [],
+                "rules": [
+                    {"type": "WORKFLOW_STATE", "data": {"state_variable_id": "dd3391bf-c818-4eba-aac5-912618ba412f"}}
+                ],
                 "combinator": "OR",
             },
         },
@@ -172,15 +172,7 @@ def test_serialize_node__prompt_inputs__state_reference():
             "id": "e138f06e-d705-46bc-8ac4-c844b0e9131a",
             "key": "bar",
             "value": {
-                "rules": [
-                    {
-                        "type": "CONSTANT_VALUE",
-                        "data": {
-                            "type": "STRING",
-                            "value": "baz",
-                        },
-                    }
-                ],
+                "rules": [{"type": "CONSTANT_VALUE", "data": {"type": "STRING", "value": "baz"}}],
                 "combinator": "OR",
             },
         },
@@ -201,7 +193,7 @@ def test_serialize_node__prompt_inputs__state_reference():
                     "key": "foo",
                     "value": {
                         "type": "WORKFLOW_STATE",
-                        "state_variable_id": "45649791-c642-4405-aff9-a1fafd780ea1",
+                        "state_variable_id": "dd3391bf-c818-4eba-aac5-912618ba412f",
                     },
                 },
                 {

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

@@ -3,6 +3,12 @@ from datetime import datetime
 from uuid import UUID, uuid4
 from typing import Type
 
+from vellum import (
+    ReleaseEnvironment,
+    WorkflowDeploymentRelease,
+    WorkflowDeploymentReleaseWorkflowDeployment,
+    WorkflowDeploymentReleaseWorkflowVersion,
+)
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes.displayable.subworkflow_deployment_node.node import SubworkflowDeploymentNode
 from vellum_ee.workflows.display.nodes.vellum.subworkflow_deployment_node import BaseSubworkflowDeploymentNodeDisplay
@@ -29,28 +35,32 @@ def _display_class_with_node_input_ids_by_name_with_inputs_prefix(Node: Type[Sub
 
 @pytest.fixture
 def mock_fetch_deployment(mocker):
-    # Create a mock deployment response
-    mock_deployment = mocker.Mock(
-        id="test-id",
-        name="test-deployment",
-        label="Test Deployment",
-        status="ACTIVE",
-        environment="DEVELOPMENT",
+    # Create a mock deployment release response
+    deployment_release = WorkflowDeploymentRelease(
+        id=str(uuid4()),
         created=datetime.now(),
-        last_deployed_on=datetime.now(),
-        last_deployed_history_item_id=str(uuid4()),
-        input_variables=[],
-        output_variables=[],
-        description="Test deployment description",
+        environment=ReleaseEnvironment(id=str(uuid4()), name="DEVELOPMENT", label="Development"),
+        workflow_version=WorkflowDeploymentReleaseWorkflowVersion(
+            id=str(uuid4()),
+            input_variables=[],
+            output_variables=[],
+        ),
+        deployment=WorkflowDeploymentReleaseWorkflowDeployment(
+            id="test-id",
+            name="test-deployment",
+        ),
+        release_tags=[],
+        reviews=[],
     )
 
-    # Patch the create_vellum_client function to return our mock client
-    mocker.patch(
-        "vellum.client.resources.workflow_deployments.client.WorkflowDeploymentsClient.retrieve",
-        return_value=mock_deployment,
+    # Patch the retrieve_workflow_deployment_release method
+    patch_path = (
+        "vellum.client.resources.workflow_deployments.client."
+        "WorkflowDeploymentsClient.retrieve_workflow_deployment_release"
     )
+    mocker.patch(patch_path, return_value=deployment_release)
 
-    return mock_deployment
+    return deployment_release
 
 
 @pytest.mark.parametrize(

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

@@ -153,7 +153,9 @@ def test_serialize_node__prompt_inputs__mixed_values():
 
 
 def test_serialize_node__tool_calling_node__mcp_server_api_key():
-    # GIVEN a tool calling node with an mcp server
+    """Tests that MCPServer with EnvironmentVariableReference serializes as ARRAY_REFERENCE."""
+
+    # GIVEN a tool calling node with an mcp server using an environment variable for the API key
     class MyToolCallingNode(ToolCallingNode):
         functions = [
             MCPServer(
@@ -184,26 +186,70 @@ def test_serialize_node__tool_calling_node__mcp_server_api_key():
         attribute for attribute in my_tool_calling_node["attributes"] if attribute["name"] == "functions"
     )
 
+    # AND the functions attribute should be an ARRAY_REFERENCE with a DICTIONARY_REFERENCE
+    # containing the MCP server fields, with api_key_header_value as ENVIRONMENT_VARIABLE
     assert functions_attribute == {
         "id": "ff00c2d6-f99c-458b-9bcd-181f8e43b2d1",
         "name": "functions",
         "value": {
-            "type": "CONSTANT_VALUE",
-            "value": {
-                "type": "JSON",
-                "value": [
-                    {
-                        "type": "MCP_SERVER",
-                        "name": "my-mcp-server",
-                        "description": "",
-                        "url": "https://my-mcp-server.com",
-                        "authorization_type": "API_KEY",
-                        "bearer_token_value": None,
-                        "api_key_header_key": "my-api-key-header-key",
-                        "api_key_header_value": "my-api-key-header-value",
-                    }
-                ],
-            },
+            "type": "ARRAY_REFERENCE",
+            "items": [
+                {
+                    "type": "DICTIONARY_REFERENCE",
+                    "entries": [
+                        {
+                            "id": "8cff4f0a-86d9-43fd-8d0b-542c845db53e",
+                            "key": "type",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "MCP_SERVER"}},
+                        },
+                        {
+                            "id": "29203be4-407c-4056-aaa3-6e6d1249113e",
+                            "key": "name",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "my-mcp-server"}},
+                        },
+                        {
+                            "id": "69813bb2-12c1-432b-9ef3-a0071bd29149",
+                            "key": "description",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": ""}},
+                        },
+                        {
+                            "id": "83cd2f84-bf35-4996-9840-3754a79e1091",
+                            "key": "url",
+                            "value": {
+                                "type": "CONSTANT_VALUE",
+                                "value": {"type": "STRING", "value": "https://my-mcp-server.com"},
+                            },
+                        },
+                        {
+                            "id": "5ad23cee-497e-4ca3-ba51-d13f56985c75",
+                            "key": "authorization_type",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "API_KEY"}},
+                        },
+                        {
+                            "id": "45964d28-1a9b-40a1-b544-62da92d2f62a",
+                            "key": "bearer_token_value",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
+                        {
+                            "id": "c4952fa6-1ffb-4e88-863f-3d0867b19b7a",
+                            "key": "api_key_header_key",
+                            "value": {
+                                "type": "CONSTANT_VALUE",
+                                "value": {"type": "STRING", "value": "my-api-key-header-key"},
+                            },
+                        },
+                        {
+                            "id": "7dc9cdb0-d139-4a1d-8c6f-9941c6b82b62",
+                            "key": "api_key_header_value",
+                            "value": {
+                                "type": "ENVIRONMENT_VARIABLE",
+                                "environment_variable": "my-api-key-header-value",
+                            },
+                        },
+                    ],
+                    "definition": {"name": "MCPServer", "module": ["vellum", "workflows", "types", "definition"]},
+                }
+            ],
         },
     }
 
@@ -328,7 +374,7 @@ def test_serialize_tool_router_node():
                     "lhs": {
                         "lhs": {
                             "lhs": {
-                                "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
+                                "state_variable_id": "30a4672a-96dc-471a-878c-90e0f1f7f043",
                                 "type": "WORKFLOW_STATE",
                             },
                             "operator": "<",
@@ -354,7 +400,7 @@ def test_serialize_tool_router_node():
                                     },
                                     "operator": "accessField",
                                     "rhs": {
-                                        "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
+                                        "state_variable_id": "30a4672a-96dc-471a-878c-90e0f1f7f043",
                                         "type": "WORKFLOW_STATE",
                                     },
                                     "type": "BINARY_EXPRESSION",
@@ -381,7 +427,7 @@ def test_serialize_tool_router_node():
                                     },
                                     "operator": "accessField",
                                     "rhs": {
-                                        "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
+                                        "state_variable_id": "30a4672a-96dc-471a-878c-90e0f1f7f043",
                                         "type": "WORKFLOW_STATE",
                                     },
                                     "type": "BINARY_EXPRESSION",
@@ -453,6 +499,83 @@ def test_serialize_function_node():
     assert display_data["color"] == "purple"
 
 
+def test_serialize_inline_prompt_node__mcp_server_not_serialized():
+    """
+    Test that MCP servers in inline prompt node functions are not serialized as function blocks.
+    MCP servers should be skipped during serialization (they return None from _generate_function_tools).
+    """
+
+    # GIVEN an MCP server
+    mcp_server = MCPServer(
+        name="my-mcp-server",
+        url="https://my-mcp-server.com",
+        authorization_type=AuthorizationType.API_KEY,
+        api_key_header_key="my-api-key-header-key",
+        api_key_header_value=EnvironmentVariableReference(name="my-api-key-header-value"),
+    )
+
+    # AND a regular function
+    def my_function(arg1: str) -> str:
+        return f"Result: {arg1}"
+
+    # AND an inline prompt node with both MCP server and regular function
+    class MyInlinePromptNode(InlinePromptNode):
+        ml_model = "gpt-4o-mini"
+        blocks = []
+        functions = [mcp_server, my_function]
+        prompt_inputs = None
+        parameters = PromptParameters()
+
+    # WHEN we serialize the inline prompt node
+    inline_prompt_node_display_class = get_node_display_class(MyInlinePromptNode)
+    inline_prompt_node_display = inline_prompt_node_display_class()
+
+    class Workflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = MyInlinePromptNode
+
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    display_context = workflow_display.display_context
+
+    serialized_node = inline_prompt_node_display.serialize(display_context)
+
+    # THEN the serialized node should have exec_config with prompt_template_block_data
+    assert isinstance(serialized_node, dict)
+    assert serialized_node["type"] == "PROMPT"
+    assert isinstance(serialized_node["data"], dict)
+    assert "exec_config" in serialized_node["data"]
+    exec_config = serialized_node["data"]["exec_config"]
+    assert isinstance(exec_config, dict)
+    assert "prompt_template_block_data" in exec_config
+    prompt_template_block_data = exec_config["prompt_template_block_data"]
+    assert isinstance(prompt_template_block_data, dict)
+    blocks = prompt_template_block_data["blocks"]
+    assert isinstance(blocks, list)
+
+    # AND there should be only one FUNCTION_DEFINITION block (from the regular function)
+    # MCP server should not be serialized
+    function_blocks = [
+        block for block in blocks if isinstance(block, dict) and block.get("block_type") == "FUNCTION_DEFINITION"
+    ]
+    assert len(function_blocks) == 1
+    assert function_blocks == [
+        {
+            "id": "b6b7b0c2-78b6-498f-8f97-fdaa81e082ec",
+            "block_type": "FUNCTION_DEFINITION",
+            "properties": {
+                "function_name": "my_function",
+                "function_description": None,
+                "function_parameters": {
+                    "type": "object",
+                    "properties": {"arg1": {"type": "string"}},
+                    "required": ["arg1"],
+                },
+                "function_forced": None,
+                "function_strict": None,
+            },
+        }
+    ]
+
+
 def test_serialize_node__tool_calling_node__subworkflow_with_parent_input_reference():
     """
     Test that a tool calling node with a subworkflow that references parent inputs serializes correctly
@@ -594,7 +717,7 @@ def test_serialize_tool_prompt_node_with_inline_workflow():
                         "operator": "concat",
                         "rhs": {
                             "type": "WORKFLOW_STATE",
-                            "state_variable_id": "7a1caaf5-99df-487a-8b2d-6512df2d871a",
+                            "state_variable_id": "c1d692a7-3d87-4283-8d9c-daee82c61854",
                         },
                     },
                 }
@@ -621,7 +744,7 @@ def test_serialize_tool_prompt_node_with_workflow_deployment(vellum_client):
             input_variables=[],
             output_variables=[],
         ),
-        deployment=WorkflowDeploymentReleaseWorkflowDeployment(name="test-name"),
+        deployment=WorkflowDeploymentReleaseWorkflowDeployment(id="test-deployment-id", name="test-name"),
         description="test-description",
         release_tags=[
             ReleaseReleaseTag(