vellum-ai 1.3.2__tar.gz → 1.3.4__tar.gz

{vellum_ai-1.3.2 → vellum_ai-1.3.4}/PKG-INFO

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

{vellum_ai-1.3.2 → vellum_ai-1.3.4}/ee/vellum_cli/__init__.py

@@ -5,6 +5,7 @@ import click
 from vellum_cli.aliased_group import ClickAliasedGroup
 from vellum_cli.image_push import image_push_command
 from vellum_cli.init import init_command
+from vellum_cli.move import move_command
 from vellum_cli.ping import ping_command
 from vellum_cli.pull import pull_command
 from vellum_cli.push import push_command
@@ -376,6 +377,26 @@ def image_push(
     image_push_command(image, tag, workspace, source)
 
 
+@workflows.command(name="move")
+@click.argument("old_module", required=True)
+@click.argument("new_module", required=True)
+@click.option("--workspace", type=str, help="The specific Workspace config to use when moving")
+def workflows_move(
+    old_module: str,
+    new_module: str,
+    workspace: Optional[str],
+) -> None:
+    """
+    Move/rename a Workflow module. Updates both the filesystem structure and configuration.
+    """
+
+    move_command(
+        old_module=old_module,
+        new_module=new_module,
+        workspace=workspace,
+    )
+
+
 @workflows.command(name="init")
 @click.argument("template_name", required=False)
 @click.option(

vellum_ai-1.3.4/ee/vellum_cli/move.py

@@ -0,0 +1,56 @@
+import os
+import shutil
+from typing import Optional
+
+from dotenv import load_dotenv
+
+from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, load_vellum_cli_config
+from vellum_cli.logger import load_cli_logger
+from vellum_cli.push import module_exists
+
+
+def move_command(
+    old_module: str,
+    new_module: str,
+    workspace: Optional[str] = None,
+) -> None:
+    load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
+    logger = load_cli_logger()
+    config = load_vellum_cli_config()
+
+    if not module_exists(old_module):
+        raise ValueError(f"Module '{old_module}' does not exist in the filesystem.")
+
+    if module_exists(new_module):
+        raise ValueError(f"Module '{new_module}' already exists. Cannot move to existing module.")
+
+    matching_configs = [w for w in config.workflows if w.module == old_module]
+
+    if workspace:
+        matching_configs = [w for w in matching_configs if w.workspace == workspace]
+    else:
+        matching_configs = [w for w in matching_configs if w.workspace == DEFAULT_WORKSPACE_CONFIG.name]
+
+    if not matching_configs:
+        workspace_msg = (
+            f" in workspace '{workspace}'" if workspace else f" in workspace '{DEFAULT_WORKSPACE_CONFIG.name}'"
+        )
+        raise ValueError(f"No workflow configuration found for module '{old_module}'{workspace_msg}.")
+
+    logger.info(f"Moving module from '{old_module}' to '{new_module}'...")
+
+    old_path = os.path.join(os.getcwd(), *old_module.split("."))
+    new_path = os.path.join(os.getcwd(), *new_module.split("."))
+
+    os.makedirs(os.path.dirname(new_path), exist_ok=True)
+
+    shutil.move(old_path, new_path)
+    logger.info(f"Moved filesystem directory from '{old_path}' to '{new_path}'")
+
+    for workflow_config in matching_configs:
+        workflow_config.module = new_module
+        logger.info(f"Updated workflow configuration: {workflow_config.workflow_sandbox_id}")
+
+    config.save()
+    logger.info("Updated vellum.lock.json file.")
+    logger.info(f"Successfully moved module from '{old_module}' to '{new_module}'")

vellum_ai-1.3.4/ee/vellum_cli/tests/test_move.py

@@ -0,0 +1,154 @@
+import os
+
+from click.testing import CliRunner
+
+from vellum_cli import main
+from vellum_cli.config import load_vellum_cli_config
+
+
+def test_move__happy_path(mock_module):
+    """
+    Test that the move command successfully moves a module and updates configuration.
+    """
+
+    temp_dir = mock_module.temp_dir
+    old_module = mock_module.module
+    new_module = "examples.new.workflow"
+
+    old_module_dir = os.path.join(temp_dir, *old_module.split("."))
+    os.makedirs(old_module_dir, exist_ok=True)
+    with open(os.path.join(old_module_dir, "workflow.py"), "w") as f:
+        f.write("# test workflow")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["workflows", "move", old_module, new_module])
+
+    assert result.exit_code == 0
+
+    assert not os.path.exists(old_module_dir)
+
+    new_module_dir = os.path.join(temp_dir, *new_module.split("."))
+    assert os.path.exists(new_module_dir)
+    assert os.path.exists(os.path.join(new_module_dir, "workflow.py"))
+
+    config = load_vellum_cli_config()
+    workflow_config = next((w for w in config.workflows if w.module == new_module), None)
+    assert workflow_config is not None
+    assert workflow_config.workflow_sandbox_id == mock_module.workflow_sandbox_id
+
+
+def test_move__old_module_not_exists(mock_module):
+    """
+    Test that the move command fails when the old module doesn't exist.
+    """
+
+    old_module = "nonexistent.module"
+    new_module = "examples.new.workflow"
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["workflows", "move", old_module, new_module])
+
+    assert result.exit_code != 0
+    assert "does not exist in the filesystem" in str(result.exception)
+
+
+def test_move__new_module_already_exists(mock_module):
+    """
+    Test that the move command fails when the new module already exists.
+    """
+
+    temp_dir = mock_module.temp_dir
+    old_module = mock_module.module
+    new_module = "examples.existing.workflow"
+
+    old_module_dir = os.path.join(temp_dir, *old_module.split("."))
+    new_module_dir = os.path.join(temp_dir, *new_module.split("."))
+    os.makedirs(old_module_dir, exist_ok=True)
+    os.makedirs(new_module_dir, exist_ok=True)
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["workflows", "move", old_module, new_module])
+
+    assert result.exit_code != 0
+    assert "already exists" in str(result.exception)
+
+
+def test_move__no_workflow_config_found(mock_module):
+    """
+    Test that the move command fails when no workflow config is found.
+    """
+
+    temp_dir = mock_module.temp_dir
+    old_module = "examples.unconfigured.workflow"
+    new_module = "examples.new.workflow"
+
+    old_module_dir = os.path.join(temp_dir, *old_module.split("."))
+    os.makedirs(old_module_dir, exist_ok=True)
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["workflows", "move", old_module, new_module])
+
+    assert result.exit_code != 0
+    assert "No workflow configuration found" in str(result.exception)
+
+
+def test_move__with_workspace_filter(mock_module):
+    """
+    Test that the move command works with workspace filtering.
+    """
+
+    temp_dir = mock_module.temp_dir
+    old_module = mock_module.module
+    new_module = "examples.new.workflow"
+    workspace = "default"
+
+    old_module_dir = os.path.join(temp_dir, *old_module.split("."))
+    os.makedirs(old_module_dir, exist_ok=True)
+    with open(os.path.join(old_module_dir, "workflow.py"), "w") as f:
+        f.write("# test workflow")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["workflows", "move", old_module, new_module, "--workspace", workspace])
+
+    assert result.exit_code == 0
+
+    config = load_vellum_cli_config()
+    workflow_config = next((w for w in config.workflows if w.module == new_module), None)
+    assert workflow_config is not None
+    assert workflow_config.workspace == workspace
+
+
+def test_move__preserves_workflow_metadata(mock_module):
+    """
+    Test that the move command preserves all workflow metadata except the module name.
+    """
+
+    temp_dir = mock_module.temp_dir
+    old_module = mock_module.module
+    new_module = "examples.new.workflow"
+
+    config = load_vellum_cli_config()
+    original_config = next((w for w in config.workflows if w.module == old_module), None)
+    assert original_config is not None
+    original_config.container_image_name = "test-image"
+    original_config.container_image_tag = "v1.0"
+    original_config.ignore = "sandbox.py"
+    config.save()
+
+    old_module_dir = os.path.join(temp_dir, *old_module.split("."))
+    os.makedirs(old_module_dir, exist_ok=True)
+    with open(os.path.join(old_module_dir, "workflow.py"), "w") as f:
+        f.write("# test workflow")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["workflows", "move", old_module, new_module])
+
+    assert result.exit_code == 0
+
+    config = load_vellum_cli_config()
+    workflow_config = next((w for w in config.workflows if w.module == new_module), None)
+    assert workflow_config is not None
+    assert workflow_config.workflow_sandbox_id == mock_module.workflow_sandbox_id
+    assert workflow_config.container_image_name == "test-image"
+    assert workflow_config.container_image_tag == "v1.0"
+    assert workflow_config.ignore == "sandbox.py"

{vellum_ai-1.3.2 → vellum_ai-1.3.4}/ee/vellum_ee/workflows/display/base.py

@@ -56,6 +56,7 @@ class StateValueDisplay:
 @dataclass
 class EdgeDisplay:
     id: UUID
+    z_index: Optional[int] = None
 
 
 @dataclass

{vellum_ai-1.3.2 → vellum_ai-1.3.4}/ee/vellum_ee/workflows/display/editor/types.py

@@ -17,6 +17,7 @@ class NodeDisplayComment(UniversalBaseModel):
 
 class NodeDisplayData(UniversalBaseModel):
     position: NodeDisplayPosition = Field(default_factory=NodeDisplayPosition)
+    z_index: Optional[int] = None
     width: Optional[int] = None
     height: Optional[int] = None
     comment: Optional[NodeDisplayComment] = None

{vellum_ai-1.3.2 → vellum_ai-1.3.4}/ee/vellum_ee/workflows/display/nodes/base_node_display.py

@@ -431,6 +431,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
             )
             return NodeDisplayData(
                 position=explicit_value.position,
+                z_index=explicit_value.z_index,
                 width=explicit_value.width,
                 height=explicit_value.height,
                 comment=comment,

{vellum_ai-1.3.2 → vellum_ai-1.3.4}/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.2 → vellum_ai-1.3.4}/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,7 @@ 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.workflows.base_workflow_display import BaseWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -379,3 +379,53 @@ 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

{vellum_ai-1.3.2 → vellum_ai-1.3.4}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_prompt_node_serialization.py

@@ -47,17 +47,20 @@ def test_serialize_workflow():
 
     # AND its output variables should be what we expect
     output_variables = serialized_workflow["output_variables"]
-    assert len(output_variables) == 1
+    assert len(output_variables) == 2
     assert not DeepDiff(
-        [{"id": "15a0ab89-8ed4-43b9-afa2-3c0b29d4dc3e", "key": "results", "type": "JSON"}],
+        [
+            {"id": "15a0ab89-8ed4-43b9-afa2-3c0b29d4dc3e", "key": "results", "type": "JSON"},
+            {"id": "0ef1608e-1737-41cc-9b90-a8e124138f70", "key": "json", "type": "JSON"},
+        ],
         output_variables,
         ignore_order=True,
     )
 
     # AND its raw data should be what we expect
     workflow_raw_data = serialized_workflow["workflow_raw_data"]
-    assert len(workflow_raw_data["edges"]) == 2
-    assert len(workflow_raw_data["nodes"]) == 3
+    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]
@@ -240,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,
                                 }
@@ -305,7 +309,7 @@ def test_serialize_workflow():
                     },
                 }
             ],
-            "display_data": {"position": {"x": 400.0, "y": -50.0}},
+            "display_data": {"position": {"x": 400.0, "y": 75.0}},
             "base": {
                 "name": "FinalOutputNode",
                 "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -336,6 +340,14 @@ def test_serialize_workflow():
                 "target_handle_id": "46c99277-2b4b-477d-851c-ea497aef6b16",
                 "type": "DEFAULT",
             },
+            {
+                "id": "0b1a2960-4cd5-4045-844f-42b6c87487aa",
+                "source_node_id": "8450dd06-975a-41a4-a564-808ee8808fe6",
+                "source_handle_id": "d4a097ab-e22d-42f1-b6bc-2ed96856377a",
+                "target_node_id": "1f4e3b7b-6af1-42c8-ab33-05b0f01e2b62",
+                "target_handle_id": "7d94907f-c840-4ced-b813-ee3b17f2a8a9",
+                "type": "DEFAULT",
+            },
         ],
         serialized_edges,
         ignore_order=True,

{vellum_ai-1.3.2 → vellum_ai-1.3.4}/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,
                                 },

{vellum_ai-1.3.2 → vellum_ai-1.3.4}/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.2 → vellum_ai-1.3.4}/ee/vellum_ee/workflows/display/utils/expressions.py

@@ -1,4 +1,7 @@
 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
@@ -48,13 +51,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 +420,29 @@ 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)
+        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.2 → vellum_ai-1.3.4}/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.2 → vellum_ai-1.3.4}/ee/vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -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]] = []

{vellum_ai-1.3.2 → vellum_ai-1.3.4}/pyproject.toml

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

{vellum_ai-1.3.2 → vellum_ai-1.3.4}/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.2",
+            "User-Agent": "vellum-ai/1.3.4",
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "1.3.2",
+            "X-Fern-SDK-Version": "1.3.4",
             **(self.get_custom_headers() or {}),
         }
         if self._api_version is not None: