vellum-ai 0.10.7__py3-none-any.whl → 0.10.9__py3-none-any.whl

vellum_cli/pull.py

@@ -7,28 +7,38 @@ from typing import Optional
 from dotenv import load_dotenv
 
 from vellum.workflows.vellum_client import create_vellum_client
-from vellum_cli.config import load_vellum_cli_config
+from vellum_cli.config import WorkflowConfig, load_vellum_cli_config
 from vellum_cli.logger import load_cli_logger
 
 
 def pull_command(
-    module: Optional[str], legacy_module: Optional[bool] = None, include_json: Optional[bool] = None
+    module: Optional[str] = None,
+    workflow_sandbox_id: Optional[str] = None,
+    include_json: Optional[bool] = None,
+    exclude_code: Optional[bool] = None,
 ) -> None:
     load_dotenv()
     logger = load_cli_logger()
     config = load_vellum_cli_config()
 
-    if not config.workflows:
-        raise ValueError("No Workflows found in project to pull.")
-
-    if len(config.workflows) > 1 and not module:
-        raise ValueError("Multiple workflows found in project to pull. Pulling only a single workflow is supported.")
-
     workflow_config = (
-        next((w for w in config.workflows if w.module == module), None) if module else config.workflows[0]
+        next((w for w in config.workflows if w.module == module), None)
+        if module
+        else (config.workflows[0] if config.workflows else None)
     )
+    save_lock_file = False
     if workflow_config is None:
-        raise ValueError(f"No workflow config for '{module}' found in project to push.")
+        if module:
+            raise ValueError(f"No workflow config for '{module}' found in project to pull.")
+        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)
+            save_lock_file = True
+        else:
+            raise ValueError("No workflow config found in project to pull from.")
 
     if not workflow_config.workflow_sandbox_id:
         raise ValueError("No workflow sandbox ID found in project to pull from.")
@@ -36,10 +46,10 @@ def pull_command(
     logger.info(f"Pulling workflow into {workflow_config.module}")
     client = create_vellum_client()
     query_parameters = {}
-    if legacy_module:
-        query_parameters["legacyModule"] = legacy_module
     if include_json:
         query_parameters["include_json"] = include_json
+    if exclude_code:
+        query_parameters["exclude_code"] = exclude_code
 
     response = client.workflows.pull(
         workflow_config.workflow_sandbox_id,
@@ -81,6 +91,11 @@ def pull_command(
                 target.write(source.read().decode("utf-8"))
 
     if include_json:
-        logger.warning("The pulled JSON representation of the Workflow should be used for debugging purposely only. Its schema should be considered unstable and subject to change at any time.")
+        logger.warning(
+            "The pulled JSON representation of the Workflow should be used for debugging purposely only. Its schema should be considered unstable and subject to change at any time."
+        )
+
+    if save_lock_file:
+        config.save()
 
     logger.info(f"Successfully pulled Workflow into {workflow_config.module}")

vellum_cli/tests/test_pull.py

@@ -69,8 +69,33 @@ def test_pull(vellum_client, mock_module):
     pull_command(module)
 
     # THEN the workflow.py file is written to the module directory
-    assert os.path.exists(os.path.join(temp_dir, *module.split("."), "workflow.py"))
-    with open(os.path.join(temp_dir, *module.split("."), "workflow.py")) as f:
+    workflow_py = os.path.join(temp_dir, *module.split("."), "workflow.py")
+    assert os.path.exists(workflow_py)
+    with open(workflow_py) as f:
+        assert f.read() == "print('hello')"
+
+
+def test_pull__sandbox_id_with_no_config(vellum_client):
+    # GIVEN a workflow sandbox id
+    workflow_sandbox_id = "87654321-0000-0000-0000-000000000000"
+
+    # 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 sandbox id and no module
+    pull_command(workflow_sandbox_id=workflow_sandbox_id)
+    os.chdir(current_dir)
+
+    # THEN the pull api is called with exclude_code=True
+    vellum_client.workflows.pull.assert_called_once()
+    workflow_py = os.path.join(temp_dir, "workflow_87654321", "workflow.py")
+    assert os.path.exists(workflow_py)
+    with open(workflow_py) as f:
         assert f.read() == "print('hello')"
 
 
@@ -168,3 +193,21 @@ def test_pull__include_json(vellum_client, mock_module):
     vellum_client.workflows.pull.assert_called_once()
     call_args = vellum_client.workflows.pull.call_args.kwargs
     assert call_args["request_options"]["additional_query_parameters"] == {"include_json": True}
+
+
+def test_pull__exclude_code(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    _, module = mock_module
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter(
+        [zip_file_map({"workflow.py": "print('hello')", "workflow.json": "{}"})]
+    )
+
+    # WHEN the user runs the pull command
+    pull_command(module, exclude_code=True)
+
+    # THEN the pull api is called with exclude_code=True
+    vellum_client.workflows.pull.assert_called_once()
+    call_args = vellum_client.workflows.pull.call_args.kwargs
+    assert call_args["request_options"]["additional_query_parameters"] == {"exclude_code": True}

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -1,7 +1,7 @@
 from functools import cached_property
 import inspect
 from uuid import UUID
-from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, TypeVar, get_args, get_origin, cast
+from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, TypeVar, cast, get_args, get_origin
 
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node

vellum_ee/workflows/display/nodes/vellum/__init__.py

@@ -1,6 +1,7 @@
 from .api_node import BaseAPINodeDisplay
 from .code_execution_node import BaseCodeExecutionNodeDisplay
 from .conditional_node import BaseConditionalNodeDisplay
+from .error_node import BaseErrorNodeDisplay
 from .final_output_node import BaseFinalOutputNodeDisplay
 from .guardrail_node import BaseGuardrailNodeDisplay
 from .inline_prompt_node import BaseInlinePromptNodeDisplay
@@ -16,19 +17,20 @@ from .try_node import BaseTryNodeDisplay
 
 # All node display classes must be imported here to be registered in BaseNodeDisplay's node display registry
 __all__ = [
+    "BaseAPINodeDisplay",
     "BaseCodeExecutionNodeDisplay",
     "BaseConditionalNodeDisplay",
+    "BaseErrorNodeDisplay",
+    "BaseFinalOutputNodeDisplay",
     "BaseGuardrailNodeDisplay",
     "BaseInlinePromptNodeDisplay",
     "BaseInlineSubworkflowNodeDisplay",
-    "BaseAPINodeDisplay",
     "BaseMapNodeDisplay",
     "BaseMergeNodeDisplay",
     "BaseNoteNodeDisplay",
+    "BasePromptDeploymentNodeDisplay",
     "BaseSearchNodeDisplay",
     "BaseSubworkflowDeploymentNodeDisplay",
     "BaseTemplatingNodeDisplay",
-    "BasePromptDeploymentNodeDisplay",
-    "BaseFinalOutputNodeDisplay",
-    "BaseTryNodeDisplay",
+    "BaseTryNodeDisplay"
 ]

vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import ClassVar, Generic, Optional, TypeVar
+from typing import ClassVar, Dict, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes.displayable.code_execution_node import CodeExecutionNode
 from vellum.workflows.nodes.displayable.code_execution_node.utils import read_file_from_path
@@ -20,6 +20,8 @@ class BaseCodeExecutionNodeDisplay(BaseNodeVellumDisplay[_CodeExecutionNodeType]
     output_id: ClassVar[Optional[UUID]] = None
     log_output_id: ClassVar[Optional[UUID]] = None
 
+    node_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
+
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
@@ -27,6 +29,19 @@ class BaseCodeExecutionNodeDisplay(BaseNodeVellumDisplay[_CodeExecutionNodeType]
         node_id = self.node_id
 
         code = read_file_from_path(raise_if_descriptor(node.filepath))
+        code_inputs = raise_if_descriptor(node.code_inputs)
+
+        inputs = [
+            create_node_input(
+                node_id=node_id,
+                input_name=variable_name,
+                value=variable_value,
+                display_context=display_context,
+                input_id=self.node_input_ids_by_name.get(variable_name),
+            )
+            for variable_name, variable_value in code_inputs.items()
+        ]
+
         code_node_input = create_node_input(
             node_id=node_id,
             input_name="code",
@@ -41,7 +56,7 @@ class BaseCodeExecutionNodeDisplay(BaseNodeVellumDisplay[_CodeExecutionNodeType]
             display_context=display_context,
             input_id=self.runtime_input_id,
         )
-        inputs = [code_node_input, runtime_node_input]
+        inputs.extend([code_node_input, runtime_node_input])
 
         packages = raise_if_descriptor(node.packages)
 

vellum_ee/workflows/display/nodes/vellum/error_node.py

@@ -0,0 +1,49 @@
+from uuid import UUID
+from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar
+
+from vellum.workflows.nodes import ErrorNode
+from vellum.workflows.types.core import EntityInputsInterface, Json, JsonObject
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+
+_ErrorNodeType = TypeVar("_ErrorNodeType", bound=ErrorNode)
+
+class BaseErrorNodeDisplay(BaseNodeVellumDisplay[_ErrorNodeType], Generic[_ErrorNodeType]):
+    error_output_id: ClassVar[Optional[UUID]] = None
+    error_inputs_by_name: ClassVar[Dict[str, Any]] = {}
+    name: ClassVar[str] = "error-node"
+
+    def serialize(
+            self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
+    ) -> JsonObject:
+        node = self._node
+        node_id = self.node_id
+        error_source_input_id = self.node_input_ids_by_name.get("error_source_input_id")
+
+        node_inputs = [
+            create_node_input(
+                node_id=node_id,
+                input_name=variable_name,
+                value=variable_value,
+                display_context=display_context,
+                input_id=self.node_input_ids_by_name.get(variable_name),
+            )
+            for variable_name, variable_value in self.error_inputs_by_name.items()
+        ]
+
+        return {
+            "id": str(node_id),
+            "type": "ERROR",
+            "inputs": [node_input.dict() for node_input in node_inputs],
+            "data": {
+                "name": self.name,
+                "label": self.label,
+                "source_handle_id": str(self.get_source_handle_id(display_context.port_displays)),
+                "target_handle_id": str(self.get_target_handle_id()),
+                "error_source_input_id": str(error_source_input_id),
+                "error_output_id": str(self.error_output_id),
+            },
+            "display_data": self.get_display_data().dict(),
+            "definition": self.get_definition().dict(),
+        }

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py

@@ -0,0 +1,203 @@
+from unittest import mock
+
+from deepdiff import DeepDiff
+
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.basic_error_node.workflow import BasicErrorNodeWorkflow
+
+
+def test_serialize_workflow():
+    # GIVEN a Workflow with an error node
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=BasicErrorNodeWorkflow
+    )
+
+    # TODO: Support serialization of BaseNode
+    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
+    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
+        mocked_serialize.return_value = {"type": "MOCKED"}
+        serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a serialized representation of the Workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "output_variables",
+    }
+
+    # AND its input variables should be what we expect
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 1
+    assert not DeepDiff(
+        [
+            {
+                "id": "5d9edd44-b35b-4bad-ad51-ccdfe8185ff5",
+                "key": "threshold",
+                "type": "NUMBER",
+                "default": None,
+                "required": True,
+                "extensions": {"color": None},
+            }
+        ],
+        input_variables,
+        ignore_order=True,
+    )
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 1
+    assert not DeepDiff(
+        [
+            {
+                "id": "04c5c6be-f5e1-41b8-b668-39e179790d9e",
+                "key": "final_value",
+                "type": "NUMBER",
+            }
+        ],
+        output_variables,
+        ignore_order=True,
+    )
+
+    # AND its raw data should be what we expect
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    assert workflow_raw_data.keys() == {"edges", "nodes", "display_data", "definition"}
+    assert len(workflow_raw_data["edges"]) == 4
+    assert len(workflow_raw_data["nodes"]) == 5
+
+    # AND each node should be serialized correctly
+    entrypoint_node = workflow_raw_data["nodes"][0]
+    assert entrypoint_node == {
+        "id": "10e90662-e998-421d-a5c9-ec16e37a8de1",
+        "type": "ENTRYPOINT",
+        "inputs": [],
+        "data": {
+            "label": "Entrypoint Node",
+            "source_handle_id": "7d86498b-84ed-4feb-8e62-2188058c2c4e",
+        },
+        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "definition": {
+            "name": "BaseNode",
+            "module": ["vellum", "workflows", "nodes", "bases", "base"],
+            "bases": [],
+        },
+    }
+
+    error_node, error_index = next(
+        (
+            (node, index)
+            for index, node in enumerate(workflow_raw_data["nodes"])
+            if node.get("data", {}).get("label") == "Fail Node"
+        ),
+        (None, None),
+    )
+    assert not DeepDiff(
+        {
+            "id": "5cf9c5e3-0eae-4daf-8d73-8b9536258eb9",
+            "type": "ERROR",
+            "inputs": [],
+            "data": {
+                "name": "error-node",
+                "label": "Fail Node",
+                "source_handle_id": "ca17d318-a0f5-4f7c-be6c-59c9dc1dd7ed",
+                "target_handle_id": "70c19f1c-309c-4a5d-ba65-664c0bb2fedf",
+                "error_source_input_id": "None",
+                "error_output_id": "None",
+            },
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "definition": {
+                "name": "FailNode",
+                "module": ["tests", "workflows", "basic_error_node", "workflow"],
+                "bases": [
+                    {
+                        "name": "ErrorNode",
+                        "module": [
+                            "vellum",
+                            "workflows",
+                            "nodes",
+                            "core",
+                            "error_node",
+                            "node",
+                        ],
+                    }
+                ],
+            },
+        },
+        error_node,
+        ignore_order=True,
+    )
+
+    mocked_base_nodes = [
+        node
+        for i, node in enumerate(workflow_raw_data["nodes"])
+        if i != error_index and i != 0 and i != len(workflow_raw_data["nodes"]) - 1
+    ]
+
+    assert not DeepDiff(
+        [
+            {
+                "type": "MOCKED",
+            },
+            {
+                "type": "MOCKED",
+            },
+        ],
+        mocked_base_nodes,
+    )
+
+    terminal_node = workflow_raw_data["nodes"][-1]
+    assert not DeepDiff(
+        {
+            "id": "e5fff999-80c7-4cbc-9d99-06c653f3ec77",
+            "type": "TERMINAL",
+            "data": {
+                "label": "Final Output",
+                "name": "final_value",
+                "target_handle_id": "b070e9bc-e9b7-46d3-8f5b-0b646bd25cf0",
+                "output_id": "04c5c6be-f5e1-41b8-b668-39e179790d9e",
+                "output_type": "NUMBER",
+                "node_input_id": "39ff42c9-eae8-432e-ad41-e208fba77027",
+            },
+            "inputs": [
+                {
+                    "id": "39ff42c9-eae8-432e-ad41-e208fba77027",
+                    "key": "node_input",
+                    "value": {
+                        "rules": [
+                            {
+                                "type": "NODE_OUTPUT",
+                                "data": {
+                                    "node_id": "1eee9b4e-531f-45f2-a4b9-42207fac2c33",
+                                    "output_id": "c6b017a4-25e9-4296-8d81-6aa4b3dad171",
+                                },
+                            }
+                        ],
+                        "combinator": "OR",
+                    },
+                }
+            ],
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "definition": {
+                "name": "FinalOutputNode",
+                "module": [
+                    "vellum",
+                    "workflows",
+                    "nodes",
+                    "displayable",
+                    "final_output_node",
+                    "node",
+                ],
+                "bases": [
+                    {
+                        "name": "BaseNode",
+                        "module": ["vellum", "workflows", "nodes", "bases", "base"],
+                        "bases": [],
+                    }
+                ],
+            },
+        },
+        terminal_node,
+    )

vellum/workflows/events/utils.py

@@ -1,5 +0,0 @@
-from vellum.workflows.events.workflow import WorkflowEvent
-
-
-def is_terminal_event(event: WorkflowEvent) -> bool:
-    return event.name in {"workflow.execution.fulfilled", "workflow.execution.rejected", "workflow.execution.paused"}