vellum-ai 0.13.10__py3-none-any.whl → 0.13.11__py3-none-any.whl

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.13.10",
+            "X-Fern-SDK-Version": "0.13.11",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/workflows/errors/types.py

@@ -61,3 +61,24 @@ def workflow_event_error_to_workflow_error(error: WorkflowEventError) -> Workflo
         message=error.message,
         code=_WORKFLOW_EVENT_ERROR_CODE_TO_WORKFLOW_ERROR_CODE.get(error.code, WorkflowErrorCode.INTERNAL_ERROR),
     )
+
+
+_WORKFLOW_ERROR_CODE_TO_VELLUM_ERROR_CODE: Dict[WorkflowErrorCode, VellumErrorCodeEnum] = {
+    WorkflowErrorCode.INVALID_WORKFLOW: "INVALID_REQUEST",
+    WorkflowErrorCode.INVALID_INPUTS: "INVALID_INPUTS",
+    WorkflowErrorCode.INVALID_OUTPUTS: "INVALID_REQUEST",
+    WorkflowErrorCode.INVALID_STATE: "INVALID_REQUEST",
+    WorkflowErrorCode.INVALID_TEMPLATE: "INVALID_INPUTS",
+    WorkflowErrorCode.INTERNAL_ERROR: "INTERNAL_SERVER_ERROR",
+    WorkflowErrorCode.NODE_EXECUTION: "USER_DEFINED_ERROR",
+    WorkflowErrorCode.PROVIDER_ERROR: "PROVIDER_ERROR",
+    WorkflowErrorCode.USER_DEFINED_ERROR: "USER_DEFINED_ERROR",
+    WorkflowErrorCode.WORKFLOW_CANCELLED: "REQUEST_TIMEOUT",
+}
+
+
+def workflow_error_to_vellum_error(error: WorkflowError) -> VellumError:
+    return VellumError(
+        message=error.message,
+        code=_WORKFLOW_ERROR_CODE_TO_VELLUM_ERROR_CODE.get(error.code, "INTERNAL_SERVER_ERROR"),
+    )

vellum/workflows/nodes/displayable/bases/base_prompt_node/node.py

@@ -2,6 +2,7 @@ from abc import abstractmethod
 from typing import ClassVar, Generator, Generic, Iterator, List, Optional, Union
 
 from vellum import AdHocExecutePromptEvent, ExecutePromptEvent, PromptOutput
+from vellum.client.core.api_error import ApiError
 from vellum.core import RequestOptions
 from vellum.workflows.errors.types import WorkflowErrorCode, vellum_error_to_workflow_error
 from vellum.workflows.exceptions import NodeException
@@ -36,7 +37,19 @@ class BasePromptNode(BaseNode, Generic[StateType]):
             )
 
     def _process_prompt_event_stream(self) -> Generator[BaseOutput, None, Optional[List[PromptOutput]]]:
-        prompt_event_stream = self._get_prompt_event_stream()
+        try:
+            prompt_event_stream = self._get_prompt_event_stream()
+        except ApiError as e:
+            if e.status_code and e.status_code >= 400 and e.status_code < 500 and isinstance(e.body, dict):
+                raise NodeException(
+                    message=e.body.get("detail", "Failed to execute prompt"),
+                    code=WorkflowErrorCode.INVALID_INPUTS,
+                ) from e
+
+            raise NodeException(
+                message="Failed to execute prompt",
+                code=WorkflowErrorCode.INTERNAL_ERROR,
+            ) from e
 
         outputs: Optional[List[PromptOutput]] = None
         for event in prompt_event_stream:

vellum/workflows/nodes/displayable/bases/tests/test_utils.py

@@ -1,8 +1,11 @@
 import pytest
 import enum
 
+from pydantic import BaseModel
+
 from vellum.client.types.chat_history_vellum_value import ChatHistoryVellumValue
 from vellum.client.types.chat_message import ChatMessage
+from vellum.client.types.error_vellum_value import ErrorVellumValue
 from vellum.client.types.json_vellum_value import JsonVellumValue
 from vellum.client.types.number_vellum_value import NumberVellumValue
 from vellum.client.types.search_result import SearchResult
@@ -10,6 +13,8 @@ from vellum.client.types.search_result_document import SearchResultDocument
 from vellum.client.types.search_results_vellum_value import SearchResultsVellumValue
 from vellum.client.types.string_vellum_value import StringVellumValue
 from vellum.client.types.string_vellum_value_request import StringVellumValueRequest
+from vellum.client.types.vellum_error import VellumError
+from vellum.workflows.errors.types import WorkflowError, WorkflowErrorCode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value, primitive_to_vellum_value_request
 
 
@@ -17,6 +22,10 @@ class MockEnum(enum.Enum):
     FOO = "foo"
 
 
+class RandomPydanticModel(BaseModel):
+    foo: str
+
+
 @pytest.mark.parametrize(
     ["value", "expected_output"],
     [
@@ -51,6 +60,15 @@ class MockEnum(enum.Enum):
         (StringVellumValue(value="hello"), StringVellumValue(value="hello")),
         (StringVellumValueRequest(value="hello"), StringVellumValueRequest(value="hello")),
         ({"foo": "bar"}, JsonVellumValue(value={"foo": "bar"})),
+        (
+            VellumError(message="hello", code="USER_DEFINED_ERROR"),
+            ErrorVellumValue(value=VellumError(message="hello", code="USER_DEFINED_ERROR")),
+        ),
+        (
+            WorkflowError(message="hello", code=WorkflowErrorCode.USER_DEFINED_ERROR),
+            ErrorVellumValue(value=VellumError(message="hello", code="USER_DEFINED_ERROR")),
+        ),
+        (RandomPydanticModel(foo="bar"), JsonVellumValue(value={"foo": "bar"})),
     ],
 )
 def test_primitive_to_vellum_value(value, expected_output):

vellum/workflows/nodes/displayable/bases/utils.py

@@ -25,8 +25,11 @@ from vellum.client.types.search_results_vellum_value import SearchResultsVellumV
 from vellum.client.types.search_results_vellum_value_request import SearchResultsVellumValueRequest
 from vellum.client.types.string_vellum_value import StringVellumValue
 from vellum.client.types.string_vellum_value_request import StringVellumValueRequest
+from vellum.client.types.vellum_error import VellumError
 from vellum.client.types.vellum_value import VellumValue
 from vellum.client.types.vellum_value_request import VellumValueRequest
+from vellum.workflows.errors.types import WorkflowError, workflow_error_to_vellum_error
+from vellum.workflows.state.encoder import DefaultStateEncoder
 
 VELLUM_VALUE_REQUEST_TUPLE = (
     StringVellumValueRequest,
@@ -63,6 +66,10 @@ def primitive_to_vellum_value(value: Any) -> VellumValue:
     ):
         search_results = cast(Union[List[SearchResultRequest], List[SearchResult]], value)
         return SearchResultsVellumValue(value=search_results)
+    elif isinstance(value, VellumError):
+        return ErrorVellumValue(value=value)
+    elif isinstance(value, WorkflowError):
+        return ErrorVellumValue(value=workflow_error_to_vellum_error(value))
     elif isinstance(
         value,
         (
@@ -88,7 +95,7 @@ def primitive_to_vellum_value(value: Any) -> VellumValue:
         return value  # type: ignore
 
     try:
-        json_value = json.dumps(value)
+        json_value = json.dumps(value, cls=DefaultStateEncoder)
     except json.JSONDecodeError:
         raise ValueError(f"Unsupported variable type: {value.__class__.__name__}")
 

vellum/workflows/nodes/displayable/inline_prompt_node/tests/test_node.py

@@ -1,7 +1,9 @@
+import pytest
 from dataclasses import dataclass
 from uuid import uuid4
 from typing import Any, Iterator, List
 
+from vellum.client.core.api_error import ApiError
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.client.types.execute_prompt_event import ExecutePromptEvent
 from vellum.client.types.fulfilled_execute_prompt_event import FulfilledExecutePromptEvent
@@ -12,6 +14,8 @@ from vellum.client.types.initiated_execute_prompt_event import InitiatedExecuteP
 from vellum.client.types.prompt_output import PromptOutput
 from vellum.client.types.prompt_request_json_input import PromptRequestJsonInput
 from vellum.client.types.string_vellum_value import StringVellumValue
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
 
 
@@ -126,3 +130,54 @@ def test_inline_prompt_node__function_definitions(vellum_adhoc_prompt_client):
     "name": "my_function"
 }"""
     )
+
+
+@pytest.mark.parametrize(
+    ["exception", "expected_code", "expected_message"],
+    [
+        (
+            ApiError(status_code=404, body={"detail": "Model not found"}),
+            WorkflowErrorCode.INVALID_INPUTS,
+            "Model not found",
+        ),
+        (
+            ApiError(status_code=404, body={"message": "Model not found"}),
+            WorkflowErrorCode.INVALID_INPUTS,
+            "Failed to execute prompt",
+        ),
+        (
+            ApiError(status_code=404, body="Model not found"),
+            WorkflowErrorCode.INTERNAL_ERROR,
+            "Failed to execute prompt",
+        ),
+        (
+            ApiError(status_code=None, body={"detail": "Model not found"}),
+            WorkflowErrorCode.INTERNAL_ERROR,
+            "Failed to execute prompt",
+        ),
+        (
+            ApiError(status_code=500, body={"detail": "Model not found"}),
+            WorkflowErrorCode.INTERNAL_ERROR,
+            "Failed to execute prompt",
+        ),
+    ],
+    ids=["404", "invalid_dict", "invalid_body", "no_status_code", "500"],
+)
+def test_inline_prompt_node__api_error__invalid_inputs_node_exception(
+    vellum_adhoc_prompt_client, exception, expected_code, expected_message
+):
+    # GIVEN a prompt node with an invalid model name
+    class MyNode(InlinePromptNode):
+        ml_model = "my-invalid-model"
+        blocks = []
+
+    # AND the adhoc prompt client raises a 4xx error
+    vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.side_effect = exception
+
+    # WHEN the node is run
+    with pytest.raises(NodeException) as e:
+        list(MyNode().run())
+
+    # THEN the node raises the correct NodeException
+    assert e.value.code == expected_code
+    assert e.value.message == expected_message

{vellum_ai-0.13.10.dist-info → vellum_ai-0.13.11.dist-info}/METADATA

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

{vellum_ai-0.13.10.dist-info → vellum_ai-0.13.11.dist-info}/RECORD

@@ -1,39 +1,38 @@
 vellum_cli/CONTRIBUTING.md,sha256=FtDC7BGxSeMnwCXAUssFsAIElXtmJE-O5Z7BpolcgvI,2935
 vellum_cli/README.md,sha256=2NudRoLzWxNKqnuVy1JuQ7DerIaxWGYkrH8kMd-asIE,90
-vellum_cli/__init__.py,sha256=L8D7nVwC63WdAe6HTy7J_S3wCvjJBHDWaBkCrw76dtU,9573
+vellum_cli/__init__.py,sha256=tRcOi6oLzT3EXORlHwGzo0Sn1-npf8Jf7mMi90NDUh0,10035
 vellum_cli/aliased_group.py,sha256=ugW498j0yv4ALJ8vS9MsO7ctDW7Jlir9j6nE_uHAP8c,3363
-vellum_cli/config.py,sha256=998IZbvkrw2avjbvs8bw6NrbEgGz5UBKRbvKAcastJg,5493
+vellum_cli/config.py,sha256=LVRB-SEJcpQYfg2QGcjKHmRSAijdSFADbS90gDY4AI8,6829
 vellum_cli/image_push.py,sha256=SJwhwWJsLjwGNezNVd_oCVpFMfPsAB3dfLWmriZZUtw,4419
 vellum_cli/logger.py,sha256=PuRFa0WCh4sAGFS5aqWB0QIYpS6nBWwPJrIXpWxugV4,1022
 vellum_cli/ping.py,sha256=lWyJw6sziXjyTopTYRdFF5hV-sYPVDdX0yVbG5fzcY4,585
 vellum_cli/pull.py,sha256=zf0y22XptUYI_hMP_4Q1CEo9s2wALsTJcCXNd-_ibd8,7551
-vellum_cli/push.py,sha256=KHakWiUwdeZff8QZSDF0l8xmCgMRo9ntan8kaLD02Lc,7677
+vellum_cli/push.py,sha256=0Q3Of25KzqDq0KTgkXQ22nAhSH4_JI424hKwKkj3RWs,8627
 vellum_cli/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum_cli/tests/conftest.py,sha256=Jv-QUjXoCDhsvVvXEjOenNqRArO_xXhtNuCYt4wiOyE,1421
+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_main.py,sha256=qDZG-aQauPwBwM6A2DIu1494n47v3pL28XakTbLGZ-k,272
 vellum_cli/tests/test_ping.py,sha256=QtbhYKMYn1DFnDyBij2mkQO32j9KOpZ5Pf0yek7k_Ao,1284
-vellum_cli/tests/test_pull.py,sha256=6gbASF6ASy5YcdWjOCt6b5K0u2VWsFegdrTWu6sEVKs,19613
-vellum_cli/tests/test_push.py,sha256=PVxKwbWHjb1QwQ0n4tTqh2Tj6yg2cOGupOSXaXl08DI,11044
+vellum_cli/tests/test_pull.py,sha256=KRvQ_fwsFH6tBU49O-KOOAzcHxdB-DPbUkdIdahjuTs,19841
+vellum_cli/tests/test_push.py,sha256=3to7bOOVoWB888ee7TIAVl2VgXCsJhP_aPcrXwF9c_c,14023
 vellum_ee/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/base.py,sha256=3ZFUYRNKL24fBqXhKpa_Dq2W1a-a86J20dmJYA3H2eY,1755
 vellum_ee/workflows/display/nodes/__init__.py,sha256=5XOcZJXYUgaLS55QgRJzyQ_W1tpeprjnYAeYVezqoGw,160
-vellum_ee/workflows/display/nodes/base_node_display.py,sha256=aucM0ta33owcDwxP-zOvCm9Nb7q7wzjDfxrnyKPV5W4,11423
-vellum_ee/workflows/display/nodes/base_node_vellum_display.py,sha256=BxA-YVUJvB36NM-Q5DNCwckeqymwLKMp9DVCaTyrPbs,2253
-vellum_ee/workflows/display/nodes/get_node_display_class.py,sha256=vyKeJAevAXvEAEtWeTEdBZXD3eJQYW_DEXLKVZ5KmYc,1325
+vellum_ee/workflows/display/nodes/base_node_display.py,sha256=AnACr-QfL4Hwc3JEN9tUDmr0c8DE9uAjP8w7NyXAfj4,15413
+vellum_ee/workflows/display/nodes/base_node_vellum_display.py,sha256=pLO0dORfRu--Ne9NgoyFT_CNjfpr5fGCsgbsMkUF5GM,2845
+vellum_ee/workflows/display/nodes/get_node_display_class.py,sha256=VpGL-tnMM9wEOkuHNxX74RHAFhz5jcVXvmBkKciOBg0,1367
 vellum_ee/workflows/display/nodes/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/nodes/tests/test_base_node_display.py,sha256=QqR3Ly0RNrXwOeLdW5nERDFt0gRPf76n1bPES6o5UN4,1093
 vellum_ee/workflows/display/nodes/types.py,sha256=St1BB6no528OyELGiyRabWao0GGw6mLhstQAvEACbGk,247
 vellum_ee/workflows/display/nodes/utils.py,sha256=sloya5TpXsnot1HURc9L51INwflRqUzHxRVnCS9Cd-4,973
-vellum_ee/workflows/display/nodes/vellum/__init__.py,sha256=_raKY9eKi_OvIFn6nGvf9xKSboKtYLHCWaWCwDQFbOc,1567
+vellum_ee/workflows/display/nodes/vellum/__init__.py,sha256=nmPLj8vkbVCS46XQqmHq8Xj8Mr36wCK_vWf26A9KIkw,1505
 vellum_ee/workflows/display/nodes/vellum/api_node.py,sha256=hoV-cUtS6H9kmRQXHd2py95GRWI_dAnnaPwvlNBkDOQ,8571
-vellum_ee/workflows/display/nodes/vellum/base_node.py,sha256=BE5bn7_SQ23r2jWEFVJcS6fZzB5fXX9uaADya9NIE6Y,4799
 vellum_ee/workflows/display/nodes/vellum/code_execution_node.py,sha256=z00Z3L0d4PsUQo4S8FRDTtOFLtjdi17TJbatNVF4nM8,4288
 vellum_ee/workflows/display/nodes/vellum/conditional_node.py,sha256=ybLIa4uclqVIy3VAQvI1ivg2tnK5Ug_1R5a69DFqL7E,11104
-vellum_ee/workflows/display/nodes/vellum/error_node.py,sha256=Tyx74dUmj_ef0slptoUXHtkjLbNd3f4hXeoEozFaFcw,2023
+vellum_ee/workflows/display/nodes/vellum/error_node.py,sha256=I1Jkp2htRINJATtv1e-zs9BrReFX842djpiVgBPHDYg,2186
 vellum_ee/workflows/display/nodes/vellum/final_output_node.py,sha256=p-PvlnxpBQ7IKskZi2A19jKAtKnSxJ8LPbGMA83VkFk,2805
 vellum_ee/workflows/display/nodes/vellum/guardrail_node.py,sha256=aYZSJTxknU4LMiQdWk9LcK6CkhdozeDEMiRxfAyUNEc,2202
 vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py,sha256=h8K183cZsU2b86WYQlqFtwMkwIHydnNr4ZaaQcWnFWo,7316
@@ -46,24 +45,25 @@ vellum_ee/workflows/display/nodes/vellum/search_node.py,sha256=TxcAGZDl_hvJ7Y1hU
 vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py,sha256=lfevlHpGEX14dEDym6qmnkw3nvzQPTB1_D2ch12B_Rk,2701
 vellum_ee/workflows/display/nodes/vellum/templating_node.py,sha256=JVIMPR3WpveOCWZubHKZkE04mavnTdb_9QY_r3XliRg,3424
 vellum_ee/workflows/display/nodes/vellum/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vellum_ee/workflows/display/nodes/vellum/tests/test_error_node.py,sha256=ulrpoYUW-5kIxfG4Lf5F2p0k_EoYKhmahEbF3P_eruM,1648
 vellum_ee/workflows/display/nodes/vellum/tests/test_try_node.py,sha256=mtzB8LJlFCHVFM4H5AanLp29gQfaVmnN4A4iaRGJHoI,2427
 vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py,sha256=3uT7Gbc0f_mQ3u8uZuCWd0mJ4GtWbz2gbUMySYaVlNE,3774
-vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=tkKILq38n3QXcd2A-RfW8DeHas0I3MII3mexIYYKMtg,5985
+vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=EoU1J7HfcszAZr7ROy_xsNhaDBRiI95-wTK-OzgBwvg,5861
 vellum_ee/workflows/display/nodes/vellum/utils.py,sha256=aBmTg9pJloG1BxhYcqzxs5JOA-IixncLVr1XiDzANnA,4313
 vellum_ee/workflows/display/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/test_vellum_workflow_display.py,sha256=h4bE187MyM0eT_elFnOxwsCeMHeuzSK0orjgAy8kV1k,4952
 vellum_ee/workflows/display/tests/workflow_serialization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py,sha256=g4R4vnZDcpeZ7jCGpqZhDx8z7Kt1dH7r9iDfCLKJ56A,2252
-vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py,sha256=jCvt_plgTaKlf5H7alWhYHihvc-Y4p635UhHCjQSbmM,7115
-vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py,sha256=hcsK36oU9UvYNOdhdKnn2nO36IPzqM4XOxLcz6od9mE,10228
-vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outputs_serialization.py,sha256=kuJek64zXtrQZPcSFENr12EST_zaypPzFh67xIr6PXU,6260
-vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py,sha256=7BRX1Qz1VF_eMqazp6FhlYcOK6sj8p-qu43l7SI4JF0,37351
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py,sha256=EenmEdBtHUFQ0OS-kE7Vboax3JnDdj-K4Qixt5oR0Po,2253
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py,sha256=pYBubuTzwNczAGFTYhaMCwvMLAcoJfCLkn5PMPbrlMI,7187
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py,sha256=9CV44pY33cL_RF9cKmNVwn5hDGEZbPIETavOrvIL3oY,10229
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outputs_serialization.py,sha256=-12ZkZb3f5gyoNASV2yeQtMo5HmNsVEo8nXwL6IC-I8,6261
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py,sha256=6th6kCwzql6lddjkTQx4Jbvvs4ChqtJwctW-B4QuBhI,37352
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py,sha256=EbVgg_3_ipTt3MOop4RARX0fmNjwqZtkhIXzx9nGw7Y,4487
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_api_node_serialization.py,sha256=bXZWxOKAVjZlbP3iLHPHGA4aPs0EguKlQqmYPNGv3cU,16391
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py,sha256=uvABhpstyxNNOz50-XJMAr3SKp8gluPp1mUtJjvRL0c,29410
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py,sha256=n3F_Eyru0DYOszBH4jplz7Mzt2FfBNxGlCkTFqvrX-M,48399
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py,sha256=QVlkczxzaKuOhwbRvVP70otPDNnJcSGDfN79j92lFyk,5534
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py,sha256=8LhQgW0uzVOhzz0AwdM-EYugVO-0mGWglxWo_lON4D8,6079
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_generic_node_serialization.py,sha256=yWDykoHUjuiVYdevcJxWmYDDmbIszpHKYBu19vqR-I8,5746
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_guardrail_node_serialization.py,sha256=mEkvKUrt8U6e9bN65QRG7Zd3KdCdoMvHm96LjGwy96k,7427
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py,sha256=CxbnZ3tB8gir9xopixjvT0hNYyjMHvT_m8BzY5kl8Dg,21497
@@ -83,7 +83,7 @@ vellum_ee/workflows/display/vellum.py,sha256=8xXRI8b8Tt661H-iZreTQTvLNEKUr4lf-Xa
 vellum_ee/workflows/display/workflows/__init__.py,sha256=kapXsC67VJcgSuiBMa86FdePG5A9kMB5Pi4Uy1O2ob4,207
 vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=x4C621UTH_SZLJYmVBiC1q-D-wadJlTiSGlKAI6CUbA,14082
 vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=kp0u8LN_2IwshLrhMImhpZx1hRyAcD5gXY-kDuuaGMQ,1269
-vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=ReEaGb0IBWYeTmEvP5ObcJlTe6RtAogOZlemHkEwBcM,16486
+vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=wij6U6-20QSDlfLMCeHjmpkWWuSA8fWYWJV17dw6L98,17187
 vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/server/virtual_file_loader.py,sha256=X_DdNK7MfyOjKWekk6YQpOSCT6klKcdjT6nVJcBH1sM,1481
 vellum_ee/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -112,7 +112,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=8nZt88C9SVwWanjLbIQMU3rzb32h5UZfFMBx3VPHB50,111887
 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=QFtPX3r8rOjA_HWWWzqeN6YXR94sGx14HQGv7kmubTA,1869
+vellum/client/core/client_wrapper.py,sha256=WdGHbGL6B-WAReI6C1XsH60HY57mtWSg-jZ-2xz4eiQ,1869
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
 vellum/client/core/http_client.py,sha256=R0pQpCppnEtxccGvXl4uJ76s7ro_65Fo_erlNNLp_AI,19228
@@ -1260,7 +1260,7 @@ vellum/workflows/emitters/base.py,sha256=D5SADKIvnbgKwIBgYm77jaqvpo1o0rz4MmuX_mu
 vellum/workflows/environment/__init__.py,sha256=wGHslgSEZ7Octe4C-hNtl84EFelNimgmWQoi7px4-uw,71
 vellum/workflows/environment/environment.py,sha256=IwcVSwA51jW1JTN_MAN22UbeAtZhiAVa_JRV4rmpAAs,219
 vellum/workflows/errors/__init__.py,sha256=tWGPu5xyAU8gRb8_bl0fL7OfU3wxQ9UH6qVwy4X4P_Q,113
-vellum/workflows/errors/types.py,sha256=mvywC2sXtpRrkDRfTx-O6aXAp0hOFyx27a1tqguKt4g,2493
+vellum/workflows/errors/types.py,sha256=-Ls41oM4HN28b78Ou13FoZnaWRuaxBgTarKgSoBlCKc,3413
 vellum/workflows/events/__init__.py,sha256=6pxxceJo2dcaRkWtkDAYlUQZ-PHBQSZytIoyuUK48Qw,759
 vellum/workflows/events/node.py,sha256=uHT6If0esgZ3nLjrjmUPTKf3qbjGhoV_x5YKpjDBDcU,5280
 vellum/workflows/events/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1339,16 +1339,16 @@ vellum/workflows/nodes/displayable/bases/__init__.py,sha256=0mWIx3qUrzllV7jqt7wN
 vellum/workflows/nodes/displayable/bases/api_node/__init__.py,sha256=1jwx4WC358CLA1jgzl_UD-rZmdMm2v9Mps39ndwCD7U,64
 vellum/workflows/nodes/displayable/bases/api_node/node.py,sha256=kKiwrz8XGJkvFjipBthxs9h1e2ndWPgmpEzZiYs0I5I,2579
 vellum/workflows/nodes/displayable/bases/base_prompt_node/__init__.py,sha256=Org3xTvgp1pA0uUXFfnJr29D3HzCey2lEdYF4zbIUgo,70
-vellum/workflows/nodes/displayable/bases/base_prompt_node/node.py,sha256=reDW1aHetmkesyq9iVrbTtesql-SFUsd2LygGLfdJPo,2258
+vellum/workflows/nodes/displayable/bases/base_prompt_node/node.py,sha256=r_gOmeGFQQfzhAc1_bmzcwUvH-Xllc93gE7miTV4rQE,2824
 vellum/workflows/nodes/displayable/bases/inline_prompt_node/__init__.py,sha256=Hl35IAoepRpE-j4cALaXVJIYTYOF3qszyVbxTj4kS1s,82
 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=j6AKAVIl_uU77a6JMGnMTyFnnK-jsfPVthxrgbltQRE,6453
 vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py,sha256=FAY6cR3rlIWsufprA6WXefOaU2J1LK0HFfQkBxCQ-yI,5181
 vellum/workflows/nodes/displayable/bases/search_node.py,sha256=3UtbqY3QO4kzfJHbmUNZGnEEfJmaoiF892u8H6TGjp8,5381
 vellum/workflows/nodes/displayable/bases/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/workflows/nodes/displayable/bases/tests/test_utils.py,sha256=uQM9RTVmC-fMQ4hd4481QZbviokshwYSxk1ee9leh3A,2406
+vellum/workflows/nodes/displayable/bases/tests/test_utils.py,sha256=eqdqbKNRWVMDPevgwLg1i6YK0g4L4bCy-7xCBN5yYZI,3156
 vellum/workflows/nodes/displayable/bases/types.py,sha256=C37B2Qh2YP7s7pUjd-EYKc2Zl1TbnCgI_mENuUSb8bo,1706
-vellum/workflows/nodes/displayable/bases/utils.py,sha256=hlbIRLAAeqRhzZPLK7WDvBzUT2b6tupaXYOF54iYvS0,4973
+vellum/workflows/nodes/displayable/bases/utils.py,sha256=ckMUenSsNkiYmSw6FmjSMHYaCk8Y8_sUjL6lkFFEqts,5412
 vellum/workflows/nodes/displayable/code_execution_node/__init__.py,sha256=0FLWMMktpzSnmBMizQglBpcPrP80fzVsoJwJgf822Cg,76
 vellum/workflows/nodes/displayable/code_execution_node/node.py,sha256=Hzc0h9sogX8DfFbJqdDLxGU2CMY1WvGosZ7taJC4sPE,8754
 vellum/workflows/nodes/displayable/code_execution_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1365,7 +1365,7 @@ vellum/workflows/nodes/displayable/guardrail_node/node.py,sha256=h5nIBzQxbXTrdTq
 vellum/workflows/nodes/displayable/inline_prompt_node/__init__.py,sha256=gSUOoEZLlrx35-tQhSAd3An8WDwBqyiQh-sIebLU9wU,74
 vellum/workflows/nodes/displayable/inline_prompt_node/node.py,sha256=UWOxIa3kizhDa5joAmEDlh0P66EaqatqxzLLbvbUybw,2294
 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=57JYbHvfZAjEGfa4CvYSpAt5yLDj0bPIr1_RoG8n9oI,4725
+vellum/workflows/nodes/displayable/inline_prompt_node/tests/test_node.py,sha256=BEKdts18soB6gaW_d3VbhDy7IFzFmC76j5lRYiTOj44,6643
 vellum/workflows/nodes/displayable/merge_node/__init__.py,sha256=J8IC08dSH7P76wKlNuxe1sn7toNGtSQdFirUbtPDEs0,60
 vellum/workflows/nodes/displayable/merge_node/node.py,sha256=nZtGGVAvY4fvGg8vwV6sTQ8_QLRnigeXt0vf2FL272A,450
 vellum/workflows/nodes/displayable/note_node/__init__.py,sha256=KWA3P4fyYJ-fOTky8qNGlcOotQ-HeHJ9AjZt6mRQmCE,58
@@ -1441,8 +1441,8 @@ vellum/workflows/vellum_client.py,sha256=ODrq_TSl-drX2aezXegf7pizpWDVJuTXH-j6528
 vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
 vellum/workflows/workflows/base.py,sha256=k0kUWWko4fHyCqLSU_1cBK_pXZpl9MXekWiG-bdOAo0,18353
 vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
-vellum_ai-0.13.10.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.13.10.dist-info/METADATA,sha256=qsbrVIIP35E7Q0tVUt2vd_NXWi9dYkUcksWgKYnCPy8,5335
-vellum_ai-0.13.10.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.13.10.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.13.10.dist-info/RECORD,,
+vellum_ai-0.13.11.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.13.11.dist-info/METADATA,sha256=NHPCibqYzo79FZeV0pFNsUB0dbOubh76GZRnzEJAc74,5335
+vellum_ai-0.13.11.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.13.11.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.13.11.dist-info/RECORD,,

vellum_cli/__init__.py

@@ -37,13 +37,15 @@ class PushGroup(ClickAliasedGroup):
 
 @main.group(invoke_without_command=True, cls=PushGroup)
 @click.pass_context
+@click.option("--workspace", type=str, help="The specific Workspace config to use when pushing")
 def push(
     ctx: click.Context,
+    workspace: Optional[str],
 ) -> None:
     """Push Resources to Vellum"""
 
     if ctx.invoked_subcommand is None:
-        push_command()
+        push_command(workspace=workspace)
 
 
 @main.group()
@@ -69,6 +71,7 @@ def workflows():
     is_flag=True,
     help="Raises an error if we detect an unexpected discrepancy in the generated artifact.",
 )
+@click.option("--workspace", type=str, help="The specific Workspace config to use when pushing")
 def workflows_push(
     module: Optional[str],
     deploy: Optional[bool],
@@ -78,6 +81,7 @@ def workflows_push(
     release_tag: Optional[List[str]],
     dry_run: Optional[bool],
     strict: Optional[bool],
+    workspace: Optional[str],
 ) -> None:
     """
     Push Workflows to Vellum. If a module is provided, only the Workflow for that module will be pushed.
@@ -93,6 +97,7 @@ def workflows_push(
         release_tags=release_tag,
         dry_run=dry_run,
         strict=strict,
+        workspace=workspace,
     )
 
 
@@ -113,6 +118,7 @@ def workflows_push(
     is_flag=True,
     help="Raises an error if we detect an unexpected discrepancy in the generated artifact.",
 )
+@click.option("--workspace", type=str, help="The specific Workspace config to use when pushing")
 def push_module(
     ctx: click.Context,
     deploy: Optional[bool],
@@ -122,6 +128,7 @@ def push_module(
     release_tag: Optional[List[str]],
     dry_run: Optional[bool],
     strict: Optional[bool],
+    workspace: Optional[str],
 ) -> None:
     """Push a specific module to Vellum"""
 
@@ -135,6 +142,7 @@ def push_module(
             release_tags=release_tag,
             dry_run=dry_run,
             strict=strict,
+            workspace=workspace,
         )
 
 

vellum_cli/config.py

@@ -13,6 +13,17 @@ LOCKFILE_PATH = "vellum.lock.json"
 PYPROJECT_TOML_PATH = "pyproject.toml"
 
 
+class WorkspaceConfig(UniversalBaseModel):
+    name: str
+    api_key: str
+
+    def merge(self, other: "WorkspaceConfig") -> "WorkspaceConfig":
+        return WorkspaceConfig(name=self.name or other.name, api_key=self.api_key or other.api_key)
+
+
+DEFAULT_WORKSPACE_CONFIG = WorkspaceConfig(name="default", api_key="VELLUM_API_KEY")
+
+
 class WorkflowDeploymentConfig(UniversalBaseModel):
     id: Optional[UUID] = None
     label: Optional[str] = None
@@ -37,6 +48,7 @@ class WorkflowConfig(UniversalBaseModel):
     deployments: List[WorkflowDeploymentConfig] = field(default_factory=list)
     container_image_name: Optional[str] = None
     container_image_tag: Optional[str] = None
+    workspace: str = DEFAULT_WORKSPACE_CONFIG.name
 
     def merge(self, other: "WorkflowConfig") -> "WorkflowConfig":
         self_deployment_by_id = {
@@ -72,12 +84,14 @@ class WorkflowConfig(UniversalBaseModel):
             deployments=merged_deployments,
             container_image_tag=self.container_image_tag or other.container_image_tag,
             container_image_name=self.container_image_name or other.container_image_name,
+            workspace=self.workspace or other.workspace,
         )
 
 
 class VellumCliConfig(UniversalBaseModel):
     version: Literal["1.0"] = "1.0"
     workflows: List[WorkflowConfig] = field(default_factory=list)
+    workspaces: List[WorkspaceConfig] = field(default_factory=list)
 
     def save(self) -> None:
         lockfile_path = os.path.join(os.getcwd(), LOCKFILE_PATH)
@@ -102,7 +116,21 @@ class VellumCliConfig(UniversalBaseModel):
             elif other_workflow:
                 merged_workflows.append(other_workflow)
 
-        return VellumCliConfig(workflows=merged_workflows, version=self.version)
+        self_workspace_by_name = {workspace.name: workspace for workspace in self.workspaces}
+        other_workspace_by_name = {workspace.name: workspace for workspace in other.workspaces}
+        all_names = sorted(set(self_workspace_by_name.keys()).union(set(other_workspace_by_name.keys())))
+        merged_workspaces = []
+        for name in all_names:
+            self_workspace = self_workspace_by_name.get(name)
+            other_workspace = other_workspace_by_name.get(name)
+            if self_workspace and other_workspace:
+                merged_workspaces.append(self_workspace.merge(other_workspace))
+            elif self_workspace:
+                merged_workspaces.append(self_workspace)
+            elif other_workspace:
+                merged_workspaces.append(other_workspace)
+
+        return VellumCliConfig(workflows=merged_workflows, workspaces=merged_workspaces, version=self.version)
 
 
 def load_vellum_cli_config(root_dir: Optional[str] = None) -> VellumCliConfig:

vellum_cli/push.py

@@ -15,7 +15,7 @@ from vellum.types import WorkflowPushDeploymentConfigRequest
 from vellum.workflows.utils.names import snake_to_title_case
 from vellum.workflows.vellum_client import create_vellum_client
 from vellum.workflows.workflows.base import BaseWorkflow
-from vellum_cli.config import WorkflowDeploymentConfig, load_vellum_cli_config
+from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, WorkflowConfig, WorkflowDeploymentConfig, load_vellum_cli_config
 from vellum_cli.logger import load_cli_logger
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
@@ -30,8 +30,9 @@ def push_command(
     release_tags: Optional[List[str]] = None,
     dry_run: Optional[bool] = None,
     strict: Optional[bool] = None,
+    workspace: Optional[str] = None,
 ) -> None:
-    load_dotenv()
+    load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
     logger = load_cli_logger()
     config = load_vellum_cli_config()
 
@@ -46,7 +47,27 @@ def push_command(
         raise ValueError(f"No workflow config for '{module}' found in project to push.")
 
     logger.info(f"Loading workflow from {workflow_config.module}")
-    client = create_vellum_client()
+    resolved_workspace = workspace or workflow_config.workspace or DEFAULT_WORKSPACE_CONFIG.name
+    workspace_config = (
+        next((w for w in config.workspaces if w.name == resolved_workspace), DEFAULT_WORKSPACE_CONFIG)
+        if workspace
+        else DEFAULT_WORKSPACE_CONFIG
+    )
+    api_key = os.getenv(workspace_config.api_key)
+    if not api_key:
+        raise ValueError(f"No API key value found in environment for workspace '{workspace_config.name}'.")
+
+    if workspace_config.name != workflow_config.workspace:
+        # We are pushing to a new workspace, so we need a new workflow config
+        workflow_config = WorkflowConfig(
+            module=workflow_config.module,
+            workspace=workspace_config.name,
+        )
+        config.workflows.append(workflow_config)
+
+    client = create_vellum_client(
+        api_key=api_key,
+    )
     sys.path.insert(0, os.getcwd())
 
     # Remove this once we could serialize using the artifact in Vembda

vellum_cli/tests/conftest.py

@@ -46,6 +46,9 @@ def mock_module(request) -> Generator[MockModuleResult, None, None]:
         }
     )
 
+    with open(os.path.join(temp_dir, ".env"), "w") as f:
+        f.write("VELLUM_API_KEY=abcdef123456")
+
     yield MockModuleResult(
         temp_dir=temp_dir,
         module=module,

vellum_cli/tests/test_pull.py

@@ -124,6 +124,7 @@ def test_pull__sandbox_id_with_no_config(vellum_client):
         lock_data = json.loads(f.read())
         assert lock_data == {
             "version": "1.0",
+            "workspaces": [],
             "workflows": [
                 {
                     "module": "workflow_87654321",
@@ -132,6 +133,7 @@ def test_pull__sandbox_id_with_no_config(vellum_client):
                     "deployments": [],
                     "container_image_tag": None,
                     "container_image_name": None,
+                    "workspace": "default",
                 }
             ],
         }
@@ -208,8 +210,10 @@ def test_pull__workflow_deployment_with_no_config(vellum_client):
                     "deployments": [],
                     "container_image_tag": None,
                     "container_image_name": None,
+                    "workspace": "default",
                 }
             ],
+            "workspaces": [],
         }
 
 
@@ -449,6 +453,7 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                 ],
                 "container_image_name": None,
                 "container_image_tag": None,
+                "workspace": "default",
             },
             {
                 "module": "workflow_87654321",
@@ -457,6 +462,7 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                 "deployments": [],
                 "container_image_name": "test",
                 "container_image_tag": "1.0",
+                "workspace": "default",
             },
         ]