vellum-ai 1.3.3__py3-none-any.whl → 1.3.5__py3-none-any.whl

vellum/workflows/utils/functions.py

@@ -135,6 +135,9 @@ def _compile_workflow_deployment_input(input_var: Any) -> dict[str, Any]:
 def compile_function_definition(function: Callable) -> FunctionDefinition:
     """
     Converts a Python function into our Vellum-native FunctionDefinition type.
+
+    Args:
+        function: The Python function to compile
     """
 
     try:
@@ -142,10 +145,18 @@ def compile_function_definition(function: Callable) -> FunctionDefinition:
     except ValueError as e:
         raise ValueError(f"Failed to get signature for function {function.__name__}: {str(e)}")
 
+    # Get inputs from the decorator if present
+    inputs = getattr(function, "__vellum_inputs__", {})
+    exclude_params = set(inputs.keys())
+
     properties = {}
     required = []
     defs: dict[str, Any] = {}
     for param in signature.parameters.values():
+        # Skip parameters that are in the exclude_params set
+        if exclude_params and param.name in exclude_params:
+            continue
+
         # Check if parameter uses Annotated type hint
         if get_origin(param.annotation) is Annotated:
             args = get_args(param.annotation)
@@ -248,3 +259,27 @@ def compile_workflow_deployment_function_definition(
         description=description,
         parameters=parameters,
     )
+
+
+def use_tool_inputs(**inputs):
+    """
+    Decorator to specify which parameters of a tool function should be provided
+    from the parent workflow inputs rather than from the LLM.
+
+    Args:
+        **inputs: Mapping of parameter names to parent input references
+
+    Example:
+        @use_tool_inputs(
+            parent_input=ParentInputs.parent_input,
+        )
+        def get_string(parent_input: str, user_query: str) -> str:
+            return f"Parent: {parent_input}, Query: {user_query}"
+    """
+
+    def decorator(func: Callable) -> Callable:
+        # Store the inputs mapping on the function for later use
+        setattr(func, "__vellum_inputs__", inputs)
+        return func
+
+    return decorator

vellum/workflows/workflows/base.py

@@ -1,3 +1,4 @@
+from dataclasses import field
 from datetime import datetime
 from functools import lru_cache
 import importlib
@@ -199,6 +200,12 @@ class _BaseWorkflowMeta(type):
         if inputs_class is not BaseInputs and inputs_class.__parent_class__ is type(None):
             inputs_class.__parent_class__ = workflow_class
 
+        # TODO: Uncomment this once we figure out why it's causing the ipython reload test to fail
+        # workflow_class.Outputs.__parent_class__ = workflow_class
+        workflow_class.__output_ids__ = {
+            ref.name: uuid4_from_hash(f"{workflow_class.__id__}|id|{ref.name}") for ref in workflow_class.Outputs
+        }
+
         return workflow_class
 
 
@@ -207,6 +214,7 @@ GraphAttribute = Union[Type[BaseNode], Graph, Set[Type[BaseNode]], Set[Graph]]
 
 class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
     __id__: UUID = uuid4_from_hash(__qualname__)
+    __output_ids__: Dict[str, UUID] = {}
     graph: ClassVar[GraphAttribute]
     unused_graphs: ClassVar[Set[GraphAttribute]]  # nodes or graphs that are defined but not used in the graph
     emitters: List[BaseWorkflowEmitter]
@@ -214,7 +222,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
     is_dynamic: ClassVar[bool] = False
 
     class Outputs(BaseOutputs):
-        pass
+        __parent_class__: Type["BaseWorkflow"] = field(init=False)
 
     WorkflowEvent = Union[  # type: ignore
         GenericWorkflowEvent,

{vellum_ai-1.3.3.dist-info → vellum_ai-1.3.5.dist-info}/METADATA

@@ -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.dist-info → vellum_ai-1.3.5.dist-info}/RECORD

@@ -39,7 +39,7 @@ vellum_ee/workflows/display/nodes/utils.py,sha256=sloya5TpXsnot1HURc9L51INwflRqU
 vellum_ee/workflows/display/nodes/vellum/__init__.py,sha256=nUIgH2s0-7IbQRNrBhLPyRNe8YIrx3Yo9HeeW-aXXFk,1668
 vellum_ee/workflows/display/nodes/vellum/api_node.py,sha256=TD23uXA4vhRvsfiq7B10XHpFnsS1CA9M7WH9jsOYNB0,9242
 vellum_ee/workflows/display/nodes/vellum/base_adornment_node.py,sha256=FHhPoGmmny4Xcxi2pm12Sk3ZOREanWEVrOWcjRhncH4,6337
-vellum_ee/workflows/display/nodes/vellum/code_execution_node.py,sha256=w8tO4EZj4Swy5e_2lIYTeoWzghj8ikMLnjML5OqevLg,4570
+vellum_ee/workflows/display/nodes/vellum/code_execution_node.py,sha256=GfwA2Tn0OpdIUiPW9tD2dyv33e3M8yhGt9VfOeN81bU,5146
 vellum_ee/workflows/display/nodes/vellum/conditional_node.py,sha256=dtO9A-rjbDEJrywwrOxwEXahqrW-S493OIDHOti9Sjs,11498
 vellum_ee/workflows/display/nodes/vellum/error_node.py,sha256=8tSb8qGVoRIELubu0qPeoDlt1LpiIqc6d9_30GWRd_k,2266
 vellum_ee/workflows/display/nodes/vellum/final_output_node.py,sha256=eeu8i01HbLDvO3KwHnAfaMrYKiEchNAdDdjufjCi6OU,2991
@@ -72,7 +72,7 @@ vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py,sha256=rHybfUAWwa0L
 vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=47fOnSCEFnY8th9m2yTYlgnoUuzgyRZdjg-SXwn--lk,4079
 vellum_ee/workflows/display/nodes/vellum/utils.py,sha256=oICunzyaXPs0tYnW5zH1r93Bx35MSH7mcD-n0DEWRok,4978
 vellum_ee/workflows/display/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum_ee/workflows/display/tests/test_base_workflow_display.py,sha256=l0ljTRBpoTJ5oKL-fNcU5qeruLw2pUvQLOJodg0yYLg,12436
+vellum_ee/workflows/display/tests/test_base_workflow_display.py,sha256=fh-F7TmzOSiWEBwhtd3Nm_Ix9v5W4LhMNhiIxT7G100,15990
 vellum_ee/workflows/display/tests/workflow_serialization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py,sha256=Y-ajeT65b5varmrZCw6L3hir4hJCFq-eO0jZfRcrs7g,1886
@@ -88,7 +88,7 @@ vellum_ee/workflows/display/tests/workflow_serialization/test_basic_default_stat
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py,sha256=03Mk8OE3kWcoZW9lNBe7v4KThCYN-pXg5Rjbkfx-jOQ,6031
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_generic_node_serialization.py,sha256=-T0cd2jx1bC0ZNtAESF78fnYD_0nOqo8zMMLwRHUTRM,6227
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_guardrail_node_serialization.py,sha256=LnZp1YXDn-AaaxiYgxrhCQeH-rLgmlu_r_lvM65EQ5w,7517
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_prompt_node_serialization.py,sha256=fze2nu8UlH36giAxu60YQzmYwOm8QK3tv_v7MYBm3WY,26306
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_prompt_node_serialization.py,sha256=9w90pk1DzX2PDRML401bZpD9hdTW53Dw5oGb58mxpiw,26358
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py,sha256=yKmOyunzt5_0cUcqhvCmV2pu81TTkshVi_uN3yt76Wc,21816
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py,sha256=W4--Ldih7mRMnfyJ9G7kdyeoKkeebSu_5d33TJQzShU,16735
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py,sha256=UnfWTfKN8DrOLpjCfWMgENJBjgNLWcRVfexbwERKY-c,8501
@@ -100,26 +100,28 @@ vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_nod
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_composio_serialization.py,sha256=oVXCjkU0G56QJmqnd_xIwF3D9bhJwALFibM2wmRhwUk,3739
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_inline_workflow_serialization.py,sha256=Sg82qV5NCzQDy-RD90hA6QaHgFHOq6ESNkbWXygsnNw,26367
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_mcp_serialization.py,sha256=QhQbijeCnFeX1i3SMjHJg2WVAEt5JEO3dhFRv-mofdA,2458
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_serialization.py,sha256=-7NvpYUZEl2T15AS9MPt_scyhKmUmS3wEPa320sgWOI,10085
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_parent_input.py,sha256=HMMa4liBACtL7vU2b-9Ui8Oltyxxb5dbAf-CziLGCes,4284
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_serialization.py,sha256=cKaxme5vUIvKa8aBU7xdeFxXF9wVZ5fW3T5Ie5vToU0,10152
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_workflow_deployment_serialization.py,sha256=XIZZr5POo2NLn2uEWm9EC3rejeBMoO4X-JtzTH6mvp4,4074
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_try_node_serialization.py,sha256=pLCyMScV88DTBXRH7jXaXOEA1GBq8NIipCUFwIAWnwI,2771
 vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py,sha256=exT7U-axwtYgFylagScflSQLJEND51qIAx2UATju6JM,6023
+vellum_ee/workflows/display/tests/workflow_serialization/test_final_output_node_map_reference_serialization.py,sha256=vl3pxUJlrYRA8zzFJ-gRm7fe-5fviLNSIsUC7imnMqk,3502
 vellum_ee/workflows/display/tests/workflow_serialization/test_web_search_node_serialization.py,sha256=vbDFBrWUPeeW7cxjNA6SXrsHlYcbOAhlQ4C45Vdnr1c,3428
 vellum_ee/workflows/display/tests/workflow_serialization/test_workflow_input_parameterization_error.py,sha256=vAdmn3YTBDpo55znbydQxsgg9ASqHcvsUPwiBR_7wfo,1461
 vellum_ee/workflows/display/types.py,sha256=cyZruu4sXAdHjwuFc7dydM4DcFNf-pp_CmulXItxac4,3679
 vellum_ee/workflows/display/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/utils/auto_layout.py,sha256=f4GiLn_LazweupfqTpubcdtdfE_vrOcmZudSsnYIY9E,3906
-vellum_ee/workflows/display/utils/events.py,sha256=KC684ZWODHRGgJtik_KjGXIjM5x4_58_1dGuco_9O8c,1748
+vellum_ee/workflows/display/utils/events.py,sha256=MEG2BT6GgAzkbv1dMaFpov5OShtaAZeAb1-g3xDFsAM,1826
 vellum_ee/workflows/display/utils/exceptions.py,sha256=LSwwxCYNxFkf5XMUcFkaZKpQ13OSrI7y_bpEUwbKVk0,169
-vellum_ee/workflows/display/utils/expressions.py,sha256=yPrUIEQKL5fDwUC9NpVO9yjL6emTOdFJw8VjoZgZ0E0,16942
+vellum_ee/workflows/display/utils/expressions.py,sha256=ci6vOzfCwqVrSRI9A-HR7t4bz_reuS2kzr71SHYIloI,19118
 vellum_ee/workflows/display/utils/registry.py,sha256=1qXiBTdsnro6FeCX0FGBEK7CIf6wa--Jt50iZ_nEp_M,3460
 vellum_ee/workflows/display/utils/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/utils/tests/test_auto_layout.py,sha256=vfXI769418s9vda5Gb5NFBH747WMOwSgHRXeLCTLVm8,2356
-vellum_ee/workflows/display/utils/tests/test_events.py,sha256=BHRq2Or_z1ri7EVWFdVbhDxSWRsCVI87VeqqJN8YnIQ,3995
-vellum_ee/workflows/display/utils/vellum.py,sha256=mtoXmSYwR7rvrq-d6CzCW_auaJXTct0Mi1F0xpRCiNQ,5627
+vellum_ee/workflows/display/utils/tests/test_events.py,sha256=Qze6wEmFJx23_sKQhX-i329apWgMeS9zTptWlRca6Ko,4528
+vellum_ee/workflows/display/utils/vellum.py,sha256=sZwU0KdmZZTKWW62SyxJTl2tC8tN6p_BpZ-lDoinV-U,5670
 vellum_ee/workflows/display/vellum.py,sha256=J2mdJZ1sdLW535DDUkq_Vm8Z572vhuxHxVZF9deKSdk,391
 vellum_ee/workflows/display/workflows/__init__.py,sha256=JTB9ObEV3l4gGGdtfBHwVJtTTKC22uj-a-XjTVwXCyA,148
-vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=_s7mYr56vdavU-CR1FtOQKQlcSpepVgJia84w2OEOUk,43042
+vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=5PCZY37qMA3z1Qpu35Ot0koF5-iH4XVctlLzEMJ26OQ,43991
 vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=gxz76AeCqgAZ9D2lZeTiZzxY9eMgn3qOSfVgiqYcOh8,2028
 vellum_ee/workflows/display/workflows/tests/test_workflow_display.py,sha256=pb7BTH-ivRnya1LQU3j-MApWk_m8POpPNOdD0oEK82A,37847
 vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -153,7 +155,7 @@ vellum/client/README.md,sha256=b6XKeYBBbhQx0v1sHWfM0gIJeJhUFF-aqL2ig7ADa08,5564
 vellum/client/__init__.py,sha256=T5Ht_w-Mk_9nzGqdadhQB8V20M0vYj7am06ut0A3P1o,73401
 vellum/client/core/__init__.py,sha256=lTcqUPXcx4112yLDd70RAPeqq6tu3eFMe1pKOqkW9JQ,1562
 vellum/client/core/api_error.py,sha256=44vPoTyWN59gonCIZMdzw7M1uspygiLnr3GNFOoVL2Q,614
-vellum/client/core/client_wrapper.py,sha256=8my-Ir86AssDOMCGyePyF0hcUkPe6DsVFwW880WD_c8,2840
+vellum/client/core/client_wrapper.py,sha256=Mh9pKr9Xxrby4VuKOk1M3W1PdfGhP8jsj32r0WIAdN8,2840
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=d4NNbX8XvXP32z8KpK2Xovv33nFfruIrpz0QWxlgpZk,2663
 vellum/client/core/force_multipart.py,sha256=awxh5MtcRYe74ehY8U76jzv6fYM_w_D3Rur7KQQzSDk,429
@@ -420,7 +422,7 @@ vellum/client/types/function_call_prompt_block.py,sha256=YcuS0ndY-avXWIcG5UE9mNa
 vellum/client/types/function_call_request.py,sha256=5nPW3gx-l9UN-9puJuDaEZKIkUJUd0glcRLaF0j1x_k,684
 vellum/client/types/function_call_vellum_value.py,sha256=ITXP69WEE3yELyDylbpmLgIxnh4FWdNI14KkIX8boKA,723
 vellum/client/types/function_call_vellum_value_request.py,sha256=9Xz4AZDQhtrUCICXdlq0IS9hpFr0bFHE6haf3g3Bn9c,752
-vellum/client/types/function_definition.py,sha256=cdDwCBQcW5iXgHF7ZBRsLYEitx7oyDLNEwV2CVSWDZA,1657
+vellum/client/types/function_definition.py,sha256=G4mQtdrJzC3I9984YfiyNmvozppEyWqt66NdZBqrgKc,1840
 vellum/client/types/generate_options_request.py,sha256=2ZXU_zNraXRqcUwky8GqDAnhR10qb2NmcrIUaBFeF6w,746
 vellum/client/types/generate_request.py,sha256=R_l2CmLeEF-GVIZHyV93Cbt6BBh2WklTslsyNbI3Vfo,1250
 vellum/client/types/generate_response.py,sha256=_7F1PqE1aQLIBYqKLgw3UDQx7etR15sZgW-VzHfaaWA,1099
@@ -635,13 +637,13 @@ vellum/client/types/rich_text_child_block.py,sha256=JCC-6FLdxHG--gE-L6lJtRNI0pmn
 vellum/client/types/rich_text_prompt_block.py,sha256=Te2qnJOiCgqV0gaiui8mVjnATX7mFNJjHqXDX4fd5jA,1000
 vellum/client/types/sandbox_scenario.py,sha256=gaHbNzufZhibRGoo42Ryo6xZNRo30lLU39oSDS6mqIs,808
 vellum/client/types/scenario_input.py,sha256=BxKRtG7u9boXbavG0kD-_-pFVk19LPzNFPZyfw64Lfo,967
-vellum/client/types/scenario_input_audio_variable_value.py,sha256=n5dF5IxRPsE-x4MiH-2hiiOlZSBFD8WbY8cTzOmKyHs,656
+vellum/client/types/scenario_input_audio_variable_value.py,sha256=YsJ9vgdZtZ2KJkB-zK6xSQpg4LzRahQMdnIQVnqLTUw,680
 vellum/client/types/scenario_input_chat_history_variable_value.py,sha256=FzXccgtyH9iTyGJJ6XcOM0TmzCMaPKWr1FmGtFCvxsQ,799
-vellum/client/types/scenario_input_document_variable_value.py,sha256=yMQtXkCuaQ-A30YJFQzgiFZx5LKk-MBi4zniO_xqxDs,674
-vellum/client/types/scenario_input_image_variable_value.py,sha256=0Hjj9e0qTAw8DO5Yob4oR9JN_d3SVTImyjpCOX3pRSc,656
+vellum/client/types/scenario_input_document_variable_value.py,sha256=xg-OOAbs7uegG9_O5fV3WhcUsM8hGXFeVL88-FJovec,698
+vellum/client/types/scenario_input_image_variable_value.py,sha256=o3rwhVnXp3XCOKYSgLRzUXDo7k9WVIvqRCgL1hjyO18,680
 vellum/client/types/scenario_input_json_variable_value.py,sha256=07_Fj9UOzMGvz1H0j1awmMz0Q6ErliwjgKVU_iUyZb8,716
 vellum/client/types/scenario_input_string_variable_value.py,sha256=PI_pk-WOXVRBSW6PZKDZW11svxd2r9E6bCpl0zC0MTM,717
-vellum/client/types/scenario_input_video_variable_value.py,sha256=IoTqaDZFEgdKUI9ACg-UAUdVFGsl6x76ky_c9GbJQB8,656
+vellum/client/types/scenario_input_video_variable_value.py,sha256=6AxOsLtHNZv-1UyBue6fMs5_vLFCXidx7e_9kP0hHrs,680
 vellum/client/types/search_filters_request.py,sha256=Moxewgl8KwjIy-VBistrchXQAqfsn5fkqO5R2UzLlI8,1276
 vellum/client/types/search_node_result.py,sha256=2DbnKCus81-fj56IoGZqIL6Tw6hjf5jH0O2P_AXah8o,713
 vellum/client/types/search_node_result_data.py,sha256=rXs0R0OG2pjTcE6bMp5lN6_cUusFiVUDwkC3pVo5l6o,810
@@ -1729,14 +1731,14 @@ vellum/workflows/errors/__init__.py,sha256=tWGPu5xyAU8gRb8_bl0fL7OfU3wxQ9UH6qVwy
 vellum/workflows/errors/types.py,sha256=nUWuniEfrhdtb-_2GzoDGlYnSJ_yuNUGjVkaKLNr-rM,4049
 vellum/workflows/events/__init__.py,sha256=V4mh766fyA70WvHelm9kfVZGrUgEKcJ9tJt8EepfQYU,832
 vellum/workflows/events/context.py,sha256=vCfMIPmz4j9Om36rRWa35A_JU_VccWWS52_mZkkqxak,3345
-vellum/workflows/events/node.py,sha256=2R5Yk86v6DNNfgPiCCQjrQ9JYKc71xKNokW2f2EFUUU,5883
+vellum/workflows/events/node.py,sha256=yHVd-rX2E3qc2XLnZr0fW6uq4ZCMm34mnY2tzYceyOg,5884
 vellum/workflows/events/relational_threads.py,sha256=zmLrBCBYpdpQV0snKH3HleST-_hWAMy2LIT0xScfzi4,1516
 vellum/workflows/events/stream.py,sha256=xhXJTZirFi0xad5neAQNogrIQ4h47fpnKbVC3vCM5Js,889
 vellum/workflows/events/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/events/tests/test_basic_workflow.py,sha256=Pj6orHsXz37jWC5FARi0Sx2Gjf99Owri2Cvr6Chb79k,1765
-vellum/workflows/events/tests/test_event.py,sha256=6R3RuN_lMJ9ooHQ9-HE1p0j2xL0FtGiRxkQ7PucFbD0,18450
+vellum/workflows/events/tests/test_event.py,sha256=V26TNmbo2aL4sDvcY3nPxCWgcoS4ejF-7VuesMeD04U,18451
 vellum/workflows/events/types.py,sha256=mVrqAH9Hs9SpXm08Hcxdyap_ImQPwGsxRr56rSNMP34,5043
-vellum/workflows/events/workflow.py,sha256=XW9tpAmDb9BPj2RLYpdbN-SJz_UhwqKBTaMJkZM2ahM,8996
+vellum/workflows/events/workflow.py,sha256=kLSWFXiDZH0TELWoDjQ_kHVomFnw8MVVUPDGIzgyosw,8997
 vellum/workflows/exceptions.py,sha256=NiBiR3ggfmPxBVqD-H1SqmjI-7mIn0EStSN1BqApvCM,1213
 vellum/workflows/expressions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/expressions/accessor.py,sha256=3lu1-_-dBfZdJvtX-q66jbmRAZtqIfdsh_3_JNuzg1E,4462
@@ -1797,7 +1799,7 @@ vellum/workflows/integrations/tests/test_mcp_service.py,sha256=q_DYrDkIqI4sQBNgI
 vellum/workflows/logging.py,sha256=_a217XogktV4Ncz6xKFz7WfYmZAzkfVRVuC0rWob8ls,437
 vellum/workflows/nodes/__init__.py,sha256=zymtc3_iW2rFmMR-sayTLuN6ZsAw8VnJweWPsjQk2-Q,1197
 vellum/workflows/nodes/bases/__init__.py,sha256=cniHuz_RXdJ4TQgD8CBzoiKDiPxg62ErdVpCbWICX64,58
-vellum/workflows/nodes/bases/base.py,sha256=9vOPHP6K6oK36bEby1KFRXpMevMwq6slwuSRhC-bWFc,20752
+vellum/workflows/nodes/bases/base.py,sha256=424pCSFZtXqMpAYX01nQCP3Ej2s16_ynWY3VtxUwTdo,20497
 vellum/workflows/nodes/bases/base_adornment_node.py,sha256=hrgzuTetM4ynPd9YGHoK8Vwwn4XITi3aZZ_OCnQrq4Y,3433
 vellum/workflows/nodes/bases/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/bases/tests/test_base_adornment_node.py,sha256=fXZI9KqpS4XMBrBnIEkK3foHaBVvyHwYcQWWDKay7ic,1148
@@ -1810,7 +1812,7 @@ vellum/workflows/nodes/core/inline_subworkflow_node/node.py,sha256=TCmO0wPbt7kc8
 vellum/workflows/nodes/core/inline_subworkflow_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/core/inline_subworkflow_node/tests/test_node.py,sha256=RK2g1h2ib-ruQZ9A2_2L-B9WBdHV44WZj75rkDNL0cE,5766
 vellum/workflows/nodes/core/map_node/__init__.py,sha256=MXpZYmGfhsMJHqqlpd64WiJRtbAtAMQz-_3fCU_cLV0,56
-vellum/workflows/nodes/core/map_node/node.py,sha256=gyOUJw2EhuI5FYT0782moEFlMKG1ajk3ukTiDCtFlxw,9678
+vellum/workflows/nodes/core/map_node/node.py,sha256=kSyQmIWk4v-KSt4WBf3d-0_QueKYkjtrEmQPGTbMryw,10054
 vellum/workflows/nodes/core/map_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/core/map_node/tests/test_node.py,sha256=Xc2xZY5ShSy-bsIQe41JbvIjq3TE95duS-ygaELRVkk,9320
 vellum/workflows/nodes/core/retry_node/__init__.py,sha256=lN2bIy5a3Uzhs_FYCrooADyYU6ZGShtvLKFWpelwPvo,60
@@ -1837,7 +1839,7 @@ vellum/workflows/nodes/displayable/bases/api_node/tests/test_node.py,sha256=5C59
 vellum/workflows/nodes/displayable/bases/base_prompt_node/__init__.py,sha256=Org3xTvgp1pA0uUXFfnJr29D3HzCey2lEdYF4zbIUgo,70
 vellum/workflows/nodes/displayable/bases/base_prompt_node/node.py,sha256=ea20icDM1HB942wkH-XtXNSNCBDcjeOiN3vowkHL4fs,4477
 vellum/workflows/nodes/displayable/bases/inline_prompt_node/__init__.py,sha256=Hl35IAoepRpE-j4cALaXVJIYTYOF3qszyVbxTj4kS1s,82
-vellum/workflows/nodes/displayable/bases/inline_prompt_node/node.py,sha256=cW0fgmferJVm6LctqC5kCFu4Qg9udMQEb0IISLS7CTo,13795
+vellum/workflows/nodes/displayable/bases/inline_prompt_node/node.py,sha256=wtJjaV0tJ5svO0shxFA-lInDQB25wfv24VUvAOYLOZ8,13793
 vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/test_inline_prompt_node.py,sha256=Hk_u2IxLIeeqL_s0RTgoyL5QGYwY9VllKT8z5_JHiDU,24956
 vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py,sha256=0a40fkkZkFMmZN0CsWf6EP_y1H6x36EGa3WcfVNyOsM,9797
@@ -1861,7 +1863,7 @@ vellum/workflows/nodes/displayable/final_output_node/node.py,sha256=6SMaGeBlHQ5r
 vellum/workflows/nodes/displayable/final_output_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/final_output_node/tests/test_node.py,sha256=E6LQ74qZjY4Xi4avx2qdOCgGhF8pEcNLBh8cqYRkzMI,709
 vellum/workflows/nodes/displayable/guardrail_node/__init__.py,sha256=Ab5eXmOoBhyV4dMWdzh32HLUmnPIBEK_zFCT38C4Fng,68
-vellum/workflows/nodes/displayable/guardrail_node/node.py,sha256=GRWRDCVfPbL03hGbyZP4FAUtAjg_8K3XrU0xJ_lIFeA,5820
+vellum/workflows/nodes/displayable/guardrail_node/node.py,sha256=axYUojar_kdB3gi4LG3g9euJ8VkOxNtiFxJNI46v-SQ,5869
 vellum/workflows/nodes/displayable/guardrail_node/test_node.py,sha256=SAGv6hSFcBwQkudn1VxtaKNsXSXWWELl3eK05zM6tS0,5410
 vellum/workflows/nodes/displayable/guardrail_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/guardrail_node/tests/test_node.py,sha256=X2pd6TI8miYxIa7rgvs1pHTEreyWcf77EyR0_Jsa700,2055
@@ -1891,13 +1893,13 @@ vellum/workflows/nodes/displayable/tests/test_search_node_error_handling.py,sha2
 vellum/workflows/nodes/displayable/tests/test_search_node_wth_text_output.py,sha256=VepO5z1277c1y5N6LLIC31nnWD1aak2m5oPFplfJHHs,6935
 vellum/workflows/nodes/displayable/tests/test_text_prompt_deployment_node.py,sha256=Bjv-wZyFgNaVZb9KEMMZd9lFoLzbPEPjEMpANizMZw4,2413
 vellum/workflows/nodes/displayable/tool_calling_node/__init__.py,sha256=3n0-ysmFKsr40CVxPthc0rfJgqVJeZuUEsCmYudLVRg,117
-vellum/workflows/nodes/displayable/tool_calling_node/node.py,sha256=diAm1hJ4ZywdRmwZAYhBq_TjmCsUt_qc6ew7tywV1rs,7840
+vellum/workflows/nodes/displayable/tool_calling_node/node.py,sha256=KxKU8JWFr-XUdwYelEtUTfAh3UaN3FTm2wK-nioXu9U,8065
 vellum/workflows/nodes/displayable/tool_calling_node/state.py,sha256=CcBVb_YtwfSSka4ze678k6-qwmzMSfjfVP8_Y95feSo,302
 vellum/workflows/nodes/displayable/tool_calling_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/tool_calling_node/tests/test_composio_service.py,sha256=in1fbEz5x1tx3uKv9YXdvOncsHucNL8Ro6Go7lBuuOQ,8962
 vellum/workflows/nodes/displayable/tool_calling_node/tests/test_node.py,sha256=GZoeybB9uM7ai8sBLAtUMHrMVgh-WrJDWrKZci6feDs,11892
 vellum/workflows/nodes/displayable/tool_calling_node/tests/test_utils.py,sha256=SIu5GCj4tIE4fz-cAcdULtQfqZIhrcc3Doo6TWLXBws,8804
-vellum/workflows/nodes/displayable/tool_calling_node/utils.py,sha256=wymqUCytUHiWKJxJDvMc7cLdJ69zfxNrHXAmXkFSUiQ,23189
+vellum/workflows/nodes/displayable/tool_calling_node/utils.py,sha256=VC0TenXfktfWHd4ZYQFqiXaZvkhizLJFIK3nQvKoVG0,23780
 vellum/workflows/nodes/displayable/web_search_node/__init__.py,sha256=8FOnEP-n-U68cvxTlJW9wphIAGHq5aqjzLM-DoSSXnU,61
 vellum/workflows/nodes/displayable/web_search_node/node.py,sha256=NQYux2bOtuBF5E4tn-fXi5y3btURPRrNqMSM9MAZYI4,5091
 vellum/workflows/nodes/displayable/web_search_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1912,7 +1914,7 @@ vellum/workflows/nodes/tests/test_mocks.py,sha256=mfPvrs75PKcsNsbJLQAN6PDFoVqs9T
 vellum/workflows/nodes/tests/test_utils.py,sha256=BUugAHx2C9YuCwTlsTXV1Glxca0kW3St6T9o_QFatSU,5649
 vellum/workflows/nodes/utils.py,sha256=wCvf8K5qruT5GwtvnHcQ-LMllktTD8aaFmAGpKQy--c,10720
 vellum/workflows/outputs/__init__.py,sha256=AyZ4pRh_ACQIGvkf0byJO46EDnSix1ZCAXfvh-ms1QE,94
-vellum/workflows/outputs/base.py,sha256=PUn0zhGzYCSZL34JXtXg9zALlXS_cqxZldLilPxDzb8,9614
+vellum/workflows/outputs/base.py,sha256=XLt2WnjOhHBC0tT4-1FIUHLVdI1J5NXME6UZxKCKwOI,9741
 vellum/workflows/ports/__init__.py,sha256=bZuMt-R7z5bKwpu4uPW7LlJeePOQWmCcDSXe5frUY5g,101
 vellum/workflows/ports/node_ports.py,sha256=SM9uLAaoaE1HwR-Uqwf2v5zZK5iFnphKs6mE5Ls7ldE,2877
 vellum/workflows/ports/port.py,sha256=PYhmzEHgJyjUjSFkPIQ38cNIKpcXrhYiZlj7nZC5yCk,3989
@@ -1925,24 +1927,24 @@ vellum/workflows/references/external_input.py,sha256=c_4SojTpykCSbGS1Pjmx9FfquyY
 vellum/workflows/references/input.py,sha256=3INu-TLTi4dziWmva6LO3WvgDlPzsjayUx61cVvqLJA,325
 vellum/workflows/references/lazy.py,sha256=jgUYmgt-yAybzPf_R-74MzdU8VuNwMYI8EQqrj9lVR0,2948
 vellum/workflows/references/node.py,sha256=LP854wDVs-9I_aZ7-nkbwXqL2H7W2_3LED2e9FixNS8,1418
-vellum/workflows/references/output.py,sha256=Odpjqnw2uY6lbmt49sUwDclBPZMndYxgtFAGn1iKj8k,3387
+vellum/workflows/references/output.py,sha256=utYoYPVfsPCTFPXcAc0zTVn2aZcFwL_-y1_WLiKMhHA,3329
 vellum/workflows/references/state_value.py,sha256=bInUF0A3Pt4-zhA0f6LdSuyv8tz7n5QRkHAEn4gsmqI,711
 vellum/workflows/references/tests/test_lazy.py,sha256=0s50-LizMTlSTBQahpK0fg_xqCucA8YTp6QmIMqPvMk,919
 vellum/workflows/references/vellum_secret.py,sha256=Od4d19a5yletWMqNfJR5d_mZQUkVcFzj29mE-T9J7yE,480
 vellum/workflows/references/workflow_input.py,sha256=W3rOK1EPd2gYHb04WJwmNm1CUSdvZ9LKrs8RMKxACBs,1751
 vellum/workflows/resolvers/__init__.py,sha256=eH6hTvZO4IciDaf_cf7aM2vs-DkBDyJPycOQevJxQnI,82
 vellum/workflows/resolvers/base.py,sha256=wrQiSC02Bw4-dBwgFjJIHsjpe-4xz4rUJs_1RdErKA0,1164
-vellum/workflows/resolvers/resolver.py,sha256=0Y9ArmN_2Gu62TpgYjFvSzhFMMFy7oGWOJh9zg3jaFk,4085
-vellum/workflows/resolvers/tests/test_resolver.py,sha256=JX6-zjIQShNvt4XYmub6zmY0iUf6z90KZynqvEd33qA,5202
+vellum/workflows/resolvers/resolver.py,sha256=yK-oY0HDsFJcjlNKAm3vpsPKRIFerIh59FcTwuEN5RY,4839
+vellum/workflows/resolvers/tests/test_resolver.py,sha256=jXkJBb9SwtoH__bBN-ECohpyD0aTIB9ErEvtFhuTMQM,9750
 vellum/workflows/resolvers/types.py,sha256=Hndhlk69g6EKLh_LYg5ILepW5U_h_BYNllfzhS9k8p4,237
 vellum/workflows/runner/__init__.py,sha256=i1iG5sAhtpdsrlvwgH6B-m49JsINkiWyPWs8vyT-bqM,72
-vellum/workflows/runner/runner.py,sha256=kuW28ozlBi2thQRhgrKigxXtAcOYoGp0M6lYucYsrMw,40321
+vellum/workflows/runner/runner.py,sha256=UJPCSddeczBcGu06j63KGM2IbbQ5xISTOwP-tFDptaI,40341
 vellum/workflows/sandbox.py,sha256=jwlFFQjHDwmbVoBah_Q3i8K_BrzOt-F6TXFauiyVyIk,3021
 vellum/workflows/state/__init__.py,sha256=yUUdR-_Vl7UiixNDYQZ-GEM_kJI9dnOia75TtuNEsnE,60
 vellum/workflows/state/base.py,sha256=m9fCqbZn21GshCVCjJTD1dPZEQjFrsMXqlg7tM9fIwM,24283
 vellum/workflows/state/context.py,sha256=khM30U1iDNts5Xp8LXa_WfpkITNITexrDUUFJ5wZ2W4,8445
 vellum/workflows/state/delta.py,sha256=7h8wR10lRCm15SykaPj-gSEvvsMjCwYLPsOx3nsvBQg,440
-vellum/workflows/state/encoder.py,sha256=HdNlabmBOcFSeY_dgn4LNtQEugyzsw3p4mvn2ChC1Io,3380
+vellum/workflows/state/encoder.py,sha256=EynuS9aCt9Neb-H6HRCinEVZX5olCzME03W1TSXfpxs,1961
 vellum/workflows/state/store.py,sha256=uVe-oN73KwGV6M6YLhwZMMUQhzTQomsVfVnb8V91gVo,1147
 vellum/workflows/state/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/state/tests/test_state.py,sha256=zEVFIY2any41X2BA5Us_qqKpzH5HRqmyrUJ04GTO0pU,7484
@@ -1960,7 +1962,7 @@ vellum/workflows/types/tests/test_definition.py,sha256=rvDYjdJ1rvAv0qHBN7i7s-_WA
 vellum/workflows/types/tests/test_utils.py,sha256=UnZog59tR577mVwqZRqqWn2fScoOU1H6up0EzS8zYhw,2536
 vellum/workflows/types/utils.py,sha256=mTctHITBybpt4855x32oCKALBEcMNLn-9cCmfEKgJHQ,6498
 vellum/workflows/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/workflows/utils/functions.py,sha256=_mJ7vGxHLIByPbfMpDVKsWVEUGlCof1Bc2aZefF56Ns,8967
+vellum/workflows/utils/functions.py,sha256=6WRRMb_XbxtvhUKOJq5ZChy0KKvlBaQCBiPhvecXT7I,10029
 vellum/workflows/utils/hmac.py,sha256=JJCczc6pyV6DuE1Oa0QVfYPUN_of3zEYmGFib3OZnrE,1135
 vellum/workflows/utils/names.py,sha256=QLUqfJ1tmSEeUwBKTTiv_Qk3QGbInC2RSmlXfGXc8Wo,380
 vellum/workflows/utils/pydantic_schema.py,sha256=eR_bBtY-T0pttJP-ARwagSdCOnwPUtiT3cegm2lzDTQ,1310
@@ -1974,13 +1976,13 @@ vellum/workflows/utils/vellum_variables.py,sha256=YHLNiQGWDNssGH1FQoG9Z1jUFZ-zYe
 vellum/workflows/utils/zip.py,sha256=HVg_YZLmBOTXKaDV3Xhaf3V6sYnfqqZXQ8CpuafkbPY,1181
 vellum/workflows/vellum_client.py,sha256=xkfoucodxNK5JR2-lbRqZx3xzDgExWkP6kySrpi_Ubc,1079
 vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
-vellum/workflows/workflows/base.py,sha256=PT7Hveate5A4Nt1UQoS6wwFehxgZQ5fuYjS1vLnlkl0,28008
+vellum/workflows/workflows/base.py,sha256=nEFaKuLQahpEAxi-3VmlIbxEcTaKJ3_1G3WNVATVvzA,28466
 vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
 vellum/workflows/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/workflows/tests/test_base_workflow.py,sha256=ptMntHzVyy8ZuzNgeTuk7hREgKQ5UBdgq8VJFSGaW4Y,20832
 vellum/workflows/workflows/tests/test_context.py,sha256=VJBUcyWVtMa_lE5KxdhgMu0WYNYnUQUDvTF7qm89hJ0,2333
-vellum_ai-1.3.3.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-1.3.3.dist-info/METADATA,sha256=bNuEdxi3O1ua5f1Q56poP2dIfUnw_dHq58lIVX6pMds,5547
-vellum_ai-1.3.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-1.3.3.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-1.3.3.dist-info/RECORD,,
+vellum_ai-1.3.5.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-1.3.5.dist-info/METADATA,sha256=cGN1uqwggZ4itH8f0IlQACy6h-uM2e-IspLFVz3DwKU,5547
+vellum_ai-1.3.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-1.3.5.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-1.3.5.dist-info/RECORD,,

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_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_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_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_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"},
                         },
                     },
                 },