vellum-ai 0.10.6__py3-none-any.whl → 0.10.8__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/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/final_output_node.py

@@ -1,14 +1,16 @@
 from uuid import UUID
 from typing import Any, ClassVar, Generic, Optional, TypeVar
 
+from vellum.workflows.nodes.core.map_node.node import MapNode
+from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
+from vellum.workflows.references.output import OutputReference
+from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.utils import to_kebab_case
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
-from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
-from vellum.workflows.types.core import JsonObject
 
 _FinalOutputNodeType = TypeVar("_FinalOutputNodeType", bound=FinalOutputNode)
 

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

@@ -1,59 +1,32 @@
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, List, Optional, Type, TypeVar
+from typing import Collection, Dict, Generic, List, Optional, TypeVar, cast
 
-from vellum import VellumVariable
 from vellum.workflows.nodes import MapNode
 from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 _MapNodeType = TypeVar("_MapNodeType", bound=MapNode)
 
 
 class BaseMapNodeDisplay(BaseNodeVellumDisplay[_MapNodeType], Generic[_MapNodeType]):
-    workflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
-
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
 
-        workflow_inputs: List[VellumVariable] = []
         subworkflow = raise_if_descriptor(node.subworkflow)
-        for descriptor in subworkflow.get_inputs_class():
-            # In WaC it's always 'all_items'
-            # In Vellum it's always 'items'
-            variable_name = descriptor.name if descriptor.name != "all_items" else "items"
-            variable_id = str(
-                self.workflow_input_ids_by_name.get(variable_name)
-                or uuid4_from_hash(f"{self.node_id}|{variable_name}")
-            )
-            workflow_inputs.append(
-                VellumVariable(
-                    id=variable_id,
-                    key=variable_name,
-                    type=infer_vellum_variable_type(descriptor),
-                )
-            )
-
-        items_workflow_input = next(input for input in workflow_inputs if input.key == "items")
-        item_workflow_input = next(input for input in workflow_inputs if input.key == "item")
-        index_workflow_input = next(input for input in workflow_inputs if input.key == "index")
-
-        workflow_outputs = self._generate_workflow_outputs(node)
 
         items_node_input = create_node_input(
             node_id=node_id,
             input_name="items",
             value=node.items,
             display_context=display_context,
-            input_id=UUID(items_workflow_input.id),
+            input_id=self.node_input_ids_by_name.get("items"),
         )
         node_inputs = [items_node_input]
 
@@ -63,6 +36,19 @@ class BaseMapNodeDisplay(BaseNodeVellumDisplay[_MapNodeType], Generic[_MapNodeTy
         )
         serialized_subworkflow = subworkflow_display.serialize()
 
+        renamed_input_variables = []
+        for input_variable in cast(List[Dict[str, str]], serialized_subworkflow["input_variables"]):
+            if input_variable["key"] == "all_items":
+                renamed_item = { **input_variable, "key": "items" }
+                renamed_input_variables.append(renamed_item)
+            else:
+                renamed_input_variables.append(input_variable)
+
+        # Note: This must match the items input ID for the map node's node input
+        items_workflow_input_id = next(input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "items")
+        item_workflow_input_id = next(input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "item")
+        index_workflow_input_id = next(input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "index")
+
         return {
             "id": str(node_id),
             "type": "MAP",
@@ -74,27 +60,13 @@ class BaseMapNodeDisplay(BaseNodeVellumDisplay[_MapNodeType], Generic[_MapNodeTy
                 "target_handle_id": str(self.get_target_handle_id()),
                 "variant": "INLINE",
                 "workflow_raw_data": serialized_subworkflow["workflow_raw_data"],
-                "input_variables": [workflow_input.dict() for workflow_input in workflow_inputs],
-                "output_variables": [workflow_output.dict() for workflow_output in workflow_outputs],
+                "input_variables": cast(JsonObject, renamed_input_variables),
+                "output_variables": serialized_subworkflow["output_variables"],
                 "concurrency": raise_if_descriptor(node.concurrency),
-                "items_input_id": str(items_workflow_input.id),
-                "item_input_id": str(item_workflow_input.id),
-                "index_input_id": str(index_workflow_input.id),
+                "items_input_id": items_workflow_input_id,
+                "item_input_id": item_workflow_input_id,
+                "index_input_id": index_workflow_input_id,
             },
             "display_data": self.get_display_data().dict(),
             "definition": self.get_definition().dict(),
         }
-
-    def _generate_workflow_outputs(
-        self,
-        node: Type[MapNode],
-    ) -> List[VellumVariable]:
-        workflow_outputs: List[VellumVariable] = []
-        for output_descriptor in raise_if_descriptor(node.subworkflow).Outputs:  # type: ignore[union-attr]
-            node_output_display = self.get_node_output_display(output_descriptor)
-            output_type = infer_vellum_variable_type(output_descriptor)
-            workflow_outputs.append(
-                VellumVariable(id=str(node_output_display.id), key=node_output_display.name, type=output_type)
-            )
-
-        return workflow_outputs

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

@@ -13,30 +13,20 @@ from vellum.workflows.expressions.does_not_equal import DoesNotEqualExpression
 from vellum.workflows.expressions.ends_with import EndsWithExpression
 from vellum.workflows.expressions.equals import EqualsExpression
 from vellum.workflows.expressions.greater_than import GreaterThanExpression
-from vellum.workflows.expressions.greater_than_or_equal_to import (
-    GreaterThanOrEqualToExpression,
-)
+from vellum.workflows.expressions.greater_than_or_equal_to import GreaterThanOrEqualToExpression
 from vellum.workflows.expressions.in_ import InExpression
 from vellum.workflows.expressions.is_not_null import IsNotNullExpression
 from vellum.workflows.expressions.is_null import IsNullExpression
 from vellum.workflows.expressions.less_than import LessThanExpression
-from vellum.workflows.expressions.less_than_or_equal_to import (
-    LessThanOrEqualToExpression,
-)
+from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualToExpression
 from vellum.workflows.expressions.not_between import NotBetweenExpression
 from vellum.workflows.expressions.not_in import NotInExpression
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import (
-    BaseNodeVellumDisplay,
-)
+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 vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 from tests.workflows.basic_conditional_node.workflow import CategoryWorkflow
-from tests.workflows.basic_conditional_node.workflow_with_only_one_conditional_node import (
-    create_simple_workflow,
-)
+from tests.workflows.basic_conditional_node.workflow_with_only_one_conditional_node import create_simple_workflow
 
 
 def test_serialize_workflow():
@@ -965,7 +955,6 @@ def test_conditional_node_serialize_all_operators_with_lhs_and_rhs(
 
     # AND the conditional node should be what we expect
     conditional_node = workflow_raw_data["nodes"][1]
-    print(conditional_node)
     assert not DeepDiff(
         {
             "id": "a9143814-6bb0-4cb3-a817-4fc076417121",

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

@@ -228,15 +228,15 @@ def test_serialize_workflow():
                     },
                 },
                 "input_variables": [
-                    {"id": "0fe9dff5-7595-4ff3-a420-416f3fc6f809", "key": "item", "type": "JSON"},
-                    {"id": "27dfe8ba-8be7-4de9-a930-2ec3e375b149", "key": "index", "type": "NUMBER"},
-                    {"id": "4c0a109f-599e-4d04-a396-a51474fc2996", "key": "items", "type": "JSON"},
+                    {"id": "b29bb546-9bc8-4136-857d-8c7a464ba9d4", "key": "item", "type": "JSON", "required": True, "default": None, "extensions": {"color": None}},
+                    {"id": "17e7ca49-668f-450d-a792-e1f97d13db67", "key": "index", "type": "NUMBER", "required": True, "default": None, "extensions": {"color": None}},
+                    {"id": "d6fc6c7a-235f-4b98-86f3-e258d1198f93", "key": "items", "type": "JSON", "required": True, "default": None, "extensions": {"color": None}},
                 ],
-                "output_variables": [{"id": "56b34d4a-d3a6-4bdb-80a4-dbf644791274", "key": "count", "type": "NUMBER"}],
+                "output_variables": [{"id": "2a957315-fae0-4366-8a35-f0b315c5eade", "key": "count", "type": "NUMBER"}],
                 "concurrency": None,
-                "items_input_id": "4c0a109f-599e-4d04-a396-a51474fc2996",
-                "item_input_id": "0fe9dff5-7595-4ff3-a420-416f3fc6f809",
-                "index_input_id": "27dfe8ba-8be7-4de9-a930-2ec3e375b149",
+                "items_input_id": "d6fc6c7a-235f-4b98-86f3-e258d1198f93",
+                "item_input_id": "b29bb546-9bc8-4136-857d-8c7a464ba9d4",
+                "index_input_id": "17e7ca49-668f-450d-a792-e1f97d13db67",
             },
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "definition": {
@@ -288,7 +288,7 @@ def test_serialize_workflow():
                                 "type": "NODE_OUTPUT",
                                 "data": {
                                     "node_id": "bf83099a-40df-4445-b90d-1f6f1067ebe3",
-                                    "output_id": "56b34d4a-d3a6-4bdb-80a4-dbf644791274",
+                                    "output_id": "2a957315-fae0-4366-8a35-f0b315c5eade",
                                 },
                             }
                         ],
@@ -323,6 +323,9 @@ def test_serialize_workflow():
             },
         },
         workflow_raw_data["nodes"][2],
+        # TODO: Fix output ID not referencing map node workflow output
+        # https://app.shortcut.com/vellum/story/5667/fix-output-display-reference-on-map-nodes
+        exclude_regex_paths=r"root\['inputs'\]\[0\]\['value'\]\['rules'\]\[0\]\['data'\]\['output_id'\]",
     )
 
     # AND each edge should be serialized correctly

vellum_ee/workflows/display/utils/vellum.py

@@ -4,7 +4,9 @@ import typing
 from typing import Any, List, Union, cast
 
 from vellum import ChatMessage, SearchResult, SearchResultRequest, VellumVariableType
+from vellum.client.types.array_vellum_value import ArrayVellumValue
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.nodes.core.map_node.node import MapNodeItemType
 from vellum.workflows.references import OutputReference, WorkflowInputReference
 from vellum.workflows.references.execution_count import ExecutionCountReference
 from vellum.workflows.references.node import NodeReference
@@ -45,9 +47,8 @@ def infer_vellum_variable_type(value: Any) -> VellumVariableType:
                 raise ValueError(
                     f"Expected NodeReference {descriptor.name} to have an instance pointing to a descriptor"
                 )
-
             descriptor = descriptor.instance
-
+        
         inferred_type = primitive_type_to_vellum_variable_type(descriptor)
     else:
         vellum_variable_value = primitive_to_vellum_value(value)