vellum-ai 1.1.2__py3-none-any.whl → 1.1.4__py3-none-any.whl

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

@@ -42,7 +42,7 @@ class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Gene
                 script_filepath=filepath,
             )
             if not file_code:
-                raise NodeValidationError(f"Filepath '{filepath}' does not exist", node.__name__)
+                display_context.add_error(NodeValidationError(f"Filepath '{filepath}' does not exist", node.__name__))
             code_value = file_code
         else:
             code_value = ""

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

@@ -197,7 +197,61 @@ def test_serialize_node__with_non_exist_code_input_path_with_dry_run():
                     "display_data": {"position": {"x": 0.0, "y": -50.0}},
                     "base": None,
                     "definition": None,
-                }
+                },
+                {
+                    "id": "ac90c0ce-f393-438c-a24f-e5e9a9286182",
+                    "type": "CODE_EXECUTION",
+                    "inputs": [
+                        {
+                            "id": "a0b9d6f6-ce59-4075-8db6-866781bc73ef",
+                            "key": "code",
+                            "value": {
+                                "rules": [{"type": "CONSTANT_VALUE", "data": {"type": "JSON", "value": None}}],
+                                "combinator": "OR",
+                            },
+                        },
+                        {
+                            "id": "58598cc8-aa8a-4b4f-99fb-09f6815b6c01",
+                            "key": "runtime",
+                            "value": {
+                                "rules": [
+                                    {"type": "CONSTANT_VALUE", "data": {"type": "STRING", "value": "PYTHON_3_11_6"}}
+                                ],
+                                "combinator": "OR",
+                            },
+                        },
+                    ],
+                    "data": {
+                        "label": "My Node",
+                        "error_output_id": None,
+                        "source_handle_id": "7afa3858-f50c-4116-936a-a401e3b2c60f",
+                        "target_handle_id": "3a39ea63-9f86-4891-a902-0216a7190720",
+                        "code_input_id": "a0b9d6f6-ce59-4075-8db6-866781bc73ef",
+                        "runtime_input_id": "58598cc8-aa8a-4b4f-99fb-09f6815b6c01",
+                        "output_type": "STRING",
+                        "packages": [],
+                        "output_id": "00b2120e-b642-46e4-8276-5f3c69d8a6cb",
+                        "log_output_id": "47e3eeca-4bf8-492e-b8ac-28c7d389c886",
+                    },
+                    "display_data": {"position": {"x": 200.0, "y": -50.0}},
+                    "base": {
+                        "name": "CodeExecutionNode",
+                        "module": ["vellum", "workflows", "nodes", "displayable", "code_execution_node", "node"],
+                    },
+                    "definition": {
+                        "name": "MyNode",
+                        "module": [
+                            "vellum_ee",
+                            "workflows",
+                            "display",
+                            "nodes",
+                            "vellum",
+                            "tests",
+                            "test_code_execution_node",
+                        ],
+                    },
+                    "ports": [{"id": "7afa3858-f50c-4116-936a-a401e3b2c60f", "name": "default", "type": "DEFAULT"}],
+                },
             ],
             "edges": [
                 {

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

@@ -1,9 +1,13 @@
+from vellum.client.types.prompt_parameters import PromptParameters
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
 from vellum.workflows.nodes.displayable.tool_calling_node.node import ToolCallingNode
+from vellum.workflows.nodes.displayable.tool_calling_node.state import ToolCallingState
+from vellum.workflows.nodes.displayable.tool_calling_node.utils import create_router_node, create_tool_prompt_node
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.types.definition import AuthorizationType, EnvironmentVariableReference, MCPServer
+from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
@@ -161,57 +165,170 @@ def test_serialize_node__tool_calling_node__mcp_server_api_key():
         "id": "6c0f7d4f-3c8a-4201-b588-8398d3c97480",
         "name": "functions",
         "value": {
-            "type": "ARRAY_REFERENCE",
-            "items": [
-                {
-                    "type": "DICTIONARY_REFERENCE",
-                    "entries": [
-                        {
-                            "id": "bcf2713b-19fc-4b4b-8ff5-b45c8e63c665",
-                            "key": "type",
-                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "MCP_SERVER"}},
-                        },
-                        {
-                            "id": "4e00439e-ce6f-4e0a-be4c-0fc05990ec44",
-                            "key": "name",
-                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "my-mcp-server"}},
-                        },
-                        {
-                            "id": "9a3dcca9-4595-4efb-ada6-c011721f7018",
-                            "key": "url",
-                            "value": {
-                                "type": "CONSTANT_VALUE",
-                                "value": {"type": "STRING", "value": "https://my-mcp-server.com"},
+            "type": "CONSTANT_VALUE",
+            "value": {
+                "type": "JSON",
+                "value": [
+                    {
+                        "type": "MCP_SERVER",
+                        "name": "my-mcp-server",
+                        "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",
+                    }
+                ],
+            },
+        },
+    }
+
+
+def test_serialize_tool_router_node():
+    """
+    Test that the tool router node created by create_router_node serializes successfully.
+    """
+
+    # GIVEN a simple function for tool calling
+    def my_function(arg1: str) -> str:
+        return f"Result: {arg1}"
+
+    tool_prompt_node = create_tool_prompt_node(
+        ml_model="gpt-4o-mini",
+        blocks=[],
+        functions=[my_function],
+        prompt_inputs=None,
+        parameters=PromptParameters(),
+    )
+
+    # WHEN we create a router node using create_router_node
+    router_node = create_router_node(
+        functions=[my_function],
+        tool_prompt_node=tool_prompt_node,
+    )
+
+    router_node_display_class = get_node_display_class(router_node)
+    router_node_display = router_node_display_class()
+
+    class Workflow(BaseWorkflow[BaseInputs, ToolCallingState]):
+        graph = tool_prompt_node >> router_node
+
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    display_context = workflow_display.display_context
+
+    # WHEN we serialize the router node
+    serialized_router_node = router_node_display.serialize(display_context)
+
+    # THEN the router node should serialize to the exact expected structure
+    assert serialized_router_node == {
+        "adornments": None,
+        "attributes": [
+            {
+                "id": "cd208919-c66b-451b-a739-bcf7d3451dea",
+                "name": "prompt_outputs",
+                "value": {
+                    "node_id": "19e664fc-3b57-48d2-b47a-b143b475406a",
+                    "node_output_id": "c2a5a7f7-a234-45dc-adee-bc6fc0bd28dd",
+                    "type": "NODE_OUTPUT",
+                },
+            }
+        ],
+        "base": {
+            "module": ["vellum", "workflows", "nodes", "displayable", "tool_calling_node", "utils"],
+            "name": "RouterNode",
+        },
+        "definition": {
+            "module": ["vellum", "workflows", "nodes", "displayable", "tool_calling_node", "utils"],
+            "name": "RouterNode",
+        },
+        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "id": "690c66e1-1e18-4984-b695-84beb0157541",
+        "label": "RouterNode",
+        "outputs": [],
+        "ports": [
+            {
+                "expression": {
+                    "lhs": {
+                        "lhs": {
+                            "lhs": {
+                                "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
+                                "type": "WORKFLOW_STATE",
                             },
+                            "operator": "<",
+                            "rhs": {
+                                "lhs": {
+                                    "node_id": "19e664fc-3b57-48d2-b47a-b143b475406a",
+                                    "node_output_id": "c2a5a7f7-a234-45dc-adee-bc6fc0bd28dd",
+                                    "type": "NODE_OUTPUT",
+                                },
+                                "operator": "length",
+                                "type": "UNARY_EXPRESSION",
+                            },
+                            "type": "BINARY_EXPRESSION",
                         },
-                        {
-                            "id": "21c3a6ce-7607-42a4-92b8-9d7c4061edce",
-                            "key": "authorization_type",
-                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "API_KEY"}},
-                        },
-                        {
-                            "id": "dcf8e8f0-84d2-4ffb-b0e6-218c569015a4",
-                            "key": "bearer_token_value",
-                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
-                        },
-                        {
-                            "id": "fe93f516-42c0-40cf-b476-50bd04857d5f",
-                            "key": "api_key_header_key",
-                            "value": {
-                                "type": "CONSTANT_VALUE",
-                                "value": {"type": "STRING", "value": "my-api-key-header-key"},
+                        "operator": "and",
+                        "rhs": {
+                            "lhs": {
+                                "lhs": {
+                                    "lhs": {
+                                        "node_id": "19e664fc-3b57-48d2-b47a-b143b475406a",
+                                        "node_output_id": "c2a5a7f7-a234-45dc-adee-bc6fc0bd28dd",
+                                        "type": "NODE_OUTPUT",
+                                    },
+                                    "operator": "accessField",
+                                    "rhs": {
+                                        "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
+                                        "type": "WORKFLOW_STATE",
+                                    },
+                                    "type": "BINARY_EXPRESSION",
+                                },
+                                "operator": "accessField",
+                                "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "type"}},
+                                "type": "BINARY_EXPRESSION",
                             },
+                            "operator": "=",
+                            "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "FUNCTION_CALL"}},
+                            "type": "BINARY_EXPRESSION",
                         },
-                        {
-                            "id": "750ed8ca-3bf3-46be-96f0-2c92d19d2084",
-                            "key": "api_key_header_value",
-                            "value": {
-                                "type": "ENVIRONMENT_VARIABLE",
-                                "environment_variable": "my-api-key-header-value",
+                        "type": "BINARY_EXPRESSION",
+                    },
+                    "operator": "and",
+                    "rhs": {
+                        "lhs": {
+                            "lhs": {
+                                "lhs": {
+                                    "lhs": {
+                                        "node_id": "19e664fc-3b57-48d2-b47a-b143b475406a",
+                                        "node_output_id": "c2a5a7f7-a234-45dc-adee-bc6fc0bd28dd",
+                                        "type": "NODE_OUTPUT",
+                                    },
+                                    "operator": "accessField",
+                                    "rhs": {
+                                        "state_variable_id": "0dd7f5a1-1d73-4153-9191-ca828ace4920",
+                                        "type": "WORKFLOW_STATE",
+                                    },
+                                    "type": "BINARY_EXPRESSION",
+                                },
+                                "operator": "accessField",
+                                "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "value"}},
+                                "type": "BINARY_EXPRESSION",
                             },
+                            "operator": "accessField",
+                            "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "name"}},
+                            "type": "BINARY_EXPRESSION",
                         },
-                    ],
-                }
-            ],
-        },
+                        "operator": "=",
+                        "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "my_function"}},
+                        "type": "BINARY_EXPRESSION",
+                    },
+                    "type": "BINARY_EXPRESSION",
+                },
+                "id": "afb4b09d-659b-459e-9a28-cf73ba6e0574",
+                "name": "my_function",
+                "type": "IF",
+            },
+            {"expression": None, "id": "4ecd916e-b5d0-407e-aab4-35551c76d02c", "name": "default", "type": "ELSE"},
+        ],
+        "trigger": {"id": "73a96f44-c2dd-40cc-96f6-49b9f914b166", "merge_behavior": "AWAIT_ATTRIBUTES"},
+        "type": "GENERIC",
     }

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

@@ -45,54 +45,20 @@ def test_serialize_workflow():
         "id": "20adf593-c4f0-4c67-8e36-37eb66f28f66",
         "name": "functions",
         "value": {
-            "type": "ARRAY_REFERENCE",
-            "items": [
-                {
-                    "type": "DICTIONARY_REFERENCE",
-                    "entries": [
-                        {
-                            "id": "bcf2713b-19fc-4b4b-8ff5-b45c8e63c665",
-                            "key": "type",
-                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "MCP_SERVER"}},
-                        },
-                        {
-                            "id": "5a3aef6f-b8a1-4f37-8688-b513da42a35a",
-                            "key": "name",
-                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "github"}},
-                        },
-                        {
-                            "id": "641decf3-79c5-4ef9-9fc8-06570d8a69af",
-                            "key": "url",
-                            "value": {
-                                "type": "CONSTANT_VALUE",
-                                "value": {"type": "STRING", "value": "https://api.githubcopilot.com/mcp/"},
-                            },
-                        },
-                        {
-                            "id": "801a74ca-7966-4ac3-b1b5-bebb71a7de07",
-                            "key": "authorization_type",
-                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "BEARER_TOKEN"}},
-                        },
-                        {
-                            "id": "fcd70e2f-0fb2-4011-a73b-90d5d7643be4",
-                            "key": "bearer_token_value",
-                            "value": {
-                                "type": "ENVIRONMENT_VARIABLE",
-                                "environment_variable": "GITHUB_PERSONAL_ACCESS_TOKEN",
-                            },
-                        },
-                        {
-                            "id": "b2fa2900-0e09-44ff-99db-c5399fd76d28",
-                            "key": "api_key_header_key",
-                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
-                        },
-                        {
-                            "id": "6ab23414-5f1b-49c1-a0bc-891bbba9124c",
-                            "key": "api_key_header_value",
-                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
-                        },
-                    ],
-                }
-            ],
+            "type": "CONSTANT_VALUE",
+            "value": {
+                "type": "JSON",
+                "value": [
+                    {
+                        "type": "MCP_SERVER",
+                        "name": "github",
+                        "url": "https://api.githubcopilot.com/mcp/",
+                        "authorization_type": "BEARER_TOKEN",
+                        "bearer_token_value": "GITHUB_PERSONAL_ACCESS_TOKEN",
+                        "api_key_header_key": None,
+                        "api_key_header_value": None,
+                    }
+                ],
+            },
         },
     }

vellum_ee/workflows/display/types.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Dict, Tuple, Type
+from typing import TYPE_CHECKING, Dict, Iterator, List, Tuple, Type
 
 from vellum.client import Vellum as VellumClient
 from vellum.workflows.descriptors.base import BaseDescriptor
@@ -51,3 +51,16 @@ class WorkflowDisplayContext:
     workflow_output_displays: WorkflowOutputDisplays = field(default_factory=dict)
     edge_displays: EdgeDisplays = field(default_factory=dict)
     port_displays: PortDisplays = field(default_factory=dict)
+    _errors: List[Exception] = field(default_factory=list)
+    _dry_run: bool = False
+
+    def add_error(self, error: Exception) -> None:
+        if self._dry_run:
+            self._errors.append(error)
+            return
+
+        raise error
+
+    @property
+    def errors(self) -> Iterator[Exception]:
+        return iter(self._errors)

vellum_ee/workflows/display/utils/expressions.py

@@ -266,10 +266,19 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
         }
 
     if isinstance(value, EnvironmentVariableReference):
-        return {
-            "type": "ENVIRONMENT_VARIABLE",
-            "environment_variable": value.name,
-        }
+        if value.serialize_as_constant:
+            return {
+                "type": "CONSTANT_VALUE",
+                "value": {
+                    "type": "STRING",
+                    "value": value.name,
+                },
+            }
+        else:
+            return {
+                "type": "ENVIRONMENT_VARIABLE",
+                "environment_variable": value.name,
+            }
 
     if isinstance(value, ExecutionCountReference):
         node_class_display = display_context.global_node_displays[value.node_class]

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -6,7 +6,7 @@ import inspect
 import logging
 import os
 from uuid import UUID
-from typing import Any, Dict, ForwardRef, Generic, Iterator, List, Optional, Tuple, Type, TypeVar, Union, cast, get_args
+from typing import Any, Dict, ForwardRef, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast, get_args
 
 from vellum.client import Vellum as VellumClient
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
@@ -96,8 +96,6 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
     # Used to explicitly specify display data for a workflow's ports.
     port_displays: PortDisplays = {}
 
-    _errors: List[Exception]
-
     _serialized_files: List[str]
 
     _dry_run: bool
@@ -116,7 +114,6 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             if self._parent_display_context
             else create_vellum_client()
         )
-        self._errors = []
         self._serialized_files = []
         self._dry_run = dry_run
 
@@ -196,7 +193,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             try:
                 serialized_node = node_display.serialize(self.display_context)
             except (NotImplementedError, NodeValidationError) as e:
-                self.add_error(e)
+                self.display_context.add_error(e)
                 continue
 
             serialized_nodes[node_display.node_id] = serialized_node
@@ -315,7 +312,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         # If there are terminal nodes with no workflow output reference,
         # raise a serialization error
         if len(unreferenced_final_output_node_outputs) > 0:
-            self.add_error(
+            self.display_context.add_error(
                 ValueError("Unable to serialize terminal nodes that are not referenced by workflow outputs.")
             )
 
@@ -375,7 +372,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             try:
                 self._apply_auto_layout(nodes_dict_list, edges)
             except Exception as e:
-                self.add_error(e)
+                self.display_context.add_error(e)
 
         return {
             "workflow_raw_data": {
@@ -446,17 +443,6 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         """Can be overridden as a class attribute to specify a custom workflow id."""
         return self._workflow.__id__
 
-    def add_error(self, error: Exception) -> None:
-        if self._dry_run:
-            self._errors.append(error)
-            return
-
-        raise error
-
-    @property
-    def errors(self) -> Iterator[Exception]:
-        return iter(self._errors)
-
     def _enrich_global_node_output_displays(
         self,
         node: Type[BaseNode],
@@ -606,6 +592,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             edge_displays=edge_displays,
             port_displays=port_displays,
             workflow_display_class=self.__class__,
+            _dry_run=self._dry_run,
         )
 
     def _generate_workflow_meta_display(self) -> WorkflowMetaDisplay:
@@ -896,7 +883,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
 
         return WorkflowSerializationResult(
             exec_config=exec_config,
-            errors=[str(error) for error in workflow_display.errors],
+            errors=[str(error) for error in workflow_display.display_context.errors],
         )
 
     def _gather_additional_module_files(self, module_path: str) -> Dict[str, str]: