vellum-ai 0.14.15__py3-none-any.whl → 0.14.16__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.14.15",
+            "X-Fern-SDK-Version": "0.14.16",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/client/resources/document_indexes/client.py

@@ -10,7 +10,6 @@ from json.decoder import JSONDecodeError
 from ...core.api_error import ApiError
 from ...types.document_index_indexing_config_request import DocumentIndexIndexingConfigRequest
 from ...types.entity_status import EntityStatus
-from ...types.environment_enum import EnvironmentEnum
 from ...types.document_index_read import DocumentIndexRead
 from ...core.serialization import convert_and_respect_annotation_metadata
 from ...core.jsonable_encoder import jsonable_encoder
@@ -108,7 +107,6 @@ class DocumentIndexesClient:
         name: str,
         indexing_config: DocumentIndexIndexingConfigRequest,
         status: typing.Optional[EntityStatus] = OMIT,
-        environment: typing.Optional[EnvironmentEnum] = OMIT,
         copy_documents_from_index_id: typing.Optional[str] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> DocumentIndexRead:
@@ -131,13 +129,6 @@ class DocumentIndexesClient:
             * `ACTIVE` - Active
             * `ARCHIVED` - Archived
 
-        environment : typing.Optional[EnvironmentEnum]
-            The environment this document index is used in
-
-            * `DEVELOPMENT` - Development
-            * `STAGING` - Staging
-            * `PRODUCTION` - Production
-
         copy_documents_from_index_id : typing.Optional[str]
             Optionally specify the id of a document index from which you'd like to copy and re-index its documents into this newly created index
 
@@ -191,7 +182,6 @@ class DocumentIndexesClient:
                 "label": label,
                 "name": name,
                 "status": status,
-                "environment": environment,
                 "indexing_config": convert_and_respect_annotation_metadata(
                     object_=indexing_config, annotation=DocumentIndexIndexingConfigRequest, direction="write"
                 ),
@@ -268,7 +258,6 @@ class DocumentIndexesClient:
         *,
         label: str,
         status: typing.Optional[EntityStatus] = OMIT,
-        environment: typing.Optional[EnvironmentEnum] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> DocumentIndexRead:
         """
@@ -288,13 +277,6 @@ class DocumentIndexesClient:
             * `ACTIVE` - Active
             * `ARCHIVED` - Archived
 
-        environment : typing.Optional[EnvironmentEnum]
-            The environment this document index is used in
-
-            * `DEVELOPMENT` - Development
-            * `STAGING` - Staging
-            * `PRODUCTION` - Production
-
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
 
@@ -322,7 +304,6 @@ class DocumentIndexesClient:
             json={
                 "label": label,
                 "status": status,
-                "environment": environment,
             },
             request_options=request_options,
             omit=OMIT,
@@ -388,7 +369,6 @@ class DocumentIndexesClient:
         *,
         label: typing.Optional[str] = OMIT,
         status: typing.Optional[EntityStatus] = OMIT,
-        environment: typing.Optional[EnvironmentEnum] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> DocumentIndexRead:
         """
@@ -408,13 +388,6 @@ class DocumentIndexesClient:
             * `ACTIVE` - Active
             * `ARCHIVED` - Archived
 
-        environment : typing.Optional[EnvironmentEnum]
-            The environment this document index is used in
-
-            * `DEVELOPMENT` - Development
-            * `STAGING` - Staging
-            * `PRODUCTION` - Production
-
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
 
@@ -441,7 +414,6 @@ class DocumentIndexesClient:
             json={
                 "label": label,
                 "status": status,
-                "environment": environment,
             },
             request_options=request_options,
             omit=OMIT,
@@ -651,7 +623,6 @@ class AsyncDocumentIndexesClient:
         name: str,
         indexing_config: DocumentIndexIndexingConfigRequest,
         status: typing.Optional[EntityStatus] = OMIT,
-        environment: typing.Optional[EnvironmentEnum] = OMIT,
         copy_documents_from_index_id: typing.Optional[str] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> DocumentIndexRead:
@@ -674,13 +645,6 @@ class AsyncDocumentIndexesClient:
             * `ACTIVE` - Active
             * `ARCHIVED` - Archived
 
-        environment : typing.Optional[EnvironmentEnum]
-            The environment this document index is used in
-
-            * `DEVELOPMENT` - Development
-            * `STAGING` - Staging
-            * `PRODUCTION` - Production
-
         copy_documents_from_index_id : typing.Optional[str]
             Optionally specify the id of a document index from which you'd like to copy and re-index its documents into this newly created index
 
@@ -742,7 +706,6 @@ class AsyncDocumentIndexesClient:
                 "label": label,
                 "name": name,
                 "status": status,
-                "environment": environment,
                 "indexing_config": convert_and_respect_annotation_metadata(
                     object_=indexing_config, annotation=DocumentIndexIndexingConfigRequest, direction="write"
                 ),
@@ -827,7 +790,6 @@ class AsyncDocumentIndexesClient:
         *,
         label: str,
         status: typing.Optional[EntityStatus] = OMIT,
-        environment: typing.Optional[EnvironmentEnum] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> DocumentIndexRead:
         """
@@ -847,13 +809,6 @@ class AsyncDocumentIndexesClient:
             * `ACTIVE` - Active
             * `ARCHIVED` - Archived
 
-        environment : typing.Optional[EnvironmentEnum]
-            The environment this document index is used in
-
-            * `DEVELOPMENT` - Development
-            * `STAGING` - Staging
-            * `PRODUCTION` - Production
-
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
 
@@ -889,7 +844,6 @@ class AsyncDocumentIndexesClient:
             json={
                 "label": label,
                 "status": status,
-                "environment": environment,
             },
             request_options=request_options,
             omit=OMIT,
@@ -963,7 +917,6 @@ class AsyncDocumentIndexesClient:
         *,
         label: typing.Optional[str] = OMIT,
         status: typing.Optional[EntityStatus] = OMIT,
-        environment: typing.Optional[EnvironmentEnum] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> DocumentIndexRead:
         """
@@ -983,13 +936,6 @@ class AsyncDocumentIndexesClient:
             * `ACTIVE` - Active
             * `ARCHIVED` - Archived
 
-        environment : typing.Optional[EnvironmentEnum]
-            The environment this document index is used in
-
-            * `DEVELOPMENT` - Development
-            * `STAGING` - Staging
-            * `PRODUCTION` - Production
-
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
 
@@ -1024,7 +970,6 @@ class AsyncDocumentIndexesClient:
             json={
                 "label": label,
                 "status": status,
-                "environment": environment,
             },
             request_options=request_options,
             omit=OMIT,

vellum/client/types/document_index_read.py

@@ -5,7 +5,6 @@ import datetime as dt
 import pydantic
 import typing
 from .entity_status import EntityStatus
-from .environment_enum import EnvironmentEnum
 from .document_index_indexing_config import DocumentIndexIndexingConfig
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
 
@@ -31,15 +30,6 @@ class DocumentIndexRead(UniversalBaseModel):
     * `ARCHIVED` - Archived
     """
 
-    environment: typing.Optional[EnvironmentEnum] = pydantic.Field(default=None)
-    """
-    The environment this document index is used in
-    
-    * `DEVELOPMENT` - Development
-    * `STAGING` - Staging
-    * `PRODUCTION` - Production
-    """
-
     indexing_config: DocumentIndexIndexingConfig
 
     if IS_PYDANTIC_V2:

vellum/plugins/pydantic.py

@@ -1,17 +1,27 @@
 from functools import lru_cache
-from typing import Any, Dict, Literal, Optional, Tuple, Union
+from typing import Any, Dict, Literal, NamedTuple, Optional, Tuple, Union
+from typing_extensions import TypeAlias
 
 from pydantic.fields import FieldInfo
 from pydantic.plugin import (
     PydanticPluginProtocol,
-    SchemaKind,
-    SchemaTypePath,
     ValidateJsonHandlerProtocol,
     ValidatePythonHandlerProtocol,
     ValidateStringsHandlerProtocol,
 )
 from pydantic_core import CoreSchema
 
+# Redefined manually instead of imported from pydantic to support versions < 2.5
+SchemaKind: TypeAlias = Literal["BaseModel", "TypeAdapter", "dataclass", "create_model", "validate_call"]
+
+
+# Redefined manually instead of imported from pydantic to support versions < 2.5
+class SchemaTypePath(NamedTuple):
+    """Path defining where `schema_type` was defined, or where `TypeAdapter` was called."""
+
+    module: str
+    name: str
+
 
 @lru_cache(maxsize=1)
 def import_base_descriptor():
@@ -81,7 +91,7 @@ class VellumPydanticPlugin(PydanticPluginProtocol):
         self,
         schema: CoreSchema,
         schema_type: Any,
-        schema_type_path: SchemaTypePath,
+        schema_type_path: SchemaTypePath,  # type: ignore
         schema_kind: SchemaKind,
         config: Any,
         plugin_settings: Dict[str, Any],

vellum/workflows/vellum_client.py

@@ -10,9 +10,13 @@ def create_vellum_client(api_key: Optional[str] = None) -> Vellum:
 
     return Vellum(
         api_key=api_key,
-        environment=VellumEnvironment(
-            default=os.getenv("VELLUM_DEFAULT_API_URL", os.getenv("VELLUM_API_URL", "https://api.vellum.ai")),
-            documents=os.getenv("VELLUM_DOCUMENTS_API_URL", os.getenv("VELLUM_API_URL", "https://documents.vellum.ai")),
-            predict=os.getenv("VELLUM_PREDICT_API_URL", os.getenv("VELLUM_API_URL", "https://predict.vellum.ai")),
-        ),
+        environment=create_vellum_environment(),
+    )
+
+
+def create_vellum_environment() -> VellumEnvironment:
+    return VellumEnvironment(
+        default=os.getenv("VELLUM_DEFAULT_API_URL", os.getenv("VELLUM_API_URL", "https://api.vellum.ai")),
+        documents=os.getenv("VELLUM_DOCUMENTS_API_URL", os.getenv("VELLUM_API_URL", "https://documents.vellum.ai")),
+        predict=os.getenv("VELLUM_PREDICT_API_URL", os.getenv("VELLUM_API_URL", "https://predict.vellum.ai")),
     )

{vellum_ai-0.14.15.dist-info → vellum_ai-0.14.16.dist-info}/METADATA

@@ -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.dist-info → vellum_ai-0.14.16.dist-info}/RECORD

@@ -3,7 +3,7 @@ vellum_cli/README.md,sha256=2NudRoLzWxNKqnuVy1JuQ7DerIaxWGYkrH8kMd-asIE,90
 vellum_cli/__init__.py,sha256=7aO9XFnaEVRiVshn86cFudebFUccT-gV8xIARJWqKYo,12257
 vellum_cli/aliased_group.py,sha256=ugW498j0yv4ALJ8vS9MsO7ctDW7Jlir9j6nE_uHAP8c,3363
 vellum_cli/config.py,sha256=aKnhvM5B8QdPA4cQC5Sqg7ImP-vNcVdSkZmk_OBpQTw,9309
-vellum_cli/image_push.py,sha256=4auU15Pb6c8DTGvT-AQ5HHXXrvIvEDs6L02d4OvJYI8,5199
+vellum_cli/image_push.py,sha256=8DDvRDJEZ-FukUCqGW1827bg1ybF4xBbx9WyqWYQE-g,6816
 vellum_cli/init.py,sha256=WpnMXPItPmh0f0bBGIer3p-e5gu8DUGwSArT_FuoMEw,5093
 vellum_cli/logger.py,sha256=PuRFa0WCh4sAGFS5aqWB0QIYpS6nBWwPJrIXpWxugV4,1022
 vellum_cli/ping.py,sha256=lWyJw6sziXjyTopTYRdFF5hV-sYPVDdX0yVbG5fzcY4,585
@@ -12,6 +12,7 @@ vellum_cli/push.py,sha256=xjTNbLwOVFNU3kpBrm56Bk5QkSRrJ9z86qceghCzfIA,9655
 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=i3lJuW8nFRwL1M1OF6752IZYvGAFgKmkB2hd_rYlsmg,2028
 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=QtbhYKMYn1DFnDyBij2mkQO32j9KOpZ5Pf0yek7k_Ao,1284
@@ -61,11 +62,11 @@ vellum_ee/workflows/display/tests/test_vellum_workflow_display.py,sha256=VD-4USi
 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=A1-tIpC5KIKG9JA_rkd1nLS8zUG3Kb4QiVdvb3boFxE,2509
-vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py,sha256=mFmxXfnUPrwndaBurW9E-VSBUQjF3cGv3JpRuNmwWh8,15475
-vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py,sha256=1cszL6N6FNGVm61MOa7AEiHnF0QjZWqDQuPOp4yiG94,18277
-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=FD94h96w5wWYlLVC8_tqr2aUHPyO2UQh4f03jDX4re4,39783
-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/generic_nodes/test_adornments_serialization.py,sha256=2ZDTmpiU2CcZe4ukxoQ79SJx8V8kDRQHgUxP5Vtesus,15605
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py,sha256=lScpUiz8fvUmJK18Cpygs9xQrfHQlkSIlmxapGYTS0g,18770
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outputs_serialization.py,sha256=KgxSN8L5AWPL8EpQfSZiTqK428nAIElcWrV-sR17528,6456
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py,sha256=yS8GpjbH0objjcNYK73j8rDHbhqinIbG7CcoNjnolBg,40611
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py,sha256=_-ADur28OkVy4YEgMHkkVYqMbS1j3GcOtX9MTYmmvf0,4632
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_api_node_serialization.py,sha256=IRazH2QR6F8RGqNemEnHueyj5DtEa6rFTYhT16S4jI8,15917
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py,sha256=V__y7uu-dy6TJjPeu4UDvaoO2yYwBRbPiW9uJdzWRx4,29828
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py,sha256=R8DW1DUb0DOSLtnF2E1HaCTmtpG-ski0LfcM2WeLVNo,47672
@@ -123,7 +124,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=tKtdM1_GqmGq1gpi9ydWD_T-MM7fPn8QdHh8ww19cNI,117564
 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=x6cZlx3EA5bHNlGRnX0klOVCrqkvqWv_oUW4xjYEehU,1869
+vellum/client/core/client_wrapper.py,sha256=0XmYpWbtLyEE04hBfhFVCncXeuj0BKKSGnpz2iz-nyA,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
@@ -150,7 +151,7 @@ vellum/client/resources/deployments/types/__init__.py,sha256=29GVdoLOJsADSSSqZwb
 vellum/client/resources/deployments/types/deployments_list_request_status.py,sha256=CxlQD16KZXme7x31YYCe_3aAgEueutDTeJo5A4Au-aU,174
 vellum/client/resources/deployments/types/list_deployment_release_tags_request_source.py,sha256=hRGgWMYZL9uKCmD_2dU8-u9RCPUUGItpNn1tUY-NXKY,180
 vellum/client/resources/document_indexes/__init__.py,sha256=YpOl_9IV7xOlH4OmusQxtAJB11kxQfCSMDyT1_UD0oM,165
-vellum/client/resources/document_indexes/client.py,sha256=XDrceFkbX1MtrR4hoFR4b-UIC2PFS1AAlKPCWR6OcZQ,37520
+vellum/client/resources/document_indexes/client.py,sha256=UcznU0NyvdNBpV4UCsTqG3Ejj6w4dK4UBmgrTWzZFtw,35438
 vellum/client/resources/document_indexes/types/__init__.py,sha256=IoFqKHN_VBdEhC7VL8_6Jbatrn0e0zuYEJAJUahcUR0,196
 vellum/client/resources/document_indexes/types/document_indexes_list_request_status.py,sha256=sfUEB0cvOSmlE2iITqnMVyHv05Zy2fWP4QjCIYqMg0M,178
 vellum/client/resources/documents/__init__.py,sha256=FTtvy8EDg9nNNg9WCatVgKTRYV8-_v1roeGPAKoa_pw,65
@@ -266,7 +267,7 @@ vellum/client/types/document_index_chunking.py,sha256=TU0Y7z0Xacm3dhzEDuDIG3ZKJC
 vellum/client/types/document_index_chunking_request.py,sha256=g9BKCsHKg5kzjG7YYeMNQ_5R8TXLeSgumJlMXoSfBcs,435
 vellum/client/types/document_index_indexing_config.py,sha256=xL1pCzUOkw5sSie1OrBpasE3bVnv0UyZBn7uZztbhbs,781
 vellum/client/types/document_index_indexing_config_request.py,sha256=Wt-ys1o_acHNyLU0c1laG2PVT7rgCfwO54f5nudAxk4,832
-vellum/client/types/document_index_read.py,sha256=cXL115A4h-TFiGc29tAkXb1pkuK0RzIquyOu1Pv7Jug,1469
+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
@@ -719,7 +720,7 @@ vellum/evaluations/utils/env.py,sha256=Xj_nxsoU5ox06EOTjRopR4lrigQI6Le6qbWGltYoE
 vellum/evaluations/utils/exceptions.py,sha256=dXMAkzqbHV_AP5FjjbegPlfUE0zQDlpA3qOsoOJUxfg,49
 vellum/evaluations/utils/paginator.py,sha256=rEED_BJAXAM6tM1yMwHePNzszjq_tTq4NbQvi1jWQ_Q,697
 vellum/plugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/plugins/pydantic.py,sha256=BMUwC4OkxtNf6hpLwb0T4MBUKgkGy55T_5Ww7GHpHYo,3068
+vellum/plugins/pydantic.py,sha256=XpQQs_o8DQY_5gFFDzJzaXahUqR_l5TEI4b5ZoH9ob8,3528
 vellum/plugins/utils.py,sha256=cPmxE9R2CK1bki2jKE8rB-G9zMf2pzHjSPDHFPXwd3Q,878
 vellum/plugins/vellum_mypy.py,sha256=QTuMSq6PiZW1dyTUZ5Bf1d4XkgFj0TKAgZLP8f4UgL4,27914
 vellum/prompts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1505,15 +1506,15 @@ vellum/workflows/utils/tests/test_uuids.py,sha256=i77ABQ0M3S-aFLzDXHJq_yr5FPkJEW
 vellum/workflows/utils/tests/test_vellum_variables.py,sha256=6H-BpmbIEmVRO75QQ3Rfy4bEUMMP2qwGzx2Gp1uXbfw,879
 vellum/workflows/utils/uuids.py,sha256=DFzPv9RCvsKhvdTEIQyfSek2A31D6S_QcmeLPbgrgTY,739
 vellum/workflows/utils/vellum_variables.py,sha256=fC2aSLvlS31D15dOWu43LBRR0QsgUKNXBiCUvvaLXSs,3231
-vellum/workflows/vellum_client.py,sha256=ODrq_TSl-drX2aezXegf7pizpWDVJuTXH-j6528t75s,683
+vellum/workflows/vellum_client.py,sha256=GxOy3dX6A04xiY69vPv1S4YGuQ_TMxwHi6WRMimQBBE,762
 vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
 vellum/workflows/workflows/base.py,sha256=TSS2BHC8LAi-N5GdEa75BeChwzwTzL7yldFnTlLINro,22665
 vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
 vellum/workflows/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/workflows/tests/test_base_workflow.py,sha256=NRteiICyJvDM5zrtUfq2fZoXcGQVaWC9xmNlLLVW0cU,7979
 vellum/workflows/workflows/tests/test_context.py,sha256=VJBUcyWVtMa_lE5KxdhgMu0WYNYnUQUDvTF7qm89hJ0,2333
-vellum_ai-0.14.15.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.14.15.dist-info/METADATA,sha256=8TM16vJdsQpUXVDgasscuRmbLESHxIgmMj8QlKoLUQU,5408
-vellum_ai-0.14.15.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.14.15.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.14.15.dist-info/RECORD,,
+vellum_ai-0.14.16.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.14.16.dist-info/METADATA,sha256=yxCtedPJiwtcov54unW-4jIJErky44J7Mhs3BEedjW8,5408
+vellum_ai-0.14.16.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.14.16.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.14.16.dist-info/RECORD,,

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_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