vellum-ai 0.13.18__py3-none-any.whl → 0.13.20__py3-none-any.whl

vellum_cli/config.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from dataclasses import field
 import json
 import os
@@ -88,6 +89,50 @@ class WorkflowConfig(UniversalBaseModel):
         )
 
 
+def merge_workflows_by_sandbox_id(
+    workflows: List[WorkflowConfig], other_workflows: List[WorkflowConfig]
+) -> List[WorkflowConfig]:
+    merged_workflows: List[WorkflowConfig] = []
+    for self_workflow in workflows:
+        if self_workflow.workflow_sandbox_id is None:
+            # If the user defines a workflow in the pyproject.toml without a sandbox_id,
+            # we merge the workflow with one of the ones in the lockfile.
+            other_workflow = next(
+                (
+                    other_workflow
+                    for other_workflow in other_workflows
+                    if self_workflow.workspace == other_workflow.workspace
+                ),
+                None,
+            )
+            if other_workflow is not None:
+                merged_workflows.append(self_workflow.merge(other_workflow))
+            else:
+                merged_workflows.append(self_workflow)
+        else:
+            # If the user defines a workflow in the pyproject.toml with a sandbox_id,
+            # we merge the workflow with one of the ones in the lockfile with the same sandbox_id.
+            other_workflow = next(
+                (
+                    other_workflow
+                    for other_workflow in other_workflows
+                    if self_workflow.workflow_sandbox_id == other_workflow.workflow_sandbox_id
+                ),
+                None,
+            )
+            if other_workflow is not None:
+                merged_workflows.append(self_workflow.merge(other_workflow))
+            else:
+                merged_workflows.append(self_workflow)
+
+    workflow_sandbox_ids_so_far = {workflow.workflow_sandbox_id for workflow in merged_workflows}
+    for other_workflow in other_workflows:
+        if other_workflow.workflow_sandbox_id not in workflow_sandbox_ids_so_far:
+            merged_workflows.append(other_workflow)
+
+    return merged_workflows
+
+
 class VellumCliConfig(UniversalBaseModel):
     version: Literal["1.0"] = "1.0"
     workflows: List[WorkflowConfig] = field(default_factory=list)
@@ -97,24 +142,31 @@ class VellumCliConfig(UniversalBaseModel):
         lockfile_path = os.path.join(os.getcwd(), LOCKFILE_PATH)
         with open(lockfile_path, "w") as f:
             json.dump(self.model_dump(), f, indent=2, cls=DefaultStateEncoder)
+            # Makes sure the file ends with a newline, consistent with most autoformatters
+            f.write("\n")
 
     def merge(self, other: "VellumCliConfig") -> "VellumCliConfig":
         if other.version != self.version:
             raise ValueError("Lockfile version mismatch")
 
-        self_workflow_by_module = {workflow.module: workflow for workflow in self.workflows}
-        other_workflow_by_module = {workflow.module: workflow for workflow in other.workflows}
-        all_modules = sorted(set(self_workflow_by_module.keys()).union(set(other_workflow_by_module.keys())))
+        self_workflows_by_module = self.get_workflows_by_module_mapping()
+        other_workflows_by_module = other.get_workflows_by_module_mapping()
+        all_modules = sorted(set(self_workflows_by_module.keys()).union(set(other_workflows_by_module.keys())))
         merged_workflows = []
         for module in all_modules:
-            self_workflow = self_workflow_by_module.get(module)
-            other_workflow = other_workflow_by_module.get(module)
-            if self_workflow and other_workflow:
-                merged_workflows.append(self_workflow.merge(other_workflow))
-            elif self_workflow:
-                merged_workflows.append(self_workflow)
-            elif other_workflow:
-                merged_workflows.append(other_workflow)
+            self_workflows = self_workflows_by_module.get(module)
+            other_workflows = other_workflows_by_module.get(module)
+            if self_workflows and other_workflows:
+                merged_workflows.extend(
+                    merge_workflows_by_sandbox_id(
+                        self_workflows,
+                        other_workflows,
+                    )
+                )
+            elif self_workflows:
+                merged_workflows.extend(self_workflows)
+            elif other_workflows:
+                merged_workflows.extend(other_workflows)
 
         self_workspace_by_name = {workspace.name: workspace for workspace in self.workspaces}
         other_workspace_by_name = {workspace.name: workspace for workspace in other.workspaces}
@@ -132,6 +184,12 @@ class VellumCliConfig(UniversalBaseModel):
 
         return VellumCliConfig(workflows=merged_workflows, workspaces=merged_workspaces, version=self.version)
 
+    def get_workflows_by_module_mapping(self) -> Dict[str, List[WorkflowConfig]]:
+        workflows_by_module = defaultdict(list)
+        for workflow in self.workflows:
+            workflows_by_module[workflow.module].append(workflow)
+        return workflows_by_module
+
 
 def load_vellum_cli_config(root_dir: Optional[str] = None) -> VellumCliConfig:
     if root_dir is None:

vellum_cli/tests/test_pull.py

@@ -428,7 +428,7 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                 "workflows": [
                     {
                         "module": module,
-                        "workflow_sandbox_id": "0edc07cd-45b9-43e8-99bc-1f181972a857",
+                        "workflow_sandbox_id": workflow_sandbox_id,
                         "ignore": "tests/*",
                         "deployments": [
                             {
@@ -646,3 +646,44 @@ def test_pull__module_not_in_config(vellum_client, mock_module):
                 "workspace": "default",
             }
         ]
+
+
+def test_pull__multiple_instances_of_same_module__keep_when_pulling_another_module(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    module = mock_module.module
+    temp_dir = mock_module.temp_dir
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
+
+    # AND the vellum lock file has two instances of some other module
+    lock_data = {
+        "workflows": [
+            {
+                "module": "some_other_module",
+                "workflow_sandbox_id": str(uuid4()),
+                "workspace": "default",
+            },
+            {
+                "module": "some_other_module",
+                "workflow_sandbox_id": str(uuid4()),
+                "workspace": "other",
+            },
+        ]
+    }
+    lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    with open(lock_json, "w") as f:
+        json.dump(lock_data, f)
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+
+    # WHEN the user runs the pull command on the new module
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", module, "--workflow-sandbox-id", workflow_sandbox_id])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0, (result.output, result.exception)
+
+    # AND the lockfile should have all three entries
+    with open(lock_json) as f:
+        lock_data = json.load(f)
+        assert len(lock_data["workflows"]) == 3

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

@@ -18,7 +18,6 @@ from .try_node import BaseTryNodeDisplay
 
 # All node display classes must be imported here to be registered in BaseNodeDisplay's node display registry
 __all__ = [
-    "BaseRetryNodeDisplay",
     "BaseAPINodeDisplay",
     "BaseCodeExecutionNodeDisplay",
     "BaseConditionalNodeDisplay",
@@ -31,6 +30,7 @@ __all__ = [
     "BaseMergeNodeDisplay",
     "BaseNoteNodeDisplay",
     "BasePromptDeploymentNodeDisplay",
+    "BaseRetryNodeDisplay",
     "BaseSearchNodeDisplay",
     "BaseSubworkflowDeploymentNodeDisplay",
     "BaseTemplatingNodeDisplay",

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

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import ClassVar, Dict, Generic, Optional, TypeVar
+from typing import Generic, Optional, TypeVar
 
 from vellum.workflows.nodes import SubworkflowDeploymentNode
 from vellum.workflows.types.core import JsonObject
@@ -15,7 +15,6 @@ _SubworkflowDeploymentNodeType = TypeVar("_SubworkflowDeploymentNodeType", bound
 class BaseSubworkflowDeploymentNodeDisplay(
     BaseNodeVellumDisplay[_SubworkflowDeploymentNodeType], Generic[_SubworkflowDeploymentNodeType]
 ):
-    subworkflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
@@ -30,7 +29,7 @@ class BaseSubworkflowDeploymentNodeDisplay(
                 input_name=variable_name,
                 value=variable_value,
                 display_context=display_context,
-                input_id=self.subworkflow_input_ids_by_name.get(variable_name),
+                input_id=self.node_input_ids_by_name.get(variable_name),
             )
             for variable_name, variable_value in subworkflow_inputs.items()
         ]

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

@@ -27,7 +27,7 @@ def create_node_input(
     input_name: str,
     value: Any,
     display_context: WorkflowDisplayContext,
-    input_id: Union[Optional[UUID], Optional[str]],
+    input_id: Union[Optional[UUID], Optional[str]] = None,
     pointer_type: Optional[Type[NodeInputValuePointerRule]] = ConstantValuePointer,
 ) -> NodeInput:
     input_id = str(input_id) if input_id else str(uuid4_from_hash(f"{node_id}|{input_name}"))

vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py

@@ -80,6 +80,11 @@ def test_serialize_node__retry(serialize_node):
                         "module": ["vellum", "workflows", "nodes", "core", "retry_node", "node"],
                     },
                     "attributes": [
+                        {
+                            "id": "8a07dc58-3fed-41d4-8ca6-31ee0bb86c61",
+                            "name": "delay",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
                         {
                             "id": "f388e93b-8c68-4f54-8577-bbd0c9091557",
                             "name": "max_attempts",

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

@@ -210,11 +210,11 @@ def test_serialize_workflow(vellum_client):
                     "target_handle_id": "06853542-e1a1-4a00-bd1e-4ac40f347b32",
                     "output_id": "9a37bf7d-484e-4725-903e-f3254df38a0a",
                     "output_type": "JSON",
-                    "node_input_id": "b544fc9a-2747-47e9-b28b-3e88d87b0f95",
+                    "node_input_id": "b49f0f85-37fc-4686-81a7-287c06634661",
                 },
                 "inputs": [
                     {
-                        "id": "b544fc9a-2747-47e9-b28b-3e88d87b0f95",
+                        "id": "b49f0f85-37fc-4686-81a7-287c06634661",
                         "key": "node_input",
                         "value": {
                             "rules": [
@@ -233,14 +233,7 @@ def test_serialize_workflow(vellum_client):
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
                 "base": {
                     "name": "FinalOutputNode",
-                    "module": [
-                        "vellum",
-                        "workflows",
-                        "nodes",
-                        "displayable",
-                        "final_output_node",
-                        "node",
-                    ],
+                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
                 },
                 "definition": None,
             },
@@ -253,11 +246,11 @@ def test_serialize_workflow(vellum_client):
                     "target_handle_id": "80d0894f-642e-4d2e-b43a-f236e7bedb3c",
                     "output_id": "5090e96d-5787-4a08-bf58-129101cf2548",
                     "output_type": "JSON",
-                    "node_input_id": "6828b56e-80b7-4699-b9dd-fd1f0820732e",
+                    "node_input_id": "5e892e5b-0004-4a04-bd2e-1ea9e0e5d3f9",
                 },
                 "inputs": [
                     {
-                        "id": "6828b56e-80b7-4699-b9dd-fd1f0820732e",
+                        "id": "5e892e5b-0004-4a04-bd2e-1ea9e0e5d3f9",
                         "key": "node_input",
                         "value": {
                             "rules": [
@@ -276,14 +269,7 @@ def test_serialize_workflow(vellum_client):
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
                 "base": {
                     "name": "FinalOutputNode",
-                    "module": [
-                        "vellum",
-                        "workflows",
-                        "nodes",
-                        "displayable",
-                        "final_output_node",
-                        "node",
-                    ],
+                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
                 },
                 "definition": None,
             },
@@ -296,11 +282,11 @@ def test_serialize_workflow(vellum_client):
                     "target_handle_id": "0c98c306-b519-40d7-8b05-321b1dfd7f11",
                     "output_id": "44ea8d75-e2a8-4627-85b1-8504b65d25c9",
                     "output_type": "NUMBER",
-                    "node_input_id": "500ff745-344f-4094-9425-48c4b40b7a5d",
+                    "node_input_id": "14345321-7e6b-4e2a-918a-7a5b0064f047",
                 },
                 "inputs": [
                     {
-                        "id": "500ff745-344f-4094-9425-48c4b40b7a5d",
+                        "id": "14345321-7e6b-4e2a-918a-7a5b0064f047",
                         "key": "node_input",
                         "value": {
                             "rules": [
@@ -319,14 +305,7 @@ def test_serialize_workflow(vellum_client):
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
                 "base": {
                     "name": "FinalOutputNode",
-                    "module": [
-                        "vellum",
-                        "workflows",
-                        "nodes",
-                        "displayable",
-                        "final_output_node",
-                        "node",
-                    ],
+                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
                 },
                 "definition": None,
             },

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

@@ -125,11 +125,11 @@ def test_serialize_workflow_with_filepath():
                     "target_handle_id": "30fb0f4a-61c3-49de-a0aa-7dfdcee6ea07",
                     "output_id": "f6a3e3e0-f83f-4491-8b7a-b20fddd7160c",
                     "output_type": "NUMBER",
-                    "node_input_id": "ae302487-ff2a-457a-81ed-9e0348e91833",
+                    "node_input_id": "529bdd20-0985-4c99-87dc-590907937c1d",
                 },
                 "inputs": [
                     {
-                        "id": "ae302487-ff2a-457a-81ed-9e0348e91833",
+                        "id": "529bdd20-0985-4c99-87dc-590907937c1d",
                         "key": "node_input",
                         "value": {
                             "rules": [
@@ -161,11 +161,11 @@ def test_serialize_workflow_with_filepath():
                     "target_handle_id": "1e126004-9de7-42c0-b1e1-87f9eb0642e2",
                     "output_id": "1cee930f-342f-421c-89fc-ff212b3764bb",
                     "output_type": "STRING",
-                    "node_input_id": "c6593516-ffc5-49a8-8a65-1038cccec3f8",
+                    "node_input_id": "09501b65-d9b3-4920-81d4-96f93c840667",
                 },
                 "inputs": [
                     {
-                        "id": "c6593516-ffc5-49a8-8a65-1038cccec3f8",
+                        "id": "09501b65-d9b3-4920-81d4-96f93c840667",
                         "key": "node_input",
                         "value": {
                             "rules": [
@@ -358,11 +358,11 @@ def test_serialize_workflow_with_code():
                     "target_handle_id": "de8f2cc2-8c32-4782-87d5-4eb5afcd42e3",
                     "output_id": "283d6849-f3ed-4beb-b261-cf70f90e8d10",
                     "output_type": "NUMBER",
-                    "node_input_id": "b38ba7a8-0b2a-4146-8d58-9fa0bcba8cd5",
+                    "node_input_id": "101d5222-8578-413a-be80-b1acf7559a0d",
                 },
                 "inputs": [
                     {
-                        "id": "b38ba7a8-0b2a-4146-8d58-9fa0bcba8cd5",
+                        "id": "101d5222-8578-413a-be80-b1acf7559a0d",
                         "key": "node_input",
                         "value": {
                             "rules": [
@@ -394,11 +394,11 @@ def test_serialize_workflow_with_code():
                     "target_handle_id": "6b7d7f2c-5cc8-4005-9e66-cdb2c97b1998",
                     "output_id": "4c136180-050b-4422-a7a4-2a1c6729042c",
                     "output_type": "STRING",
-                    "node_input_id": "76d49710-1ed0-4105-a1d7-9190c0408558",
+                    "node_input_id": "de7ea5c4-da2f-4cfd-9192-0257d0f9ba53",
                 },
                 "inputs": [
                     {
-                        "id": "76d49710-1ed0-4105-a1d7-9190c0408558",
+                        "id": "de7ea5c4-da2f-4cfd-9192-0257d0f9ba53",
                         "key": "node_input",
                         "value": {
                             "rules": [
@@ -601,29 +601,17 @@ def test_serialize_workflow__try_wrapped():
             {
                 "id": "af4fc1ef-7701-43df-b5e7-4f354f707db2",
                 "type": "TERMINAL",
-                "base": {
-                    "module": [
-                        "vellum",
-                        "workflows",
-                        "nodes",
-                        "displayable",
-                        "final_output_node",
-                        "node",
-                    ],
-                    "name": "FinalOutputNode",
-                },
-                "definition": None,
                 "data": {
                     "label": "Final Output",
                     "name": "log",
                     "target_handle_id": "d243df8d-46f6-4928-ac31-7c775c5d73a9",
                     "output_id": "5fbd27a0-9831-49c7-93c8-9c2a28c78696",
                     "output_type": "STRING",
-                    "node_input_id": "569e1c05-53fa-413b-abf0-0353eaa44208",
+                    "node_input_id": "a6ffaacf-1284-4e70-907e-042bb281e6fa",
                 },
                 "inputs": [
                     {
-                        "id": "569e1c05-53fa-413b-abf0-0353eaa44208",
+                        "id": "a6ffaacf-1284-4e70-907e-042bb281e6fa",
                         "key": "node_input",
                         "value": {
                             "rules": [
@@ -640,33 +628,26 @@ def test_serialize_workflow__try_wrapped():
                     }
                 ],
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
-            },
-            {
-                "id": "4cbfa5f7-fc12-4ab2-81cb-168c5caef4f0",
-                "type": "TERMINAL",
                 "base": {
-                    "module": [
-                        "vellum",
-                        "workflows",
-                        "nodes",
-                        "displayable",
-                        "final_output_node",
-                        "node",
-                    ],
                     "name": "FinalOutputNode",
+                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
                 },
                 "definition": None,
+            },
+            {
+                "id": "4cbfa5f7-fc12-4ab2-81cb-168c5caef4f0",
+                "type": "TERMINAL",
                 "data": {
                     "label": "Final Output",
                     "name": "result",
                     "target_handle_id": "9c43709e-25cb-4548-b840-3fcf6a1c9f3e",
                     "output_id": "400f9ffe-e700-4204-a810-e06123565947",
                     "output_type": "NUMBER",
-                    "node_input_id": "3d6cb7ef-985f-48f1-ad23-de49be60666a",
+                    "node_input_id": "2cbb7351-a3e8-4a6b-b33e-204b04da7ad6",
                 },
                 "inputs": [
                     {
-                        "id": "3d6cb7ef-985f-48f1-ad23-de49be60666a",
+                        "id": "2cbb7351-a3e8-4a6b-b33e-204b04da7ad6",
                         "key": "node_input",
                         "value": {
                             "rules": [
@@ -683,6 +664,11 @@ def test_serialize_workflow__try_wrapped():
                     }
                 ],
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "base": {
+                    "name": "FinalOutputNode",
+                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
+                },
+                "definition": None,
             },
         ],
         sorted(final_output_nodes, key=lambda x: x["id"], reverse=True),