vellum-ai 0.14.33__py3-none-any.whl → 0.14.34__py3-none-any.whl

vellum/client/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.14.33",
+            "X-Fern-SDK-Version": "0.14.34",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/client/types/workflow_deployment_release.py

@@ -1,6 +1,8 @@
 # This file was auto-generated by Fern from our API Definition.
 
+from __future__ import annotations
 from ..core.pydantic_utilities import UniversalBaseModel
+from .array_vellum_value import ArrayVellumValue
 import datetime as dt
 from .release_environment import ReleaseEnvironment
 import typing
@@ -10,6 +12,7 @@ from .release_release_tag import ReleaseReleaseTag
 from .slim_release_review import SlimReleaseReview
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import pydantic
+from ..core.pydantic_utilities import update_forward_refs
 
 
 class WorkflowDeploymentRelease(UniversalBaseModel):
@@ -30,3 +33,6 @@ class WorkflowDeploymentRelease(UniversalBaseModel):
             frozen = True
             smart_union = True
             extra = pydantic.Extra.allow
+
+
+update_forward_refs(ArrayVellumValue, WorkflowDeploymentRelease=WorkflowDeploymentRelease)

vellum/client/types/workflow_deployment_release_workflow_version.py

@@ -1,13 +1,19 @@
 # This file was auto-generated by Fern from our API Definition.
 
+from __future__ import annotations
 from ..core.pydantic_utilities import UniversalBaseModel
-from ..core.pydantic_utilities import IS_PYDANTIC_V2
+from .array_vellum_value import ArrayVellumValue
 import typing
+from .vellum_variable import VellumVariable
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import pydantic
+from ..core.pydantic_utilities import update_forward_refs
 
 
 class WorkflowDeploymentReleaseWorkflowVersion(UniversalBaseModel):
     id: str
+    input_variables: typing.List[VellumVariable]
+    output_variables: typing.List[VellumVariable]
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
@@ -17,3 +23,6 @@ class WorkflowDeploymentReleaseWorkflowVersion(UniversalBaseModel):
             frozen = True
             smart_union = True
             extra = pydantic.Extra.allow
+
+
+update_forward_refs(ArrayVellumValue, WorkflowDeploymentReleaseWorkflowVersion=WorkflowDeploymentReleaseWorkflowVersion)

vellum/workflows/nodes/core/inline_subworkflow_node/node.py

@@ -75,9 +75,7 @@ class InlineSubworkflowNode(
         with execution_context(parent_context=get_parent_context()):
             subworkflow = self.subworkflow(
                 parent_state=self.state,
-                context=WorkflowContext(
-                    vellum_client=self._context.vellum_client, generated_files=self._context.generated_files
-                ),
+                context=WorkflowContext.create_from(self._context),
             )
             subworkflow_stream = subworkflow.stream(
                 inputs=self._compile_subworkflow_inputs(),

vellum/workflows/nodes/core/map_node/node.py

@@ -171,9 +171,7 @@ class MapNode(BaseAdornmentNode[StateType], Generic[StateType, MapNodeItemType])
             self._run_subworkflow(item=item, index=index)
 
     def _run_subworkflow(self, *, item: MapNodeItemType, index: int) -> None:
-        context = WorkflowContext(
-            vellum_client=self._context.vellum_client, generated_files=self._context.generated_files
-        )
+        context = WorkflowContext.create_from(self._context)
         subworkflow = self.subworkflow(
             parent_state=self.state,
             context=context,

vellum/workflows/nodes/core/retry_node/node.py

@@ -45,9 +45,7 @@ class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
             with execution_context(parent_context=parent_context):
                 subworkflow = self.subworkflow(
                     parent_state=self.state,
-                    context=WorkflowContext(
-                        vellum_client=self._context.vellum_client, generated_files=self._context.generated_files
-                    ),
+                    context=WorkflowContext.create_from(self._context),
                 )
                 subworkflow_stream = subworkflow.stream(
                     inputs=self.SubworkflowInputs(attempt_number=attempt_number),

vellum/workflows/nodes/core/try_node/node.py

@@ -32,9 +32,7 @@ class TryNode(BaseAdornmentNode[StateType], Generic[StateType]):
         with execution_context(parent_context=parent_context):
             subworkflow = self.subworkflow(
                 parent_state=self.state,
-                context=WorkflowContext(
-                    vellum_client=self._context.vellum_client, generated_files=self._context.generated_files
-                ),
+                context=WorkflowContext.create_from(self._context),
             )
             subworkflow_stream = subworkflow.stream(
                 event_filter=all_workflow_event_filter,

vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py

@@ -1,5 +1,6 @@
 import pytest
 import os
+import re
 from typing import Any, List, Union
 
 from pydantic import BaseModel
@@ -653,7 +654,9 @@ def main():
         node.run()
 
     # AND the error should contain the execution error details
-    assert exc_info.value.message == "Expected an output of type 'int | float', but received 'str'"
+    assert re.match(
+        r"Expected an output of type '(int \| float|float \| int)', but received 'str'", exc_info.value.message
+    )
 
 
 def test_run_node__chat_history_output_type():

vellum/workflows/state/context.py

@@ -74,3 +74,7 @@ class WorkflowContext:
 
     def _get_all_node_output_mocks(self) -> List[MockNodeExecution]:
         return [mock for mocks in self._node_output_mocks_map.values() for mock in mocks]
+
+    @classmethod
+    def create_from(cls, context):
+        return cls(vellum_client=context.vellum_client, generated_files=context.generated_files)

{vellum_ai-0.14.33.dist-info → vellum_ai-0.14.34.dist-info}/METADATA

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

{vellum_ai-0.14.33.dist-info → vellum_ai-0.14.34.dist-info}/RECORD

@@ -7,7 +7,7 @@ vellum_cli/image_push.py,sha256=8DDvRDJEZ-FukUCqGW1827bg1ybF4xBbx9WyqWYQE-g,6816
 vellum_cli/init.py,sha256=WpnMXPItPmh0f0bBGIer3p-e5gu8DUGwSArT_FuoMEw,5093
 vellum_cli/logger.py,sha256=PuRFa0WCh4sAGFS5aqWB0QIYpS6nBWwPJrIXpWxugV4,1022
 vellum_cli/ping.py,sha256=p_BCCRjgPhng6JktuECtkDQLbhopt6JpmrtGoLnLJT8,1161
-vellum_cli/pull.py,sha256=OfP2FIChfYzi7ksWObnvmkibln0JubMQftr1TZL-Bxk,10061
+vellum_cli/pull.py,sha256=YTo5cVCcp7RjS9lHednOuud4rW4bH7jqV3GMdbic_Uk,12002
 vellum_cli/push.py,sha256=xjTNbLwOVFNU3kpBrm56Bk5QkSRrJ9z86qceghCzfIA,9655
 vellum_cli/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_cli/tests/conftest.py,sha256=AFYZryKA2qnUuCPBxBKmHLFoPiE0WhBFFej9tNwSHdc,1526
@@ -16,7 +16,7 @@ vellum_cli/tests/test_image_push.py,sha256=i3lJuW8nFRwL1M1OF6752IZYvGAFgKmkB2hd_
 vellum_cli/tests/test_init.py,sha256=8UOc_ThfouR4ja5cCl_URuLk7ohr9JXfCnG4yka1OUQ,18754
 vellum_cli/tests/test_main.py,sha256=qDZG-aQauPwBwM6A2DIu1494n47v3pL28XakTbLGZ-k,272
 vellum_cli/tests/test_ping.py,sha256=3ucVRThEmTadlV9LrJdCCrr1Ofj3rOjG6ue0BNR2UC0,2523
-vellum_cli/tests/test_pull.py,sha256=UmgQwtv8J7U2zkbTvOPbxi6ozvAec1Fo1b2E5Fc9gWU,34625
+vellum_cli/tests/test_pull.py,sha256=fci5idbcAsnvfJcmP8KlULX4w4gkjLtXkn3iV4wsLZA,46472
 vellum_cli/tests/test_push.py,sha256=zDv_Q1hbXtLwmTJDPRAvwDjbuHC09uNRYOy4FQujUow,23476
 vellum_ee/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -121,14 +121,14 @@ vellum_ee/workflows/tests/local_workflow/nodes/final_output.py,sha256=ZX7zBv87zi
 vellum_ee/workflows/tests/local_workflow/nodes/templating_node.py,sha256=NQwFN61QkHfI3Vssz-B0NKGfupK8PU0FDSAIAhYBLi0,325
 vellum_ee/workflows/tests/local_workflow/workflow.py,sha256=A4qOzOPNwePYxWbcAgIPLsmrVS_aVEZEc-wULSv787Q,393
 vellum_ee/workflows/tests/test_display_meta.py,sha256=C25dErwghPNXio49pvSRxyOuc96srH6eYEwTAWdE2zY,2258
-vellum_ee/workflows/tests/test_server.py,sha256=Q9XD93jMwYjZOkcwMPNM2F4BWvcY71wxkYp8HgMgGFU,4435
+vellum_ee/workflows/tests/test_server.py,sha256=SsOkS6sGO7uGC4mxvk4iv8AtcXs058P9hgFHzTWmpII,14519
 vellum_ee/workflows/tests/test_virtual_files.py,sha256=TJEcMR0v2S8CkloXNmCHA0QW0K6pYNGaIjraJz7sFvY,2762
 vellum/__init__.py,sha256=Ur02Hk_LtwXW45f-LBBtBhRX1xJRSq05ytaQUfgIGmc,40970
 vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=ki-TDOmYzC0FePN7swDyE6UpHFQV_4dK7lqy4h-3s1Y,118148
 vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
 vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-vellum/client/core/client_wrapper.py,sha256=jPo42cS6hlGPneWIZtcMfaVmCbt4TEQwCPgoWxxCTUw,1869
+vellum/client/core/client_wrapper.py,sha256=-wKm98fWVAP4wszQZhj8UDobt3IotWN6IYciYhNcNck,1869
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
 vellum/client/core/http_client.py,sha256=R0pQpCppnEtxccGvXl4uJ76s7ro_65Fo_erlNNLp_AI,19228
@@ -698,8 +698,8 @@ vellum/client/types/workflow_deployment_event_executions_response.py,sha256=x7mZ
 vellum/client/types/workflow_deployment_history_item.py,sha256=4WUPzcthBvEZ7iaisKfEg0soUtHjcTEnL_VUVaKpTyw,1420
 vellum/client/types/workflow_deployment_parent_context.py,sha256=QNyPj2o-jauaC48KrRjCWan1IKIbDgyuxLxURmhXsXM,2347
 vellum/client/types/workflow_deployment_read.py,sha256=tp1WaojTVE_dz1oiZ97h8ixMbIWDgy2yRu08A7wPMpw,2363
-vellum/client/types/workflow_deployment_release.py,sha256=pLkRtkCIEUUZuYxdfXELddVKkljqt42OG7fUA8HXt4M,1209
-vellum/client/types/workflow_deployment_release_workflow_version.py,sha256=kTZbVa6KAtAD_nCfea3QlQrZTH2AYP9tbaOKimHp9jg,581
+vellum/client/types/workflow_deployment_release.py,sha256=UshzYZE5Fv1hAwztaMTD1Fs0JNQqbeNus0IrwXLkKL0,1444
+vellum/client/types/workflow_deployment_release_workflow_version.py,sha256=NP3FoxLpgMUIK1OAPf_ei58mxE9F7BLAuw7q55CIcT8,989
 vellum/client/types/workflow_error.py,sha256=EQajkEmLS64T0wYm0goHQl0rT7Lguurk8pLwkhjsgAI,282
 vellum/client/types/workflow_event_display_context.py,sha256=tnO9lgIJKnLtuS6xum_QilI83LjOmnWCLtnSHLn1oNo,929
 vellum/client/types/workflow_event_error.py,sha256=HIewu_kh3KNPpWegAQArvAGHCp-cBIXqlUAAc_dBZhc,687
@@ -1523,22 +1523,22 @@ vellum/workflows/nodes/core/__init__.py,sha256=5zDMCmyt1v0HTJzlUBwq3U9L825yZGZhT
 vellum/workflows/nodes/core/error_node/__init__.py,sha256=g7RRnlHhqu4qByfLjBwCunmgGA8dI5gNsjS3h6TwlSI,60
 vellum/workflows/nodes/core/error_node/node.py,sha256=MFHU5vITYSK-L9CuMZ49In2ZeNLWnhZD0f8r5dWvb5Y,1270
 vellum/workflows/nodes/core/inline_subworkflow_node/__init__.py,sha256=nKNEH1QTl-1PcvmYoqSWEl0-t6gAur8GLTXHzklRQfM,84
-vellum/workflows/nodes/core/inline_subworkflow_node/node.py,sha256=gumdZs50XGpnFvKxde-WAu6NqH7ImwRf3gK6c7B46W0,6696
+vellum/workflows/nodes/core/inline_subworkflow_node/node.py,sha256=rgcjc3gaCEX9uSfkLSErHjSnNQEeqREVZk-7TEr9hUo,6595
 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=kUqwcRMMxjTHALbwGUXCJT_aJBrHS1bkg49oL8R0JO8,4337
 vellum/workflows/nodes/core/map_node/__init__.py,sha256=MXpZYmGfhsMJHqqlpd64WiJRtbAtAMQz-_3fCU_cLV0,56
-vellum/workflows/nodes/core/map_node/node.py,sha256=8AOHvV1TR8Tr4IyDjdIk6Jzi86KJpg_wS9gp30tDRLQ,9312
+vellum/workflows/nodes/core/map_node/node.py,sha256=RER4mLOXSe1BSpjEYCoJq83ZeydQ5Q5uKCZAsqhlOW8,9227
 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=uMR0AyIFn539LqTKHdwuBswnx1i-PHyqPpgtYrnmYMY,3496
 vellum/workflows/nodes/core/retry_node/__init__.py,sha256=lN2bIy5a3Uzhs_FYCrooADyYU6ZGShtvLKFWpelwPvo,60
-vellum/workflows/nodes/core/retry_node/node.py,sha256=UmTxbxom2VPBoJLk75gXdMa300JJ-1rkW0rh2MeF3jE,5354
+vellum/workflows/nodes/core/retry_node/node.py,sha256=at7RjwUmlBeUv-tHvqeOhCAxkyuSw47ySmIQKC0fJf8,5245
 vellum/workflows/nodes/core/retry_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/core/retry_node/tests/test_node.py,sha256=RM_OHwxrHwyxvlQQBJPqVBxpedFuWQ9h2-Xa3kP75sc,4399
 vellum/workflows/nodes/core/templating_node/__init__.py,sha256=GmyuYo81_A1_Bz6id69ozVFS6FKiuDsZTiA3I6MaL2U,70
 vellum/workflows/nodes/core/templating_node/node.py,sha256=iqBmr2i-f-BqhisNQJiDfewjol0ur7-XpupLStyMJsg,3731
 vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py,sha256=nXkgGDBg4wP36NwykdMEVWwx_xjv8oGT2rYkwuCB_VU,10075
 vellum/workflows/nodes/core/try_node/__init__.py,sha256=JVD4DrldTIqFQQFrubs9KtWCCc0YCAc7Fzol5ZWIWeM,56
-vellum/workflows/nodes/core/try_node/node.py,sha256=-fHfcgI9RFQobNyiugP2dow4ahqY7iDd1LObJw5aB8k,4530
+vellum/workflows/nodes/core/try_node/node.py,sha256=2r2I70j30IyZPvn3Q3zP2VgLbx3WQ1DdOr8NUlXjcq0,4429
 vellum/workflows/nodes/core/try_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/core/try_node/tests/test_node.py,sha256=h6eUc3SggvhzBWlOD0PrPUlkoCSQHwjqYn81VkxSIxU,4948
 vellum/workflows/nodes/displayable/__init__.py,sha256=6F_4DlSwvHuilWnIalp8iDjjDXl0Nmz4QzJV2PYe5RI,1023
@@ -1567,7 +1567,7 @@ vellum/workflows/nodes/displayable/code_execution_node/node.py,sha256=Ko_Dy17Ajf
 vellum/workflows/nodes/displayable/code_execution_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/code_execution_node/tests/fixtures/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/code_execution_node/tests/fixtures/main.py,sha256=5QsbmkzSlSbcbWTG_JmIqcP-JNJzOPTKxGzdHos19W4,79
-vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py,sha256=xAaoOfQHQUlp0iKlig87t0aT2cJM-8PxiTb1QDg8VmY,24641
+vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py,sha256=AzgTK2YSvVj7zr6gWZfz0-YGf1cVQ9DiSx9fe5BR4uE,24690
 vellum/workflows/nodes/displayable/code_execution_node/utils.py,sha256=G-sc7yOL5g6rLk99X8HAbXNcLxRaqpju9IXq1iUwnQI,4470
 vellum/workflows/nodes/displayable/conditional_node/__init__.py,sha256=AS_EIqFdU1F9t8aLmbZU-rLh9ry6LCJ0uj0D8F0L5Uw,72
 vellum/workflows/nodes/displayable/conditional_node/node.py,sha256=Qjfl33gZ3JEgxBA1EgzSUebboGvsARthIxxcQyvx5Gg,1152
@@ -1638,7 +1638,7 @@ vellum/workflows/runner/runner.py,sha256=ww4fjZJBENkB5HJxdj92kTz7k_EyifCeAreupy5
 vellum/workflows/sandbox.py,sha256=GVJzVjMuYzOBnSrboB0_6MMRZWBluAyQ2o7syeaeBd0,2235
 vellum/workflows/state/__init__.py,sha256=yUUdR-_Vl7UiixNDYQZ-GEM_kJI9dnOia75TtuNEsnE,60
 vellum/workflows/state/base.py,sha256=Vkhneko3VlQrPsMLU1PYSzXU_W1u7_AraJsghiv5O-4,15512
-vellum/workflows/state/context.py,sha256=B-sGK9GpZPh7bgCPb8PJCrl_evod3Ru7rHs_RofisrU,2926
+vellum/workflows/state/context.py,sha256=KOAI1wEGn8dGmhmAemJaf4SZbitP3jpIBcwKfznQaRE,3076
 vellum/workflows/state/encoder.py,sha256=TnOQojc5lTQ83g9QbpA4UCqShJvutmTMxbpKt-9gNe4,1911
 vellum/workflows/state/store.py,sha256=uVe-oN73KwGV6M6YLhwZMMUQhzTQomsVfVnb8V91gVo,1147
 vellum/workflows/state/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1670,8 +1670,8 @@ vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnad
 vellum/workflows/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/workflows/tests/test_base_workflow.py,sha256=tCxrV3QBHL8wfdEO3bvKteDdw32xBlUl1_WxkAwaONw,8344
 vellum/workflows/workflows/tests/test_context.py,sha256=VJBUcyWVtMa_lE5KxdhgMu0WYNYnUQUDvTF7qm89hJ0,2333
-vellum_ai-0.14.33.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.14.33.dist-info/METADATA,sha256=8j_YxkXJ0HIBjUZszx0257w_iMhp8KxdW3Y3pUmabMs,5484
-vellum_ai-0.14.33.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.14.33.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.14.33.dist-info/RECORD,,
+vellum_ai-0.14.34.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.14.34.dist-info/METADATA,sha256=1DTJPYOhbXBtorWcrH9e992hHsfsTsvKtleo96NqcaY,5484
+vellum_ai-0.14.34.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.14.34.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.14.34.dist-info/RECORD,,

vellum_cli/pull.py

@@ -2,6 +2,7 @@ import io
 import json
 import os
 from pathlib import Path
+from uuid import UUID
 import zipfile
 from typing import Optional
 
@@ -10,8 +11,9 @@ from pydash import snake_case
 
 from vellum.client.core.api_error import ApiError
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.utils.uuid import is_valid_uuid
 from vellum.workflows.vellum_client import create_vellum_client
-from vellum_cli.config import VellumCliConfig, WorkflowConfig, load_vellum_cli_config
+from vellum_cli.config import VellumCliConfig, WorkflowConfig, WorkflowDeploymentConfig, load_vellum_cli_config
 from vellum_cli.logger import load_cli_logger
 
 ERROR_LOG_FILE_NAME = "error.log"
@@ -31,6 +33,7 @@ class RunnerConfig(UniversalBaseModel):
 class PullContentsMetadata(UniversalBaseModel):
     label: Optional[str] = None
     runner_config: Optional[RunnerConfig] = None
+    deployment_id: Optional[UUID] = None
     deployment_name: Optional[str] = None
 
 
@@ -79,6 +82,29 @@ def _resolve_workflow_config(
             pk=workflow_config.workflow_sandbox_id,
         )
     elif workflow_deployment:
+        if is_valid_uuid(workflow_deployment):
+            # name may also be a valid UUID
+            workflow_config = next(
+                (
+                    w
+                    for w in config.workflows
+                    if w.deployments
+                    and (
+                        str(w.deployments[0].id) == workflow_deployment or w.deployments[0].name == workflow_deployment
+                    )
+                ),
+                None,
+            )
+        else:
+            workflow_config = next(
+                (w for w in config.workflows if w.deployments and w.deployments[0].name == workflow_deployment), None
+            )
+        if workflow_config:
+            return WorkflowConfigResolutionResult(
+                workflow_config=workflow_config,
+                pk=workflow_deployment,
+            )
+
         workflow_config = WorkflowConfig(
             module="",
         )
@@ -177,6 +203,24 @@ def pull_command(
                 if not workflow_config.module and pull_contents_metadata.label:
                     workflow_config.module = snake_case(pull_contents_metadata.label)
 
+                # Save or update the deployment info when pulling with --workflow-deployment
+                if workflow_deployment:
+                    workflow_deployment_id = pull_contents_metadata.deployment_id
+                    existing_deployment = next(
+                        (d for d in workflow_config.deployments if d.id == workflow_deployment_id), None
+                    )
+
+                    if existing_deployment:
+                        if pull_contents_metadata.label:
+                            existing_deployment.label = pull_contents_metadata.label
+                    else:
+                        deployment_config = WorkflowDeploymentConfig(
+                            id=workflow_deployment_id,
+                            label=pull_contents_metadata.label,
+                            name=pull_contents_metadata.deployment_name,
+                        )
+                        workflow_config.deployments.append(deployment_config)
+
             if not workflow_config.module:
                 raise ValueError(f"Failed to resolve a module name for Workflow {pk}")
 

vellum_cli/tests/test_pull.py

@@ -337,6 +337,7 @@ def test_pull__sandbox_id_with_other_workflow_configured(vellum_client, mock_mod
 def test_pull__workflow_deployment_with_no_config(vellum_client):
     # GIVEN a workflow deployment
     workflow_deployment = "my-deployment"
+    deployment_id = str(uuid4())
 
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter(
@@ -344,7 +345,9 @@ def test_pull__workflow_deployment_with_no_config(vellum_client):
             _zip_file_map(
                 {
                     "workflow.py": "print('hello')",
-                    "metadata.json": json.dumps({"deployment_name": workflow_deployment, "label": "Some Label"}),
+                    "metadata.json": json.dumps(
+                        {"deployment_id": deployment_id, "deployment_name": workflow_deployment, "label": "Some Label"}
+                    ),
                 }
             )
         ]
@@ -382,7 +385,15 @@ def test_pull__workflow_deployment_with_no_config(vellum_client):
                     "module": "my_deployment",
                     "workflow_sandbox_id": None,
                     "ignore": None,
-                    "deployments": [],
+                    "deployments": [
+                        {
+                            "id": deployment_id,
+                            "label": "Some Label",
+                            "name": "my-deployment",
+                            "description": None,
+                            "release_tags": None,
+                        }
+                    ],
                     "container_image_tag": None,
                     "container_image_name": None,
                     "workspace": "default",
@@ -937,3 +948,313 @@ def test_pull__unauthorized_error_path(vellum_client):
     # THEN the command returns an error
     assert result.exit_code == 1
     assert str(result.exception) == "Please make sure your `VELLUM_API_KEY` environment variable is set correctly."
+
+
+@pytest.mark.parametrize(
+    "workflow_deployment",
+    [
+        "test-workflow-deployment-id",
+        str(uuid4()),
+    ],
+)
+def test_pull__workflow_deployment_adds_deployment_to_config(vellum_client, workflow_deployment):
+    # GIVEN a workflow deployment ID
+    deployment_id = str(uuid4())  # config will always use the deployment_id return from the API
+    deployment_name = "Test Deployment"
+    deployment_label = "Test Label"
+
+    # AND the workflow pull API call returns a zip file with metadata
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {
+                    "workflow.py": "print('hello')",
+                    "metadata.json": json.dumps(
+                        {
+                            "deployment_id": deployment_id,
+                            "deployment_name": deployment_name,
+                            "label": deployment_label,
+                        }
+                    ),
+                }
+            )
+        ]
+    )
+
+    # AND we are currently in a new directory
+    current_dir = os.getcwd()
+    temp_dir = tempfile.mkdtemp()
+    os.chdir(temp_dir)
+
+    # WHEN the user runs the pull command with the workflow deployment
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-deployment", workflow_deployment])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the deployment is saved in the config
+    vellum_lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    assert os.path.exists(vellum_lock_json)
+    with open(vellum_lock_json) as f:
+        lock_data = json.loads(f.read())
+        assert len(lock_data["workflows"]) == 1
+        assert len(lock_data["workflows"][0]["deployments"]) == 1
+        deployment = lock_data["workflows"][0]["deployments"][0]
+        assert str(deployment["id"]) == deployment_id
+        assert deployment["name"] == deployment_name
+        assert deployment["label"] == deployment_label
+
+    os.chdir(current_dir)
+
+
+def test_pull__workflow_deployment_name_is_uuid(vellum_client):
+    # GIVEN a workflow deployment name that is a valid UUID
+    deployment_id = str(uuid4())
+    deployment_name = str(uuid4())
+
+    # AND an existing configuration with this deployment
+    current_dir = os.getcwd()
+    temp_dir = tempfile.mkdtemp()
+    os.chdir(temp_dir)
+
+    # Create initial config with a deployment
+    vellum_lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    with open(vellum_lock_json, "w") as f:
+        json.dump(
+            {
+                "version": "1.0",
+                "workflows": [
+                    {
+                        "module": "test_workflow",
+                        "workflow_sandbox_id": None,
+                        "ignore": None,
+                        "deployments": [
+                            {
+                                "id": deployment_id,
+                                "label": "Test Label",
+                                "name": deployment_name,
+                                "description": None,
+                                "release_tags": None,
+                            }
+                        ],
+                        "container_image_name": None,
+                        "container_image_tag": None,
+                        "workspace": "default",
+                        "target_directory": None,
+                    }
+                ],
+                "workspaces": [],
+            },
+            f,
+        )
+
+    # AND the workflow pull API call returns a zip file with updated metadata
+    updated_label = "Updated Label"
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {
+                    "workflow.py": "print('hello')",
+                    "metadata.json": json.dumps(
+                        {
+                            "deployment_id": deployment_id,
+                            "deployment_name": deployment_name,
+                            "label": updated_label,
+                        }
+                    ),
+                }
+            )
+        ]
+    )
+
+    # WHEN the user runs the pull command with the workflow deployment
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-deployment", deployment_name])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the deployment info is updated in the config
+    with open(vellum_lock_json) as f:
+        lock_data = json.loads(f.read())
+        assert len(lock_data["workflows"]) == 1
+        assert len(lock_data["workflows"][0]["deployments"]) == 1
+        deployment = lock_data["workflows"][0]["deployments"][0]
+        assert str(deployment["id"]) == deployment_id
+        assert deployment["name"] == deployment_name
+        assert deployment["label"] == updated_label
+
+    os.chdir(current_dir)
+
+
+@pytest.mark.parametrize("get_identifier", [(lambda d: d), (lambda d: "Test Name")])
+def test_pull__workflow_deployment_updates_existing_deployment(vellum_client, get_identifier):
+    """
+    This test is to ensure that the deployment info is updated in the config
+    when the user runs the pull command with the workflow deployment
+
+    get_identifier is a function that returns the identifier to use for the deployment
+    it can be the deployment_id or the deployment_name
+    """
+    # GIVEN a workflow deployment id and name
+    deployment_id = str(uuid4())
+    deployment_name = "Test Name"
+
+    # AND an existing configuration with this deployment
+    current_dir = os.getcwd()
+    temp_dir = tempfile.mkdtemp()
+    os.chdir(temp_dir)
+
+    # Create initial config with a deployment
+    vellum_lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    with open(vellum_lock_json, "w") as f:
+        json.dump(
+            {
+                "version": "1.0",
+                "workflows": [
+                    {
+                        "module": "test_workflow",
+                        "workflow_sandbox_id": None,
+                        "ignore": None,
+                        "deployments": [
+                            {
+                                "id": deployment_id,
+                                "label": "Test Label",
+                                "name": deployment_name,
+                                "description": None,
+                                "release_tags": None,
+                            }
+                        ],
+                        "container_image_name": None,
+                        "container_image_tag": None,
+                        "workspace": "default",
+                        "target_directory": None,
+                    }
+                ],
+                "workspaces": [],
+            },
+            f,
+        )
+
+    # AND the workflow pull API call returns a zip file with updated metadata
+    updated_label = "Updated Label"
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {
+                    "workflow.py": "print('hello')",
+                    "metadata.json": json.dumps(
+                        {
+                            "deployment_id": deployment_id,
+                            "deployment_name": deployment_name,
+                            "label": updated_label,
+                        }
+                    ),
+                }
+            )
+        ]
+    )
+
+    # WHEN the user runs the pull command with the workflow deployment
+    runner = CliRunner()
+    identifier_to_use = get_identifier(deployment_id)
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-deployment", identifier_to_use])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the deployment info is updated in the config
+    with open(vellum_lock_json) as f:
+        lock_data = json.loads(f.read())
+        assert len(lock_data["workflows"]) == 1
+        assert len(lock_data["workflows"][0]["deployments"]) == 1
+        deployment = lock_data["workflows"][0]["deployments"][0]
+        assert str(deployment["id"]) == deployment_id
+        assert deployment["name"] == deployment_name
+        assert deployment["label"] == updated_label
+
+    os.chdir(current_dir)
+
+
+def test_pull__workflow_deployment_with_name_and_id(vellum_client):
+    """
+    This test is to ensure that pulling with id and name will not add a new deployment to the config
+    """
+    # GIVEN a workflow deployment ID
+    deployment_id = str(uuid4())  # config will always use the deployment_id return from the API
+    deployment_name = "Test Deployment"
+    deployment_label = "Test Label"
+
+    # AND the workflow pull API call returns a zip file with metadata
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {
+                    "workflow.py": "print('hello')",
+                    "metadata.json": json.dumps(
+                        {
+                            "deployment_id": deployment_id,
+                            "deployment_name": deployment_name,
+                            "label": deployment_label,
+                        }
+                    ),
+                }
+            )
+        ]
+    )
+
+    # AND we are currently in a new directory
+    current_dir = os.getcwd()
+    temp_dir = tempfile.mkdtemp()
+    os.chdir(temp_dir)
+
+    # WHEN the user runs the pull command with the workflow deployment
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-deployment", deployment_id])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the deployment is saved in the config
+    vellum_lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    assert os.path.exists(vellum_lock_json)
+    with open(vellum_lock_json) as f:
+        lock_data = json.loads(f.read())
+        assert len(lock_data["workflows"]) == 1
+        assert len(lock_data["workflows"][0]["deployments"]) == 1
+        deployment = lock_data["workflows"][0]["deployments"][0]
+        assert str(deployment["id"]) == deployment_id
+        assert deployment["name"] == deployment_name
+        assert deployment["label"] == deployment_label
+
+    os.chdir(current_dir)
+
+    # AND pull with name will not add a new deployment to the config
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {
+                    "workflow.py": "print('hello')",
+                    "metadata.json": json.dumps(
+                        {
+                            "deployment_id": deployment_id,
+                            "deployment_name": deployment_name,
+                            "label": deployment_label,
+                        }
+                    ),
+                }
+            )
+        ]
+    )
+
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-deployment", deployment_name])
+    assert result.exit_code == 0
+    with open(vellum_lock_json) as f:
+        lock_data = json.loads(f.read())
+        assert len(lock_data["workflows"][0]["deployments"]) == 1
+        assert lock_data["workflows"][0]["deployments"][0]["id"] == deployment_id
+        assert lock_data["workflows"][0]["deployments"][0]["name"] == deployment_name
+        assert lock_data["workflows"][0]["deployments"][0]["label"] == deployment_label
+
+    os.chdir(current_dir)

vellum_ee/workflows/tests/test_server.py

@@ -12,8 +12,7 @@ from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
 
 
 def test_load_workflow_event_display_context():
-    # DEPRECATED: Use `vellum.workflows.events.workflow.WorkflowEventDisplayContext` instead. Will be removed in 0.15.0
-    from vellum_ee.workflows.display.types import WorkflowEventDisplayContext
+    from vellum.workflows.events.workflow import WorkflowEventDisplayContext
 
     # We are actually just ensuring there are no circular dependencies when
     # our Workflow Server imports this class.
@@ -140,3 +139,358 @@ class CodeExecutionNode(BaseCodeExecutionNode[BaseState, int]):
 
     # AND we get the code execution result
     assert event.body.outputs == {"final_output": 5.0}
+
+
+def test_load_from_module__simple_code_execution_node_with_try(
+    vellum_client,
+):
+    # GIVEN a simple Python script
+    py_code = """def main() -> int:
+    print("Hello")
+    return 1
+"""
+
+    # AND a workflow module with only a code execution node
+    files = {
+        "__init__.py": "",
+        "workflow.py": """
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.state import BaseState
+
+from .nodes.code_execution_node import CodeExecutionNode
+
+class Workflow(BaseWorkflow):
+    graph = CodeExecutionNode
+
+    class Outputs(BaseWorkflow.Outputs):
+       final_output = CodeExecutionNode.Outputs.result
+""",
+        "nodes/__init__.py": """
+from .code_execution_node import CodeExecutionNode
+
+__all__ = ["CodeExecutionNode"]
+""",
+        "nodes/code_execution_node/__init__.py": """
+from typing import Union
+
+from vellum.workflows.nodes.core import TryNode
+from vellum.workflows.nodes.displayable import CodeExecutionNode as BaseCodeExecutionNode
+from vellum.workflows.state import BaseState
+
+@TryNode.wrap()
+class CodeExecutionNode(BaseCodeExecutionNode[BaseState, Union[float, int]]):
+    filepath = "./script.py"
+    code_inputs = {}
+    runtime = "PYTHON_3_11_6"
+    packages = []
+""",
+        "nodes/code_execution_node/script.py": py_code,
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    finder = VirtualFileFinder(files, namespace)
+    sys.meta_path.append(finder)
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=1),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN the workflow is loaded
+    Workflow = BaseWorkflow.load_from_module(namespace)
+    workflow = Workflow(context=WorkflowContext(generated_files=files))
+
+    # THEN the workflow is successfully initialized
+    assert workflow
+
+    # WHEN we run the workflow
+    event = workflow.run()
+
+    # THEN the execution is fulfilled
+    assert event.name == "workflow.execution.fulfilled"
+
+    # AND we get the code execution result
+    assert event.body.outputs == {"final_output": 1.0}
+
+
+def test_load_from_module__simple_code_execution_node_with_retry(
+    vellum_client,
+):
+    # GIVEN a simple Python script
+    py_code = """def main() -> int:
+    print("Hello")
+    return 1
+"""
+
+    # AND a workflow module with only a code execution node wrapped with RetryNode
+    files = {
+        "__init__.py": "",
+        "workflow.py": """
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.state import BaseState
+
+from .nodes.code_execution_node import CodeExecutionNode
+
+class Workflow(BaseWorkflow):
+    graph = CodeExecutionNode
+
+    class Outputs(BaseWorkflow.Outputs):
+       final_output = CodeExecutionNode.Outputs.result
+""",
+        "nodes/__init__.py": """
+from .code_execution_node import CodeExecutionNode
+
+__all__ = ["CodeExecutionNode"]
+""",
+        "nodes/code_execution_node/__init__.py": """
+from typing import Union
+
+from vellum.workflows.nodes.core.retry_node.node import RetryNode
+from vellum.workflows.nodes.displayable import CodeExecutionNode as BaseCodeExecutionNode
+from vellum.workflows.state import BaseState
+
+@RetryNode.wrap(max_attempts=3)
+class CodeExecutionNode(BaseCodeExecutionNode[BaseState, Union[float, int]]):
+    filepath = "./script.py"
+    code_inputs = {}
+    runtime = "PYTHON_3_11_6"
+    packages = []
+""",
+        "nodes/code_execution_node/script.py": py_code,
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    finder = VirtualFileFinder(files, namespace)
+    sys.meta_path.append(finder)
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=1),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN the workflow is loaded
+    Workflow = BaseWorkflow.load_from_module(namespace)
+    workflow = Workflow(context=WorkflowContext(generated_files=files))
+
+    # THEN the workflow is successfully initialized
+    assert workflow
+
+    # WHEN we run the workflow
+    event = workflow.run()
+
+    # THEN the execution is fulfilled
+    assert event.name == "workflow.execution.fulfilled"
+
+    # AND we get the code execution result
+    assert event.body.outputs == {"final_output": 1.0}
+
+
+def test_load_from_module__code_execution_within_subworkflow(
+    vellum_client,
+):
+    # GIVEN a simple Python script
+    py_code = """def main() -> int:
+    print("Hello")
+    return 1
+"""
+
+    # AND a workflow module with a subworkflow containing a code execution node
+    files = {
+        "__init__.py": "",
+        "workflow.py": """
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.state import BaseState
+
+from .nodes.subworkflow import Subworkflow
+
+class Workflow(BaseWorkflow):
+    graph = Subworkflow
+
+    class Outputs(BaseWorkflow.Outputs):
+       final_output = Subworkflow.Outputs.result
+""",
+        "nodes/__init__.py": """
+from .subworkflow import Subworkflow
+
+__all__ = ["Subworkflow"]
+""",
+        "nodes/subworkflow/__init__.py": """
+from vellum.workflows.nodes.displayable import InlineSubworkflowNode
+
+from .workflow import SubworkflowWorkflow
+
+class Subworkflow(InlineSubworkflowNode):
+    subworkflow = SubworkflowWorkflow
+""",
+        "nodes/subworkflow/nodes/__init__.py": """
+from .code_execution_node import CodeExecutionNode
+
+__all__ = ["CodeExecutionNode"]
+""",
+        "nodes/subworkflow/nodes/code_execution_node/__init__.py": """
+from typing import Union
+
+from vellum.workflows.nodes.displayable import CodeExecutionNode as BaseCodeExecutionNode
+from vellum.workflows.state import BaseState
+
+class CodeExecutionNode(BaseCodeExecutionNode[BaseState, Union[float, int]]):
+    filepath = "./script.py"
+    code_inputs = {}
+    runtime = "PYTHON_3_11_6"
+    packages = []
+""",
+        "nodes/subworkflow/nodes/code_execution_node/script.py": py_code,
+        "nodes/subworkflow/workflow.py": """
+from vellum.workflows import BaseWorkflow
+
+from .nodes.code_execution_node import CodeExecutionNode
+
+class SubworkflowWorkflow(BaseWorkflow):
+    graph = CodeExecutionNode
+
+    class Outputs(BaseWorkflow.Outputs):
+        result = CodeExecutionNode.Outputs.result
+""",
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    finder = VirtualFileFinder(files, namespace)
+    sys.meta_path.append(finder)
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=1),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN the workflow is loaded
+    Workflow = BaseWorkflow.load_from_module(namespace)
+    workflow = Workflow(context=WorkflowContext(generated_files=files))
+
+    # THEN the workflow is successfully initialized
+    assert workflow
+
+    # WHEN we run the workflow
+    event = workflow.run()
+
+    # THEN the execution is fulfilled
+    assert event.name == "workflow.execution.fulfilled"
+
+    # AND we get the code execution result from the subworkflow
+    assert event.body.outputs == {"final_output": 1.0}
+
+
+def test_load_from_module__code_execution_within_map_node(
+    vellum_client,
+):
+    # GIVEN a simple Python script
+    py_code = """def main() -> int:
+    print("Hello")
+    return 1
+"""
+
+    # AND a workflow module with a map node containing a code execution node
+    files = {
+        "__init__.py": "",
+        "workflow.py": """
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.state import BaseState
+
+from .nodes.map_node import MapNode
+
+class Workflow(BaseWorkflow):
+    graph = MapNode
+    
+    class Outputs(BaseWorkflow.Outputs):  # noqa: W293
+       results = MapNode.Outputs.final_output
+""",
+        "nodes/__init__.py": """
+from .map_node import MapNode
+
+__all__ = ["MapNode"]
+""",
+        "nodes/map_node/__init__.py": """
+from vellum.workflows.nodes.core.map_node import MapNode as BaseMapNode
+
+from .workflow import MapNodeWorkflow
+
+class MapNode(BaseMapNode):
+    items = ["foo", "bar", "baz"]
+    subworkflow = MapNodeWorkflow
+    max_concurrency = 4
+""",
+        "nodes/map_node/nodes/__init__.py": """
+from .code_execution_node import CodeExecutionNode
+
+__all__ = ["CodeExecutionNode"]
+""",
+        "nodes/map_node/nodes/code_execution_node/__init__.py": """
+from typing import Union
+
+from vellum.workflows.nodes.displayable import CodeExecutionNode as BaseCodeExecutionNode
+from vellum.workflows.state import BaseState
+
+class CodeExecutionNode(BaseCodeExecutionNode[BaseState, Union[float, int]]):
+    filepath = "./script.py"
+    code_inputs = {}
+    runtime = "PYTHON_3_11_6"
+    packages = []
+""",
+        "nodes/map_node/nodes/code_execution_node/script.py": py_code,
+        "nodes/map_node/workflow.py": """
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.state import BaseState
+
+from .nodes.code_execution_node import CodeExecutionNode
+
+class MapNodeInputs:
+    item: str
+    index: int
+
+class MapNodeWorkflow(BaseWorkflow):
+    graph = CodeExecutionNode
+    
+    class Outputs(BaseWorkflow.Outputs):  # noqa: W293
+        final_output = CodeExecutionNode.Outputs.result
+""",
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    finder = VirtualFileFinder(files, namespace)
+    sys.meta_path.append(finder)
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=3),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN the workflow is loaded
+    Workflow = BaseWorkflow.load_from_module(namespace)
+    workflow = Workflow(context=WorkflowContext(generated_files=files))
+
+    # THEN the workflow is successfully initialized
+    assert workflow
+
+    # WHEN we run the workflow
+    event = workflow.run()
+
+    # THEN the execution is fulfilled
+    assert event.name == "workflow.execution.fulfilled"
+
+    # AND we get the map node results as a list
+    assert event.body.outputs == {"results": [1.0, 1.0, 1.0]}