vellum-ai 0.14.15__tar.gz → 0.14.16__tar.gz

{vellum_ai-0.14.15 → vellum_ai-0.14.16}/PKG-INFO

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

{vellum_ai-0.14.15 → vellum_ai-0.14.16}/ee/vellum_cli/image_push.py

@@ -8,7 +8,7 @@ import docker
 from docker import DockerClient
 from dotenv import load_dotenv
 
-from vellum.workflows.vellum_client import create_vellum_client
+from vellum.workflows.vellum_client import create_vellum_client, create_vellum_environment
 from vellum_cli.logger import load_cli_logger
 
 _SUPPORTED_ARCHITECTURE = "amd64"
@@ -19,6 +19,21 @@ def image_push_command(image: str, tags: Optional[List[str]] = None) -> None:
     logger = load_cli_logger()
     vellum_client = create_vellum_client()
 
+    # Check if we are self hosted by looking at our base url
+    api_url = create_vellum_environment().default
+    is_self_hosted = not api_url.endswith("vellum.ai") and "localhost:" not in api_url and "127.0.0.1" not in api_url
+    if is_self_hosted:
+        logger.info("Self hosted install detected...")
+
+    if is_self_hosted and "/" not in image:
+        logger.error(
+            "For adding images to your self hosted install you must include the "
+            "repository address in the provided image name. You must also have "
+            "already pushed the image to the docker repository that your self "
+            "hosted install is using."
+        )
+        exit(1)
+
     # We're using docker python SDK here instead of subprocess since it connects to the docker host directly
     # instead of using the command line so it seemed like it would possibly be a little more robust since
     # it might avoid peoples' wonky paths, unfortunately it doesn't support the manifest command which we need for
@@ -29,59 +44,78 @@ def image_push_command(image: str, tags: Optional[List[str]] = None) -> None:
     docker_client = docker.from_env()
     check_architecture(docker_client, image, logger)
 
-    logger.info("Authenticating...")
-    auth = vellum_client.container_images.docker_service_token()
-
-    docker_client.login(
-        username="oauth2accesstoken",
-        password=auth.access_token,
-        registry=auth.repository,
-    )
-
     repo_split = image.split("/")
     tag_split = repo_split[-1].split(":")
     image_name = tag_split[0]
     main_tag = tag_split[1] if len(tag_split) > 1 else "latest"
-
     all_tags = [main_tag, *(tags or [])]
-    for tag in all_tags:
-        vellum_image_name = f"{auth.repository}/{image_name}:{tag}"
-
-        docker_client.api.tag(image, vellum_image_name)
-
-        push_result = docker_client.images.push(repository=vellum_image_name, stream=True)
-
-        # Here were trying to mime the output you would get from a normal docker push, which
-        # the python sdk makes as hard as possible.
-        for raw_line in push_result:
-            try:
-                for sub_line in raw_line.decode("utf-8").split("\r\n"):
-                    line = json.loads(sub_line)
-                    error_message = line.get("errorDetail", {}).get("message")
-                    status = line.get("status")
-                    id = line.get("id", "")
-
-                    if error_message:
-                        logger.error(error_message)
-                        exit(1)
-                    elif status == "Waiting":
-                        continue
-                    elif status:
-                        logger.info(f"{id}{': ' if id else ''}{status}")
-                    else:
-                        logger.info(line)
-            except Exception:
-                continue
 
-    result = subprocess.run(
+    if is_self_hosted:
+        # Include the repo for self hosted since we skip pushing it to our internal repo and expect
+        # the user the push it themselves and provide us with the repo name.
+        image_name = image.split(":")[0]
+
+        logger.info("Checking if image already exists in repository...")
+        exists_result = subprocess.run(
+            ["docker", "manifest", "inspect", image],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+
+        if exists_result.returncode != 0:
+            logger.error(
+                "Image does not exist in repository. Push the image to the "
+                "repository your self hosted install is using and try again."
+            )
+            exit(1)
+    else:
+        logger.info("Authenticating...")
+        auth = vellum_client.container_images.docker_service_token()
+
+        docker_client.login(
+            username="oauth2accesstoken",
+            password=auth.access_token,
+            registry=auth.repository,
+        )
+
+        for tag in all_tags:
+            vellum_image_name = f"{auth.repository}/{image_name}:{tag}"
+
+            docker_client.api.tag(image, vellum_image_name)
+
+            push_result = docker_client.images.push(repository=vellum_image_name, stream=True)
+
+            # Here were trying to mime the output you would get from a normal docker push, which
+            # the python sdk makes as hard as possible.
+            for raw_line in push_result:
+                try:
+                    for sub_line in raw_line.decode("utf-8").split("\r\n"):
+                        line = json.loads(sub_line)
+                        error_message = line.get("errorDetail", {}).get("message")
+                        status = line.get("status")
+                        id = line.get("id", "")
+
+                        if error_message:
+                            logger.error(error_message)
+                            exit(1)
+                        elif status == "Waiting":
+                            continue
+                        elif status:
+                            logger.info(f"{id}{': ' if id else ''}{status}")
+                        else:
+                            logger.info(line)
+                except Exception:
+                    continue
+
+    inspect_result = subprocess.run(
         ["docker", "inspect", "--format='{{index .RepoDigests 0}}'", image],
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
     )
 
     sha = ""
-    if result.returncode == 0:
-        match = re.search(r"sha256[^']*", result.stdout.decode("utf-8"))
+    if inspect_result.returncode == 0:
+        match = re.search(r"sha256[^']*", inspect_result.stdout.decode("utf-8"))
         if match and match.group(0):
             sha = match.group(0)
 

vellum_ai-0.14.16/ee/vellum_cli/tests/test_image_push.py

@@ -0,0 +1,56 @@
+import subprocess
+from unittest.mock import MagicMock, patch
+
+from click.testing import CliRunner
+
+from vellum_cli import main as cli_main
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    # GIVEN a self hosted vellum api URL env var
+    monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
+
+    # Mock Docker client
+    mock_docker_client = MagicMock()
+    mock_docker_from_env.return_value = mock_docker_client
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"manifest"),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"sha256:hellosha"),
+    ]
+
+    # WHEN the user runs the image push command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myrepo.net/myimage:latest"])
+
+    # THEN the command exits successfully
+    assert result.exit_code == 0, result.output
+
+    # AND gives the success message
+    assert "Image successfully pushed" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push__self_hosted_blocks_repo(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    # GIVEN a self hosted vellum api URL env var
+    monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
+
+    # Mock Docker client
+    mock_docker_client = MagicMock()
+    mock_docker_from_env.return_value = mock_docker_client
+
+    # WHEN the user runs the image push command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage"])
+
+    # THEN the command exits unsuccessfully
+    assert result.exit_code == 1, result.output
+
+    # AND gives the error message for self hosted installs not including the repo
+    assert "For adding images to your self hosted install you must include" in result.output

{vellum_ai-0.14.15 → vellum_ai-0.14.16}/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py

@@ -21,20 +21,18 @@ class Inputs(BaseInputs):
     input: str
 
 
-@RetryNode.wrap(max_attempts=3)
-class InnerRetryGenericNode(BaseNode):
-    input = Inputs.input
-
-    class Outputs(BaseOutputs):
-        output: str
-
+def test_serialize_node__retry(serialize_node):
+    @RetryNode.wrap(max_attempts=3)
+    class InnerRetryGenericNode(BaseNode):
+        input = Inputs.input
 
-@BaseRetryNodeDisplay.wrap(max_attempts=3)
-class InnerRetryGenericNodeDisplay(BaseNodeDisplay[InnerRetryGenericNode]):
-    pass
+        class Outputs(BaseOutputs):
+            output: str
 
+    @BaseRetryNodeDisplay.wrap(max_attempts=3)
+    class InnerRetryGenericNodeDisplay(BaseNodeDisplay[InnerRetryGenericNode]):
+        pass
 
-def test_serialize_node__retry(serialize_node):
     input_id = uuid4()
     serialized_node = serialize_node(
         node_class=InnerRetryGenericNode,
@@ -49,8 +47,8 @@ def test_serialize_node__retry(serialize_node):
     )
     assert not DeepDiff(
         {
-            "id": "f2a95e79-7d4b-47ad-b986-4f648297ec65",
-            "label": "InnerRetryGenericNode",
+            "id": "188b50aa-e518-4b7b-a5e0-e2585fb1d7b5",
+            "label": "test_serialize_node__retry.<locals>.InnerRetryGenericNode",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -66,8 +64,8 @@ def test_serialize_node__retry(serialize_node):
                     "test_adornments_serialization",
                 ],
             },
-            "trigger": {"id": "af9ba01c-4cde-4632-9aa1-7673b42e7bd8", "merge_behavior": "AWAIT_ATTRIBUTES"},
-            "ports": [{"id": "c2ecc6c0-f353-4495-9b93-a61a47248556", "name": "default", "type": "DEFAULT"}],
+            "trigger": {"id": "d38a83bf-23d1-4f9d-a875-a08dc27cf397", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "ports": [{"id": "078650c9-f775-4cd0-a08c-23af9983a361", "name": "default", "type": "DEFAULT"}],
             "adornments": [
                 {
                     "id": "5be7d260-74f7-4734-b31b-a46a94539586",
@@ -102,13 +100,13 @@ def test_serialize_node__retry(serialize_node):
             ],
             "attributes": [
                 {
-                    "id": "c363daa7-9482-4c0e-aee8-faa080602ee3",
+                    "id": "278df25e-58b5-43c3-b346-cf6444d893a5",
                     "name": "input",
                     "value": {"type": "WORKFLOW_INPUT", "input_variable_id": str(input_id)},
                 }
             ],
             "outputs": [
-                {"id": "8aaf6cd8-3fa5-4f17-a60f-ec7da5ec6498", "name": "output", "type": "STRING", "value": None}
+                {"id": "dc89dc0d-c0bd-47fd-88aa-ec7b262aa2f1", "name": "output", "type": "STRING", "value": None}
             ],
         },
         serialized_node,
@@ -136,20 +134,18 @@ def test_serialize_node__retry__no_display():
     assert exec_config is not None
 
 
-@TryNode.wrap()
-class InnerTryGenericNode(BaseNode):
-    input = Inputs.input
-
-    class Outputs(BaseOutputs):
-        output: str
-
+def test_serialize_node__try(serialize_node):
+    @TryNode.wrap()
+    class InnerTryGenericNode(BaseNode):
+        input = Inputs.input
 
-@BaseTryNodeDisplay.wrap()
-class InnerTryGenericNodeDisplay(BaseNodeDisplay[InnerTryGenericNode]):
-    pass
+        class Outputs(BaseOutputs):
+            output: str
 
+    @BaseTryNodeDisplay.wrap()
+    class InnerTryGenericNodeDisplay(BaseNodeDisplay[InnerTryGenericNode]):
+        pass
 
-def test_serialize_node__try(serialize_node):
     input_id = uuid4()
     serialized_node = serialize_node(
         base_class=BaseNodeVellumDisplay,
@@ -163,7 +159,7 @@ def test_serialize_node__try(serialize_node):
     assert not DeepDiff(
         {
             "id": str(InnerTryGenericNode.__wrapped_node__.__id__),
-            "label": "InnerTryGenericNode",
+            "label": "test_serialize_node__try.<locals>.InnerTryGenericNode",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -179,8 +175,8 @@ def test_serialize_node__try(serialize_node):
                     "test_adornments_serialization",
                 ],
             },
-            "trigger": {"id": "741f7f75-e921-47a9-8c05-9e66640d0866", "merge_behavior": "AWAIT_ATTRIBUTES"},
-            "ports": [{"id": "1b8f8ab5-a656-4015-926c-80655bbd9cb8", "name": "default", "type": "DEFAULT"}],
+            "trigger": {"id": "16bc1522-c408-47ad-9a22-0ef136384abf", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "ports": [{"id": "8d25f244-4b12-4f8b-b202-8948698679a0", "name": "default", "type": "DEFAULT"}],
             "adornments": [
                 {
                     "id": "3344083c-a32c-4a32-920b-0fb5093448fa",
@@ -200,13 +196,13 @@ def test_serialize_node__try(serialize_node):
             ],
             "attributes": [
                 {
-                    "id": "4d8b4c2c-4f92-4c7a-abf0-b9c88a15a790",
+                    "id": "51aa0077-4060-496b-8e2e-e79d56ee6a32",
                     "name": "input",
                     "value": {"type": "WORKFLOW_INPUT", "input_variable_id": str(input_id)},
                 }
             ],
             "outputs": [
-                {"id": "63ba929b-bf79-44ee-bd1f-d259dbe8d48e", "name": "output", "type": "STRING", "value": None}
+                {"id": "ce9f8b86-6d26-4c03-8bfa-a31aa2cd97f1", "name": "output", "type": "STRING", "value": None}
             ],
         },
         serialized_node,