vellum-ai 0.14.65__py3-none-any.whl → 0.14.66__py3-none-any.whl

vellum/client/README.md

@@ -20,7 +20,7 @@ pip install vellum-ai
 
 ## Reference
 
-A full reference for this library is available [here](./reference.md).
+A full reference for this library is available [here](https://github.com/vellum-ai/vellum-python-sdks/blob/HEAD/./reference.md).
 
 ## Usage
 

vellum/client/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.14.65",
+            "X-Fern-SDK-Version": "0.14.66",
         }
         headers["X-API-KEY"] = self.api_key
         return headers

vellum/client/types/document_read.py

@@ -32,7 +32,6 @@ class DocumentRead(UniversalBaseModel):
     """
 
     original_file_url: typing.Optional[str] = None
-    processed_file_url: typing.Optional[str] = None
     document_to_document_indexes: typing.List[DocumentDocumentToDocumentIndex]
     metadata: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
     """

vellum/client/types/folder_entity_prompt_sandbox_data.py

@@ -14,6 +14,7 @@ class FolderEntityPromptSandboxData(UniversalBaseModel):
     created: dt.datetime
     modified: dt.datetime
     status: EntityStatus
+    description: typing.Optional[str] = None
     last_deployed_on: typing.Optional[dt.datetime] = None
 
     if IS_PYDANTIC_V2:

vellum/client/types/folder_entity_workflow_sandbox_data.py

@@ -14,6 +14,7 @@ class FolderEntityWorkflowSandboxData(UniversalBaseModel):
     created: dt.datetime
     modified: dt.datetime
     status: EntityStatus
+    description: typing.Optional[str] = None
     last_deployed_on: typing.Optional[dt.datetime] = None
 
     if IS_PYDANTIC_V2:

vellum/workflows/expressions/accessor.py

@@ -35,20 +35,37 @@ class AccessorExpression(BaseDescriptor[Any]):
             if isinstance(self._field, int):
                 raise InvalidExpressionException("Cannot access field by index on a dataclass")
 
-            return getattr(base, self._field)
+            try:
+                return getattr(base, self._field)
+            except AttributeError:
+                raise InvalidExpressionException(f"Field '{self._field}' not found on dataclass {type(base).__name__}")
 
         if isinstance(base, BaseModel):
             if isinstance(self._field, int):
                 raise InvalidExpressionException("Cannot access field by index on a BaseModel")
 
-            return getattr(base, self._field)
+            try:
+                return getattr(base, self._field)
+            except AttributeError:
+                raise InvalidExpressionException(f"Field '{self._field}' not found on BaseModel {type(base).__name__}")
 
         if isinstance(base, Mapping):
-            return base[self._field]
+            try:
+                return base[self._field]
+            except KeyError:
+                raise InvalidExpressionException(f"Key '{self._field}' not found in mapping")
 
         if isinstance(base, Sequence):
-            index = int(self._field)
-            return base[index]
+            try:
+                index = int(self._field)
+                return base[index]
+            except (IndexError, ValueError):
+                if isinstance(self._field, int) or (isinstance(self._field, str) and self._field.lstrip("-").isdigit()):
+                    raise InvalidExpressionException(
+                        f"Index {self._field} is out of bounds for sequence of length {len(base)}"
+                    )
+                else:
+                    raise InvalidExpressionException(f"Invalid index '{self._field}' for sequence access")
 
         raise InvalidExpressionException(f"Cannot get field {self._field} from {base}")
 

vellum/workflows/expressions/tests/test_accessor.py

@@ -0,0 +1,189 @@
+import pytest
+from dataclasses import dataclass
+
+from pydantic import BaseModel
+
+from vellum.workflows.descriptors.exceptions import InvalidExpressionException
+from vellum.workflows.expressions.accessor import AccessorExpression
+from vellum.workflows.references.constant import ConstantValueReference
+from vellum.workflows.state.base import BaseState
+
+
+@dataclass
+class TestDataclass:
+    name: str
+    value: int
+
+
+class TestBaseModel(BaseModel):
+    name: str
+    value: int
+
+
+class TestState(BaseState):
+    pass
+
+
+def test_accessor_expression_dict_valid_key():
+    state = TestState()
+    base_ref = ConstantValueReference({"name": "test", "value": 42})
+    accessor = AccessorExpression(base=base_ref, field="name")
+
+    result = accessor.resolve(state)
+
+    assert result == "test"
+
+
+def test_accessor_expression_dict_invalid_key():
+    state = TestState()
+    base_ref = ConstantValueReference({"name": "test", "value": 42})
+    accessor = AccessorExpression(base=base_ref, field="missing_key")
+
+    with pytest.raises(InvalidExpressionException) as exc_info:
+        accessor.resolve(state)
+
+    assert "Key 'missing_key' not found in mapping" in str(exc_info.value)
+
+
+def test_accessor_expression_list_valid_index():
+    state = TestState()
+    base_ref = ConstantValueReference(["first", "second", "third"])
+    accessor = AccessorExpression(base=base_ref, field=1)
+
+    result = accessor.resolve(state)
+
+    assert result == "second"
+
+
+def test_accessor_expression_list_invalid_index():
+    state = TestState()
+    base_ref = ConstantValueReference(["first", "second"])
+    accessor = AccessorExpression(base=base_ref, field=5)
+
+    with pytest.raises(InvalidExpressionException) as exc_info:
+        accessor.resolve(state)
+
+    assert "Index 5 is out of bounds for sequence of length 2" in str(exc_info.value)
+
+
+def test_accessor_expression_list_negative_index():
+    state = TestState()
+    base_ref = ConstantValueReference(["first", "second", "third"])
+    accessor = AccessorExpression(base=base_ref, field=-1)
+
+    result = accessor.resolve(state)
+
+    assert result == "third"
+
+
+def test_accessor_expression_list_invalid_negative_index():
+    state = TestState()
+    base_ref = ConstantValueReference(["first", "second"])
+    accessor = AccessorExpression(base=base_ref, field=-5)
+
+    with pytest.raises(InvalidExpressionException) as exc_info:
+        accessor.resolve(state)
+
+    assert "Index -5 is out of bounds for sequence of length 2" in str(exc_info.value)
+
+
+def test_accessor_expression_list_string_index():
+    state = TestState()
+    base_ref = ConstantValueReference(["first", "second", "third"])
+    accessor = AccessorExpression(base=base_ref, field="1")
+
+    result = accessor.resolve(state)
+
+    assert result == "second"
+
+
+def test_accessor_expression_list_invalid_string_index():
+    state = TestState()
+    base_ref = ConstantValueReference(["first", "second"])
+    accessor = AccessorExpression(base=base_ref, field="invalid")
+
+    with pytest.raises(InvalidExpressionException) as exc_info:
+        accessor.resolve(state)
+
+    assert "Invalid index 'invalid' for sequence access" in str(exc_info.value)
+
+
+def test_accessor_expression_dataclass_valid_field():
+    state = TestState()
+    test_obj = TestDataclass(name="test", value=42)
+    base_ref = ConstantValueReference(test_obj)
+    accessor = AccessorExpression(base=base_ref, field="name")
+
+    result = accessor.resolve(state)
+
+    assert result == "test"
+
+
+def test_accessor_expression_dataclass_invalid_field():
+    state = TestState()
+    test_obj = TestDataclass(name="test", value=42)
+    base_ref = ConstantValueReference(test_obj)
+    accessor = AccessorExpression(base=base_ref, field="missing_field")
+
+    with pytest.raises(InvalidExpressionException) as exc_info:
+        accessor.resolve(state)
+
+    assert "Field 'missing_field' not found on dataclass TestDataclass" in str(exc_info.value)
+
+
+def test_accessor_expression_basemodel_valid_field():
+    state = TestState()
+    test_obj = TestBaseModel(name="test", value=42)
+    base_ref = ConstantValueReference(test_obj)
+    accessor = AccessorExpression(base=base_ref, field="name")
+
+    result = accessor.resolve(state)
+
+    assert result == "test"
+
+
+def test_accessor_expression_basemodel_invalid_field():
+    state = TestState()
+    test_obj = TestBaseModel(name="test", value=42)
+    base_ref = ConstantValueReference(test_obj)
+    accessor = AccessorExpression(base=base_ref, field="missing_field")
+
+    with pytest.raises(InvalidExpressionException) as exc_info:
+        accessor.resolve(state)
+
+    assert "Field 'missing_field' not found on BaseModel TestBaseModel" in str(exc_info.value)
+
+
+def test_accessor_expression_dataclass_index_access():
+    state = TestState()
+    test_obj = TestDataclass(name="test", value=42)
+    base_ref = ConstantValueReference(test_obj)
+    accessor = AccessorExpression(base=base_ref, field=0)
+
+    with pytest.raises(InvalidExpressionException) as exc_info:
+        accessor.resolve(state)
+
+    assert "Cannot access field by index on a dataclass" in str(exc_info.value)
+
+
+def test_accessor_expression_basemodel_index_access():
+    state = TestState()
+    test_obj = TestBaseModel(name="test", value=42)
+    base_ref = ConstantValueReference(test_obj)
+    accessor = AccessorExpression(base=base_ref, field=0)
+
+    with pytest.raises(InvalidExpressionException) as exc_info:
+        accessor.resolve(state)
+
+    assert "Cannot access field by index on a BaseModel" in str(exc_info.value)
+
+
+def test_accessor_expression_unsupported_type():
+    state = TestState()
+    base_ref = ConstantValueReference(42)
+    accessor = AccessorExpression(base=base_ref, field="field")
+
+    with pytest.raises(InvalidExpressionException) as exc_info:
+        accessor.resolve(state)
+
+    assert "Cannot get field field from 42" in str(exc_info.value)

vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/test_inline_prompt_node.py

@@ -192,7 +192,30 @@ def test_validation_with_extra_variables(vellum_adhoc_prompt_client):
     ]
 
 
-def test_inline_prompt_node__json_output(vellum_adhoc_prompt_client):
+@pytest.mark.parametrize(
+    "custom_parameters,test_description",
+    [
+        (
+            {
+                "json_mode": False,
+                "json_schema": {
+                    "name": "get_result",
+                    "schema": {
+                        "type": "object",
+                        "required": ["result"],
+                        "properties": {"result": {"type": "string", "description": ""}},
+                    },
+                },
+            },
+            "with json_schema configured",
+        ),
+        (
+            {},
+            "without json_mode or json_schema configured",
+        ),
+    ],
+)
+def test_inline_prompt_node__json_output(vellum_adhoc_prompt_client, custom_parameters, test_description):
     """Confirm that InlinePromptNodes output the expected JSON when run."""
 
     # GIVEN a node that subclasses InlinePromptNode
@@ -214,17 +237,7 @@ def test_inline_prompt_node__json_output(vellum_adhoc_prompt_client):
             frequency_penalty=0.0,
             presence_penalty=0.0,
             logit_bias=None,
-            custom_parameters={
-                "json_mode": False,
-                "json_schema": {
-                    "name": "get_result",
-                    "schema": {
-                        "type": "object",
-                        "required": ["result"],
-                        "properties": {"result": {"type": "string", "description": ""}},
-                    },
-                },
-            },
+            custom_parameters=custom_parameters,
         )
 
     # AND a known JSON response from invoking an inline prompt
@@ -284,17 +297,7 @@ def test_inline_prompt_node__json_output(vellum_adhoc_prompt_client):
             frequency_penalty=0.0,
             presence_penalty=0.0,
             logit_bias=None,
-            custom_parameters={
-                "json_mode": False,
-                "json_schema": {
-                    "name": "get_result",
-                    "schema": {
-                        "type": "object",
-                        "required": ["result"],
-                        "properties": {"result": {"type": "string", "description": ""}},
-                    },
-                },
-            },
+            custom_parameters=custom_parameters,
         ),
         request_options=mock.ANY,
         settings=None,

vellum/workflows/nodes/displayable/inline_prompt_node/node.py

@@ -48,26 +48,16 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
         string_outputs = []
         json_output = None
 
-        should_parse_json = False
-        if hasattr(self, "parameters"):
-            custom_params = self.parameters.custom_parameters
-            if custom_params and isinstance(custom_params, dict):
-                json_schema = custom_params.get("json_schema", {})
-                if (isinstance(json_schema, dict) and "schema" in json_schema) or custom_params.get("json_mode", {}):
-                    should_parse_json = True
-
         for output in outputs:
             if output.value is None:
                 continue
 
             if output.type == "STRING":
                 string_outputs.append(output.value)
-                if should_parse_json:
-                    try:
-                        parsed_json = json.loads(output.value)
-                        json_output = parsed_json
-                    except (json.JSONDecodeError, TypeError):
-                        pass
+                try:
+                    json_output = json.loads(output.value)
+                except (json.JSONDecodeError, TypeError):
+                    pass
             elif output.type == "JSON":
                 string_outputs.append(json.dumps(output.value, indent=4))
                 json_output = output.value

vellum/workflows/nodes/experimental/tool_calling_node/utils.py

@@ -15,7 +15,6 @@ from vellum.client.types.string_chat_message_content import StringChatMessageCon
 from vellum.client.types.variable_prompt_block import VariablePromptBlock
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
-from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.displayable.code_execution_node.node import CodeExecutionNode
 from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
@@ -143,24 +142,16 @@ def create_function_node(
         packages: Optional list of packages to install for code execution (only used for regular functions)
         runtime: The runtime to use for code execution (default: "PYTHON_3_11_6")
     """
-
     if is_workflow_class(function):
-
+        # Create a class-level wrapper that calls the original function
         def execute_function(self) -> BaseNode.Outputs:
             outputs = self.state.meta.node_outputs.get(tool_router_node.Outputs.text)
-            # first parse into json
+
             outputs = json.loads(outputs)
             arguments = outputs["arguments"]
 
-            # Dynamically define an Inputs subclass of BaseInputs
-            Inputs = type(
-                "Inputs",
-                (BaseInputs,),
-                {"__annotations__": {k: type(v) for k, v in arguments.items()}},
-            )
-
-            # Create an instance with arguments
-            inputs_instance = Inputs(**arguments)
+            # Call the function based on its type
+            inputs_instance = function.get_inputs_class()(**arguments)
 
             workflow = function()
             terminal_event = workflow.run(

{vellum_ai-0.14.65.dist-info → vellum_ai-0.14.66.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.14.65
+Version: 0.14.66
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -21,7 +21,7 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Typing :: Typed
 Requires-Dist: Jinja2 (>=3.1.0,<4.0.0)
-Requires-Dist: click (>=8.1.0,<9.0.0)
+Requires-Dist: click (>=8.1.7,<9.0.0)
 Requires-Dist: docker (>=7.1.0,<8.0.0)
 Requires-Dist: httpx (>=0.21.2)
 Requires-Dist: openai (>=1.0.0,<2.0.0)

{vellum_ai-0.14.65.dist-info → vellum_ai-0.14.66.dist-info}/RECORD

@@ -3,20 +3,21 @@ vellum_cli/README.md,sha256=2NudRoLzWxNKqnuVy1JuQ7DerIaxWGYkrH8kMd-asIE,90
 vellum_cli/__init__.py,sha256=2_6oGoVcLFUh4L63Kz4SBL4Y6XevJ70oYbg7BJ3cb5Q,12569
 vellum_cli/aliased_group.py,sha256=ugW498j0yv4ALJ8vS9MsO7ctDW7Jlir9j6nE_uHAP8c,3363
 vellum_cli/config.py,sha256=v5BmZ-t_v4Jmqd7KVuQMZF2pRI-rbMspSkVYXIRoTmI,9448
-vellum_cli/image_push.py,sha256=skFXf25ixMOX1yfcyAtii-RivYYv-_hsv-Z-bVB6m5Q,7380
+vellum_cli/image_push.py,sha256=mjK3Fj3_MIhIGmurYVz_OC-hCoAef5LqDb0OppKJIJc,9909
 vellum_cli/init.py,sha256=WpnMXPItPmh0f0bBGIer3p-e5gu8DUGwSArT_FuoMEw,5093
 vellum_cli/logger.py,sha256=dcM_OmgqXLo93vDYswO5ylyUQQcTfnA5GTd5tbIt3wM,1446
 vellum_cli/ping.py,sha256=p_BCCRjgPhng6JktuECtkDQLbhopt6JpmrtGoLnLJT8,1161
-vellum_cli/pull.py,sha256=M50yXzA_35N35gk1Y8KjLbXrzdRG86--XFQvEukxGtA,13371
+vellum_cli/pull.py,sha256=udYyPlJ6VKDdh78rApNJOZgxHl82fcV6iGnRPSdX1LY,14750
 vellum_cli/push.py,sha256=wxRlFu2mYW9SvwODYxwajri1mDQ2be0n-9i0d9QAc30,10194
 vellum_cli/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_cli/tests/conftest.py,sha256=AFYZryKA2qnUuCPBxBKmHLFoPiE0WhBFFej9tNwSHdc,1526
 vellum_cli/tests/test_config.py,sha256=uvKGDc8BoVyT9_H0Z-g8469zVxomn6Oi3Zj-vK7O_wU,2631
 vellum_cli/tests/test_image_push.py,sha256=QM-JlR_aJappvwbCLteQZZf76sd7SE1sRj3armvFK-I,5706
+vellum_cli/tests/test_image_push_error_handling.py,sha256=_Wjfkn1orI2K4Ahzqz4u8T13or7NOX01K4BtcTuTIOM,7107
 vellum_cli/tests/test_init.py,sha256=8UOc_ThfouR4ja5cCl_URuLk7ohr9JXfCnG4yka1OUQ,18754
 vellum_cli/tests/test_main.py,sha256=qDZG-aQauPwBwM6A2DIu1494n47v3pL28XakTbLGZ-k,272
 vellum_cli/tests/test_ping.py,sha256=3ucVRThEmTadlV9LrJdCCrr1Ofj3rOjG6ue0BNR2UC0,2523
-vellum_cli/tests/test_pull.py,sha256=7HRAhIdkVW5mR2VckEaNDjp4rt-MlIxOWMMI2XNUPE8,49814
+vellum_cli/tests/test_pull.py,sha256=91aOAYNzYcJuIrTwRvGX_Qa1_UOmdCfVAfvZVfwCKZ4,49902
 vellum_cli/tests/test_push.py,sha256=I8XICg3pUb3yxAFLXziVHHf5CRm354LO-uUfwtca3bU,33897
 vellum_ee/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -130,11 +131,11 @@ vellum_ee/workflows/tests/test_display_meta.py,sha256=DIzjNbwK1-4mlttPML6NskQ4rP
 vellum_ee/workflows/tests/test_server.py,sha256=SsOkS6sGO7uGC4mxvk4iv8AtcXs058P9hgFHzTWmpII,14519
 vellum_ee/workflows/tests/test_virtual_files.py,sha256=TJEcMR0v2S8CkloXNmCHA0QW0K6pYNGaIjraJz7sFvY,2762
 vellum/__init__.py,sha256=U4MVBzBx8I8-rOG_feluSz5U8_hD969e85F4cRgmVv4,42104
-vellum/client/README.md,sha256=qmaVIP42MnxAu8jV7u-CsgVFfs3-pHQODrXdZdFxtaw,4749
+vellum/client/README.md,sha256=CuGUYnaE0Imt0KqQ4sIPaUghCjLHkF3DdEvZWu14-8s,4807
 vellum/client/__init__.py,sha256=AYopGv2ZRVn3zsU8_km6KOvEHDbXiTPCVuYVI7bWvdA,120166
 vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
 vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-vellum/client/core/client_wrapper.py,sha256=955wHqaBcVMDsbO3pyETLDv2SnobfTnBs5g2chqn4cE,1869
+vellum/client/core/client_wrapper.py,sha256=gluL9LJQaydz9m87P7bWACAysxbElOZWqF6JrbjOjDw,1869
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=d4NNbX8XvXP32z8KpK2Xovv33nFfruIrpz0QWxlgpZk,2663
 vellum/client/core/http_client.py,sha256=Z77OIxIbL4OAB2IDqjRq_sYa5yNYAWfmdhdCSSvh6Y4,19552
@@ -288,7 +289,7 @@ vellum/client/types/document_index_indexing_config_request.py,sha256=Wt-ys1o_acH
 vellum/client/types/document_index_read.py,sha256=ePngiRszr65HLl9D0_FUdhAdMe84nRwyM3cKbr8rFpg,1177
 vellum/client/types/document_processing_state.py,sha256=ISlurj7jQzwHzxPzDZTqeAIgSIIGMBBPgcOSoe04pTU,211
 vellum/client/types/document_prompt_block.py,sha256=sgFxN48PILFuuF2KUIwks6PbJ3XH6sCE_8ydLEE_doU,1019
-vellum/client/types/document_read.py,sha256=6nwEvVvVe-6y2vtPNYB7KtcFoaydH2ow-WhCmCAvMQ8,1713
+vellum/client/types/document_read.py,sha256=U3YcFLF9exeZYAhLErEdAc_vVTtWZ8efXDKmmcabshg,1661
 vellum/client/types/document_status.py,sha256=GD_TSoFmZUBJnPl-chAmaQFzQ2_TYO3PSqi3-9QfEHE,122
 vellum/client/types/document_vellum_value.py,sha256=a8WQhyntwy80iN9j8L9F5v6Jmq1L4j0ETJo9c9VGabs,768
 vellum/client/types/document_vellum_value_request.py,sha256=utpoRMMVhMIsa4S4ZOaOr2lX76BgrOVolPxCwy9-pUw,797
@@ -329,11 +330,11 @@ vellum/client/types/folder_entity_document_index_data.py,sha256=UceusLf7dLYDHq4Z
 vellum/client/types/folder_entity_folder.py,sha256=1GWfyHdenyAI2GXiouIbnFhPK4ADqZGFh-6fpp_C6-U,792
 vellum/client/types/folder_entity_folder_data.py,sha256=JLYkhvRY63-ojNY84lfRTaUg25KjslSE-iNAC2NYCTI,674
 vellum/client/types/folder_entity_prompt_sandbox.py,sha256=7CGVcK5maoqO7CC7sFfi6F5X0QWdHVbEygqLyci_VDY,839
-vellum/client/types/folder_entity_prompt_sandbox_data.py,sha256=4u-t-p2-h7Z1SacNOQnEnY3tt-au5WB0hIDrB1HOeKk,781
+vellum/client/types/folder_entity_prompt_sandbox_data.py,sha256=LMTtOshCdp0MpYHdKtrtBSdpXw9B9id5BfIpkg9SsY4,826
 vellum/client/types/folder_entity_test_suite.py,sha256=M5sb7ntUX2VsCyD0AYnz7dalt7G1ejOYQZvsUU33q2w,811
 vellum/client/types/folder_entity_test_suite_data.py,sha256=RGITy3Pip6NeIb77YGavDWgRr8GLFDB9sXrTU_Dm2OI,719
 vellum/client/types/folder_entity_workflow_sandbox.py,sha256=gFmkwKZGh5O9-9y4w6VGK1IRsKLf7NZ4SCC-pgucm8M,853
-vellum/client/types/folder_entity_workflow_sandbox_data.py,sha256=MW0nLAvetgnsm6KWiGcYcKB9llR-UfosSNz0MtokN38,783
+vellum/client/types/folder_entity_workflow_sandbox_data.py,sha256=4_kzQSDrFXVJCSF1st73jo0LxKOicn_gE82XAk-QKJI,828
 vellum/client/types/fulfilled_ad_hoc_execute_prompt_event.py,sha256=19rFjVrzobaaAQhZlP_WGV9f_Rwrz4EPRXbT2aYkaJw,1016
 vellum/client/types/fulfilled_execute_prompt_event.py,sha256=E-iBwlMWvwQyRfzqPlgtX8KurE3IYsTRd5vWjXtbOmk,994
 vellum/client/types/fulfilled_execute_prompt_response.py,sha256=nKhDk2ZloCuE0uijX5XOE63_cq8PBo4UWs4hK4e3jUE,1227
@@ -1505,7 +1506,7 @@ vellum/workflows/events/types.py,sha256=LwgFlMRbptJvdPtPO1POUtGtbhGw7BSuvgHxNSgS
 vellum/workflows/events/workflow.py,sha256=i9JSCANjAhf5uc57gYspdII2V2OLItbc0BfT8yB9mF0,7728
 vellum/workflows/exceptions.py,sha256=NiBiR3ggfmPxBVqD-H1SqmjI-7mIn0EStSN1BqApvCM,1213
 vellum/workflows/expressions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/workflows/expressions/accessor.py,sha256=ItZF7fMLzVTqsdAiaXb5SiDupXmX0X9xbIus1W6hRds,1870
+vellum/workflows/expressions/accessor.py,sha256=1kY0sYfednQ_u0kh1Df2WTsDREVIOzU7UuZc5jvj5xw,2861
 vellum/workflows/expressions/and_.py,sha256=I7lNqrUM3-m_5hmjjiMhaHhJtKcLj39kEFVWPDOqwfo,916
 vellum/workflows/expressions/begins_with.py,sha256=FnWsQXbENm0ZwkfEP7dR8Qx4_MMrzj6C1yqAV2KaNHw,1123
 vellum/workflows/expressions/between.py,sha256=dVeddT6YA91eOAlE1Utg7C7gnCiYE7WP-dg17yXUeAY,1492
@@ -1535,6 +1536,7 @@ vellum/workflows/expressions/not_in.py,sha256=pFvwkFPsn3WJw61ssFgM2U1dqWEeglfz4F
 vellum/workflows/expressions/or_.py,sha256=s-8YdMSSCDS2yijR38kguwok3iqmDMMgDYKV93b4O4s,914
 vellum/workflows/expressions/parse_json.py,sha256=xsk6j3HF7bU1yF6fwt5P9Ugcyd5D9ZXrdng11FRilUI,1088
 vellum/workflows/expressions/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vellum/workflows/expressions/tests/test_accessor.py,sha256=ijmGjS1pipdurAgYGmfV8Gcle4HIJ5lQHVG18xHjNiQ,5839
 vellum/workflows/expressions/tests/test_expressions.py,sha256=3b6k8xs-CItBBw95NygFLUNoNPKxI4VA1GyWbkMtqyI,11623
 vellum/workflows/expressions/tests/test_parse_json.py,sha256=zpB_qE5_EwWQL7ULQUJm0o1PRSfWZdAqZNW6Ah13oJE,1059
 vellum/workflows/graph/__init__.py,sha256=3sHlay5d_-uD7j3QJXiGl0WHFZZ_QScRvgyDhN2GhHY,74
@@ -1589,7 +1591,7 @@ vellum/workflows/nodes/displayable/bases/inline_prompt_node/__init__.py,sha256=H
 vellum/workflows/nodes/displayable/bases/inline_prompt_node/constants.py,sha256=fnjiRWLoRlC4Puo5oQcpZD5Hd-EesxsAo9l5tGAkpZQ,270
 vellum/workflows/nodes/displayable/bases/inline_prompt_node/node.py,sha256=cD2RzOX9WE6xTKgm09dsEw4xHATZbOjeGyYCSdl8fjU,11785
 vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/test_inline_prompt_node.py,sha256=5CNag1_aEFZbCL0nrOC5e1L-t90-4rp2xDwh0h52hVI,21407
+vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/test_inline_prompt_node.py,sha256=terQtPqLGiKKGtU7zbkM3XkbVddLQI87bBj31M0zaRE,21356
 vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py,sha256=0a40fkkZkFMmZN0CsWf6EP_y1H6x36EGa3WcfVNyOsM,9797
 vellum/workflows/nodes/displayable/bases/search_node.py,sha256=3UtbqY3QO4kzfJHbmUNZGnEEfJmaoiF892u8H6TGjp8,5381
 vellum/workflows/nodes/displayable/bases/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1616,7 +1618,7 @@ vellum/workflows/nodes/displayable/guardrail_node/test_node.py,sha256=SAGv6hSFcB
 vellum/workflows/nodes/displayable/guardrail_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/guardrail_node/tests/test_node.py,sha256=X2pd6TI8miYxIa7rgvs1pHTEreyWcf77EyR0_Jsa700,2055
 vellum/workflows/nodes/displayable/inline_prompt_node/__init__.py,sha256=gSUOoEZLlrx35-tQhSAd3An8WDwBqyiQh-sIebLU9wU,74
-vellum/workflows/nodes/displayable/inline_prompt_node/node.py,sha256=8RXZqWMzViUjFfbpmcy1gkSsKnEpci8BGwsuPYv4xMQ,3380
+vellum/workflows/nodes/displayable/inline_prompt_node/node.py,sha256=ST_V1Z3-7lzZ__Ntfm6iry70u9D3GD2a2vxsl2d4Q_c,2843
 vellum/workflows/nodes/displayable/inline_prompt_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/inline_prompt_node/tests/test_node.py,sha256=bBHs90mV5SZ3rJPAL0wx4WWyawUA406LgMPOdvpZC_A,10923
 vellum/workflows/nodes/displayable/merge_node/__init__.py,sha256=J8IC08dSH7P76wKlNuxe1sn7toNGtSQdFirUbtPDEs0,60
@@ -1646,7 +1648,7 @@ vellum/workflows/nodes/experimental/openai_chat_completion_node/node.py,sha256=c
 vellum/workflows/nodes/experimental/tool_calling_node/__init__.py,sha256=S7OzT3I4cyOU5Beoz87nPwCejCMP2FsHBFL8OcVmxJ4,118
 vellum/workflows/nodes/experimental/tool_calling_node/node.py,sha256=FkhaJccpCbx2be_IZ5V2v6Lo-jPJ0WgSC5tveLvAW4A,5774
 vellum/workflows/nodes/experimental/tool_calling_node/tests/test_node.py,sha256=sxG26mOwt4N36RLoPJ-ngginPqC5qFzD_kGj9izdCFI,1833
-vellum/workflows/nodes/experimental/tool_calling_node/utils.py,sha256=DjFn33XuomM-EiojCFzTlOsWMdF_uYyHzNayBok6X-k,10055
+vellum/workflows/nodes/experimental/tool_calling_node/utils.py,sha256=vB-tdN-44iVrTFvwGLVZW3dv62imIp5fPo0FYCW0l3M,9815
 vellum/workflows/nodes/mocks.py,sha256=a1FjWEIocseMfjzM-i8DNozpUsaW0IONRpZmXBoWlyc,10455
 vellum/workflows/nodes/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/tests/test_mocks.py,sha256=mfPvrs75PKcsNsbJLQAN6PDFoVqs9TmQxpdyFKDdO60,7837
@@ -1712,8 +1714,8 @@ vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnad
 vellum/workflows/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/workflows/tests/test_base_workflow.py,sha256=fROqff6AZpCIzaSwOKSdtYy4XR0UZQ6ejxL3RJOSJVs,20447
 vellum/workflows/workflows/tests/test_context.py,sha256=VJBUcyWVtMa_lE5KxdhgMu0WYNYnUQUDvTF7qm89hJ0,2333
-vellum_ai-0.14.65.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.14.65.dist-info/METADATA,sha256=kYXF0TPnW90jthw0xtqtK8eTzOQo-bRqiD_Tmssv7dU,5556
-vellum_ai-0.14.65.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.14.65.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.14.65.dist-info/RECORD,,
+vellum_ai-0.14.66.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.14.66.dist-info/METADATA,sha256=bkr1Oj0n0HAwEiTLSb4lZEiDQQM_mtpz1uXZ1Lt7VQM,5556
+vellum_ai-0.14.66.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.14.66.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.14.66.dist-info/RECORD,,

vellum_cli/image_push.py

@@ -9,9 +9,10 @@ import docker
 from docker import DockerClient
 from dotenv import load_dotenv
 
+from vellum.client.core.api_error import ApiError
 from vellum.workflows.vellum_client import create_vellum_client, create_vellum_environment
 from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, load_vellum_cli_config
-from vellum_cli.logger import load_cli_logger
+from vellum_cli.logger import handle_cli_error, load_cli_logger
 
 _SUPPORTED_ARCHITECTURE = "amd64"
 
@@ -82,7 +83,34 @@ def image_push_command(image: str, tags: Optional[List[str]] = None, workspace:
             exit(1)
     else:
         logger.info("Authenticating...")
-        auth = vellum_client.container_images.docker_service_token()
+        try:
+            auth = vellum_client.container_images.docker_service_token()
+        except ApiError as e:
+            if e.status_code == 401 or e.status_code == 403:
+                handle_cli_error(
+                    logger,
+                    title="Authentication failed",
+                    message="Unable to authenticate with Vellum API. Please check your API key.",
+                    suggestion="Make sure your VELLUM_API_KEY environment variable is set correctly.",
+                )
+                return
+            elif e.status_code == 500:
+                handle_cli_error(
+                    logger,
+                    title="Server error",
+                    message="The Vellum API failed with an unexpected server error.",
+                    suggestion="Please try again later and contact support if the problem persists.",
+                )
+                return
+            else:
+                handle_cli_error(
+                    logger,
+                    title="API request failed",
+                    message=f"Failed to get Docker service token from Vellum API (HTTP {e.status_code}).",
+                    suggestion="Please check your configuration and try again. If the problem persists, "
+                    "contact support.",
+                )
+                return
 
         docker_client.login(
             username="oauth2accesstoken",
@@ -143,11 +171,38 @@ def image_push_command(image: str, tags: Optional[List[str]] = None, workspace:
 
     logger.info(f"Updating Vellum metadata and validating image works in our system with image digest: {sha}...")
 
-    vellum_client.container_images.push_container_image(
-        name=image_name,
-        sha=sha,
-        tags=all_tags,
-    )
+    try:
+        vellum_client.container_images.push_container_image(
+            name=image_name,
+            sha=sha,
+            tags=all_tags,
+        )
+    except ApiError as e:
+        if e.status_code == 401 or e.status_code == 403:
+            handle_cli_error(
+                logger,
+                title="Authentication failed",
+                message="Unable to push container image metadata to Vellum API. Please check your API key.",
+                suggestion="Make sure your VELLUM_API_KEY environment variable is set correctly.",
+            )
+            return
+        elif e.status_code == 500:
+            handle_cli_error(
+                logger,
+                title="Server error",
+                message="The Vellum API failed with an unexpected server error while pushing container image metadata.",
+                suggestion="Please try again later and contact support if the problem persists.",
+            )
+            return
+        else:
+            handle_cli_error(
+                logger,
+                title="API request failed",
+                message=f"Failed to push container image metadata to Vellum API (HTTP {e.status_code}).",
+                suggestion="Please check your configuration and try again. If the problem persists, contact support.",
+            )
+            return
+
     logger.info(f"Image successfully pushed as {image_name} to vellum with tags: {all_tags}.")
 
 

vellum_cli/pull.py

@@ -20,7 +20,7 @@ from vellum_cli.config import (
     WorkflowDeploymentConfig,
     load_vellum_cli_config,
 )
-from vellum_cli.logger import load_cli_logger
+from vellum_cli.logger import handle_cli_error, load_cli_logger
 
 ERROR_LOG_FILE_NAME = "error.log"
 METADATA_FILE_NAME = "metadata.json"
@@ -201,16 +201,42 @@ def pull_command(
         zip_bytes = b"".join(response)
     except ApiError as e:
         if e.status_code == 401 or e.status_code == 403:
-            raise Exception("Please make sure your `VELLUM_API_KEY` environment variable is set correctly.")
+            handle_cli_error(
+                logger,
+                title="Authentication failed",
+                message="Unable to authenticate with the Vellum API.",
+                suggestion="Please make sure your `VELLUM_API_KEY` environment variable is set correctly and that you have access to this workflow.",  # noqa: E501
+            )
+
+        if e.status_code == 404:
+            handle_cli_error(
+                logger,
+                title="Workflow not found",
+                message=f"The workflow with ID '{pk}' could not be found.",
+                suggestion="Please verify the workflow ID is correct and that you have access to it in your workspace.",
+            )
 
         if e.status_code == 500:
-            raise Exception(
-                "The Pull API failed with an unexpected error. Please try again later and contact support if the problem persists."  # noqa: E501
+            handle_cli_error(
+                logger,
+                title="Server error occurred",
+                message="The Vellum API encountered an internal server error while processing your request.",
+                suggestion="Please try again in a few moments. If the problem persists, contact Vellum support with the workflow ID and timestamp.",  # noqa: E501
+            )
+
+        if e.status_code == 502 or e.status_code == 503 or e.status_code == 504:
+            handle_cli_error(
+                logger,
+                title="Service temporarily unavailable",
+                message="The Vellum API is temporarily unavailable or experiencing high load.",
+                suggestion="Please wait a moment and try again. If the issue continues, check the Vellum status page or contact support.",  # noqa: E501
             )
 
-        # TODO: We should return an Origin header in to validate this case
-        raise Exception(
-            "The API we tried to pull is invalid. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        handle_cli_error(
+            logger,
+            title="API request failed",
+            message=f"The API request failed with status code {e.status_code}.",
+            suggestion="Please verify your `VELLUM_API_URL` environment variable is set correctly and try again.",
         )
 
     zip_buffer = io.BytesIO(zip_bytes)
@@ -299,8 +325,11 @@ def pull_command(
                         logger.info(f"Writing to {target_file}...")
                         target.write(content)
     except zipfile.BadZipFile:
-        raise Exception(
-            "The API we tried to pull from returned an invalid zip file. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        handle_cli_error(
+            logger,
+            title="Invalid response format",
+            message="The API returned an invalid zip file format.",
+            suggestion="Please verify your `VELLUM_API_URL` environment variable is set correctly and try again.",
         )
 
     if include_json:

vellum_cli/tests/test_image_push_error_handling.py

@@ -0,0 +1,184 @@
+import subprocess
+from unittest.mock import patch
+from uuid import uuid4
+
+from click.testing import CliRunner
+
+from vellum.client.core.api_error import ApiError
+from vellum_cli import main as cli_main
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_docker_service_token_401_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+    ]
+
+    vellum_client.container_images.docker_service_token.side_effect = ApiError(status_code=401, body="Unauthorized")
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Authentication failed" in result.output
+    assert "VELLUM_API_KEY" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_docker_service_token_500_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+    ]
+
+    vellum_client.container_images.docker_service_token.side_effect = ApiError(
+        status_code=500, body="Internal Server Error"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Server error" in result.output
+    assert "try again later" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_docker_service_token_other_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+    ]
+
+    vellum_client.container_images.docker_service_token.side_effect = ApiError(
+        status_code=429, body="Too Many Requests"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "API request failed" in result.output
+    assert "HTTP 429" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_container_image_401_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+    mock_docker_client.images.push.return_value = ["pushed"]
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b'[{"RepoDigests": ["test-repo@sha256:abcd1234"]}]'),
+    ]
+
+    vellum_client.container_images.docker_service_token.return_value = type(
+        "obj", (object,), {"access_token": "345678mnopqr", "organization_id": str(uuid4()), "repository": "myrepo.net"}
+    )()
+
+    vellum_client.container_images.push_container_image.side_effect = ApiError(status_code=401, body="Unauthorized")
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Authentication failed" in result.output
+    assert "VELLUM_API_KEY" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_container_image_500_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+    mock_docker_client.images.push.return_value = ["pushed"]
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b'[{"RepoDigests": ["test-repo@sha256:abcd1234"]}]'),
+    ]
+
+    vellum_client.container_images.docker_service_token.return_value = type(
+        "obj", (object,), {"access_token": "345678mnopqr", "organization_id": str(uuid4()), "repository": "myrepo.net"}
+    )()
+
+    vellum_client.container_images.push_container_image.side_effect = ApiError(
+        status_code=500, body="Internal Server Error"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Server error" in result.output
+    assert "try again later" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_container_image_other_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+    mock_docker_client.images.push.return_value = ["pushed"]
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b'[{"RepoDigests": ["test-repo@sha256:abcd1234"]}]'),
+    ]
+
+    vellum_client.container_images.docker_service_token.return_value = type(
+        "obj", (object,), {"access_token": "345678mnopqr", "organization_id": str(uuid4()), "repository": "myrepo.net"}
+    )()
+
+    vellum_client.container_images.push_container_image.side_effect = ApiError(
+        status_code=429, body="Too Many Requests"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "API request failed" in result.output
+    assert "HTTP 429" in result.output

vellum_cli/tests/test_pull.py

@@ -890,6 +890,7 @@ def test_pull__module_name_from_deployment_name(vellum_client):
         assert f.read() == "print('hello')"
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__invalid_zip_file(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -903,11 +904,12 @@ def test_pull__invalid_zip_file(vellum_client):
     # THEN the command returns an error
     assert result.exit_code == 1
     assert (
-        str(result.exception)
-        == "The API we tried to pull from returned an invalid zip file. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        "Invalid response format" in result.output
+        or "Please verify your `VELLUM_API_URL` environment variable is set correctly" in result.output
     )
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__json_decode_error(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -927,11 +929,12 @@ def test_pull__json_decode_error(vellum_client):
     # THEN the command returns an error
     assert result.exit_code == 1
     assert (
-        str(result.exception)
-        == "The API we tried to pull is invalid. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        "API request failed" in result.output
+        or "Please verify your `VELLUM_API_URL` environment variable is set correctly" in result.output
     )
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__unauthorized_error_path(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -948,9 +951,10 @@ 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."
+    assert "Please make sure your `VELLUM_API_KEY` environment variable is set correctly" in result.output
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__unexpected_error_path(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -968,9 +972,8 @@ def test_pull__unexpected_error_path(vellum_client):
     # THEN the command returns an error
     assert result.exit_code == 1
     assert (
-        str(result.exception)
-        == """The Pull API failed with an unexpected error. \
-Please try again later and contact support if the problem persists."""
+        "Server error occurred" in result.output
+        or "Please try again in a few moments. If the problem persists, contact Vellum support" in result.output
     )