vellum-ai 0.14.32__tar.gz → 0.14.34__tar.gz

{vellum_ai-0.14.32 → vellum_ai-0.14.34}/PKG-INFO

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

{vellum_ai-0.14.32 → vellum_ai-0.14.34}/ee/vellum_cli/ping.py

@@ -14,7 +14,12 @@ def ping_command():
     try:
         workspace = client.workspaces.workspace_identity()
         organization = client.organizations.organization_identity()
-    except ApiError:
+    except ApiError as e:
+        # If user did not provide an API key, we will get a 403 error
+        if e.status_code == 401 or e.status_code == 403:
+            raise Exception(
+                "Please make sure your `VELLUM_API_KEY` environment variable is set correctly."  # noqa: E501
+            ) from e
         raise Exception(
             "The API we tried to ping returned an invalid response. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
         )

{vellum_ai-0.14.32 → vellum_ai-0.14.34}/ee/vellum_cli/pull.py

@@ -2,6 +2,7 @@ import io
 import json
 import os
 from pathlib import Path
+from uuid import UUID
 import zipfile
 from typing import Optional
 
@@ -10,8 +11,9 @@ from pydash import snake_case
 
 from vellum.client.core.api_error import ApiError
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.utils.uuid import is_valid_uuid
 from vellum.workflows.vellum_client import create_vellum_client
-from vellum_cli.config import VellumCliConfig, WorkflowConfig, load_vellum_cli_config
+from vellum_cli.config import VellumCliConfig, WorkflowConfig, WorkflowDeploymentConfig, load_vellum_cli_config
 from vellum_cli.logger import load_cli_logger
 
 ERROR_LOG_FILE_NAME = "error.log"
@@ -31,6 +33,7 @@ class RunnerConfig(UniversalBaseModel):
 class PullContentsMetadata(UniversalBaseModel):
     label: Optional[str] = None
     runner_config: Optional[RunnerConfig] = None
+    deployment_id: Optional[UUID] = None
     deployment_name: Optional[str] = None
 
 
@@ -79,6 +82,29 @@ def _resolve_workflow_config(
             pk=workflow_config.workflow_sandbox_id,
         )
     elif workflow_deployment:
+        if is_valid_uuid(workflow_deployment):
+            # name may also be a valid UUID
+            workflow_config = next(
+                (
+                    w
+                    for w in config.workflows
+                    if w.deployments
+                    and (
+                        str(w.deployments[0].id) == workflow_deployment or w.deployments[0].name == workflow_deployment
+                    )
+                ),
+                None,
+            )
+        else:
+            workflow_config = next(
+                (w for w in config.workflows if w.deployments and w.deployments[0].name == workflow_deployment), None
+            )
+        if workflow_config:
+            return WorkflowConfigResolutionResult(
+                workflow_config=workflow_config,
+                pk=workflow_deployment,
+            )
+
         workflow_config = WorkflowConfig(
             module="",
         )
@@ -177,6 +203,24 @@ def pull_command(
                 if not workflow_config.module and pull_contents_metadata.label:
                     workflow_config.module = snake_case(pull_contents_metadata.label)
 
+                # Save or update the deployment info when pulling with --workflow-deployment
+                if workflow_deployment:
+                    workflow_deployment_id = pull_contents_metadata.deployment_id
+                    existing_deployment = next(
+                        (d for d in workflow_config.deployments if d.id == workflow_deployment_id), None
+                    )
+
+                    if existing_deployment:
+                        if pull_contents_metadata.label:
+                            existing_deployment.label = pull_contents_metadata.label
+                    else:
+                        deployment_config = WorkflowDeploymentConfig(
+                            id=workflow_deployment_id,
+                            label=pull_contents_metadata.label,
+                            name=pull_contents_metadata.deployment_name,
+                        )
+                        workflow_config.deployments.append(deployment_config)
+
             if not workflow_config.module:
                 raise ValueError(f"Failed to resolve a module name for Workflow {pk}")
 

{vellum_ai-0.14.32 → vellum_ai-0.14.34}/ee/vellum_cli/tests/test_ping.py

@@ -2,6 +2,7 @@ from datetime import datetime
 
 from click.testing import CliRunner
 
+from vellum.client.core.api_error import ApiError
 from vellum.client.types.organization_read import OrganizationRead
 from vellum.client.types.workspace_read import WorkspaceRead
 from vellum_cli import main as cli_main
@@ -51,9 +52,6 @@ def test_ping__error_path(vellum_client):
     # GIVEN a cli
     runner = CliRunner()
 
-    # Mock the workspace_identity method to raise ApiError directly
-    from vellum.client.core.api_error import ApiError
-
     vellum_client.workspaces.workspace_identity.side_effect = ApiError(status_code=400, body="Invalid JSON")
 
     # WHEN the user runs the ping command
@@ -65,3 +63,20 @@ def test_ping__error_path(vellum_client):
         str(result.exception)
         == "The API we tried to ping returned an invalid response. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
     )
+
+
+def test_ping__unauthorized_error_path(vellum_client):
+    # GIVEN a cli
+    runner = CliRunner()
+
+    # GIVEN an unauthorized error with the error message from the API
+    vellum_client.workspaces.workspace_identity.side_effect = ApiError(
+        status_code=403, body={"detail": "Authentication credentials were not provided."}
+    )
+
+    # WHEN the user runs the ping command
+    result = runner.invoke(cli_main, ["ping"])
+
+    # THEN the command returns an error
+    assert result.exit_code == 1
+    assert str(result.exception) == "Please make sure your `VELLUM_API_KEY` environment variable is set correctly."

{vellum_ai-0.14.32 → vellum_ai-0.14.34}/ee/vellum_cli/tests/test_pull.py

@@ -337,6 +337,7 @@ def test_pull__sandbox_id_with_other_workflow_configured(vellum_client, mock_mod
 def test_pull__workflow_deployment_with_no_config(vellum_client):
     # GIVEN a workflow deployment
     workflow_deployment = "my-deployment"
+    deployment_id = str(uuid4())
 
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter(
@@ -344,7 +345,9 @@ def test_pull__workflow_deployment_with_no_config(vellum_client):
             _zip_file_map(
                 {
                     "workflow.py": "print('hello')",
-                    "metadata.json": json.dumps({"deployment_name": workflow_deployment, "label": "Some Label"}),
+                    "metadata.json": json.dumps(
+                        {"deployment_id": deployment_id, "deployment_name": workflow_deployment, "label": "Some Label"}
+                    ),
                 }
             )
         ]
@@ -382,7 +385,15 @@ def test_pull__workflow_deployment_with_no_config(vellum_client):
                     "module": "my_deployment",
                     "workflow_sandbox_id": None,
                     "ignore": None,
-                    "deployments": [],
+                    "deployments": [
+                        {
+                            "id": deployment_id,
+                            "label": "Some Label",
+                            "name": "my-deployment",
+                            "description": None,
+                            "release_tags": None,
+                        }
+                    ],
                     "container_image_tag": None,
                     "container_image_name": None,
                     "workspace": "default",
@@ -937,3 +948,313 @@ def test_pull__unauthorized_error_path(vellum_client):
     # THEN the command returns an error
     assert result.exit_code == 1
     assert str(result.exception) == "Please make sure your `VELLUM_API_KEY` environment variable is set correctly."
+
+
+@pytest.mark.parametrize(
+    "workflow_deployment",
+    [
+        "test-workflow-deployment-id",
+        str(uuid4()),
+    ],
+)
+def test_pull__workflow_deployment_adds_deployment_to_config(vellum_client, workflow_deployment):
+    # GIVEN a workflow deployment ID
+    deployment_id = str(uuid4())  # config will always use the deployment_id return from the API
+    deployment_name = "Test Deployment"
+    deployment_label = "Test Label"
+
+    # AND the workflow pull API call returns a zip file with metadata
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {
+                    "workflow.py": "print('hello')",
+                    "metadata.json": json.dumps(
+                        {
+                            "deployment_id": deployment_id,
+                            "deployment_name": deployment_name,
+                            "label": deployment_label,
+                        }
+                    ),
+                }
+            )
+        ]
+    )
+
+    # 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])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the deployment is saved in the config
+    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 len(lock_data["workflows"]) == 1
+        assert len(lock_data["workflows"][0]["deployments"]) == 1
+        deployment = lock_data["workflows"][0]["deployments"][0]
+        assert str(deployment["id"]) == deployment_id
+        assert deployment["name"] == deployment_name
+        assert deployment["label"] == deployment_label
+
+    os.chdir(current_dir)
+
+
+def test_pull__workflow_deployment_name_is_uuid(vellum_client):
+    # GIVEN a workflow deployment name that is a valid UUID
+    deployment_id = str(uuid4())
+    deployment_name = str(uuid4())
+
+    # AND an existing configuration with this deployment
+    current_dir = os.getcwd()
+    temp_dir = tempfile.mkdtemp()
+    os.chdir(temp_dir)
+
+    # Create initial config with a deployment
+    vellum_lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    with open(vellum_lock_json, "w") as f:
+        json.dump(
+            {
+                "version": "1.0",
+                "workflows": [
+                    {
+                        "module": "test_workflow",
+                        "workflow_sandbox_id": None,
+                        "ignore": None,
+                        "deployments": [
+                            {
+                                "id": deployment_id,
+                                "label": "Test Label",
+                                "name": deployment_name,
+                                "description": None,
+                                "release_tags": None,
+                            }
+                        ],
+                        "container_image_name": None,
+                        "container_image_tag": None,
+                        "workspace": "default",
+                        "target_directory": None,
+                    }
+                ],
+                "workspaces": [],
+            },
+            f,
+        )
+
+    # AND the workflow pull API call returns a zip file with updated metadata
+    updated_label = "Updated Label"
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {
+                    "workflow.py": "print('hello')",
+                    "metadata.json": json.dumps(
+                        {
+                            "deployment_id": deployment_id,
+                            "deployment_name": deployment_name,
+                            "label": updated_label,
+                        }
+                    ),
+                }
+            )
+        ]
+    )
+
+    # WHEN the user runs the pull command with the workflow deployment
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-deployment", deployment_name])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the deployment info is updated in the config
+    with open(vellum_lock_json) as f:
+        lock_data = json.loads(f.read())
+        assert len(lock_data["workflows"]) == 1
+        assert len(lock_data["workflows"][0]["deployments"]) == 1
+        deployment = lock_data["workflows"][0]["deployments"][0]
+        assert str(deployment["id"]) == deployment_id
+        assert deployment["name"] == deployment_name
+        assert deployment["label"] == updated_label
+
+    os.chdir(current_dir)
+
+
+@pytest.mark.parametrize("get_identifier", [(lambda d: d), (lambda d: "Test Name")])
+def test_pull__workflow_deployment_updates_existing_deployment(vellum_client, get_identifier):
+    """
+    This test is to ensure that the deployment info is updated in the config
+    when the user runs the pull command with the workflow deployment
+
+    get_identifier is a function that returns the identifier to use for the deployment
+    it can be the deployment_id or the deployment_name
+    """
+    # GIVEN a workflow deployment id and name
+    deployment_id = str(uuid4())
+    deployment_name = "Test Name"
+
+    # AND an existing configuration with this deployment
+    current_dir = os.getcwd()
+    temp_dir = tempfile.mkdtemp()
+    os.chdir(temp_dir)
+
+    # Create initial config with a deployment
+    vellum_lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    with open(vellum_lock_json, "w") as f:
+        json.dump(
+            {
+                "version": "1.0",
+                "workflows": [
+                    {
+                        "module": "test_workflow",
+                        "workflow_sandbox_id": None,
+                        "ignore": None,
+                        "deployments": [
+                            {
+                                "id": deployment_id,
+                                "label": "Test Label",
+                                "name": deployment_name,
+                                "description": None,
+                                "release_tags": None,
+                            }
+                        ],
+                        "container_image_name": None,
+                        "container_image_tag": None,
+                        "workspace": "default",
+                        "target_directory": None,
+                    }
+                ],
+                "workspaces": [],
+            },
+            f,
+        )
+
+    # AND the workflow pull API call returns a zip file with updated metadata
+    updated_label = "Updated Label"
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {
+                    "workflow.py": "print('hello')",
+                    "metadata.json": json.dumps(
+                        {
+                            "deployment_id": deployment_id,
+                            "deployment_name": deployment_name,
+                            "label": updated_label,
+                        }
+                    ),
+                }
+            )
+        ]
+    )
+
+    # WHEN the user runs the pull command with the workflow deployment
+    runner = CliRunner()
+    identifier_to_use = get_identifier(deployment_id)
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-deployment", identifier_to_use])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the deployment info is updated in the config
+    with open(vellum_lock_json) as f:
+        lock_data = json.loads(f.read())
+        assert len(lock_data["workflows"]) == 1
+        assert len(lock_data["workflows"][0]["deployments"]) == 1
+        deployment = lock_data["workflows"][0]["deployments"][0]
+        assert str(deployment["id"]) == deployment_id
+        assert deployment["name"] == deployment_name
+        assert deployment["label"] == updated_label
+
+    os.chdir(current_dir)
+
+
+def test_pull__workflow_deployment_with_name_and_id(vellum_client):
+    """
+    This test is to ensure that pulling with id and name will not add a new deployment to the config
+    """
+    # GIVEN a workflow deployment ID
+    deployment_id = str(uuid4())  # config will always use the deployment_id return from the API
+    deployment_name = "Test Deployment"
+    deployment_label = "Test Label"
+
+    # AND the workflow pull API call returns a zip file with metadata
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {
+                    "workflow.py": "print('hello')",
+                    "metadata.json": json.dumps(
+                        {
+                            "deployment_id": deployment_id,
+                            "deployment_name": deployment_name,
+                            "label": deployment_label,
+                        }
+                    ),
+                }
+            )
+        ]
+    )
+
+    # 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", deployment_id])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the deployment is saved in the config
+    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 len(lock_data["workflows"]) == 1
+        assert len(lock_data["workflows"][0]["deployments"]) == 1
+        deployment = lock_data["workflows"][0]["deployments"][0]
+        assert str(deployment["id"]) == deployment_id
+        assert deployment["name"] == deployment_name
+        assert deployment["label"] == deployment_label
+
+    os.chdir(current_dir)
+
+    # AND pull with name will not add a new deployment to the config
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {
+                    "workflow.py": "print('hello')",
+                    "metadata.json": json.dumps(
+                        {
+                            "deployment_id": deployment_id,
+                            "deployment_name": deployment_name,
+                            "label": deployment_label,
+                        }
+                    ),
+                }
+            )
+        ]
+    )
+
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-deployment", deployment_name])
+    assert result.exit_code == 0
+    with open(vellum_lock_json) as f:
+        lock_data = json.loads(f.read())
+        assert len(lock_data["workflows"][0]["deployments"]) == 1
+        assert lock_data["workflows"][0]["deployments"][0]["id"] == deployment_id
+        assert lock_data["workflows"][0]["deployments"][0]["name"] == deployment_name
+        assert lock_data["workflows"][0]["deployments"][0]["label"] == deployment_label
+
+    os.chdir(current_dir)

{vellum_ai-0.14.32 → vellum_ai-0.14.34}/ee/vellum_ee/workflows/display/nodes/vellum/guardrail_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 GuardrailNode
 from vellum.workflows.types.core import JsonObject
@@ -12,8 +12,6 @@ _GuardrailNodeType = TypeVar("_GuardrailNodeType", bound=GuardrailNode)
 
 
 class BaseGuardrailNodeDisplay(BaseNodeVellumDisplay[_GuardrailNodeType], Generic[_GuardrailNodeType]):
-    metric_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
-
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
@@ -27,7 +25,8 @@ class BaseGuardrailNodeDisplay(BaseNodeVellumDisplay[_GuardrailNodeType], Generi
                 input_name=variable_name,
                 value=variable_value,
                 display_context=display_context,
-                input_id=self.metric_input_ids_by_name.get(variable_name),
+                input_id=self.node_input_ids_by_name.get(f"{GuardrailNode.metric_inputs.name}.{variable_name}")
+                or self.node_input_ids_by_name.get(variable_name),
             )
             for variable_name, variable_value in metric_inputs.items()
         ]

{vellum_ai-0.14.32 → vellum_ai-0.14.34}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_api_node_serialization.py

@@ -7,7 +7,7 @@ from vellum import WorkspaceSecretRead
 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_vellum_api_node.workflow import SimpleAPIWorkflow
+from tests.workflows.basic_api_node.workflow import SimpleAPIWorkflow
 
 
 def test_serialize_workflow(vellum_client):
@@ -99,7 +99,7 @@ def test_serialize_workflow(vellum_client):
                     "id": "c3a17ceb-e201-4025-b18a-9162aac7705e",
                     "key": "body",
                     "value": {
-                        "rules": [{"type": "CONSTANT_VALUE", "data": {"type": "JSON", "value": None}}],
+                        "rules": [{"type": "CONSTANT_VALUE", "data": {"type": "JSON", "value": {"key": "value"}}}],
                         "combinator": "OR",
                     },
                 },
@@ -193,7 +193,7 @@ def test_serialize_workflow(vellum_client):
             },
             "definition": {
                 "name": "SimpleAPINode",
-                "module": ["tests", "workflows", "basic_vellum_api_node", "workflow"],
+                "module": ["tests", "workflows", "basic_api_node", "workflow"],
             },
         },
         api_node,
@@ -373,7 +373,7 @@ def test_serialize_workflow(vellum_client):
         "module": [
             "tests",
             "workflows",
-            "basic_vellum_api_node",
+            "basic_api_node",
             "workflow",
         ],
     }

{vellum_ai-0.14.32 → vellum_ai-0.14.34}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_guardrail_node_serialization.py

@@ -73,7 +73,7 @@ def test_serialize_workflow():
         "type": "METRIC",
         "inputs": [
             {
-                "id": "3ec00ee3-c068-4d41-9488-87b7778a649e",
+                "id": "e5e9757e-27db-4200-a650-2f7bb2cd61f5",
                 "key": "expected",
                 "value": {
                     "rules": [
@@ -86,7 +86,7 @@ def test_serialize_workflow():
                 },
             },
             {
-                "id": "a43ce7ba-1685-4977-a34d-65580c42853f",
+                "id": "f5cd9336-e8e9-429e-b7bc-785402e31bfb",
                 "key": "actual",
                 "value": {
                     "rules": [