vellum-ai 1.3.3__tar.gz → 1.3.5__tar.gz

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 1.3.3
+Version: 1.3.5
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/ee/vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -1,9 +1,9 @@
 import inspect
+import os
 from uuid import UUID
 from typing import ClassVar, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes.displayable.code_execution_node import CodeExecutionNode
-from vellum.workflows.nodes.displayable.code_execution_node.utils import read_file_from_path
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
 from vellum_ee.workflows.display.exceptions import NodeValidationError
@@ -11,10 +11,26 @@ from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 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.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.expressions import virtual_open
 
 _CodeExecutionNodeType = TypeVar("_CodeExecutionNodeType", bound=CodeExecutionNode)
 
 
+def _read_file_from_path_with_virtual_support(node_filepath: str, script_filepath: str) -> Optional[str]:
+    """
+    Read a file using virtual_open which handles VirtualFileFinder instances.
+    """
+    node_filepath_dir = os.path.dirname(node_filepath)
+    normalized_script_filepath = script_filepath.lstrip("./")
+    full_filepath = os.path.join(node_filepath_dir, normalized_script_filepath)
+
+    try:
+        with virtual_open(full_filepath, "r") as file:
+            return file.read()
+    except (FileNotFoundError, IsADirectoryError):
+        return None
+
+
 class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Generic[_CodeExecutionNodeType]):
     output_id: ClassVar[Optional[UUID]] = None
     log_output_id: ClassVar[Optional[UUID]] = None
@@ -37,7 +53,7 @@ class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Gene
             code_value = raw_code
         elif filepath:
             node_file_path = inspect.getfile(node)
-            file_code = read_file_from_path(
+            file_code = _read_file_from_path_with_virtual_support(
                 node_filepath=node_file_path,
                 script_filepath=filepath,
             )

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/ee/vellum_ee/workflows/display/tests/test_base_workflow_display.py

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Dict
+from typing import Any, Dict, List, cast
 
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes import BaseNode, InlineSubworkflowNode
@@ -8,7 +8,9 @@ from vellum.workflows.ports.port import Port
 from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.state import BaseState
 from vellum.workflows.workflows.base import BaseWorkflow
-from vellum_ee.workflows.display.base import WorkflowInputsDisplay
+from vellum_ee.workflows.display.base import EdgeDisplay, WorkflowInputsDisplay
+from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
+from vellum_ee.workflows.display.nodes import BaseNodeDisplay
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -379,3 +381,98 @@ def test_global_propagation_deep_nested_subworkflows():
     inner_global_names = {ref.name for ref in inner_display.display_context.global_workflow_input_displays.keys()}
 
     assert inner_global_names == {"middle_param", "inner_param", "root_param"}
+
+
+def test_serialize_workflow_with_edge_display_data():
+    """
+    Tests that edges with z_index values serialize display_data correctly.
+    """
+
+    # GIVEN a workflow with connected nodes
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class EndNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            final: str
+
+    class TestWorkflow(BaseWorkflow):
+        graph = StartNode >> EndNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_result = EndNode.Outputs.final
+
+    class TestWorkflowDisplay(BaseWorkflowDisplay[TestWorkflow]):
+        edge_displays = {
+            (StartNode.Ports.default, EndNode): EdgeDisplay(id=UUID("12345678-1234-5678-1234-567812345678"), z_index=5)
+        }
+
+    # WHEN we serialize the workflow with the custom display
+    display = get_workflow_display(
+        base_display_class=TestWorkflowDisplay,
+        workflow_class=TestWorkflow,
+    )
+    serialized_workflow = display.serialize()
+
+    # THEN the edge should include display_data with z_index
+    workflow_raw_data = cast(Dict[str, Any], serialized_workflow["workflow_raw_data"])
+    edges = cast(List[Dict[str, Any]], workflow_raw_data["edges"])
+
+    edge_with_display_data = None
+    for edge in edges:
+        if edge["id"] == "12345678-1234-5678-1234-567812345678":
+            edge_with_display_data = edge
+            break
+
+    assert edge_with_display_data is not None, "Edge with custom UUID not found"
+    assert edge_with_display_data["display_data"] == {"z_index": 5}
+
+    assert edge_with_display_data["type"] == "DEFAULT"
+    assert "source_node_id" in edge_with_display_data
+    assert "target_node_id" in edge_with_display_data
+
+
+def test_serialize_workflow_with_node_display_data():
+    """
+    Tests that nodes with z_index values serialize display_data correctly.
+    """
+
+    # GIVEN a workflow with a node that has custom display data
+    class TestNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class TestWorkflow(BaseWorkflow):
+        graph = TestNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_result = TestNode.Outputs.result
+
+    class TestNodeDisplay(BaseNodeDisplay[TestNode]):
+        display_data = NodeDisplayData(position=NodeDisplayPosition(x=100, y=200), z_index=10, width=300, height=150)
+
+    class TestWorkflowDisplay(BaseWorkflowDisplay[TestWorkflow]):
+        pass
+
+    # WHEN we serialize the workflow with the custom node display
+    display = get_workflow_display(
+        base_display_class=TestWorkflowDisplay,
+        workflow_class=TestWorkflow,
+    )
+    serialized_workflow = display.serialize()
+
+    # THEN the node should include display_data with z_index
+    workflow_raw_data = cast(Dict[str, Any], serialized_workflow["workflow_raw_data"])
+    nodes = cast(List[Dict[str, Any]], workflow_raw_data["nodes"])
+
+    test_node = None
+    for node in nodes:
+        if node.get("type") == "GENERIC":
+            definition = node.get("definition")
+            if isinstance(definition, dict) and definition.get("name") == "TestNode":
+                test_node = node
+                break
+
+    assert test_node is not None, "TestNode not found in serialized nodes"
+    assert test_node["display_data"] == {"position": {"x": 100, "y": 200}, "z_index": 10, "width": 300, "height": 150}

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_prompt_node_serialization.py

@@ -243,6 +243,7 @@ def test_serialize_workflow():
                                     "name": "favorite_noun",
                                     "description": "Returns the favorite noun of the user",
                                     "parameters": {},
+                                    "inputs": None,
                                     "forced": None,
                                     "strict": None,
                                 }

vellum_ai-1.3.5/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_parent_input.py

@@ -0,0 +1,85 @@
+from deepdiff import DeepDiff
+
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.basic_tool_calling_node_parent_input.workflow import BasicToolCallingNodeParentInputWorkflow
+
+
+def test_serialize_workflow():
+    # GIVEN a Workflow that uses a generic node
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=BasicToolCallingNodeParentInputWorkflow)
+
+    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",
+        "state_variables",
+        "output_variables",
+    }
+
+    # AND its input variables should be what we expect
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 1
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 2
+    assert not DeepDiff(
+        [
+            {"id": "e2e36cfc-cf24-42fd-ba8f-cce39c53d47b", "key": "text", "type": "STRING"},
+            {"id": "08ca9519-e421-47e8-a42d-44f49f6aab16", "key": "chat_history", "type": "CHAT_HISTORY"},
+        ],
+        output_variables,
+        ignore_order=True,
+    )
+
+    # AND its raw data should be what we expect
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    tool_calling_node = workflow_raw_data["nodes"][2]
+
+    attributes = tool_calling_node["attributes"]
+    function_attributes = next(attribute for attribute in attributes if attribute["name"] == "functions")
+    assert function_attributes == {
+        "id": "cec9f5f2-7bb0-42e4-9c56-f215f07c5569",
+        "name": "functions",
+        "value": {
+            "type": "CONSTANT_VALUE",
+            "value": {
+                "type": "JSON",
+                "value": [
+                    {
+                        "type": "CODE_EXECUTION",
+                        "name": "get_string",
+                        "description": "\n    Get a string with the parent input, dummy input, and the populated input.\n    ",  # noqa: E501
+                        "definition": {
+                            "state": None,
+                            "cache_config": None,
+                            "name": "get_string",
+                            "description": "\n    Get a string with the parent input, dummy input, and the populated input.\n    ",  # noqa: E501
+                            "parameters": {
+                                "type": "object",
+                                "properties": {"populated_input": {"type": "string"}},
+                                "required": ["populated_input"],
+                            },
+                            "inputs": {
+                                "parent_input": {
+                                    "type": "WORKFLOW_INPUT",
+                                    "input_variable_id": "4bf1f0e7-76c6-4204-9f8c-bd9c3b73a8db",
+                                },
+                                "dummy_input": {
+                                    "type": "NODE_OUTPUT",
+                                    "node_id": "8e89ae10-a709-45ec-89f8-242f92e4a83f",
+                                    "node_output_id": "0cd09f0a-c142-4d5d-acc7-b93cd30ca58d",
+                                },
+                            },
+                            "forced": None,
+                            "strict": None,
+                        },
+                        "src": 'from vellum.workflows.utils.functions import use_tool_inputs\n\nfrom .inputs import ParentInputs\nfrom .nodes.dummy_node import DummyNode\n\n\n@use_tool_inputs(\n    parent_input=ParentInputs.parent_input,\n    dummy_input=DummyNode.Outputs.text,\n)\ndef get_string(parent_input: str, dummy_input: str, populated_input: str) -> str:\n    """\n    Get a string with the parent input, dummy input, and the populated input.\n    """\n    return f"This is the parent input: {parent_input}, this is the dummy input: {dummy_input}, and this is the populated input: {populated_input}"  # noqa: E501\n',  # noqa: E501
+                    }
+                ],
+            },
+        },
+    }

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_serialization.py

@@ -152,6 +152,7 @@ def test_serialize_workflow():
                                         },
                                         "required": ["location", "unit"],
                                     },
+                                    "inputs": None,
                                     "forced": None,
                                     "strict": None,
                                 },
@@ -194,7 +195,7 @@ def test_serialize_workflow():
                             "frequency_penalty": 0.0,
                             "presence_penalty": 0.0,
                             "logit_bias": None,
-                            "custom_parameters": None,
+                            "custom_parameters": {"mode": "initial"},
                         },
                     },
                 },

vellum_ai-1.3.5/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_final_output_node_map_reference_serialization.py

@@ -0,0 +1,88 @@
+from typing import List
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes import MapNode
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.displayable import FinalOutputNode
+from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.state import BaseState
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+class TestInputs(BaseInputs):
+    items: List[str]
+
+
+class TestIteration(BaseNode):
+    item = MapNode.SubworkflowInputs.item
+    index = MapNode.SubworkflowInputs.index
+
+    class Outputs(BaseOutputs):
+        processed: str
+
+    def run(self) -> Outputs:
+        return self.Outputs(processed=f"processed_{self.item}_{self.index}")
+
+
+class TestIterationSubworkflow(BaseWorkflow[MapNode.SubworkflowInputs, BaseState]):
+    graph = TestIteration
+
+    class Outputs(BaseOutputs):
+        processed = TestIteration.Outputs.processed
+
+
+class TestMapNode(MapNode):
+    items = TestInputs.items
+    subworkflow = TestIterationSubworkflow
+
+
+class TestFinalOutputNode(FinalOutputNode[BaseState, List[str]]):
+    class Outputs(FinalOutputNode.Outputs):
+        value = TestMapNode.Outputs.processed
+
+
+class TestWorkflowWithFinalOutputReferencingMap(BaseWorkflow[TestInputs, BaseState]):
+    graph = TestMapNode >> TestFinalOutputNode
+
+    class Outputs(BaseOutputs):
+        final_result = TestFinalOutputNode.Outputs.value
+
+
+def test_serialize_workflow__final_output_node_referencing_map_node():
+    """
+    Test that final output nodes referencing map node outputs have correct outputs structure.
+
+    This test verifies that when a FinalOutputNode references a MapNode output,
+    the serialized output contains proper NODE_OUTPUT references instead of None values.
+    This addresses the Agent Builder issue where final outputs showed value=None in the UI.
+    """
+    workflow_display = get_workflow_display(workflow_class=TestWorkflowWithFinalOutputReferencingMap)
+
+    # WHEN we serialize it
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the final output node should have the correct outputs structure
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    map_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "MAP")
+    final_output_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "TERMINAL")
+
+    # AND the map node's subworkflow should have the one output variable
+    output_variable = next(iter(map_node["data"]["output_variables"]))
+    map_node_output_id = output_variable["id"]
+
+    # AND the final output node should have an outputs array with proper structure
+    assert "outputs" in final_output_node
+    outputs = final_output_node["outputs"]
+    assert len(outputs) == 1
+
+    output = outputs[0]
+    # AND the output should have the correct structure with NODE_OUTPUT reference instead of None
+    assert output["name"] == "value"
+    assert output["type"] == "JSON"
+
+    # AND the value should be a NODE_OUTPUT reference, not None
+    assert output["value"] is not None, f"Expected NODE_OUTPUT reference but got None. Full output: {output}"
+    assert output["value"]["type"] == "NODE_OUTPUT", f"Expected NODE_OUTPUT type but got {output['value']['type']}"
+    assert "node_id" in output["value"], f"Missing node_id in output value: {output['value']}"
+    assert output["value"]["node_output_id"] == map_node_output_id

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/ee/vellum_ee/workflows/display/utils/events.py

@@ -36,6 +36,7 @@ def event_enricher(event: WorkflowExecutionInitiatedEvent) -> WorkflowExecutionI
         dry_run=True,
     )
     register_workflow_display_context(event.span_id, workflow_display.display_context)
+    event.body.display_context = workflow_display.get_event_display_context()
 
     if event.body.workflow_definition.is_dynamic or _should_mark_workflow_dynamic(event):
         register_workflow_display_class(workflow_definition, workflow_display.__class__)

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/ee/vellum_ee/workflows/display/utils/expressions.py

@@ -1,8 +1,12 @@
 from dataclasses import asdict, is_dataclass
+import inspect
+from io import StringIO
+import sys
 from typing import TYPE_CHECKING, Any, Dict, List, cast
 
 from pydantic import BaseModel
 
+from vellum.client.types.function_definition import FunctionDefinition
 from vellum.client.types.logical_operator import LogicalOperator
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.expressions.accessor import AccessorExpression
@@ -48,13 +52,31 @@ from vellum.workflows.references.workflow_input import WorkflowInputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.definition import DeploymentDefinition
 from vellum.workflows.types.generics import is_workflow_class
+from vellum.workflows.utils.functions import compile_function_definition
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
+from vellum_ee.workflows.server.virtual_file_loader import VirtualFileLoader
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 
+def virtual_open(file_path: str, mode: str = "r"):
+    """
+    Open a file, checking VirtualFileFinder instances first before falling back to regular open().
+    """
+    for finder in sys.meta_path:
+        if hasattr(finder, "loader") and isinstance(finder.loader, VirtualFileLoader):
+            namespace = finder.loader.namespace
+            if file_path.startswith(namespace + "/"):
+                relative_path = file_path[len(namespace) + 1 :]
+                content = finder.loader._get_code(relative_path)
+                if content is not None:
+                    return StringIO(content)
+
+    return open(file_path, mode)
+
+
 def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> LogicalOperator:
     if isinstance(descriptor, EqualsExpression):
         return "="
@@ -399,6 +421,40 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
         dict_value = value.model_dump()
         return serialize_value(display_context, dict_value)
 
+    if callable(value):
+        function_definition = compile_function_definition(value)
+        inputs = getattr(value, "__vellum_inputs__", {})
+
+        if inputs:
+            serialized_inputs = {}
+            for param_name, input_ref in inputs.items():
+                serialized_inputs[param_name] = serialize_value(display_context, input_ref)
+
+            model_data = function_definition.model_dump()
+            model_data["inputs"] = serialized_inputs
+            function_definition = FunctionDefinition.model_validate(model_data)
+
+        source_path = inspect.getsourcefile(value)
+        if source_path is not None:
+            with virtual_open(source_path) as f:
+                source_code = f.read()
+        else:
+            source_code = f"Source code not available for {value.__name__}"
+
+        return {
+            "type": "CONSTANT_VALUE",
+            "value": {
+                "type": "JSON",
+                "value": {
+                    "type": "CODE_EXECUTION",
+                    "name": function_definition.name,
+                    "description": function_definition.description,
+                    "definition": function_definition.model_dump(),
+                    "src": source_code,
+                },
+            },
+        }
+
     if not isinstance(value, BaseDescriptor):
         vellum_value = primitive_to_vellum_value(value)
         return {

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/ee/vellum_ee/workflows/display/utils/tests/test_events.py

@@ -67,9 +67,14 @@ def test_event_enricher_static_workflow(is_dynamic: bool, expected_config: Optio
     # WHEN the event_enricher is called with mocked dependencies
     event_enricher(event)
 
-    # AND workflow_version_exec_config is set to the expected config
+    # THEN workflow_version_exec_config is set to the expected config
     assert event.body.workflow_version_exec_config == expected_config
 
+    assert event.body.display_context is not None
+    assert hasattr(event.body.display_context, "node_displays")
+    assert hasattr(event.body.display_context, "workflow_inputs")
+    assert hasattr(event.body.display_context, "workflow_outputs")
+
 
 def test_event_enricher_marks_subworkflow_deployment_as_dynamic():
     """Test that event_enricher treats subworkflow deployments as dynamic."""
@@ -109,3 +114,8 @@ def test_event_enricher_marks_subworkflow_deployment_as_dynamic():
 
     assert hasattr(enriched_event.body, "workflow_version_exec_config")
     assert enriched_event.body.workflow_version_exec_config is not None
+
+    assert enriched_event.body.display_context is not None
+    assert hasattr(enriched_event.body.display_context, "node_displays")
+    assert hasattr(enriched_event.body.display_context, "workflow_inputs")
+    assert hasattr(enriched_event.body.display_context, "workflow_outputs")

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/ee/vellum_ee/workflows/display/utils/vellum.py

@@ -102,7 +102,9 @@ def create_node_input_value_pointer_rule(
     if isinstance(value, OutputReference):
         if value not in display_context.global_node_output_displays:
             if issubclass(value.outputs_class, BaseNode.Outputs):
-                raise ValueError(f"Reference to node '{value.outputs_class._node_class.__name__}' not found in graph.")
+                raise ValueError(
+                    f"Reference to node '{value.outputs_class.__parent_class__.__name__}' not found in graph."
+                )
 
             raise ValueError(f"Reference to outputs '{value.outputs_class.__qualname__}' is invalid.")
 

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/ee/vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -294,7 +294,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                     )
 
             elif isinstance(workflow_output.instance, OutputReference):
-                terminal_node_id = workflow_output.instance.outputs_class._node_class.__id__
+                terminal_node_id = workflow_output.instance.outputs_class.__parent_class__.__id__
                 serialized_terminal_node = serialized_nodes.get(terminal_node_id)
                 if serialized_terminal_node and isinstance(serialized_terminal_node["data"], dict):
                     serialized_terminal_node["data"]["name"] = workflow_output_display.name
@@ -329,16 +329,18 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 continue
 
             target_node_display = self.display_context.node_displays[unadorned_target_node]
-            edges.append(
-                {
-                    "id": str(entrypoint_display.edge_display.id),
-                    "source_node_id": str(entrypoint_node_id),
-                    "source_handle_id": str(entrypoint_node_source_handle_id),
-                    "target_node_id": str(target_node_display.node_id),
-                    "target_handle_id": str(target_node_display.get_trigger_id()),
-                    "type": "DEFAULT",
-                }
-            )
+            entrypoint_edge_dict: Dict[str, Json] = {
+                "id": str(entrypoint_display.edge_display.id),
+                "source_node_id": str(entrypoint_node_id),
+                "source_handle_id": str(entrypoint_node_source_handle_id),
+                "target_node_id": str(target_node_display.node_id),
+                "target_handle_id": str(target_node_display.get_trigger_id()),
+                "type": "DEFAULT",
+            }
+            display_data = self._serialize_edge_display_data(entrypoint_display.edge_display)
+            if display_data is not None:
+                entrypoint_edge_dict["display_data"] = display_data
+            edges.append(entrypoint_edge_dict)
 
         for (source_node_port, target_node), edge_display in self.display_context.edge_displays.items():
             unadorned_source_node_port = get_unadorned_port(source_node_port)
@@ -353,18 +355,20 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             source_node_port_display = self.display_context.port_displays[unadorned_source_node_port]
             target_node_display = self.display_context.node_displays[unadorned_target_node]
 
-            edges.append(
-                {
-                    "id": str(edge_display.id),
-                    "source_node_id": str(source_node_port_display.node_id),
-                    "source_handle_id": str(source_node_port_display.id),
-                    "target_node_id": str(target_node_display.node_id),
-                    "target_handle_id": str(
-                        target_node_display.get_target_handle_id_by_source_node_id(source_node_port_display.node_id)
-                    ),
-                    "type": "DEFAULT",
-                }
-            )
+            regular_edge_dict: Dict[str, Json] = {
+                "id": str(edge_display.id),
+                "source_node_id": str(source_node_port_display.node_id),
+                "source_handle_id": str(source_node_port_display.id),
+                "target_node_id": str(target_node_display.node_id),
+                "target_handle_id": str(
+                    target_node_display.get_target_handle_id_by_source_node_id(source_node_port_display.node_id)
+                ),
+                "type": "DEFAULT",
+            }
+            display_data = self._serialize_edge_display_data(edge_display)
+            if display_data is not None:
+                regular_edge_dict["display_data"] = display_data
+            edges.append(regular_edge_dict)
 
         edges.extend(synthetic_output_edges)
 
@@ -405,6 +409,12 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             "output_variables": output_variables,
         }
 
+    def _serialize_edge_display_data(self, edge_display: EdgeDisplay) -> Optional[JsonObject]:
+        """Serialize edge display data, returning None if no display data is present."""
+        if edge_display.z_index is not None:
+            return {"z_index": edge_display.z_index}
+        return None
+
     def _apply_auto_layout(self, nodes_dict_list: List[Dict[str, Any]], edges: List[Json]) -> None:
         """Apply auto-layout to nodes that are all positioned at (0,0)."""
         nodes_for_layout: List[Tuple[str, NodeDisplayData]] = []
@@ -589,7 +599,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
 
             workflow_output_display = self.output_displays.get(workflow_output)
             workflow_output_displays[workflow_output] = (
-                workflow_output_display or self._generate_workflow_output_display(workflow_output)
+                workflow_output_display or self._generate_workflow_output_display(workflow_output, self._workflow)
             )
 
         return WorkflowDisplayContext(
@@ -678,9 +688,13 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
 
         return EntrypointDisplay(id=entrypoint_id, edge_display=edge_display)
 
-    def _generate_workflow_output_display(self, output: BaseDescriptor) -> WorkflowOutputDisplay:
-        output_id = uuid4_from_hash(f"{self.workflow_id}|id|{output.name}")
-
+    def _generate_workflow_output_display(
+        self, output: OutputReference, workflow_class: Type[BaseWorkflow]
+    ) -> WorkflowOutputDisplay:
+        # TODO: use the output.id field instead once we add `__parent_class__` to BaseWorkflow.Outputs
+        output_id = workflow_class.__output_ids__.get(output.name) or uuid4_from_hash(
+            f"{self.workflow_id}|id|{output.name}"
+        )
         return WorkflowOutputDisplay(id=output_id, name=output.name)
 
     def __init_subclass__(cls, **kwargs: Any) -> None:

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-ai"
 
 [tool.poetry]
 name = "vellum-ai"
-version = "1.3.3"
+version = "1.3.5"
 description = ""
 readme = "README.md"
 authors = []

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/src/vellum/client/core/client_wrapper.py

@@ -27,10 +27,10 @@ class BaseClientWrapper:
 
     def get_headers(self) -> typing.Dict[str, str]:
         headers: typing.Dict[str, str] = {
-            "User-Agent": "vellum-ai/1.3.3",
+            "User-Agent": "vellum-ai/1.3.5",
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "1.3.3",
+            "X-Fern-SDK-Version": "1.3.5",
             **(self.get_custom_headers() or {}),
         }
         if self._api_version is not None:

{vellum_ai-1.3.3 → vellum_ai-1.3.5}/src/vellum/client/types/function_definition.py

@@ -30,6 +30,11 @@ class FunctionDefinition(UniversalBaseModel):
     An OpenAPI specification of parameters that are supported by this function.
     """
 
+    inputs: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+    """
+    Optional user defined input mappings for this function.
+    """
+
     forced: typing.Optional[bool] = pydantic.Field(default=None)
     """
     Set this option to true to force the model to return a function call of this function.