vellum-ai 0.14.30__py3-none-any.whl → 0.14.31__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.30",
+            "X-Fern-SDK-Version": "0.14.31",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/workflows/nodes/displayable/code_execution_node/node.py

@@ -227,7 +227,8 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
             )
 
         root = inspect.getfile(self.__class__)
-        code = read_file_from_path(node_filepath=root, script_filepath=self.filepath)
+
+        code = read_file_from_path(node_filepath=root, script_filepath=self.filepath, context=self._context)
         if not code:
             raise NodeException(
                 message=f"Filepath '{self.filepath}' does not exist",

vellum/workflows/nodes/displayable/code_execution_node/utils.py

@@ -1,19 +1,33 @@
 import io
 import os
-from typing import Any, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 from pydantic import BaseModel
 
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.utils import cast_to_output_type
+from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.types.core import EntityInputsInterface
 
 
-def read_file_from_path(node_filepath: str, script_filepath: str) -> Union[str, None]:
+def read_file_from_path(
+    node_filepath: str, script_filepath: str, context: Optional[WorkflowContext] = None
+) -> Union[str, None]:
     node_filepath_dir = os.path.dirname(node_filepath)
     full_filepath = os.path.join(node_filepath_dir, script_filepath)
 
+    # If dynamic file loader is present, try and read the code from there
+    if context and context.generated_files:
+        # Strip out namespace
+        normalized_path = os.path.normpath(full_filepath)
+        stripped_node_filepath = "/".join(normalized_path.split("/")[1:])
+
+        code = context.generated_files.get(stripped_node_filepath, None)
+        if code is not None:
+            return code
+
+    # Default logic for reading from filesystem
     try:
         with open(full_filepath) as file:
             return file.read()

vellum/workflows/state/context.py

@@ -19,6 +19,7 @@ class WorkflowContext:
         *,
         vellum_client: Optional[Vellum] = None,
         execution_context: Optional[ExecutionContext] = None,
+        generated_files: Optional[dict[str, str]] = None,
     ):
         self._vellum_client = vellum_client
         self._event_queue: Optional[Queue["WorkflowEvent"]] = None
@@ -26,6 +27,7 @@ class WorkflowContext:
         self._execution_context = get_execution_context()
         if not self._execution_context.parent_context and execution_context:
             self._execution_context = execution_context
+        self._generated_files = generated_files
 
     @cached_property
     def vellum_client(self) -> Vellum:
@@ -38,6 +40,10 @@ class WorkflowContext:
     def execution_context(self) -> ExecutionContext:
         return self._execution_context
 
+    @cached_property
+    def generated_files(self) -> Optional[dict[str, str]]:
+        return self._generated_files
+
     @cached_property
     def node_output_mocks_map(self) -> Dict[Type[BaseOutputs], List[MockNodeExecution]]:
         return self._node_output_mocks_map

{vellum_ai-0.14.30.dist-info → vellum_ai-0.14.31.dist-info}/METADATA

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

{vellum_ai-0.14.30.dist-info → vellum_ai-0.14.31.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=lWyJw6sziXjyTopTYRdFF5hV-sYPVDdX0yVbG5fzcY4,585
-vellum_cli/pull.py,sha256=GTYNJDzEeJ3Z48i6YjL2sF2ERFt4czOR7T6lEoQ_J6Q,9176
+vellum_cli/pull.py,sha256=Cz4_LWeaMNRGw4qvKpUBFnOpY6PXzJJ-9d7NU90S9Jc,9629
 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=QtbhYKMYn1DFnDyBij2mkQO32j9KOpZ5Pf0yek7k_Ao,1284
-vellum_cli/tests/test_pull.py,sha256=qIoft5Pu49Os6760rgtOG80PNYW0ERdF10N1h0380sA,32089
+vellum_cli/tests/test_pull.py,sha256=asWPuOOeGY1Tj_JGu574A3-pnNKCDY0HqjaSWI0BXss,32824
 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=Ll4o9gg0Q4r8uX6Kt8LWgIz0u2zLwPiZxl3TuqoZpxg,4707
+vellum_ee/workflows/tests/test_server.py,sha256=Q9XD93jMwYjZOkcwMPNM2F4BWvcY71wxkYp8HgMgGFU,4435
 vellum_ee/workflows/tests/test_virtual_files.py,sha256=TJEcMR0v2S8CkloXNmCHA0QW0K6pYNGaIjraJz7sFvY,2762
 vellum/__init__.py,sha256=88-79I29hBTQvR1uH_BOCGMWuj2a4Nx82R_8KIESg28,40470
 vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=Jv9sI5BNFo2OYA9px_aREFSIp655ryC3eaZSRI6yH1k,117826
 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=n2emxJFfx0tF0nSGCah08qytcpcFnvJBgH2UN0WLCFA,1869
+vellum/client/core/client_wrapper.py,sha256=PGwx4aPEP3yRHu1laYbSTTdEDpTOYg6DZU5rt8u-QZo,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
@@ -1543,12 +1543,12 @@ vellum/workflows/nodes/displayable/bases/tests/test_utils.py,sha256=eqdqbKNRWVMD
 vellum/workflows/nodes/displayable/bases/types.py,sha256=C37B2Qh2YP7s7pUjd-EYKc2Zl1TbnCgI_mENuUSb8bo,1706
 vellum/workflows/nodes/displayable/bases/utils.py,sha256=ckMUenSsNkiYmSw6FmjSMHYaCk8Y8_sUjL6lkFFEqts,5412
 vellum/workflows/nodes/displayable/code_execution_node/__init__.py,sha256=0FLWMMktpzSnmBMizQglBpcPrP80fzVsoJwJgf822Cg,76
-vellum/workflows/nodes/displayable/code_execution_node/node.py,sha256=oUBfghBMe6JzrPm5e3tGNuBqILwnaBfnCnOSUX-HfOE,9552
+vellum/workflows/nodes/displayable/code_execution_node/node.py,sha256=Ko_Dy17AjfSx2A4u5Xno5R0KH2p5akEHq8L0rQkySGs,9576
 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/utils.py,sha256=PI0IQysC3uASv4nof23O4gIWpoNl3tRleb1q417bfTw,3896
+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
 vellum/workflows/nodes/displayable/conftest.py,sha256=tD_WIiw5WjFqnzgnGLtEZDaMj2XhQ1DptnBTKYeBbI0,5705
@@ -1618,7 +1618,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=yePVr4CCTQn5bjo1697JOO24fKFQpVNzooL07xL4gL0,2702
+vellum/workflows/state/context.py,sha256=B-sGK9GpZPh7bgCPb8PJCrl_evod3Ru7rHs_RofisrU,2926
 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
@@ -1650,8 +1650,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.30.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.14.30.dist-info/METADATA,sha256=RwawNldGcpTbIALzGk6SBEVUdlVisV1xiZTdYAdfLRs,5484
-vellum_ai-0.14.30.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.14.30.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.14.30.dist-info/RECORD,,
+vellum_ai-0.14.31.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.14.31.dist-info/METADATA,sha256=BDTX1i0ZWziXvzVjwWKd3C6uEiD3QMQol9_a636y0-s,5484
+vellum_ai-0.14.31.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.14.31.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.14.31.dist-info/RECORD,,

vellum_cli/pull.py

@@ -151,68 +151,73 @@ def pull_command(
 
     error_content = ""
 
-    with zipfile.ZipFile(zip_buffer) as zip_file:
-        if METADATA_FILE_NAME in zip_file.namelist():
-            metadata_json: Optional[dict] = None
-            with zip_file.open(METADATA_FILE_NAME) as source:
-                metadata_json = json.load(source)
-
-            pull_contents_metadata = PullContentsMetadata.model_validate(metadata_json)
-
-            if pull_contents_metadata.runner_config:
-                workflow_config.container_image_name = pull_contents_metadata.runner_config.container_image_name
-                workflow_config.container_image_tag = pull_contents_metadata.runner_config.container_image_tag
-                if workflow_config.container_image_name and not workflow_config.container_image_tag:
-                    workflow_config.container_image_tag = "latest"
-            if not workflow_config.module and workflow_deployment and pull_contents_metadata.deployment_name:
-                workflow_config.module = snake_case(pull_contents_metadata.deployment_name)
-            if not workflow_config.module and pull_contents_metadata.label:
-                workflow_config.module = snake_case(pull_contents_metadata.label)
-
-        if not workflow_config.module:
-            raise ValueError(f"Failed to resolve a module name for Workflow {pk}")
-
-        # Use target_directory if provided, otherwise use current working directory
-        base_dir = os.path.join(os.getcwd(), target_directory) if target_directory else os.getcwd()
-        target_dir = os.path.join(base_dir, *workflow_config.module.split("."))
-        workflow_config.target_directory = target_dir if target_directory else None
-
-        # Delete files in target_dir that aren't in the zip file
-        if os.path.exists(target_dir):
-            ignore_patterns = (
-                workflow_config.ignore
-                if isinstance(workflow_config.ignore, list)
-                else [workflow_config.ignore] if isinstance(workflow_config.ignore, str) else []
-            )
-            existing_files = []
-            for root, _, files in os.walk(target_dir):
-                for file in files:
-                    rel_path = os.path.relpath(os.path.join(root, file), target_dir)
-                    existing_files.append(rel_path)
-
-            for file in existing_files:
-                if any(Path(file).match(ignore_pattern) for ignore_pattern in ignore_patterns):
-                    continue
-
-                if file not in zip_file.namelist():
-                    file_path = os.path.join(target_dir, file)
-                    logger.info(f"Deleting {file_path}...")
-                    os.remove(file_path)
-
-        for file_name in zip_file.namelist():
-            with zip_file.open(file_name) as source:
-                content = source.read().decode("utf-8")
-                if file_name == ERROR_LOG_FILE_NAME:
-                    error_content = content
-                    continue
-                if file_name == METADATA_FILE_NAME:
-                    continue
-
-                target_file = os.path.join(target_dir, file_name)
-                os.makedirs(os.path.dirname(target_file), exist_ok=True)
-                with open(target_file, "w") as target:
-                    logger.info(f"Writing to {target_file}...")
-                    target.write(content)
+    try:
+        with zipfile.ZipFile(zip_buffer) as zip_file:
+            if METADATA_FILE_NAME in zip_file.namelist():
+                metadata_json: Optional[dict] = None
+                with zip_file.open(METADATA_FILE_NAME) as source:
+                    metadata_json = json.load(source)
+
+                pull_contents_metadata = PullContentsMetadata.model_validate(metadata_json)
+
+                if pull_contents_metadata.runner_config:
+                    workflow_config.container_image_name = pull_contents_metadata.runner_config.container_image_name
+                    workflow_config.container_image_tag = pull_contents_metadata.runner_config.container_image_tag
+                    if workflow_config.container_image_name and not workflow_config.container_image_tag:
+                        workflow_config.container_image_tag = "latest"
+                if not workflow_config.module and workflow_deployment and pull_contents_metadata.deployment_name:
+                    workflow_config.module = snake_case(pull_contents_metadata.deployment_name)
+                if not workflow_config.module and pull_contents_metadata.label:
+                    workflow_config.module = snake_case(pull_contents_metadata.label)
+
+            if not workflow_config.module:
+                raise ValueError(f"Failed to resolve a module name for Workflow {pk}")
+
+            # Use target_directory if provided, otherwise use current working directory
+            base_dir = os.path.join(os.getcwd(), target_directory) if target_directory else os.getcwd()
+            target_dir = os.path.join(base_dir, *workflow_config.module.split("."))
+            workflow_config.target_directory = target_dir if target_directory else None
+
+            # Delete files in target_dir that aren't in the zip file
+            if os.path.exists(target_dir):
+                ignore_patterns = (
+                    workflow_config.ignore
+                    if isinstance(workflow_config.ignore, list)
+                    else [workflow_config.ignore] if isinstance(workflow_config.ignore, str) else []
+                )
+                existing_files = []
+                for root, _, files in os.walk(target_dir):
+                    for file in files:
+                        rel_path = os.path.relpath(os.path.join(root, file), target_dir)
+                        existing_files.append(rel_path)
+
+                for file in existing_files:
+                    if any(Path(file).match(ignore_pattern) for ignore_pattern in ignore_patterns):
+                        continue
+
+                    if file not in zip_file.namelist():
+                        file_path = os.path.join(target_dir, file)
+                        logger.info(f"Deleting {file_path}...")
+                        os.remove(file_path)
+
+            for file_name in zip_file.namelist():
+                with zip_file.open(file_name) as source:
+                    content = source.read().decode("utf-8")
+                    if file_name == ERROR_LOG_FILE_NAME:
+                        error_content = content
+                        continue
+                    if file_name == METADATA_FILE_NAME:
+                        continue
+
+                    target_file = os.path.join(target_dir, file_name)
+                    os.makedirs(os.path.dirname(target_file), exist_ok=True)
+                    with open(target_file, "w") as target:
+                        logger.info(f"Writing to {target_file}...")
+                        target.write(content)
+    except zipfile.BadZipFile:
+        raise Exception(
+            "The API we tried to pull from returned an invalid zip file. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        )
 
     if include_json:
         logger.warning(

vellum_cli/tests/test_pull.py

@@ -875,3 +875,21 @@ def test_pull__module_name_from_deployment_name(vellum_client):
     assert os.path.exists(workflow_py)
     with open(workflow_py) as f:
         assert f.read() == "print('hello')"
+
+
+def test_pull__invalid_zip_file(vellum_client):
+    workflow_deployment = "test-workflow-deployment-id"
+
+    # GIVEN a workflow pull API call returns an invalid zip file
+    vellum_client.workflows.pull.return_value = iter([b"invalid"])  # Return bytes instead of string
+
+    # WHEN the user runs the pull command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-deployment", workflow_deployment])
+
+    # THEN the command returns an error
+    assert result.exit_code == 1
+    assert (
+        str(result.exception)
+        == "The API we tried to pull from returned an invalid zip file. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+    )

vellum_ee/workflows/tests/test_server.py

@@ -1,4 +1,3 @@
-import pytest
 import sys
 from uuid import uuid4
 from typing import Type, cast
@@ -8,14 +7,10 @@ from vellum.client.types.code_executor_response import CodeExecutorResponse
 from vellum.client.types.number_vellum_value import NumberVellumValue
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes import BaseNode
+from vellum.workflows.state.context import WorkflowContext
 from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
 
 
-@pytest.fixture
-def mock_open(mocker):
-    return mocker.patch("vellum.workflows.nodes.displayable.code_execution_node.utils.open")
-
-
 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
@@ -80,7 +75,6 @@ class StartNode(BaseNode):
 
 
 def test_load_from_module__ts_code_in_file_loader(
-    mock_open,
     vellum_client,
 ):
     # GIVEN typescript code
@@ -91,19 +85,22 @@ def test_load_from_module__ts_code_in_file_loader(
     # AND a workflow module with only a code execution node
     files = {
         "__init__.py": "",
-        "workflow.py": """\
+        "workflow.py": """
 from vellum.workflows import BaseWorkflow
 from .nodes.code_execution_node import CodeExecutionNode
 
 class Workflow(BaseWorkflow):
     graph = CodeExecutionNode
+
+    class Outputs(BaseWorkflow.Outputs):
+       final_output = CodeExecutionNode.Outputs.result
 """,
-        "nodes/__init__.py": """\
+        "nodes/__init__.py": """
 from .code_execution_node import CodeExecutionNode
 
 __all__ = ["CodeExecutionNode"]
 """,
-        "nodes/code_execution_node.py": """\
+        "nodes/code_execution_node/__init__.py": """
 from typing import Any
 
 from vellum.workflows.nodes.displayable import CodeExecutionNode as BaseCodeExecutionNode
@@ -121,10 +118,8 @@ class CodeExecutionNode(BaseCodeExecutionNode[BaseState, int]):
     namespace = str(uuid4())
 
     # AND the virtual file loader is registered
-    sys.meta_path.append(VirtualFileFinder(files, namespace))
-
-    # AND the open function returns our file content
-    mock_open.return_value.__enter__.return_value.read.return_value = ts_code
+    finder = VirtualFileFinder(files, namespace)
+    sys.meta_path.append(finder)
 
     # AND we know what the Code Execution Node will respond with
     mock_code_execution = CodeExecutorResponse(
@@ -135,16 +130,13 @@ class CodeExecutionNode(BaseCodeExecutionNode[BaseState, int]):
 
     # WHEN the workflow is loaded
     Workflow = BaseWorkflow.load_from_module(namespace)
-    workflow = Workflow()
+    workflow = Workflow(context=WorkflowContext(generated_files=files))
 
     # THEN the workflow is successfully initialized
     assert workflow
 
     event = workflow.run()
-    assert event.name == "workflow.execution.fulfilled", event.model_dump_json()
-
-    # AND we pass in the correct file path to the open function
-    assert mock_open.call_args[0][0] == f"{namespace}/nodes/./script.ts"
+    assert event.name == "workflow.execution.fulfilled"
 
-    # AND we invoke the Code Execution Node with the correct code
-    assert vellum_client.execute_code.call_args.kwargs["code"] == ts_code
+    # AND we get the code execution result
+    assert event.body.outputs == {"final_output": 5.0}