databricks-sdk 0.29.0__py3-none-any.whl → 0.31.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of databricks-sdk might be problematic. Click here for more details.

Files changed (30) hide show
  1. databricks/sdk/__init__.py +89 -21
  2. databricks/sdk/config.py +61 -75
  3. databricks/sdk/core.py +16 -9
  4. databricks/sdk/credentials_provider.py +15 -15
  5. databricks/sdk/data_plane.py +65 -0
  6. databricks/sdk/errors/overrides.py +8 -0
  7. databricks/sdk/errors/platform.py +5 -0
  8. databricks/sdk/mixins/files.py +12 -4
  9. databricks/sdk/service/apps.py +977 -0
  10. databricks/sdk/service/billing.py +602 -218
  11. databricks/sdk/service/catalog.py +324 -34
  12. databricks/sdk/service/compute.py +766 -81
  13. databricks/sdk/service/dashboards.py +628 -18
  14. databricks/sdk/service/iam.py +99 -88
  15. databricks/sdk/service/jobs.py +332 -23
  16. databricks/sdk/service/marketplace.py +2 -122
  17. databricks/sdk/service/oauth2.py +127 -70
  18. databricks/sdk/service/pipelines.py +72 -52
  19. databricks/sdk/service/serving.py +303 -750
  20. databricks/sdk/service/settings.py +423 -4
  21. databricks/sdk/service/sharing.py +235 -25
  22. databricks/sdk/service/sql.py +2328 -544
  23. databricks/sdk/useragent.py +151 -0
  24. databricks/sdk/version.py +1 -1
  25. {databricks_sdk-0.29.0.dist-info → databricks_sdk-0.31.0.dist-info}/METADATA +36 -16
  26. {databricks_sdk-0.29.0.dist-info → databricks_sdk-0.31.0.dist-info}/RECORD +30 -27
  27. {databricks_sdk-0.29.0.dist-info → databricks_sdk-0.31.0.dist-info}/WHEEL +1 -1
  28. {databricks_sdk-0.29.0.dist-info → databricks_sdk-0.31.0.dist-info}/LICENSE +0 -0
  29. {databricks_sdk-0.29.0.dist-info → databricks_sdk-0.31.0.dist-info}/NOTICE +0 -0
  30. {databricks_sdk-0.29.0.dist-info → databricks_sdk-0.31.0.dist-info}/top_level.txt +0 -0
@@ -10,6 +10,9 @@ from datetime import timedelta
10
10
  from enum import Enum
11
11
  from typing import Any, BinaryIO, Callable, Dict, Iterator, List, Optional
12
12
 
13
+ import requests
14
+
15
+ from ..data_plane import DataPlaneService
13
16
  from ..errors import OperationFailed
14
17
  from ._internal import Wait, _enum, _from_dict, _repeated_dict
15
18
 
@@ -22,19 +25,29 @@ from databricks.sdk.service import oauth2
22
25
 
23
26
  @dataclass
24
27
  class Ai21LabsConfig:
25
- ai21labs_api_key: str
26
- """The Databricks secret key reference for an AI21Labs API key."""
28
+ ai21labs_api_key: Optional[str] = None
29
+ """The Databricks secret key reference for an AI21 Labs API key. If you prefer to paste your API
30
+ key directly, see `ai21labs_api_key_plaintext`. You must provide an API key using one of the
31
+ following fields: `ai21labs_api_key` or `ai21labs_api_key_plaintext`."""
32
+
33
+ ai21labs_api_key_plaintext: Optional[str] = None
34
+ """An AI21 Labs API key provided as a plaintext string. If you prefer to reference your key using
35
+ Databricks Secrets, see `ai21labs_api_key`. You must provide an API key using one of the
36
+ following fields: `ai21labs_api_key` or `ai21labs_api_key_plaintext`."""
27
37
 
28
38
  def as_dict(self) -> dict:
29
39
  """Serializes the Ai21LabsConfig into a dictionary suitable for use as a JSON request body."""
30
40
  body = {}
31
41
  if self.ai21labs_api_key is not None: body['ai21labs_api_key'] = self.ai21labs_api_key
42
+ if self.ai21labs_api_key_plaintext is not None:
43
+ body['ai21labs_api_key_plaintext'] = self.ai21labs_api_key_plaintext
32
44
  return body
33
45
 
34
46
  @classmethod
35
47
  def from_dict(cls, d: Dict[str, any]) -> Ai21LabsConfig:
36
48
  """Deserializes the Ai21LabsConfig from a dictionary."""
37
- return cls(ai21labs_api_key=d.get('ai21labs_api_key', None))
49
+ return cls(ai21labs_api_key=d.get('ai21labs_api_key', None),
50
+ ai21labs_api_key_plaintext=d.get('ai21labs_api_key_plaintext', None))
38
51
 
39
52
 
40
53
  @dataclass
@@ -42,24 +55,44 @@ class AmazonBedrockConfig:
42
55
  aws_region: str
43
56
  """The AWS region to use. Bedrock has to be enabled there."""
44
57
 
45
- aws_access_key_id: str
46
- """The Databricks secret key reference for an AWS Access Key ID with permissions to interact with
47
- Bedrock services."""
48
-
49
- aws_secret_access_key: str
50
- """The Databricks secret key reference for an AWS Secret Access Key paired with the access key ID,
51
- with permissions to interact with Bedrock services."""
52
-
53
58
  bedrock_provider: AmazonBedrockConfigBedrockProvider
54
59
  """The underlying provider in Amazon Bedrock. Supported values (case insensitive) include:
55
60
  Anthropic, Cohere, AI21Labs, Amazon."""
56
61
 
62
+ aws_access_key_id: Optional[str] = None
63
+ """The Databricks secret key reference for an AWS access key ID with permissions to interact with
64
+ Bedrock services. If you prefer to paste your API key directly, see `aws_access_key_id`. You
65
+ must provide an API key using one of the following fields: `aws_access_key_id` or
66
+ `aws_access_key_id_plaintext`."""
67
+
68
+ aws_access_key_id_plaintext: Optional[str] = None
69
+ """An AWS access key ID with permissions to interact with Bedrock services provided as a plaintext
70
+ string. If you prefer to reference your key using Databricks Secrets, see `aws_access_key_id`.
71
+ You must provide an API key using one of the following fields: `aws_access_key_id` or
72
+ `aws_access_key_id_plaintext`."""
73
+
74
+ aws_secret_access_key: Optional[str] = None
75
+ """The Databricks secret key reference for an AWS secret access key paired with the access key ID,
76
+ with permissions to interact with Bedrock services. If you prefer to paste your API key
77
+ directly, see `aws_secret_access_key_plaintext`. You must provide an API key using one of the
78
+ following fields: `aws_secret_access_key` or `aws_secret_access_key_plaintext`."""
79
+
80
+ aws_secret_access_key_plaintext: Optional[str] = None
81
+ """An AWS secret access key paired with the access key ID, with permissions to interact with
82
+ Bedrock services provided as a plaintext string. If you prefer to reference your key using
83
+ Databricks Secrets, see `aws_secret_access_key`. You must provide an API key using one of the
84
+ following fields: `aws_secret_access_key` or `aws_secret_access_key_plaintext`."""
85
+
57
86
  def as_dict(self) -> dict:
58
87
  """Serializes the AmazonBedrockConfig into a dictionary suitable for use as a JSON request body."""
59
88
  body = {}
60
89
  if self.aws_access_key_id is not None: body['aws_access_key_id'] = self.aws_access_key_id
90
+ if self.aws_access_key_id_plaintext is not None:
91
+ body['aws_access_key_id_plaintext'] = self.aws_access_key_id_plaintext
61
92
  if self.aws_region is not None: body['aws_region'] = self.aws_region
62
93
  if self.aws_secret_access_key is not None: body['aws_secret_access_key'] = self.aws_secret_access_key
94
+ if self.aws_secret_access_key_plaintext is not None:
95
+ body['aws_secret_access_key_plaintext'] = self.aws_secret_access_key_plaintext
63
96
  if self.bedrock_provider is not None: body['bedrock_provider'] = self.bedrock_provider.value
64
97
  return body
65
98
 
@@ -67,8 +100,10 @@ class AmazonBedrockConfig:
67
100
  def from_dict(cls, d: Dict[str, any]) -> AmazonBedrockConfig:
68
101
  """Deserializes the AmazonBedrockConfig from a dictionary."""
69
102
  return cls(aws_access_key_id=d.get('aws_access_key_id', None),
103
+ aws_access_key_id_plaintext=d.get('aws_access_key_id_plaintext', None),
70
104
  aws_region=d.get('aws_region', None),
71
105
  aws_secret_access_key=d.get('aws_secret_access_key', None),
106
+ aws_secret_access_key_plaintext=d.get('aws_secret_access_key_plaintext', None),
72
107
  bedrock_provider=_enum(d, 'bedrock_provider', AmazonBedrockConfigBedrockProvider))
73
108
 
74
109
 
@@ -84,249 +119,29 @@ class AmazonBedrockConfigBedrockProvider(Enum):
84
119
 
85
120
  @dataclass
86
121
  class AnthropicConfig:
87
- anthropic_api_key: str
88
- """The Databricks secret key reference for an Anthropic API key."""
122
+ anthropic_api_key: Optional[str] = None
123
+ """The Databricks secret key reference for an Anthropic API key. If you prefer to paste your API
124
+ key directly, see `anthropic_api_key_plaintext`. You must provide an API key using one of the
125
+ following fields: `anthropic_api_key` or `anthropic_api_key_plaintext`."""
126
+
127
+ anthropic_api_key_plaintext: Optional[str] = None
128
+ """The Anthropic API key provided as a plaintext string. If you prefer to reference your key using
129
+ Databricks Secrets, see `anthropic_api_key`. You must provide an API key using one of the
130
+ following fields: `anthropic_api_key` or `anthropic_api_key_plaintext`."""
89
131
 
90
132
  def as_dict(self) -> dict:
91
133
  """Serializes the AnthropicConfig into a dictionary suitable for use as a JSON request body."""
92
134
  body = {}
93
135
  if self.anthropic_api_key is not None: body['anthropic_api_key'] = self.anthropic_api_key
136
+ if self.anthropic_api_key_plaintext is not None:
137
+ body['anthropic_api_key_plaintext'] = self.anthropic_api_key_plaintext
94
138
  return body
95
139
 
96
140
  @classmethod
97
141
  def from_dict(cls, d: Dict[str, any]) -> AnthropicConfig:
98
142
  """Deserializes the AnthropicConfig from a dictionary."""
99
- return cls(anthropic_api_key=d.get('anthropic_api_key', None))
100
-
101
-
102
- @dataclass
103
- class App:
104
- name: str
105
- """The name of the app. The name must contain only lowercase alphanumeric characters and hyphens.
106
- It must be unique within the workspace."""
107
-
108
- active_deployment: Optional[AppDeployment] = None
109
- """The active deployment of the app."""
110
-
111
- create_time: Optional[str] = None
112
- """The creation time of the app. Formatted timestamp in ISO 6801."""
113
-
114
- creator: Optional[str] = None
115
- """The email of the user that created the app."""
116
-
117
- description: Optional[str] = None
118
- """The description of the app."""
119
-
120
- pending_deployment: Optional[AppDeployment] = None
121
- """The pending deployment of the app."""
122
-
123
- service_principal_id: Optional[int] = None
124
-
125
- service_principal_name: Optional[str] = None
126
-
127
- status: Optional[AppStatus] = None
128
-
129
- update_time: Optional[str] = None
130
- """The update time of the app. Formatted timestamp in ISO 6801."""
131
-
132
- updater: Optional[str] = None
133
- """The email of the user that last updated the app."""
134
-
135
- url: Optional[str] = None
136
- """The URL of the app once it is deployed."""
137
-
138
- def as_dict(self) -> dict:
139
- """Serializes the App into a dictionary suitable for use as a JSON request body."""
140
- body = {}
141
- if self.active_deployment: body['active_deployment'] = self.active_deployment.as_dict()
142
- if self.create_time is not None: body['create_time'] = self.create_time
143
- if self.creator is not None: body['creator'] = self.creator
144
- if self.description is not None: body['description'] = self.description
145
- if self.name is not None: body['name'] = self.name
146
- if self.pending_deployment: body['pending_deployment'] = self.pending_deployment.as_dict()
147
- if self.service_principal_id is not None: body['service_principal_id'] = self.service_principal_id
148
- if self.service_principal_name is not None:
149
- body['service_principal_name'] = self.service_principal_name
150
- if self.status: body['status'] = self.status.as_dict()
151
- if self.update_time is not None: body['update_time'] = self.update_time
152
- if self.updater is not None: body['updater'] = self.updater
153
- if self.url is not None: body['url'] = self.url
154
- return body
155
-
156
- @classmethod
157
- def from_dict(cls, d: Dict[str, any]) -> App:
158
- """Deserializes the App from a dictionary."""
159
- return cls(active_deployment=_from_dict(d, 'active_deployment', AppDeployment),
160
- create_time=d.get('create_time', None),
161
- creator=d.get('creator', None),
162
- description=d.get('description', None),
163
- name=d.get('name', None),
164
- pending_deployment=_from_dict(d, 'pending_deployment', AppDeployment),
165
- service_principal_id=d.get('service_principal_id', None),
166
- service_principal_name=d.get('service_principal_name', None),
167
- status=_from_dict(d, 'status', AppStatus),
168
- update_time=d.get('update_time', None),
169
- updater=d.get('updater', None),
170
- url=d.get('url', None))
171
-
172
-
173
- @dataclass
174
- class AppDeployment:
175
- source_code_path: str
176
- """The workspace file system path of the source code used to create the app deployment. This is
177
- different from `deployment_artifacts.source_code_path`, which is the path used by the deployed
178
- app. The former refers to the original source code location of the app in the workspace during
179
- deployment creation, whereas the latter provides a system generated stable snapshotted source
180
- code path used by the deployment."""
181
-
182
- mode: AppDeploymentMode
183
- """The mode of which the deployment will manage the source code."""
184
-
185
- create_time: Optional[str] = None
186
- """The creation time of the deployment. Formatted timestamp in ISO 6801."""
187
-
188
- creator: Optional[str] = None
189
- """The email of the user creates the deployment."""
190
-
191
- deployment_artifacts: Optional[AppDeploymentArtifacts] = None
192
- """The deployment artifacts for an app."""
193
-
194
- deployment_id: Optional[str] = None
195
- """The unique id of the deployment."""
196
-
197
- status: Optional[AppDeploymentStatus] = None
198
- """Status and status message of the deployment"""
199
-
200
- update_time: Optional[str] = None
201
- """The update time of the deployment. Formatted timestamp in ISO 6801."""
202
-
203
- def as_dict(self) -> dict:
204
- """Serializes the AppDeployment into a dictionary suitable for use as a JSON request body."""
205
- body = {}
206
- if self.create_time is not None: body['create_time'] = self.create_time
207
- if self.creator is not None: body['creator'] = self.creator
208
- if self.deployment_artifacts: body['deployment_artifacts'] = self.deployment_artifacts.as_dict()
209
- if self.deployment_id is not None: body['deployment_id'] = self.deployment_id
210
- if self.mode is not None: body['mode'] = self.mode.value
211
- if self.source_code_path is not None: body['source_code_path'] = self.source_code_path
212
- if self.status: body['status'] = self.status.as_dict()
213
- if self.update_time is not None: body['update_time'] = self.update_time
214
- return body
215
-
216
- @classmethod
217
- def from_dict(cls, d: Dict[str, any]) -> AppDeployment:
218
- """Deserializes the AppDeployment from a dictionary."""
219
- return cls(create_time=d.get('create_time', None),
220
- creator=d.get('creator', None),
221
- deployment_artifacts=_from_dict(d, 'deployment_artifacts', AppDeploymentArtifacts),
222
- deployment_id=d.get('deployment_id', None),
223
- mode=_enum(d, 'mode', AppDeploymentMode),
224
- source_code_path=d.get('source_code_path', None),
225
- status=_from_dict(d, 'status', AppDeploymentStatus),
226
- update_time=d.get('update_time', None))
227
-
228
-
229
- @dataclass
230
- class AppDeploymentArtifacts:
231
- source_code_path: Optional[str] = None
232
- """The snapshotted workspace file system path of the source code loaded by the deployed app."""
233
-
234
- def as_dict(self) -> dict:
235
- """Serializes the AppDeploymentArtifacts into a dictionary suitable for use as a JSON request body."""
236
- body = {}
237
- if self.source_code_path is not None: body['source_code_path'] = self.source_code_path
238
- return body
239
-
240
- @classmethod
241
- def from_dict(cls, d: Dict[str, any]) -> AppDeploymentArtifacts:
242
- """Deserializes the AppDeploymentArtifacts from a dictionary."""
243
- return cls(source_code_path=d.get('source_code_path', None))
244
-
245
-
246
- class AppDeploymentMode(Enum):
247
-
248
- AUTO_SYNC = 'AUTO_SYNC'
249
- MODE_UNSPECIFIED = 'MODE_UNSPECIFIED'
250
- SNAPSHOT = 'SNAPSHOT'
251
-
252
-
253
- class AppDeploymentState(Enum):
254
-
255
- FAILED = 'FAILED'
256
- IN_PROGRESS = 'IN_PROGRESS'
257
- STATE_UNSPECIFIED = 'STATE_UNSPECIFIED'
258
- STOPPED = 'STOPPED'
259
- SUCCEEDED = 'SUCCEEDED'
260
-
261
-
262
- @dataclass
263
- class AppDeploymentStatus:
264
- message: Optional[str] = None
265
- """Message corresponding with the deployment state."""
266
-
267
- state: Optional[AppDeploymentState] = None
268
- """State of the deployment."""
269
-
270
- def as_dict(self) -> dict:
271
- """Serializes the AppDeploymentStatus into a dictionary suitable for use as a JSON request body."""
272
- body = {}
273
- if self.message is not None: body['message'] = self.message
274
- if self.state is not None: body['state'] = self.state.value
275
- return body
276
-
277
- @classmethod
278
- def from_dict(cls, d: Dict[str, any]) -> AppDeploymentStatus:
279
- """Deserializes the AppDeploymentStatus from a dictionary."""
280
- return cls(message=d.get('message', None), state=_enum(d, 'state', AppDeploymentState))
281
-
282
-
283
- @dataclass
284
- class AppEnvironment:
285
- env: Optional[List[EnvVariable]] = None
286
-
287
- def as_dict(self) -> dict:
288
- """Serializes the AppEnvironment into a dictionary suitable for use as a JSON request body."""
289
- body = {}
290
- if self.env: body['env'] = [v.as_dict() for v in self.env]
291
- return body
292
-
293
- @classmethod
294
- def from_dict(cls, d: Dict[str, any]) -> AppEnvironment:
295
- """Deserializes the AppEnvironment from a dictionary."""
296
- return cls(env=_repeated_dict(d, 'env', EnvVariable))
297
-
298
-
299
- class AppState(Enum):
300
-
301
- CREATING = 'CREATING'
302
- DELETED = 'DELETED'
303
- DELETING = 'DELETING'
304
- ERROR = 'ERROR'
305
- IDLE = 'IDLE'
306
- RUNNING = 'RUNNING'
307
- STARTING = 'STARTING'
308
- STATE_UNSPECIFIED = 'STATE_UNSPECIFIED'
309
-
310
-
311
- @dataclass
312
- class AppStatus:
313
- message: Optional[str] = None
314
- """Message corresponding with the app state."""
315
-
316
- state: Optional[AppState] = None
317
- """State of the app."""
318
-
319
- def as_dict(self) -> dict:
320
- """Serializes the AppStatus into a dictionary suitable for use as a JSON request body."""
321
- body = {}
322
- if self.message is not None: body['message'] = self.message
323
- if self.state is not None: body['state'] = self.state.value
324
- return body
325
-
326
- @classmethod
327
- def from_dict(cls, d: Dict[str, any]) -> AppStatus:
328
- """Deserializes the AppStatus from a dictionary."""
329
- return cls(message=d.get('message', None), state=_enum(d, 'state', AppState))
143
+ return cls(anthropic_api_key=d.get('anthropic_api_key', None),
144
+ anthropic_api_key_plaintext=d.get('anthropic_api_key_plaintext', None))
330
145
 
331
146
 
332
147
  @dataclass
@@ -464,72 +279,35 @@ class ChatMessageRole(Enum):
464
279
 
465
280
  @dataclass
466
281
  class CohereConfig:
467
- cohere_api_key: str
468
- """The Databricks secret key reference for a Cohere API key."""
282
+ cohere_api_base: Optional[str] = None
283
+ """This is an optional field to provide a customized base URL for the Cohere API. If left
284
+ unspecified, the standard Cohere base URL is used."""
285
+
286
+ cohere_api_key: Optional[str] = None
287
+ """The Databricks secret key reference for a Cohere API key. If you prefer to paste your API key
288
+ directly, see `cohere_api_key_plaintext`. You must provide an API key using one of the following
289
+ fields: `cohere_api_key` or `cohere_api_key_plaintext`."""
290
+
291
+ cohere_api_key_plaintext: Optional[str] = None
292
+ """The Cohere API key provided as a plaintext string. If you prefer to reference your key using
293
+ Databricks Secrets, see `cohere_api_key`. You must provide an API key using one of the following
294
+ fields: `cohere_api_key` or `cohere_api_key_plaintext`."""
469
295
 
470
296
  def as_dict(self) -> dict:
471
297
  """Serializes the CohereConfig into a dictionary suitable for use as a JSON request body."""
472
298
  body = {}
299
+ if self.cohere_api_base is not None: body['cohere_api_base'] = self.cohere_api_base
473
300
  if self.cohere_api_key is not None: body['cohere_api_key'] = self.cohere_api_key
301
+ if self.cohere_api_key_plaintext is not None:
302
+ body['cohere_api_key_plaintext'] = self.cohere_api_key_plaintext
474
303
  return body
475
304
 
476
305
  @classmethod
477
306
  def from_dict(cls, d: Dict[str, any]) -> CohereConfig:
478
307
  """Deserializes the CohereConfig from a dictionary."""
479
- return cls(cohere_api_key=d.get('cohere_api_key', None))
480
-
481
-
482
- @dataclass
483
- class CreateAppDeploymentRequest:
484
- source_code_path: str
485
- """The workspace file system path of the source code used to create the app deployment. This is
486
- different from `deployment_artifacts.source_code_path`, which is the path used by the deployed
487
- app. The former refers to the original source code location of the app in the workspace during
488
- deployment creation, whereas the latter provides a system generated stable snapshotted source
489
- code path used by the deployment."""
490
-
491
- mode: AppDeploymentMode
492
- """The mode of which the deployment will manage the source code."""
493
-
494
- app_name: Optional[str] = None
495
- """The name of the app."""
496
-
497
- def as_dict(self) -> dict:
498
- """Serializes the CreateAppDeploymentRequest into a dictionary suitable for use as a JSON request body."""
499
- body = {}
500
- if self.app_name is not None: body['app_name'] = self.app_name
501
- if self.mode is not None: body['mode'] = self.mode.value
502
- if self.source_code_path is not None: body['source_code_path'] = self.source_code_path
503
- return body
504
-
505
- @classmethod
506
- def from_dict(cls, d: Dict[str, any]) -> CreateAppDeploymentRequest:
507
- """Deserializes the CreateAppDeploymentRequest from a dictionary."""
508
- return cls(app_name=d.get('app_name', None),
509
- mode=_enum(d, 'mode', AppDeploymentMode),
510
- source_code_path=d.get('source_code_path', None))
511
-
512
-
513
- @dataclass
514
- class CreateAppRequest:
515
- name: str
516
- """The name of the app. The name must contain only lowercase alphanumeric characters and hyphens.
517
- It must be unique within the workspace."""
518
-
519
- description: Optional[str] = None
520
- """The description of the app."""
521
-
522
- def as_dict(self) -> dict:
523
- """Serializes the CreateAppRequest into a dictionary suitable for use as a JSON request body."""
524
- body = {}
525
- if self.description is not None: body['description'] = self.description
526
- if self.name is not None: body['name'] = self.name
527
- return body
528
-
529
- @classmethod
530
- def from_dict(cls, d: Dict[str, any]) -> CreateAppRequest:
531
- """Deserializes the CreateAppRequest from a dictionary."""
532
- return cls(description=d.get('description', None), name=d.get('name', None))
308
+ return cls(cohere_api_base=d.get('cohere_api_base', None),
309
+ cohere_api_key=d.get('cohere_api_key', None),
310
+ cohere_api_key_plaintext=d.get('cohere_api_key_plaintext', None))
533
311
 
534
312
 
535
313
  @dataclass
@@ -573,19 +351,30 @@ class CreateServingEndpoint:
573
351
 
574
352
  @dataclass
575
353
  class DatabricksModelServingConfig:
576
- databricks_api_token: str
577
- """The Databricks secret key reference for a Databricks API token that corresponds to a user or
578
- service principal with Can Query access to the model serving endpoint pointed to by this
579
- external model."""
580
-
581
354
  databricks_workspace_url: str
582
355
  """The URL of the Databricks workspace containing the model serving endpoint pointed to by this
583
356
  external model."""
584
357
 
358
+ databricks_api_token: Optional[str] = None
359
+ """The Databricks secret key reference for a Databricks API token that corresponds to a user or
360
+ service principal with Can Query access to the model serving endpoint pointed to by this
361
+ external model. If you prefer to paste your API key directly, see
362
+ `databricks_api_token_plaintext`. You must provide an API key using one of the following fields:
363
+ `databricks_api_token` or `databricks_api_token_plaintext`."""
364
+
365
+ databricks_api_token_plaintext: Optional[str] = None
366
+ """The Databricks API token that corresponds to a user or service principal with Can Query access
367
+ to the model serving endpoint pointed to by this external model provided as a plaintext string.
368
+ If you prefer to reference your key using Databricks Secrets, see `databricks_api_token`. You
369
+ must provide an API key using one of the following fields: `databricks_api_token` or
370
+ `databricks_api_token_plaintext`."""
371
+
585
372
  def as_dict(self) -> dict:
586
373
  """Serializes the DatabricksModelServingConfig into a dictionary suitable for use as a JSON request body."""
587
374
  body = {}
588
375
  if self.databricks_api_token is not None: body['databricks_api_token'] = self.databricks_api_token
376
+ if self.databricks_api_token_plaintext is not None:
377
+ body['databricks_api_token_plaintext'] = self.databricks_api_token_plaintext
589
378
  if self.databricks_workspace_url is not None:
590
379
  body['databricks_workspace_url'] = self.databricks_workspace_url
591
380
  return body
@@ -594,6 +383,7 @@ class DatabricksModelServingConfig:
594
383
  def from_dict(cls, d: Dict[str, any]) -> DatabricksModelServingConfig:
595
384
  """Deserializes the DatabricksModelServingConfig from a dictionary."""
596
385
  return cls(databricks_api_token=d.get('databricks_api_token', None),
386
+ databricks_api_token_plaintext=d.get('databricks_api_token_plaintext', None),
597
387
  databricks_workspace_url=d.get('databricks_workspace_url', None))
598
388
 
599
389
 
@@ -846,6 +636,7 @@ class EndpointStateConfigUpdate(Enum):
846
636
 
847
637
  IN_PROGRESS = 'IN_PROGRESS'
848
638
  NOT_UPDATING = 'NOT_UPDATING'
639
+ UPDATE_CANCELED = 'UPDATE_CANCELED'
849
640
  UPDATE_FAILED = 'UPDATE_FAILED'
850
641
 
851
642
 
@@ -879,28 +670,6 @@ class EndpointTag:
879
670
  return cls(key=d.get('key', None), value=d.get('value', None))
880
671
 
881
672
 
882
- @dataclass
883
- class EnvVariable:
884
- name: Optional[str] = None
885
-
886
- value: Optional[str] = None
887
-
888
- value_from: Optional[str] = None
889
-
890
- def as_dict(self) -> dict:
891
- """Serializes the EnvVariable into a dictionary suitable for use as a JSON request body."""
892
- body = {}
893
- if self.name is not None: body['name'] = self.name
894
- if self.value is not None: body['value'] = self.value
895
- if self.value_from is not None: body['value_from'] = self.value_from
896
- return body
897
-
898
- @classmethod
899
- def from_dict(cls, d: Dict[str, any]) -> EnvVariable:
900
- """Deserializes the EnvVariable from a dictionary."""
901
- return cls(name=d.get('name', None), value=d.get('value', None), value_from=d.get('value_from', None))
902
-
903
-
904
673
  @dataclass
905
674
  class ExportMetricsResponse:
906
675
  contents: Optional[BinaryIO] = None
@@ -921,8 +690,8 @@ class ExportMetricsResponse:
921
690
  class ExternalModel:
922
691
  provider: ExternalModelProvider
923
692
  """The name of the provider for the external model. Currently, the supported providers are
924
- 'ai21labs', 'anthropic', 'amazon-bedrock', 'cohere', 'databricks-model-serving', 'openai', and
925
- 'palm'.","""
693
+ 'ai21labs', 'anthropic', 'amazon-bedrock', 'cohere', 'databricks-model-serving',
694
+ 'google-cloud-vertex-ai', 'openai', and 'palm'.","""
926
695
 
927
696
  name: str
928
697
  """The name of the external model."""
@@ -945,6 +714,9 @@ class ExternalModel:
945
714
  databricks_model_serving_config: Optional[DatabricksModelServingConfig] = None
946
715
  """Databricks Model Serving Config. Only required if the provider is 'databricks-model-serving'."""
947
716
 
717
+ google_cloud_vertex_ai_config: Optional[GoogleCloudVertexAiConfig] = None
718
+ """Google Cloud Vertex AI Config. Only required if the provider is 'google-cloud-vertex-ai'."""
719
+
948
720
  openai_config: Optional[OpenAiConfig] = None
949
721
  """OpenAI Config. Only required if the provider is 'openai'."""
950
722
 
@@ -960,6 +732,8 @@ class ExternalModel:
960
732
  if self.cohere_config: body['cohere_config'] = self.cohere_config.as_dict()
961
733
  if self.databricks_model_serving_config:
962
734
  body['databricks_model_serving_config'] = self.databricks_model_serving_config.as_dict()
735
+ if self.google_cloud_vertex_ai_config:
736
+ body['google_cloud_vertex_ai_config'] = self.google_cloud_vertex_ai_config.as_dict()
963
737
  if self.name is not None: body['name'] = self.name
964
738
  if self.openai_config: body['openai_config'] = self.openai_config.as_dict()
965
739
  if self.palm_config: body['palm_config'] = self.palm_config.as_dict()
@@ -976,6 +750,8 @@ class ExternalModel:
976
750
  cohere_config=_from_dict(d, 'cohere_config', CohereConfig),
977
751
  databricks_model_serving_config=_from_dict(d, 'databricks_model_serving_config',
978
752
  DatabricksModelServingConfig),
753
+ google_cloud_vertex_ai_config=_from_dict(d, 'google_cloud_vertex_ai_config',
754
+ GoogleCloudVertexAiConfig),
979
755
  name=d.get('name', None),
980
756
  openai_config=_from_dict(d, 'openai_config', OpenAiConfig),
981
757
  palm_config=_from_dict(d, 'palm_config', PaLmConfig),
@@ -985,14 +761,15 @@ class ExternalModel:
985
761
 
986
762
  class ExternalModelProvider(Enum):
987
763
  """The name of the provider for the external model. Currently, the supported providers are
988
- 'ai21labs', 'anthropic', 'amazon-bedrock', 'cohere', 'databricks-model-serving', 'openai', and
989
- 'palm'.","""
764
+ 'ai21labs', 'anthropic', 'amazon-bedrock', 'cohere', 'databricks-model-serving',
765
+ 'google-cloud-vertex-ai', 'openai', and 'palm'.","""
990
766
 
991
767
  AI21LABS = 'ai21labs'
992
768
  AMAZON_BEDROCK = 'amazon-bedrock'
993
769
  ANTHROPIC = 'anthropic'
994
770
  COHERE = 'cohere'
995
771
  DATABRICKS_MODEL_SERVING = 'databricks-model-serving'
772
+ GOOGLE_CLOUD_VERTEX_AI = 'google-cloud-vertex-ai'
996
773
  OPENAI = 'openai'
997
774
  PALM = 'palm'
998
775
 
@@ -1091,45 +868,48 @@ class GetServingEndpointPermissionLevelsResponse:
1091
868
 
1092
869
 
1093
870
  @dataclass
1094
- class ListAppDeploymentsResponse:
1095
- app_deployments: Optional[List[AppDeployment]] = None
1096
- """Deployment history of the app."""
1097
-
1098
- next_page_token: Optional[str] = None
1099
- """Pagination token to request the next page of apps."""
1100
-
1101
- def as_dict(self) -> dict:
1102
- """Serializes the ListAppDeploymentsResponse into a dictionary suitable for use as a JSON request body."""
1103
- body = {}
1104
- if self.app_deployments: body['app_deployments'] = [v.as_dict() for v in self.app_deployments]
1105
- if self.next_page_token is not None: body['next_page_token'] = self.next_page_token
1106
- return body
1107
-
1108
- @classmethod
1109
- def from_dict(cls, d: Dict[str, any]) -> ListAppDeploymentsResponse:
1110
- """Deserializes the ListAppDeploymentsResponse from a dictionary."""
1111
- return cls(app_deployments=_repeated_dict(d, 'app_deployments', AppDeployment),
1112
- next_page_token=d.get('next_page_token', None))
871
+ class GoogleCloudVertexAiConfig:
872
+ private_key: Optional[str] = None
873
+ """The Databricks secret key reference for a private key for the service account which has access
874
+ to the Google Cloud Vertex AI Service. See [Best practices for managing service account keys].
875
+ If you prefer to paste your API key directly, see `private_key_plaintext`. You must provide an
876
+ API key using one of the following fields: `private_key` or `private_key_plaintext`
877
+
878
+ [Best practices for managing service account keys]: https://cloud.google.com/iam/docs/best-practices-for-managing-service-account-keys"""
1113
879
 
880
+ private_key_plaintext: Optional[str] = None
881
+ """The private key for the service account which has access to the Google Cloud Vertex AI Service
882
+ provided as a plaintext secret. See [Best practices for managing service account keys]. If you
883
+ prefer to reference your key using Databricks Secrets, see `private_key`. You must provide an
884
+ API key using one of the following fields: `private_key` or `private_key_plaintext`.
885
+
886
+ [Best practices for managing service account keys]: https://cloud.google.com/iam/docs/best-practices-for-managing-service-account-keys"""
1114
887
 
1115
- @dataclass
1116
- class ListAppsResponse:
1117
- apps: Optional[List[App]] = None
888
+ project_id: Optional[str] = None
889
+ """This is the Google Cloud project id that the service account is associated with."""
1118
890
 
1119
- next_page_token: Optional[str] = None
1120
- """Pagination token to request the next page of apps."""
891
+ region: Optional[str] = None
892
+ """This is the region for the Google Cloud Vertex AI Service. See [supported regions] for more
893
+ details. Some models are only available in specific regions.
894
+
895
+ [supported regions]: https://cloud.google.com/vertex-ai/docs/general/locations"""
1121
896
 
1122
897
  def as_dict(self) -> dict:
1123
- """Serializes the ListAppsResponse into a dictionary suitable for use as a JSON request body."""
898
+ """Serializes the GoogleCloudVertexAiConfig into a dictionary suitable for use as a JSON request body."""
1124
899
  body = {}
1125
- if self.apps: body['apps'] = [v.as_dict() for v in self.apps]
1126
- if self.next_page_token is not None: body['next_page_token'] = self.next_page_token
900
+ if self.private_key is not None: body['private_key'] = self.private_key
901
+ if self.private_key_plaintext is not None: body['private_key_plaintext'] = self.private_key_plaintext
902
+ if self.project_id is not None: body['project_id'] = self.project_id
903
+ if self.region is not None: body['region'] = self.region
1127
904
  return body
1128
905
 
1129
906
  @classmethod
1130
- def from_dict(cls, d: Dict[str, any]) -> ListAppsResponse:
1131
- """Deserializes the ListAppsResponse from a dictionary."""
1132
- return cls(apps=_repeated_dict(d, 'apps', App), next_page_token=d.get('next_page_token', None))
907
+ def from_dict(cls, d: Dict[str, any]) -> GoogleCloudVertexAiConfig:
908
+ """Deserializes the GoogleCloudVertexAiConfig from a dictionary."""
909
+ return cls(private_key=d.get('private_key', None),
910
+ private_key_plaintext=d.get('private_key_plaintext', None),
911
+ project_id=d.get('project_id', None),
912
+ region=d.get('region', None))
1133
913
 
1134
914
 
1135
915
  @dataclass
@@ -1172,19 +952,35 @@ class OpenAiConfig:
1172
952
  """This field is only required for Azure AD OpenAI and is the Microsoft Entra Client ID."""
1173
953
 
1174
954
  microsoft_entra_client_secret: Optional[str] = None
1175
- """The Databricks secret key reference for the Microsoft Entra Client Secret that is only required
1176
- for Azure AD OpenAI."""
955
+ """The Databricks secret key reference for a client secret used for Microsoft Entra ID
956
+ authentication. If you prefer to paste your client secret directly, see
957
+ `microsoft_entra_client_secret_plaintext`. You must provide an API key using one of the
958
+ following fields: `microsoft_entra_client_secret` or `microsoft_entra_client_secret_plaintext`."""
959
+
960
+ microsoft_entra_client_secret_plaintext: Optional[str] = None
961
+ """The client secret used for Microsoft Entra ID authentication provided as a plaintext string. If
962
+ you prefer to reference your key using Databricks Secrets, see `microsoft_entra_client_secret`.
963
+ You must provide an API key using one of the following fields: `microsoft_entra_client_secret`
964
+ or `microsoft_entra_client_secret_plaintext`."""
1177
965
 
1178
966
  microsoft_entra_tenant_id: Optional[str] = None
1179
967
  """This field is only required for Azure AD OpenAI and is the Microsoft Entra Tenant ID."""
1180
968
 
1181
969
  openai_api_base: Optional[str] = None
1182
- """This is the base URL for the OpenAI API (default: "https://api.openai.com/v1"). For Azure
1183
- OpenAI, this field is required, and is the base URL for the Azure OpenAI API service provided by
1184
- Azure."""
970
+ """This is a field to provide a customized base URl for the OpenAI API. For Azure OpenAI, this
971
+ field is required, and is the base URL for the Azure OpenAI API service provided by Azure. For
972
+ other OpenAI API types, this field is optional, and if left unspecified, the standard OpenAI
973
+ base URL is used."""
1185
974
 
1186
975
  openai_api_key: Optional[str] = None
1187
- """The Databricks secret key reference for an OpenAI or Azure OpenAI API key."""
976
+ """The Databricks secret key reference for an OpenAI API key using the OpenAI or Azure service. If
977
+ you prefer to paste your API key directly, see `openai_api_key_plaintext`. You must provide an
978
+ API key using one of the following fields: `openai_api_key` or `openai_api_key_plaintext`."""
979
+
980
+ openai_api_key_plaintext: Optional[str] = None
981
+ """The OpenAI API key using the OpenAI or Azure service provided as a plaintext string. If you
982
+ prefer to reference your key using Databricks Secrets, see `openai_api_key`. You must provide an
983
+ API key using one of the following fields: `openai_api_key` or `openai_api_key_plaintext`."""
1188
984
 
1189
985
  openai_api_type: Optional[str] = None
1190
986
  """This is an optional field to specify the type of OpenAI API to use. For Azure OpenAI, this field
@@ -1210,10 +1006,14 @@ class OpenAiConfig:
1210
1006
  body['microsoft_entra_client_id'] = self.microsoft_entra_client_id
1211
1007
  if self.microsoft_entra_client_secret is not None:
1212
1008
  body['microsoft_entra_client_secret'] = self.microsoft_entra_client_secret
1009
+ if self.microsoft_entra_client_secret_plaintext is not None:
1010
+ body['microsoft_entra_client_secret_plaintext'] = self.microsoft_entra_client_secret_plaintext
1213
1011
  if self.microsoft_entra_tenant_id is not None:
1214
1012
  body['microsoft_entra_tenant_id'] = self.microsoft_entra_tenant_id
1215
1013
  if self.openai_api_base is not None: body['openai_api_base'] = self.openai_api_base
1216
1014
  if self.openai_api_key is not None: body['openai_api_key'] = self.openai_api_key
1015
+ if self.openai_api_key_plaintext is not None:
1016
+ body['openai_api_key_plaintext'] = self.openai_api_key_plaintext
1217
1017
  if self.openai_api_type is not None: body['openai_api_type'] = self.openai_api_type
1218
1018
  if self.openai_api_version is not None: body['openai_api_version'] = self.openai_api_version
1219
1019
  if self.openai_deployment_name is not None:
@@ -1226,9 +1026,12 @@ class OpenAiConfig:
1226
1026
  """Deserializes the OpenAiConfig from a dictionary."""
1227
1027
  return cls(microsoft_entra_client_id=d.get('microsoft_entra_client_id', None),
1228
1028
  microsoft_entra_client_secret=d.get('microsoft_entra_client_secret', None),
1029
+ microsoft_entra_client_secret_plaintext=d.get('microsoft_entra_client_secret_plaintext',
1030
+ None),
1229
1031
  microsoft_entra_tenant_id=d.get('microsoft_entra_tenant_id', None),
1230
1032
  openai_api_base=d.get('openai_api_base', None),
1231
1033
  openai_api_key=d.get('openai_api_key', None),
1034
+ openai_api_key_plaintext=d.get('openai_api_key_plaintext', None),
1232
1035
  openai_api_type=d.get('openai_api_type', None),
1233
1036
  openai_api_version=d.get('openai_api_version', None),
1234
1037
  openai_deployment_name=d.get('openai_deployment_name', None),
@@ -1237,19 +1040,29 @@ class OpenAiConfig:
1237
1040
 
1238
1041
  @dataclass
1239
1042
  class PaLmConfig:
1240
- palm_api_key: str
1241
- """The Databricks secret key reference for a PaLM API key."""
1043
+ palm_api_key: Optional[str] = None
1044
+ """The Databricks secret key reference for a PaLM API key. If you prefer to paste your API key
1045
+ directly, see `palm_api_key_plaintext`. You must provide an API key using one of the following
1046
+ fields: `palm_api_key` or `palm_api_key_plaintext`."""
1047
+
1048
+ palm_api_key_plaintext: Optional[str] = None
1049
+ """The PaLM API key provided as a plaintext string. If you prefer to reference your key using
1050
+ Databricks Secrets, see `palm_api_key`. You must provide an API key using one of the following
1051
+ fields: `palm_api_key` or `palm_api_key_plaintext`."""
1242
1052
 
1243
1053
  def as_dict(self) -> dict:
1244
1054
  """Serializes the PaLmConfig into a dictionary suitable for use as a JSON request body."""
1245
1055
  body = {}
1246
1056
  if self.palm_api_key is not None: body['palm_api_key'] = self.palm_api_key
1057
+ if self.palm_api_key_plaintext is not None:
1058
+ body['palm_api_key_plaintext'] = self.palm_api_key_plaintext
1247
1059
  return body
1248
1060
 
1249
1061
  @classmethod
1250
1062
  def from_dict(cls, d: Dict[str, any]) -> PaLmConfig:
1251
1063
  """Deserializes the PaLmConfig from a dictionary."""
1252
- return cls(palm_api_key=d.get('palm_api_key', None))
1064
+ return cls(palm_api_key=d.get('palm_api_key', None),
1065
+ palm_api_key_plaintext=d.get('palm_api_key_plaintext', None))
1253
1066
 
1254
1067
 
1255
1068
  @dataclass
@@ -1581,11 +1394,10 @@ class ServedEntityInput:
1581
1394
  external_model: Optional[ExternalModel] = None
1582
1395
  """The external model to be served. NOTE: Only one of external_model and (entity_name,
1583
1396
  entity_version, workload_size, workload_type, and scale_to_zero_enabled) can be specified with
1584
- the latter set being used for custom model serving for a Databricks registered model. When an
1585
- external_model is present, the served entities list can only have one served_entity object. For
1586
- an existing endpoint with external_model, it can not be updated to an endpoint without
1397
+ the latter set being used for custom model serving for a Databricks registered model. For an
1398
+ existing endpoint with external_model, it cannot be updated to an endpoint without
1587
1399
  external_model. If the endpoint is created without external_model, users cannot update it to add
1588
- external_model later."""
1400
+ external_model later. The task type of all external models within an endpoint must be the same."""
1589
1401
 
1590
1402
  instance_profile_arn: Optional[str] = None
1591
1403
  """ARN of the instance profile that the served entity uses to access AWS resources."""
@@ -2404,32 +2216,6 @@ class ServingEndpointPermissionsRequest:
2404
2216
  serving_endpoint_id=d.get('serving_endpoint_id', None))
2405
2217
 
2406
2218
 
2407
- @dataclass
2408
- class StartAppRequest:
2409
- name: Optional[str] = None
2410
- """The name of the app."""
2411
-
2412
-
2413
- @dataclass
2414
- class StopAppRequest:
2415
- name: Optional[str] = None
2416
- """The name of the app."""
2417
-
2418
-
2419
- @dataclass
2420
- class StopAppResponse:
2421
-
2422
- def as_dict(self) -> dict:
2423
- """Serializes the StopAppResponse into a dictionary suitable for use as a JSON request body."""
2424
- body = {}
2425
- return body
2426
-
2427
- @classmethod
2428
- def from_dict(cls, d: Dict[str, any]) -> StopAppResponse:
2429
- """Deserializes the StopAppResponse from a dictionary."""
2430
- return cls()
2431
-
2432
-
2433
2219
  @dataclass
2434
2220
  class TrafficConfig:
2435
2221
  routes: Optional[List[Route]] = None
@@ -2447,28 +2233,6 @@ class TrafficConfig:
2447
2233
  return cls(routes=_repeated_dict(d, 'routes', Route))
2448
2234
 
2449
2235
 
2450
- @dataclass
2451
- class UpdateAppRequest:
2452
- name: str
2453
- """The name of the app. The name must contain only lowercase alphanumeric characters and hyphens.
2454
- It must be unique within the workspace."""
2455
-
2456
- description: Optional[str] = None
2457
- """The description of the app."""
2458
-
2459
- def as_dict(self) -> dict:
2460
- """Serializes the UpdateAppRequest into a dictionary suitable for use as a JSON request body."""
2461
- body = {}
2462
- if self.description is not None: body['description'] = self.description
2463
- if self.name is not None: body['name'] = self.name
2464
- return body
2465
-
2466
- @classmethod
2467
- def from_dict(cls, d: Dict[str, any]) -> UpdateAppRequest:
2468
- """Deserializes the UpdateAppRequest from a dictionary."""
2469
- return cls(description=d.get('description', None), name=d.get('name', None))
2470
-
2471
-
2472
2236
  @dataclass
2473
2237
  class V1ResponseChoiceElement:
2474
2238
  finish_reason: Optional[str] = None
@@ -2506,333 +2270,6 @@ class V1ResponseChoiceElement:
2506
2270
  text=d.get('text', None))
2507
2271
 
2508
2272
 
2509
- class AppsAPI:
2510
- """Apps run directly on a customer’s Databricks instance, integrate with their data, use and extend
2511
- Databricks services, and enable users to interact through single sign-on."""
2512
-
2513
- def __init__(self, api_client):
2514
- self._api = api_client
2515
-
2516
- def wait_get_app_idle(self,
2517
- name: str,
2518
- timeout=timedelta(minutes=20),
2519
- callback: Optional[Callable[[App], None]] = None) -> App:
2520
- deadline = time.time() + timeout.total_seconds()
2521
- target_states = (AppState.IDLE, )
2522
- failure_states = (AppState.ERROR, )
2523
- status_message = 'polling...'
2524
- attempt = 1
2525
- while time.time() < deadline:
2526
- poll = self.get(name=name)
2527
- status = poll.status.state
2528
- status_message = f'current status: {status}'
2529
- if poll.status:
2530
- status_message = poll.status.message
2531
- if status in target_states:
2532
- return poll
2533
- if callback:
2534
- callback(poll)
2535
- if status in failure_states:
2536
- msg = f'failed to reach IDLE, got {status}: {status_message}'
2537
- raise OperationFailed(msg)
2538
- prefix = f"name={name}"
2539
- sleep = attempt
2540
- if sleep > 10:
2541
- # sleep 10s max per attempt
2542
- sleep = 10
2543
- _LOG.debug(f'{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)')
2544
- time.sleep(sleep + random.random())
2545
- attempt += 1
2546
- raise TimeoutError(f'timed out after {timeout}: {status_message}')
2547
-
2548
- def wait_get_deployment_app_succeeded(
2549
- self,
2550
- app_name: str,
2551
- deployment_id: str,
2552
- timeout=timedelta(minutes=20),
2553
- callback: Optional[Callable[[AppDeployment], None]] = None) -> AppDeployment:
2554
- deadline = time.time() + timeout.total_seconds()
2555
- target_states = (AppDeploymentState.SUCCEEDED, )
2556
- failure_states = (AppDeploymentState.FAILED, )
2557
- status_message = 'polling...'
2558
- attempt = 1
2559
- while time.time() < deadline:
2560
- poll = self.get_deployment(app_name=app_name, deployment_id=deployment_id)
2561
- status = poll.status.state
2562
- status_message = f'current status: {status}'
2563
- if poll.status:
2564
- status_message = poll.status.message
2565
- if status in target_states:
2566
- return poll
2567
- if callback:
2568
- callback(poll)
2569
- if status in failure_states:
2570
- msg = f'failed to reach SUCCEEDED, got {status}: {status_message}'
2571
- raise OperationFailed(msg)
2572
- prefix = f"app_name={app_name}, deployment_id={deployment_id}"
2573
- sleep = attempt
2574
- if sleep > 10:
2575
- # sleep 10s max per attempt
2576
- sleep = 10
2577
- _LOG.debug(f'{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)')
2578
- time.sleep(sleep + random.random())
2579
- attempt += 1
2580
- raise TimeoutError(f'timed out after {timeout}: {status_message}')
2581
-
2582
- def create(self, name: str, *, description: Optional[str] = None) -> Wait[App]:
2583
- """Create an app.
2584
-
2585
- Creates a new app.
2586
-
2587
- :param name: str
2588
- The name of the app. The name must contain only lowercase alphanumeric characters and hyphens. It
2589
- must be unique within the workspace.
2590
- :param description: str (optional)
2591
- The description of the app.
2592
-
2593
- :returns:
2594
- Long-running operation waiter for :class:`App`.
2595
- See :method:wait_get_app_idle for more details.
2596
- """
2597
- body = {}
2598
- if description is not None: body['description'] = description
2599
- if name is not None: body['name'] = name
2600
- headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
2601
-
2602
- op_response = self._api.do('POST', '/api/2.0/preview/apps', body=body, headers=headers)
2603
- return Wait(self.wait_get_app_idle, response=App.from_dict(op_response), name=op_response['name'])
2604
-
2605
- def create_and_wait(self,
2606
- name: str,
2607
- *,
2608
- description: Optional[str] = None,
2609
- timeout=timedelta(minutes=20)) -> App:
2610
- return self.create(description=description, name=name).result(timeout=timeout)
2611
-
2612
- def delete(self, name: str):
2613
- """Delete an app.
2614
-
2615
- Deletes an app.
2616
-
2617
- :param name: str
2618
- The name of the app.
2619
-
2620
-
2621
- """
2622
-
2623
- headers = {'Accept': 'application/json', }
2624
-
2625
- self._api.do('DELETE', f'/api/2.0/preview/apps/{name}', headers=headers)
2626
-
2627
- def deploy(self, app_name: str, source_code_path: str, mode: AppDeploymentMode) -> Wait[AppDeployment]:
2628
- """Create an app deployment.
2629
-
2630
- Creates an app deployment for the app with the supplied name.
2631
-
2632
- :param app_name: str
2633
- The name of the app.
2634
- :param source_code_path: str
2635
- The workspace file system path of the source code used to create the app deployment. This is
2636
- different from `deployment_artifacts.source_code_path`, which is the path used by the deployed app.
2637
- The former refers to the original source code location of the app in the workspace during deployment
2638
- creation, whereas the latter provides a system generated stable snapshotted source code path used by
2639
- the deployment.
2640
- :param mode: :class:`AppDeploymentMode`
2641
- The mode of which the deployment will manage the source code.
2642
-
2643
- :returns:
2644
- Long-running operation waiter for :class:`AppDeployment`.
2645
- See :method:wait_get_deployment_app_succeeded for more details.
2646
- """
2647
- body = {}
2648
- if mode is not None: body['mode'] = mode.value
2649
- if source_code_path is not None: body['source_code_path'] = source_code_path
2650
- headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
2651
-
2652
- op_response = self._api.do('POST',
2653
- f'/api/2.0/preview/apps/{app_name}/deployments',
2654
- body=body,
2655
- headers=headers)
2656
- return Wait(self.wait_get_deployment_app_succeeded,
2657
- response=AppDeployment.from_dict(op_response),
2658
- app_name=app_name,
2659
- deployment_id=op_response['deployment_id'])
2660
-
2661
- def deploy_and_wait(self,
2662
- app_name: str,
2663
- source_code_path: str,
2664
- mode: AppDeploymentMode,
2665
- timeout=timedelta(minutes=20)) -> AppDeployment:
2666
- return self.deploy(app_name=app_name, mode=mode,
2667
- source_code_path=source_code_path).result(timeout=timeout)
2668
-
2669
- def get(self, name: str) -> App:
2670
- """Get an app.
2671
-
2672
- Retrieves information for the app with the supplied name.
2673
-
2674
- :param name: str
2675
- The name of the app.
2676
-
2677
- :returns: :class:`App`
2678
- """
2679
-
2680
- headers = {'Accept': 'application/json', }
2681
-
2682
- res = self._api.do('GET', f'/api/2.0/preview/apps/{name}', headers=headers)
2683
- return App.from_dict(res)
2684
-
2685
- def get_deployment(self, app_name: str, deployment_id: str) -> AppDeployment:
2686
- """Get an app deployment.
2687
-
2688
- Retrieves information for the app deployment with the supplied name and deployment id.
2689
-
2690
- :param app_name: str
2691
- The name of the app.
2692
- :param deployment_id: str
2693
- The unique id of the deployment.
2694
-
2695
- :returns: :class:`AppDeployment`
2696
- """
2697
-
2698
- headers = {'Accept': 'application/json', }
2699
-
2700
- res = self._api.do('GET',
2701
- f'/api/2.0/preview/apps/{app_name}/deployments/{deployment_id}',
2702
- headers=headers)
2703
- return AppDeployment.from_dict(res)
2704
-
2705
- def get_environment(self, name: str) -> AppEnvironment:
2706
- """Get app environment.
2707
-
2708
- Retrieves app environment.
2709
-
2710
- :param name: str
2711
- The name of the app.
2712
-
2713
- :returns: :class:`AppEnvironment`
2714
- """
2715
-
2716
- headers = {'Accept': 'application/json', }
2717
-
2718
- res = self._api.do('GET', f'/api/2.0/preview/apps/{name}/environment', headers=headers)
2719
- return AppEnvironment.from_dict(res)
2720
-
2721
- def list(self, *, page_size: Optional[int] = None, page_token: Optional[str] = None) -> Iterator[App]:
2722
- """List apps.
2723
-
2724
- Lists all apps in the workspace.
2725
-
2726
- :param page_size: int (optional)
2727
- Upper bound for items returned.
2728
- :param page_token: str (optional)
2729
- Pagination token to go to the next page of apps. Requests first page if absent.
2730
-
2731
- :returns: Iterator over :class:`App`
2732
- """
2733
-
2734
- query = {}
2735
- if page_size is not None: query['page_size'] = page_size
2736
- if page_token is not None: query['page_token'] = page_token
2737
- headers = {'Accept': 'application/json', }
2738
-
2739
- while True:
2740
- json = self._api.do('GET', '/api/2.0/preview/apps', query=query, headers=headers)
2741
- if 'apps' in json:
2742
- for v in json['apps']:
2743
- yield App.from_dict(v)
2744
- if 'next_page_token' not in json or not json['next_page_token']:
2745
- return
2746
- query['page_token'] = json['next_page_token']
2747
-
2748
- def list_deployments(self,
2749
- app_name: str,
2750
- *,
2751
- page_size: Optional[int] = None,
2752
- page_token: Optional[str] = None) -> Iterator[AppDeployment]:
2753
- """List app deployments.
2754
-
2755
- Lists all app deployments for the app with the supplied name.
2756
-
2757
- :param app_name: str
2758
- The name of the app.
2759
- :param page_size: int (optional)
2760
- Upper bound for items returned.
2761
- :param page_token: str (optional)
2762
- Pagination token to go to the next page of apps. Requests first page if absent.
2763
-
2764
- :returns: Iterator over :class:`AppDeployment`
2765
- """
2766
-
2767
- query = {}
2768
- if page_size is not None: query['page_size'] = page_size
2769
- if page_token is not None: query['page_token'] = page_token
2770
- headers = {'Accept': 'application/json', }
2771
-
2772
- while True:
2773
- json = self._api.do('GET',
2774
- f'/api/2.0/preview/apps/{app_name}/deployments',
2775
- query=query,
2776
- headers=headers)
2777
- if 'app_deployments' in json:
2778
- for v in json['app_deployments']:
2779
- yield AppDeployment.from_dict(v)
2780
- if 'next_page_token' not in json or not json['next_page_token']:
2781
- return
2782
- query['page_token'] = json['next_page_token']
2783
-
2784
- def start(self, name: str) -> AppDeployment:
2785
- """Start an app.
2786
-
2787
- Start the last active deployment of the app in the workspace.
2788
-
2789
- :param name: str
2790
- The name of the app.
2791
-
2792
- :returns: :class:`AppDeployment`
2793
- """
2794
-
2795
- headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
2796
-
2797
- res = self._api.do('POST', f'/api/2.0/preview/apps/{name}/start', headers=headers)
2798
- return AppDeployment.from_dict(res)
2799
-
2800
- def stop(self, name: str):
2801
- """Stop an app.
2802
-
2803
- Stops the active deployment of the app in the workspace.
2804
-
2805
- :param name: str
2806
- The name of the app.
2807
-
2808
-
2809
- """
2810
-
2811
- headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
2812
-
2813
- self._api.do('POST', f'/api/2.0/preview/apps/{name}/stop', headers=headers)
2814
-
2815
- def update(self, name: str, *, description: Optional[str] = None) -> App:
2816
- """Update an app.
2817
-
2818
- Updates the app with the supplied name.
2819
-
2820
- :param name: str
2821
- The name of the app. The name must contain only lowercase alphanumeric characters and hyphens. It
2822
- must be unique within the workspace.
2823
- :param description: str (optional)
2824
- The description of the app.
2825
-
2826
- :returns: :class:`App`
2827
- """
2828
- body = {}
2829
- if description is not None: body['description'] = description
2830
- headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
2831
-
2832
- res = self._api.do('PATCH', f'/api/2.0/preview/apps/{name}', body=body, headers=headers)
2833
- return App.from_dict(res)
2834
-
2835
-
2836
2273
  class ServingEndpointsAPI:
2837
2274
  """The Serving Endpoints API allows you to create, update, and delete model serving endpoints.
2838
2275
 
@@ -2855,7 +2292,8 @@ class ServingEndpointsAPI:
2855
2292
  callback: Optional[Callable[[ServingEndpointDetailed], None]] = None) -> ServingEndpointDetailed:
2856
2293
  deadline = time.time() + timeout.total_seconds()
2857
2294
  target_states = (EndpointStateConfigUpdate.NOT_UPDATING, )
2858
- failure_states = (EndpointStateConfigUpdate.UPDATE_FAILED, )
2295
+ failure_states = (EndpointStateConfigUpdate.UPDATE_FAILED, EndpointStateConfigUpdate.UPDATE_CANCELED,
2296
+ )
2859
2297
  status_message = 'polling...'
2860
2298
  attempt = 1
2861
2299
  while time.time() < deadline:
@@ -3335,3 +2773,118 @@ class ServingEndpointsAPI:
3335
2773
  body=body,
3336
2774
  headers=headers)
3337
2775
  return ServingEndpointPermissions.from_dict(res)
2776
+
2777
+
2778
+ class ServingEndpointsDataPlaneAPI:
2779
+ """Serving endpoints DataPlane provides a set of operations to interact with data plane endpoints for Serving
2780
+ endpoints service."""
2781
+
2782
+ def __init__(self, api_client, control_plane):
2783
+ self._api = api_client
2784
+ self._control_plane = control_plane
2785
+ self._data_plane_service = DataPlaneService()
2786
+
2787
+ def query(self,
2788
+ name: str,
2789
+ *,
2790
+ dataframe_records: Optional[List[Any]] = None,
2791
+ dataframe_split: Optional[DataframeSplitInput] = None,
2792
+ extra_params: Optional[Dict[str, str]] = None,
2793
+ input: Optional[Any] = None,
2794
+ inputs: Optional[Any] = None,
2795
+ instances: Optional[List[Any]] = None,
2796
+ max_tokens: Optional[int] = None,
2797
+ messages: Optional[List[ChatMessage]] = None,
2798
+ n: Optional[int] = None,
2799
+ prompt: Optional[Any] = None,
2800
+ stop: Optional[List[str]] = None,
2801
+ stream: Optional[bool] = None,
2802
+ temperature: Optional[float] = None) -> QueryEndpointResponse:
2803
+ """Query a serving endpoint.
2804
+
2805
+ :param name: str
2806
+ The name of the serving endpoint. This field is required.
2807
+ :param dataframe_records: List[Any] (optional)
2808
+ Pandas Dataframe input in the records orientation.
2809
+ :param dataframe_split: :class:`DataframeSplitInput` (optional)
2810
+ Pandas Dataframe input in the split orientation.
2811
+ :param extra_params: Dict[str,str] (optional)
2812
+ The extra parameters field used ONLY for __completions, chat,__ and __embeddings external &
2813
+ foundation model__ serving endpoints. This is a map of strings and should only be used with other
2814
+ external/foundation model query fields.
2815
+ :param input: Any (optional)
2816
+ The input string (or array of strings) field used ONLY for __embeddings external & foundation
2817
+ model__ serving endpoints and is the only field (along with extra_params if needed) used by
2818
+ embeddings queries.
2819
+ :param inputs: Any (optional)
2820
+ Tensor-based input in columnar format.
2821
+ :param instances: List[Any] (optional)
2822
+ Tensor-based input in row format.
2823
+ :param max_tokens: int (optional)
2824
+ The max tokens field used ONLY for __completions__ and __chat external & foundation model__ serving
2825
+ endpoints. This is an integer and should only be used with other chat/completions query fields.
2826
+ :param messages: List[:class:`ChatMessage`] (optional)
2827
+ The messages field used ONLY for __chat external & foundation model__ serving endpoints. This is a
2828
+ map of strings and should only be used with other chat query fields.
2829
+ :param n: int (optional)
2830
+ The n (number of candidates) field used ONLY for __completions__ and __chat external & foundation
2831
+ model__ serving endpoints. This is an integer between 1 and 5 with a default of 1 and should only be
2832
+ used with other chat/completions query fields.
2833
+ :param prompt: Any (optional)
2834
+ The prompt string (or array of strings) field used ONLY for __completions external & foundation
2835
+ model__ serving endpoints and should only be used with other completions query fields.
2836
+ :param stop: List[str] (optional)
2837
+ The stop sequences field used ONLY for __completions__ and __chat external & foundation model__
2838
+ serving endpoints. This is a list of strings and should only be used with other chat/completions
2839
+ query fields.
2840
+ :param stream: bool (optional)
2841
+ The stream field used ONLY for __completions__ and __chat external & foundation model__ serving
2842
+ endpoints. This is a boolean defaulting to false and should only be used with other chat/completions
2843
+ query fields.
2844
+ :param temperature: float (optional)
2845
+ The temperature field used ONLY for __completions__ and __chat external & foundation model__ serving
2846
+ endpoints. This is a float between 0.0 and 2.0 with a default of 1.0 and should only be used with
2847
+ other chat/completions query fields.
2848
+
2849
+ :returns: :class:`QueryEndpointResponse`
2850
+ """
2851
+ body = {}
2852
+ if dataframe_records is not None: body['dataframe_records'] = [v for v in dataframe_records]
2853
+ if dataframe_split is not None: body['dataframe_split'] = dataframe_split.as_dict()
2854
+ if extra_params is not None: body['extra_params'] = extra_params
2855
+ if input is not None: body['input'] = input
2856
+ if inputs is not None: body['inputs'] = inputs
2857
+ if instances is not None: body['instances'] = [v for v in instances]
2858
+ if max_tokens is not None: body['max_tokens'] = max_tokens
2859
+ if messages is not None: body['messages'] = [v.as_dict() for v in messages]
2860
+ if n is not None: body['n'] = n
2861
+ if prompt is not None: body['prompt'] = prompt
2862
+ if stop is not None: body['stop'] = [v for v in stop]
2863
+ if stream is not None: body['stream'] = stream
2864
+ if temperature is not None: body['temperature'] = temperature
2865
+
2866
+ def info_getter():
2867
+ response = self._control_plane.get(name=name, )
2868
+ if response.data_plane_info is None:
2869
+ raise Exception("Resource does not support direct Data Plane access")
2870
+ return response.data_plane_info.query_info
2871
+
2872
+ get_params = [name, ]
2873
+ data_plane_details = self._data_plane_service.get_data_plane_details('query', get_params, info_getter,
2874
+ self._api.get_oauth_token)
2875
+ token = data_plane_details.token
2876
+
2877
+ def auth(r: requests.PreparedRequest) -> requests.PreparedRequest:
2878
+ authorization = f"{token.token_type} {token.access_token}"
2879
+ r.headers["Authorization"] = authorization
2880
+ return r
2881
+
2882
+ headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
2883
+ response_headers = ['served-model-name', ]
2884
+ res = self._api.do('POST',
2885
+ url=data_plane_details.endpoint_url,
2886
+ body=body,
2887
+ headers=headers,
2888
+ response_headers=response_headers,
2889
+ auth=auth)
2890
+ return QueryEndpointResponse.from_dict(res)