vellum-ai 0.11.8__tar.gz → 0.11.10__tar.gz

{vellum_ai-0.11.8 → vellum_ai-0.11.10}/PKG-INFO

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

{vellum_ai-0.11.8 → vellum_ai-0.11.10}/ee/vellum_cli/__init__.py

@@ -148,6 +148,11 @@ def pull(
 Should only be used for debugging purposes.""",
 )
 @click.option("--workflow-sandbox-id", type=str, help="Pull the Workflow from a specific Sandbox ID")
+@click.option(
+    "--workflow-deployment",
+    type=str,
+    help="""Pull the Workflow from a specific Deployment. Can use the name or the ID of the Deployment.""",
+)
 @click.option(
     "--exclude-code",
     is_flag=True,
@@ -158,6 +163,7 @@ def workflows_pull(
     module: Optional[str],
     include_json: Optional[bool],
     workflow_sandbox_id: Optional[str],
+    workflow_deployment: Optional[str],
     exclude_code: Optional[bool],
 ) -> None:
     """
@@ -169,6 +175,7 @@ def workflows_pull(
         module=module,
         include_json=include_json,
         workflow_sandbox_id=workflow_sandbox_id,
+        workflow_deployment=workflow_deployment,
         exclude_code=exclude_code,
     )
 

{vellum_ai-0.11.8 → vellum_ai-0.11.10}/ee/vellum_cli/pull.py

@@ -1,39 +1,84 @@
 import io
 import os
 from pathlib import Path
+from uuid import UUID
 import zipfile
-from typing import Optional
+from typing import Optional, Union
 
 from dotenv import load_dotenv
+from pydash import snake_case
 
+from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.workflows.vellum_client import create_vellum_client
 from vellum_cli.config import VellumCliConfig, WorkflowConfig, load_vellum_cli_config
 from vellum_cli.logger import load_cli_logger
 
 
-def resolve_workflow_config(
+def _is_valid_uuid(val: Union[str, UUID, None]) -> bool:
+    try:
+        UUID(str(val))
+        return True
+    except (ValueError, TypeError):
+        return False
+
+
+class WorkflowConfigResolutionResult(UniversalBaseModel):
+    workflow_config: Optional[WorkflowConfig] = None
+    pk: Optional[str] = None
+
+
+def _resolve_workflow_config(
     config: VellumCliConfig,
     module: Optional[str] = None,
     workflow_sandbox_id: Optional[str] = None,
-) -> Optional[WorkflowConfig]:
+    workflow_deployment: Optional[str] = None,
+) -> WorkflowConfigResolutionResult:
+    if workflow_sandbox_id and workflow_deployment:
+        raise ValueError("Cannot specify both workflow_sandbox_id and workflow_deployment")
+
     if module:
-        return next((w for w in config.workflows if w.module == module), None)
+        workflow_config = next((w for w in config.workflows if w.module == module), None)
+        return WorkflowConfigResolutionResult(
+            workflow_config=workflow_config,
+            pk=workflow_config.workflow_sandbox_id if workflow_config else None,
+        )
     elif workflow_sandbox_id:
         workflow_config = WorkflowConfig(
             workflow_sandbox_id=workflow_sandbox_id,
             module=f"workflow_{workflow_sandbox_id.split('-')[0]}",
         )
         config.workflows.append(workflow_config)
-        return workflow_config
+        return WorkflowConfigResolutionResult(
+            workflow_config=workflow_config,
+            pk=workflow_config.workflow_sandbox_id,
+        )
+    elif workflow_deployment:
+        module = (
+            f"workflow_{workflow_deployment.split('-')[0]}"
+            if _is_valid_uuid(workflow_deployment)
+            else snake_case(workflow_deployment)
+        )
+        workflow_config = WorkflowConfig(
+            module=module,
+        )
+        config.workflows.append(workflow_config)
+        return WorkflowConfigResolutionResult(
+            workflow_config=workflow_config,
+            pk=workflow_deployment,
+        )
     elif config.workflows:
-        return config.workflows[0]
+        return WorkflowConfigResolutionResult(
+            workflow_config=config.workflows[0],
+            pk=config.workflows[0].workflow_sandbox_id,
+        )
 
-    return None
+    return WorkflowConfigResolutionResult()
 
 
 def pull_command(
     module: Optional[str] = None,
     workflow_sandbox_id: Optional[str] = None,
+    workflow_deployment: Optional[str] = None,
     include_json: Optional[bool] = None,
     exclude_code: Optional[bool] = None,
 ) -> None:
@@ -41,17 +86,20 @@ def pull_command(
     logger = load_cli_logger()
     config = load_vellum_cli_config()
 
-    workflow_config = resolve_workflow_config(
-        config,
-        module,
-        workflow_sandbox_id,
+    workflow_config_result = _resolve_workflow_config(
+        config=config,
+        module=module,
+        workflow_sandbox_id=workflow_sandbox_id,
+        workflow_deployment=workflow_deployment,
     )
     save_lock_file = not module
 
+    workflow_config = workflow_config_result.workflow_config
     if not workflow_config:
         raise ValueError("No workflow config found in project to pull from.")
 
-    if not workflow_config.workflow_sandbox_id:
+    pk = workflow_config_result.pk
+    if not pk:
         raise ValueError("No workflow sandbox ID found in project to pull from.")
 
     logger.info(f"Pulling workflow into {workflow_config.module}")
@@ -63,7 +111,7 @@ def pull_command(
         query_parameters["exclude_code"] = exclude_code
 
     response = client.workflows.pull(
-        workflow_config.workflow_sandbox_id,
+        pk,
         request_options={"additional_query_parameters": query_parameters},
     )
 

{vellum_ai-0.11.8 → vellum_ai-0.11.10}/ee/vellum_cli/tests/test_pull.py

@@ -164,6 +164,86 @@ def test_pull__sandbox_id_with_other_workflow_configured(vellum_client, mock_mod
         assert f.read() == "print('hello')"
 
 
+def test_pull__workflow_deployment_with_no_config(vellum_client):
+    # GIVEN a workflow deployment
+    workflow_deployment = "my-deployment"
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+
+    # 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])
+    os.chdir(current_dir)
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the pull api is called with the workflow deployment
+    vellum_client.workflows.pull.assert_called_once()
+    workflow_py = os.path.join(temp_dir, "my_deployment", "workflow.py")
+    assert os.path.exists(workflow_py)
+    with open(workflow_py) as f:
+        assert f.read() == "print('hello')"
+
+    # AND the vellum.lock.json file is created
+    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 lock_data == {
+            "version": "1.0",
+            "workflows": [
+                {
+                    "module": "my_deployment",
+                    "workflow_sandbox_id": None,
+                    "ignore": None,
+                    "deployments": [],
+                }
+            ],
+        }
+
+
+def test_pull__both_workflow_sandbox_id_and_deployment(vellum_client):
+    # GIVEN a workflow sandbox id
+    workflow_sandbox_id = "87654321-0000-0000-0000-000000000000"
+
+    # AND a workflow deployment
+    workflow_deployment = "my-deployment"
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+
+    # 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-sandbox-id",
+            workflow_sandbox_id,
+            "--workflow-deployment",
+            workflow_deployment,
+        ],
+    )
+    os.chdir(current_dir)
+
+    # THEN the command returns successfully
+    assert result.exit_code == 1
+    assert "Cannot specify both workflow_sandbox_id and workflow_deployment" == str(result.exception)
+
+
 def test_pull__remove_missing_files(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
     temp_dir = mock_module.temp_dir

{vellum_ai-0.11.8 → vellum_ai-0.11.10}/ee/vellum_ee/workflows/display/nodes/base_node_vellum_display.py

@@ -30,6 +30,15 @@ class BaseNodeVellumDisplay(BaseNodeDisplay[NodeType]):
     def get_target_handle_id(self) -> UUID:
         return self._get_node_display_uuid("target_handle_id")
 
+    def get_target_handle_id_by_source_node_id(self, source_node_id: UUID) -> UUID:
+        """
+        In the vast majority of cases, nodes will only have a single target handle and can be retrieved independently
+        of the source node. However, in rare cases (such as legacy Merge nodes), this method can be overridden to
+        account for the case of retrieving one amongst multiple target handles on a node.
+        """
+
+        return self.get_target_handle_id()
+
     def get_source_handle_id(self, port_displays: Dict[Port, PortDisplay]) -> UUID:
         default_port = self._node.Ports.default
         default_port_display = port_displays[default_port]

{vellum_ai-0.11.8 → vellum_ai-0.11.10}/ee/vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -30,7 +30,10 @@ class BaseCodeExecutionNodeDisplay(BaseNodeVellumDisplay[_CodeExecutionNodeType]
         node_id = self.node_id
 
         node_file_path = inspect.getfile(node)
-        code = read_file_from_path(node_filepath=node_file_path, script_filepath=(raise_if_descriptor(node.filepath)))
+        code = read_file_from_path(
+            node_filepath=node_file_path,
+            script_filepath=(raise_if_descriptor(node.filepath)),  # type: ignore
+        )
         code_inputs = raise_if_descriptor(node.code_inputs)
 
         inputs = [

{vellum_ai-0.11.8 → vellum_ai-0.11.10}/ee/vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -12,7 +12,6 @@ from vellum.workflows.references import WorkflowInputReference
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.generics import WorkflowType
-from vellum_ee.workflows.display.nodes import BaseMergeNodeDisplay
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.types import PortDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -369,10 +368,7 @@ class VellumWorkflowDisplay(
         target_node_id = target_node_display.node_id
 
         target_handle_id: UUID
-        if isinstance(target_node_display, BaseMergeNodeDisplay):
-            target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(source_node_id)
-        else:
-            target_handle_id = target_node_display.get_target_handle_id()
+        target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(source_node_id)
 
         return self._generate_edge_display_from_source(
             source_node_id, source_handle_id, target_node_id, target_handle_id, overrides

{vellum_ai-0.11.8 → vellum_ai-0.11.10}/pyproject.toml

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

{vellum_ai-0.11.8 → vellum_ai-0.11.10}/src/vellum/client/core/client_wrapper.py

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

{vellum_ai-0.11.8 → vellum_ai-0.11.10}/src/vellum/workflows/nodes/displayable/code_execution_node/node.py

@@ -73,7 +73,8 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
     request_options: Optional[RequestOptions] = None - The request options to use for the custom script.
     """
 
-    filepath: ClassVar[str]
+    filepath: ClassVar[Optional[str]] = None
+    code: ClassVar[Optional[str]] = None
 
     code_inputs: ClassVar[EntityInputsInterface]
     runtime: CodeExecutionRuntime = "PYTHON_3_11_6"
@@ -190,6 +191,21 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
         return compiled_inputs
 
     def _resolve_code(self) -> str:
+        if self.code and self.filepath:
+            raise NodeException(
+                message="Cannot specify both `code` and `filepath` for a CodeExecutionNode",
+                code=VellumErrorCode.INVALID_INPUTS,
+            )
+
+        if self.code:
+            return self.code
+
+        if not self.filepath:
+            raise NodeException(
+                message="Must specify either `code` or `filepath` for a CodeExecutionNode",
+                code=VellumErrorCode.INVALID_INPUTS,
+            )
+
         root = inspect.getfile(self.__class__)
         code = read_file_from_path(node_filepath=root, script_filepath=self.filepath)
         if not code:

vellum_ai-0.11.10/src/vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py

@@ -0,0 +1,251 @@
+import pytest
+import os
+
+from vellum import CodeExecutorResponse, NumberVellumValue, StringInput
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.displayable.code_execution_node import CodeExecutionNode
+from vellum.workflows.references.vellum_secret import VellumSecretReference
+from vellum.workflows.state.base import BaseState, StateMeta
+
+
+def test_run_workflow__happy_path(vellum_client):
+    """Confirm that CodeExecutionNodes output the expected text and results when run."""
+
+    # GIVEN a node that subclasses CodeExecutionNode
+    class Inputs(BaseInputs):
+        word: str
+
+    class State(BaseState):
+        pass
+
+    fixture = os.path.abspath(os.path.join(__file__, "../fixtures/main.py"))
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        filepath = fixture
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "word": Inputs.word,
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=5),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(
+        state=State(
+            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
+        )
+    )
+    outputs = node.run()
+
+    # THEN the node should have produced the outputs we expect
+    assert outputs == {"result": 5, "log": "hello"}
+
+    # AND we should have invoked the Code with the expected inputs
+    vellum_client.execute_code.assert_called_once_with(
+        input_values=[
+            StringInput(name="word", value="hello"),
+        ],
+        code="""\
+def main(word: str) -> int:
+    print(word)  # noqa: T201
+    return len(word)
+""",
+        runtime="PYTHON_3_11_6",
+        output_type="NUMBER",
+        packages=[],
+        request_options=None,
+    )
+
+
+def test_run_workflow__code_attribute(vellum_client):
+    """Confirm that CodeExecutionNodes can use the `code` attribute to specify the code to execute."""
+
+    # GIVEN a node that subclasses CodeExecutionNode
+    class Inputs(BaseInputs):
+        word: str
+
+    class State(BaseState):
+        pass
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        code = """\
+def main(word: str) -> int:
+    print(word)  # noqa: T201
+    return len(word)
+"""
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "word": Inputs.word,
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=5),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(
+        state=State(
+            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
+        )
+    )
+    outputs = node.run()
+
+    # THEN the node should have produced the outputs we expect
+    assert outputs == {"result": 5, "log": "hello"}
+
+    # AND we should have invoked the Code with the expected inputs
+    vellum_client.execute_code.assert_called_once_with(
+        input_values=[
+            StringInput(name="word", value="hello"),
+        ],
+        code="""\
+def main(word: str) -> int:
+    print(word)  # noqa: T201
+    return len(word)
+""",
+        runtime="PYTHON_3_11_6",
+        output_type="NUMBER",
+        packages=[],
+        request_options=None,
+    )
+
+
+def test_run_workflow__code_and_filepath_defined(vellum_client):
+    """Confirm that CodeExecutionNodes raise an error if both `code` and `filepath` are defined."""
+
+    # GIVEN a node that subclasses CodeExecutionNode
+    class Inputs(BaseInputs):
+        word: str
+
+    class State(BaseState):
+        pass
+
+    fixture = os.path.abspath(os.path.join(__file__, "../fixtures/main.py"))
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        filepath = fixture
+        code = """\
+def main(word: str) -> int:
+    print(word)  # noqa: T201
+    return len(word)
+"""
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "word": Inputs.word,
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=5),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(
+        state=State(
+            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
+        )
+    )
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN the node should have produced the exception we expected
+    assert exc_info.value.message == "Cannot specify both `code` and `filepath` for a CodeExecutionNode"
+
+
+def test_run_workflow__code_and_filepath_not_defined(vellum_client):
+    """Confirm that CodeExecutionNodes raise an error if neither `code` nor `filepath` are defined."""
+
+    # GIVEN a node that subclasses CodeExecutionNode
+    class Inputs(BaseInputs):
+        word: str
+
+    class State(BaseState):
+        pass
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "word": Inputs.word,
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=5),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(
+        state=State(
+            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
+        )
+    )
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN the node should have produced the exception we expected
+    assert exc_info.value.message == "Must specify either `code` or `filepath` for a CodeExecutionNode"
+
+
+def test_run_workflow__vellum_secret(vellum_client):
+    """Confirm that CodeExecutionNodes can use Vellum Secrets"""
+
+    # GIVEN a node that subclasses CodeExecutionNode that references a Vellum Secret
+    class State(BaseState):
+        pass
+
+    fixture = os.path.abspath(os.path.join(__file__, "../fixtures/main.py"))
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        filepath = fixture
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "token": VellumSecretReference("OPENAI_API_KEY"),
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="",
+        output=NumberVellumValue(value=0),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(state=State())
+    outputs = node.run()
+
+    # THEN the node should have produced the outputs we expect
+    assert outputs == {"result": 0, "log": ""}
+
+    # AND we should have invoked the Code with the expected inputs
+    vellum_client.execute_code.assert_called_once_with(
+        input_values=[
+            {"name": "token", "type": "SECRET", "value": "OPENAI_API_KEY"},
+        ],
+        code="""\
+def main(word: str) -> int:
+    print(word)  # noqa: T201
+    return len(word)
+""",
+        runtime="PYTHON_3_11_6",
+        output_type="NUMBER",
+        packages=[],
+        request_options=None,
+    )