vellum-ai 0.14.30__tar.gz → 0.14.31__tar.gz

{vellum_ai-0.14.30 → vellum_ai-0.14.31}/PKG-INFO

@@ -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 → vellum_ai-0.14.31}/ee/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_ai-0.14.30 → vellum_ai-0.14.31}/ee/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_ai-0.14.30 → vellum_ai-0.14.31}/ee/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}

{vellum_ai-0.14.30 → vellum_ai-0.14.31}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vellum-ai"
-version = "0.14.30"
+version = "0.14.31"
 description = ""
 readme = "README.md"
 authors = []

{vellum_ai-0.14.30 → vellum_ai-0.14.31}/src/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_ai-0.14.30 → vellum_ai-0.14.31}/src/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_ai-0.14.30 → vellum_ai-0.14.31}/src/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_ai-0.14.30 → vellum_ai-0.14.31}/src/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