databricks-sdk 0.29.0__py3-none-any.whl → 0.30.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.

@@ -3,14 +3,20 @@
3
3
  from __future__ import annotations
4
4
 
5
5
  import logging
6
+ import random
7
+ import time
6
8
  from dataclasses import dataclass
9
+ from datetime import timedelta
7
10
  from enum import Enum
8
- from typing import Dict, Iterator, List, Optional
11
+ from typing import Callable, Dict, Iterator, List, Optional
9
12
 
10
- from ._internal import _enum, _from_dict, _repeated_dict
13
+ from ..errors import OperationFailed
14
+ from ._internal import Wait, _enum, _from_dict, _repeated_dict
11
15
 
12
16
  _LOG = logging.getLogger('databricks.sdk')
13
17
 
18
+ from databricks.sdk.service import sql
19
+
14
20
  # all definitions in this file are in alphabetical order
15
21
 
16
22
 
@@ -202,7 +208,6 @@ class Dashboard:
202
208
  class DashboardView(Enum):
203
209
 
204
210
  DASHBOARD_VIEW_BASIC = 'DASHBOARD_VIEW_BASIC'
205
- DASHBOARD_VIEW_FULL = 'DASHBOARD_VIEW_FULL'
206
211
 
207
212
 
208
213
  @dataclass
@@ -233,6 +238,244 @@ class DeleteSubscriptionResponse:
233
238
  return cls()
234
239
 
235
240
 
241
+ @dataclass
242
+ class GenieAttachment:
243
+ """Genie AI Response"""
244
+
245
+ query: Optional[QueryAttachment] = None
246
+
247
+ text: Optional[TextAttachment] = None
248
+
249
+ def as_dict(self) -> dict:
250
+ """Serializes the GenieAttachment into a dictionary suitable for use as a JSON request body."""
251
+ body = {}
252
+ if self.query: body['query'] = self.query.as_dict()
253
+ if self.text: body['text'] = self.text.as_dict()
254
+ return body
255
+
256
+ @classmethod
257
+ def from_dict(cls, d: Dict[str, any]) -> GenieAttachment:
258
+ """Deserializes the GenieAttachment from a dictionary."""
259
+ return cls(query=_from_dict(d, 'query', QueryAttachment), text=_from_dict(d, 'text', TextAttachment))
260
+
261
+
262
+ @dataclass
263
+ class GenieConversation:
264
+ id: str
265
+ """Conversation ID"""
266
+
267
+ space_id: str
268
+ """Genie space ID"""
269
+
270
+ user_id: int
271
+ """ID of the user who created the conversation"""
272
+
273
+ title: str
274
+ """Conversation title"""
275
+
276
+ created_timestamp: Optional[int] = None
277
+ """Timestamp when the message was created"""
278
+
279
+ last_updated_timestamp: Optional[int] = None
280
+ """Timestamp when the message was last updated"""
281
+
282
+ def as_dict(self) -> dict:
283
+ """Serializes the GenieConversation into a dictionary suitable for use as a JSON request body."""
284
+ body = {}
285
+ if self.created_timestamp is not None: body['created_timestamp'] = self.created_timestamp
286
+ if self.id is not None: body['id'] = self.id
287
+ if self.last_updated_timestamp is not None:
288
+ body['last_updated_timestamp'] = self.last_updated_timestamp
289
+ if self.space_id is not None: body['space_id'] = self.space_id
290
+ if self.title is not None: body['title'] = self.title
291
+ if self.user_id is not None: body['user_id'] = self.user_id
292
+ return body
293
+
294
+ @classmethod
295
+ def from_dict(cls, d: Dict[str, any]) -> GenieConversation:
296
+ """Deserializes the GenieConversation from a dictionary."""
297
+ return cls(created_timestamp=d.get('created_timestamp', None),
298
+ id=d.get('id', None),
299
+ last_updated_timestamp=d.get('last_updated_timestamp', None),
300
+ space_id=d.get('space_id', None),
301
+ title=d.get('title', None),
302
+ user_id=d.get('user_id', None))
303
+
304
+
305
+ @dataclass
306
+ class GenieCreateConversationMessageRequest:
307
+ content: str
308
+ """User message content."""
309
+
310
+ conversation_id: Optional[str] = None
311
+ """The ID associated with the conversation."""
312
+
313
+ space_id: Optional[str] = None
314
+ """The ID associated with the Genie space where the conversation is started."""
315
+
316
+ def as_dict(self) -> dict:
317
+ """Serializes the GenieCreateConversationMessageRequest into a dictionary suitable for use as a JSON request body."""
318
+ body = {}
319
+ if self.content is not None: body['content'] = self.content
320
+ if self.conversation_id is not None: body['conversation_id'] = self.conversation_id
321
+ if self.space_id is not None: body['space_id'] = self.space_id
322
+ return body
323
+
324
+ @classmethod
325
+ def from_dict(cls, d: Dict[str, any]) -> GenieCreateConversationMessageRequest:
326
+ """Deserializes the GenieCreateConversationMessageRequest from a dictionary."""
327
+ return cls(content=d.get('content', None),
328
+ conversation_id=d.get('conversation_id', None),
329
+ space_id=d.get('space_id', None))
330
+
331
+
332
+ @dataclass
333
+ class GenieGetMessageQueryResultResponse:
334
+ statement_response: Optional[sql.StatementResponse] = None
335
+ """SQL Statement Execution response. See [Get status, manifest, and result first
336
+ chunk](:method:statementexecution/getstatement) for more details."""
337
+
338
+ def as_dict(self) -> dict:
339
+ """Serializes the GenieGetMessageQueryResultResponse into a dictionary suitable for use as a JSON request body."""
340
+ body = {}
341
+ if self.statement_response: body['statement_response'] = self.statement_response.as_dict()
342
+ return body
343
+
344
+ @classmethod
345
+ def from_dict(cls, d: Dict[str, any]) -> GenieGetMessageQueryResultResponse:
346
+ """Deserializes the GenieGetMessageQueryResultResponse from a dictionary."""
347
+ return cls(statement_response=_from_dict(d, 'statement_response', sql.StatementResponse))
348
+
349
+
350
+ @dataclass
351
+ class GenieMessage:
352
+ id: str
353
+ """Message ID"""
354
+
355
+ space_id: str
356
+ """Genie space ID"""
357
+
358
+ conversation_id: str
359
+ """Conversation ID"""
360
+
361
+ content: str
362
+ """User message content"""
363
+
364
+ attachments: Optional[List[GenieAttachment]] = None
365
+ """AI produced response to the message"""
366
+
367
+ created_timestamp: Optional[int] = None
368
+ """Timestamp when the message was created"""
369
+
370
+ error: Optional[MessageError] = None
371
+ """Error message if AI failed to respond to the message"""
372
+
373
+ last_updated_timestamp: Optional[int] = None
374
+ """Timestamp when the message was last updated"""
375
+
376
+ query_result: Optional[Result] = None
377
+ """The result of SQL query if the message has a query attachment"""
378
+
379
+ status: Optional[MessageStatus] = None
380
+ """MesssageStatus. The possible values are: * `FETCHING_METADATA`: Fetching metadata from the data
381
+ sources. * `ASKING_AI`: Waiting for the LLM to respond to the users question. *
382
+ `EXECUTING_QUERY`: Executing AI provided SQL query. Get the SQL query result by calling
383
+ [getMessageQueryResult](:method:genie/getMessageQueryResult) API. **Important: The message
384
+ status will stay in the `EXECUTING_QUERY` until a client calls
385
+ [getMessageQueryResult](:method:genie/getMessageQueryResult)**. * `FAILED`: Generating a
386
+ response or the executing the query failed. Please see `error` field. * `COMPLETED`: Message
387
+ processing is completed. Results are in the `attachments` field. Get the SQL query result by
388
+ calling [getMessageQueryResult](:method:genie/getMessageQueryResult) API. * `SUBMITTED`: Message
389
+ has been submitted. * `QUERY_RESULT_EXPIRED`: SQL result is not available anymore. The user
390
+ needs to execute the query again. * `CANCELLED`: Message has been cancelled."""
391
+
392
+ user_id: Optional[int] = None
393
+ """ID of the user who created the message"""
394
+
395
+ def as_dict(self) -> dict:
396
+ """Serializes the GenieMessage into a dictionary suitable for use as a JSON request body."""
397
+ body = {}
398
+ if self.attachments: body['attachments'] = [v.as_dict() for v in self.attachments]
399
+ if self.content is not None: body['content'] = self.content
400
+ if self.conversation_id is not None: body['conversation_id'] = self.conversation_id
401
+ if self.created_timestamp is not None: body['created_timestamp'] = self.created_timestamp
402
+ if self.error: body['error'] = self.error.as_dict()
403
+ if self.id is not None: body['id'] = self.id
404
+ if self.last_updated_timestamp is not None:
405
+ body['last_updated_timestamp'] = self.last_updated_timestamp
406
+ if self.query_result: body['query_result'] = self.query_result.as_dict()
407
+ if self.space_id is not None: body['space_id'] = self.space_id
408
+ if self.status is not None: body['status'] = self.status.value
409
+ if self.user_id is not None: body['user_id'] = self.user_id
410
+ return body
411
+
412
+ @classmethod
413
+ def from_dict(cls, d: Dict[str, any]) -> GenieMessage:
414
+ """Deserializes the GenieMessage from a dictionary."""
415
+ return cls(attachments=_repeated_dict(d, 'attachments', GenieAttachment),
416
+ content=d.get('content', None),
417
+ conversation_id=d.get('conversation_id', None),
418
+ created_timestamp=d.get('created_timestamp', None),
419
+ error=_from_dict(d, 'error', MessageError),
420
+ id=d.get('id', None),
421
+ last_updated_timestamp=d.get('last_updated_timestamp', None),
422
+ query_result=_from_dict(d, 'query_result', Result),
423
+ space_id=d.get('space_id', None),
424
+ status=_enum(d, 'status', MessageStatus),
425
+ user_id=d.get('user_id', None))
426
+
427
+
428
+ @dataclass
429
+ class GenieStartConversationMessageRequest:
430
+ content: str
431
+ """The text of the message that starts the conversation."""
432
+
433
+ space_id: Optional[str] = None
434
+ """The ID associated with the Genie space where you want to start a conversation."""
435
+
436
+ def as_dict(self) -> dict:
437
+ """Serializes the GenieStartConversationMessageRequest into a dictionary suitable for use as a JSON request body."""
438
+ body = {}
439
+ if self.content is not None: body['content'] = self.content
440
+ if self.space_id is not None: body['space_id'] = self.space_id
441
+ return body
442
+
443
+ @classmethod
444
+ def from_dict(cls, d: Dict[str, any]) -> GenieStartConversationMessageRequest:
445
+ """Deserializes the GenieStartConversationMessageRequest from a dictionary."""
446
+ return cls(content=d.get('content', None), space_id=d.get('space_id', None))
447
+
448
+
449
+ @dataclass
450
+ class GenieStartConversationResponse:
451
+ message_id: str
452
+ """Message ID"""
453
+
454
+ conversation_id: str
455
+ """Conversation ID"""
456
+
457
+ conversation: Optional[GenieConversation] = None
458
+
459
+ message: Optional[GenieMessage] = None
460
+
461
+ def as_dict(self) -> dict:
462
+ """Serializes the GenieStartConversationResponse into a dictionary suitable for use as a JSON request body."""
463
+ body = {}
464
+ if self.conversation: body['conversation'] = self.conversation.as_dict()
465
+ if self.conversation_id is not None: body['conversation_id'] = self.conversation_id
466
+ if self.message: body['message'] = self.message.as_dict()
467
+ if self.message_id is not None: body['message_id'] = self.message_id
468
+ return body
469
+
470
+ @classmethod
471
+ def from_dict(cls, d: Dict[str, any]) -> GenieStartConversationResponse:
472
+ """Deserializes the GenieStartConversationResponse from a dictionary."""
473
+ return cls(conversation=_from_dict(d, 'conversation', GenieConversation),
474
+ conversation_id=d.get('conversation_id', None),
475
+ message=_from_dict(d, 'message', GenieMessage),
476
+ message_id=d.get('message_id', None))
477
+
478
+
236
479
  class LifecycleState(Enum):
237
480
 
238
481
  ACTIVE = 'ACTIVE'
@@ -305,6 +548,90 @@ class ListSubscriptionsResponse:
305
548
  subscriptions=_repeated_dict(d, 'subscriptions', Subscription))
306
549
 
307
550
 
551
+ @dataclass
552
+ class MessageError:
553
+ error: Optional[str] = None
554
+
555
+ type: Optional[MessageErrorType] = None
556
+
557
+ def as_dict(self) -> dict:
558
+ """Serializes the MessageError into a dictionary suitable for use as a JSON request body."""
559
+ body = {}
560
+ if self.error is not None: body['error'] = self.error
561
+ if self.type is not None: body['type'] = self.type.value
562
+ return body
563
+
564
+ @classmethod
565
+ def from_dict(cls, d: Dict[str, any]) -> MessageError:
566
+ """Deserializes the MessageError from a dictionary."""
567
+ return cls(error=d.get('error', None), type=_enum(d, 'type', MessageErrorType))
568
+
569
+
570
+ class MessageErrorType(Enum):
571
+
572
+ BLOCK_MULTIPLE_EXECUTIONS_EXCEPTION = 'BLOCK_MULTIPLE_EXECUTIONS_EXCEPTION'
573
+ CHAT_COMPLETION_CLIENT_EXCEPTION = 'CHAT_COMPLETION_CLIENT_EXCEPTION'
574
+ CHAT_COMPLETION_CLIENT_TIMEOUT_EXCEPTION = 'CHAT_COMPLETION_CLIENT_TIMEOUT_EXCEPTION'
575
+ CHAT_COMPLETION_NETWORK_EXCEPTION = 'CHAT_COMPLETION_NETWORK_EXCEPTION'
576
+ CONTENT_FILTER_EXCEPTION = 'CONTENT_FILTER_EXCEPTION'
577
+ CONTEXT_EXCEEDED_EXCEPTION = 'CONTEXT_EXCEEDED_EXCEPTION'
578
+ COULD_NOT_GET_UC_SCHEMA_EXCEPTION = 'COULD_NOT_GET_UC_SCHEMA_EXCEPTION'
579
+ DEPLOYMENT_NOT_FOUND_EXCEPTION = 'DEPLOYMENT_NOT_FOUND_EXCEPTION'
580
+ FUNCTIONS_NOT_AVAILABLE_EXCEPTION = 'FUNCTIONS_NOT_AVAILABLE_EXCEPTION'
581
+ FUNCTION_ARGUMENTS_INVALID_EXCEPTION = 'FUNCTION_ARGUMENTS_INVALID_EXCEPTION'
582
+ FUNCTION_ARGUMENTS_INVALID_JSON_EXCEPTION = 'FUNCTION_ARGUMENTS_INVALID_JSON_EXCEPTION'
583
+ FUNCTION_CALL_MISSING_PARAMETER_EXCEPTION = 'FUNCTION_CALL_MISSING_PARAMETER_EXCEPTION'
584
+ GENERIC_CHAT_COMPLETION_EXCEPTION = 'GENERIC_CHAT_COMPLETION_EXCEPTION'
585
+ GENERIC_CHAT_COMPLETION_SERVICE_EXCEPTION = 'GENERIC_CHAT_COMPLETION_SERVICE_EXCEPTION'
586
+ GENERIC_SQL_EXEC_API_CALL_EXCEPTION = 'GENERIC_SQL_EXEC_API_CALL_EXCEPTION'
587
+ ILLEGAL_PARAMETER_DEFINITION_EXCEPTION = 'ILLEGAL_PARAMETER_DEFINITION_EXCEPTION'
588
+ INVALID_CERTIFIED_ANSWER_FUNCTION_EXCEPTION = 'INVALID_CERTIFIED_ANSWER_FUNCTION_EXCEPTION'
589
+ INVALID_CERTIFIED_ANSWER_IDENTIFIER_EXCEPTION = 'INVALID_CERTIFIED_ANSWER_IDENTIFIER_EXCEPTION'
590
+ INVALID_CHAT_COMPLETION_JSON_EXCEPTION = 'INVALID_CHAT_COMPLETION_JSON_EXCEPTION'
591
+ INVALID_COMPLETION_REQUEST_EXCEPTION = 'INVALID_COMPLETION_REQUEST_EXCEPTION'
592
+ INVALID_FUNCTION_CALL_EXCEPTION = 'INVALID_FUNCTION_CALL_EXCEPTION'
593
+ INVALID_TABLE_IDENTIFIER_EXCEPTION = 'INVALID_TABLE_IDENTIFIER_EXCEPTION'
594
+ LOCAL_CONTEXT_EXCEEDED_EXCEPTION = 'LOCAL_CONTEXT_EXCEEDED_EXCEPTION'
595
+ MESSAGE_DELETED_WHILE_EXECUTING_EXCEPTION = 'MESSAGE_DELETED_WHILE_EXECUTING_EXCEPTION'
596
+ MESSAGE_UPDATED_WHILE_EXECUTING_EXCEPTION = 'MESSAGE_UPDATED_WHILE_EXECUTING_EXCEPTION'
597
+ NO_TABLES_TO_QUERY_EXCEPTION = 'NO_TABLES_TO_QUERY_EXCEPTION'
598
+ RATE_LIMIT_EXCEEDED_GENERIC_EXCEPTION = 'RATE_LIMIT_EXCEEDED_GENERIC_EXCEPTION'
599
+ RATE_LIMIT_EXCEEDED_SPECIFIED_WAIT_EXCEPTION = 'RATE_LIMIT_EXCEEDED_SPECIFIED_WAIT_EXCEPTION'
600
+ REPLY_PROCESS_TIMEOUT_EXCEPTION = 'REPLY_PROCESS_TIMEOUT_EXCEPTION'
601
+ RETRYABLE_PROCESSING_EXCEPTION = 'RETRYABLE_PROCESSING_EXCEPTION'
602
+ SQL_EXECUTION_EXCEPTION = 'SQL_EXECUTION_EXCEPTION'
603
+ TABLES_MISSING_EXCEPTION = 'TABLES_MISSING_EXCEPTION'
604
+ TOO_MANY_CERTIFIED_ANSWERS_EXCEPTION = 'TOO_MANY_CERTIFIED_ANSWERS_EXCEPTION'
605
+ TOO_MANY_TABLES_EXCEPTION = 'TOO_MANY_TABLES_EXCEPTION'
606
+ UNEXPECTED_REPLY_PROCESS_EXCEPTION = 'UNEXPECTED_REPLY_PROCESS_EXCEPTION'
607
+ UNKNOWN_AI_MODEL = 'UNKNOWN_AI_MODEL'
608
+ WAREHOUSE_ACCESS_MISSING_EXCEPTION = 'WAREHOUSE_ACCESS_MISSING_EXCEPTION'
609
+ WAREHOUSE_NOT_FOUND_EXCEPTION = 'WAREHOUSE_NOT_FOUND_EXCEPTION'
610
+
611
+
612
+ class MessageStatus(Enum):
613
+ """MesssageStatus. The possible values are: * `FETCHING_METADATA`: Fetching metadata from the data
614
+ sources. * `ASKING_AI`: Waiting for the LLM to respond to the users question. *
615
+ `EXECUTING_QUERY`: Executing AI provided SQL query. Get the SQL query result by calling
616
+ [getMessageQueryResult](:method:genie/getMessageQueryResult) API. **Important: The message
617
+ status will stay in the `EXECUTING_QUERY` until a client calls
618
+ [getMessageQueryResult](:method:genie/getMessageQueryResult)**. * `FAILED`: Generating a
619
+ response or the executing the query failed. Please see `error` field. * `COMPLETED`: Message
620
+ processing is completed. Results are in the `attachments` field. Get the SQL query result by
621
+ calling [getMessageQueryResult](:method:genie/getMessageQueryResult) API. * `SUBMITTED`: Message
622
+ has been submitted. * `QUERY_RESULT_EXPIRED`: SQL result is not available anymore. The user
623
+ needs to execute the query again. * `CANCELLED`: Message has been cancelled."""
624
+
625
+ ASKING_AI = 'ASKING_AI'
626
+ CANCELLED = 'CANCELLED'
627
+ COMPLETED = 'COMPLETED'
628
+ EXECUTING_QUERY = 'EXECUTING_QUERY'
629
+ FAILED = 'FAILED'
630
+ FETCHING_METADATA = 'FETCHING_METADATA'
631
+ QUERY_RESULT_EXPIRED = 'QUERY_RESULT_EXPIRED'
632
+ SUBMITTED = 'SUBMITTED'
633
+
634
+
308
635
  @dataclass
309
636
  class MigrateDashboardRequest:
310
637
  source_dashboard_id: str
@@ -392,6 +719,76 @@ class PublishedDashboard:
392
719
  warehouse_id=d.get('warehouse_id', None))
393
720
 
394
721
 
722
+ @dataclass
723
+ class QueryAttachment:
724
+ description: Optional[str] = None
725
+ """Description of the query"""
726
+
727
+ id: Optional[str] = None
728
+
729
+ instruction_id: Optional[str] = None
730
+ """If the query was created on an instruction (trusted asset) we link to the id"""
731
+
732
+ instruction_title: Optional[str] = None
733
+ """Always store the title next to the id in case the original instruction title changes or the
734
+ instruction is deleted."""
735
+
736
+ last_updated_timestamp: Optional[int] = None
737
+ """Time when the user updated the query last"""
738
+
739
+ query: Optional[str] = None
740
+ """AI generated SQL query"""
741
+
742
+ title: Optional[str] = None
743
+ """Name of the query"""
744
+
745
+ def as_dict(self) -> dict:
746
+ """Serializes the QueryAttachment into a dictionary suitable for use as a JSON request body."""
747
+ body = {}
748
+ if self.description is not None: body['description'] = self.description
749
+ if self.id is not None: body['id'] = self.id
750
+ if self.instruction_id is not None: body['instruction_id'] = self.instruction_id
751
+ if self.instruction_title is not None: body['instruction_title'] = self.instruction_title
752
+ if self.last_updated_timestamp is not None:
753
+ body['last_updated_timestamp'] = self.last_updated_timestamp
754
+ if self.query is not None: body['query'] = self.query
755
+ if self.title is not None: body['title'] = self.title
756
+ return body
757
+
758
+ @classmethod
759
+ def from_dict(cls, d: Dict[str, any]) -> QueryAttachment:
760
+ """Deserializes the QueryAttachment from a dictionary."""
761
+ return cls(description=d.get('description', None),
762
+ id=d.get('id', None),
763
+ instruction_id=d.get('instruction_id', None),
764
+ instruction_title=d.get('instruction_title', None),
765
+ last_updated_timestamp=d.get('last_updated_timestamp', None),
766
+ query=d.get('query', None),
767
+ title=d.get('title', None))
768
+
769
+
770
+ @dataclass
771
+ class Result:
772
+ row_count: Optional[int] = None
773
+ """Row count of the result"""
774
+
775
+ statement_id: Optional[str] = None
776
+ """Statement Execution API statement id. Use [Get status, manifest, and result first
777
+ chunk](:method:statementexecution/getstatement) to get the full result data."""
778
+
779
+ def as_dict(self) -> dict:
780
+ """Serializes the Result into a dictionary suitable for use as a JSON request body."""
781
+ body = {}
782
+ if self.row_count is not None: body['row_count'] = self.row_count
783
+ if self.statement_id is not None: body['statement_id'] = self.statement_id
784
+ return body
785
+
786
+ @classmethod
787
+ def from_dict(cls, d: Dict[str, any]) -> Result:
788
+ """Deserializes the Result from a dictionary."""
789
+ return cls(row_count=d.get('row_count', None), statement_id=d.get('statement_id', None))
790
+
791
+
395
792
  @dataclass
396
793
  class Schedule:
397
794
  cron_schedule: CronSchedule
@@ -565,6 +962,26 @@ class SubscriptionSubscriberUser:
565
962
  return cls(user_id=d.get('user_id', None))
566
963
 
567
964
 
965
+ @dataclass
966
+ class TextAttachment:
967
+ content: Optional[str] = None
968
+ """AI generated message"""
969
+
970
+ id: Optional[str] = None
971
+
972
+ def as_dict(self) -> dict:
973
+ """Serializes the TextAttachment into a dictionary suitable for use as a JSON request body."""
974
+ body = {}
975
+ if self.content is not None: body['content'] = self.content
976
+ if self.id is not None: body['id'] = self.id
977
+ return body
978
+
979
+ @classmethod
980
+ def from_dict(cls, d: Dict[str, any]) -> TextAttachment:
981
+ """Deserializes the TextAttachment from a dictionary."""
982
+ return cls(content=d.get('content', None), id=d.get('id', None))
983
+
984
+
568
985
  @dataclass
569
986
  class TrashDashboardResponse:
570
987
 
@@ -675,6 +1092,193 @@ class UpdateScheduleRequest:
675
1092
  schedule_id=d.get('schedule_id', None))
676
1093
 
677
1094
 
1095
+ class GenieAPI:
1096
+ """Genie provides a no-code experience for business users, powered by AI/BI. Analysts set up spaces that
1097
+ business users can use to ask questions using natural language. Genie uses data registered to Unity
1098
+ Catalog and requires at least CAN USE permission on a Pro or Serverless SQL warehouse. Also, Databricks
1099
+ Assistant must be enabled."""
1100
+
1101
+ def __init__(self, api_client):
1102
+ self._api = api_client
1103
+
1104
+ def wait_get_message_genie_completed(
1105
+ self,
1106
+ conversation_id: str,
1107
+ message_id: str,
1108
+ space_id: str,
1109
+ timeout=timedelta(minutes=20),
1110
+ callback: Optional[Callable[[GenieMessage], None]] = None) -> GenieMessage:
1111
+ deadline = time.time() + timeout.total_seconds()
1112
+ target_states = (MessageStatus.COMPLETED, )
1113
+ failure_states = (MessageStatus.FAILED, )
1114
+ status_message = 'polling...'
1115
+ attempt = 1
1116
+ while time.time() < deadline:
1117
+ poll = self.get_message(conversation_id=conversation_id, message_id=message_id, space_id=space_id)
1118
+ status = poll.status
1119
+ status_message = f'current status: {status}'
1120
+ if status in target_states:
1121
+ return poll
1122
+ if callback:
1123
+ callback(poll)
1124
+ if status in failure_states:
1125
+ msg = f'failed to reach COMPLETED, got {status}: {status_message}'
1126
+ raise OperationFailed(msg)
1127
+ prefix = f"conversation_id={conversation_id}, message_id={message_id}, space_id={space_id}"
1128
+ sleep = attempt
1129
+ if sleep > 10:
1130
+ # sleep 10s max per attempt
1131
+ sleep = 10
1132
+ _LOG.debug(f'{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)')
1133
+ time.sleep(sleep + random.random())
1134
+ attempt += 1
1135
+ raise TimeoutError(f'timed out after {timeout}: {status_message}')
1136
+
1137
+ def create_message(self, space_id: str, conversation_id: str, content: str) -> Wait[GenieMessage]:
1138
+ """Create conversation message.
1139
+
1140
+ Create new message in [conversation](:method:genie/startconversation). The AI response uses all
1141
+ previously created messages in the conversation to respond.
1142
+
1143
+ :param space_id: str
1144
+ The ID associated with the Genie space where the conversation is started.
1145
+ :param conversation_id: str
1146
+ The ID associated with the conversation.
1147
+ :param content: str
1148
+ User message content.
1149
+
1150
+ :returns:
1151
+ Long-running operation waiter for :class:`GenieMessage`.
1152
+ See :method:wait_get_message_genie_completed for more details.
1153
+ """
1154
+ body = {}
1155
+ if content is not None: body['content'] = content
1156
+ headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
1157
+
1158
+ op_response = self._api.do(
1159
+ 'POST',
1160
+ f'/api/2.0/genie/spaces/{space_id}/conversations/{conversation_id}/messages',
1161
+ body=body,
1162
+ headers=headers)
1163
+ return Wait(self.wait_get_message_genie_completed,
1164
+ response=GenieMessage.from_dict(op_response),
1165
+ conversation_id=conversation_id,
1166
+ message_id=op_response['id'],
1167
+ space_id=space_id)
1168
+
1169
+ def create_message_and_wait(self,
1170
+ space_id: str,
1171
+ conversation_id: str,
1172
+ content: str,
1173
+ timeout=timedelta(minutes=20)) -> GenieMessage:
1174
+ return self.create_message(content=content, conversation_id=conversation_id,
1175
+ space_id=space_id).result(timeout=timeout)
1176
+
1177
+ def execute_message_query(self, space_id: str, conversation_id: str,
1178
+ message_id: str) -> GenieGetMessageQueryResultResponse:
1179
+ """Execute SQL query in a conversation message.
1180
+
1181
+ Execute the SQL query in the message.
1182
+
1183
+ :param space_id: str
1184
+ Genie space ID
1185
+ :param conversation_id: str
1186
+ Conversation ID
1187
+ :param message_id: str
1188
+ Message ID
1189
+
1190
+ :returns: :class:`GenieGetMessageQueryResultResponse`
1191
+ """
1192
+
1193
+ headers = {'Accept': 'application/json', }
1194
+
1195
+ res = self._api.do(
1196
+ 'POST',
1197
+ f'/api/2.0/genie/spaces/{space_id}/conversations/{conversation_id}/messages/{message_id}/execute-query',
1198
+ headers=headers)
1199
+ return GenieGetMessageQueryResultResponse.from_dict(res)
1200
+
1201
+ def get_message(self, space_id: str, conversation_id: str, message_id: str) -> GenieMessage:
1202
+ """Get conversation message.
1203
+
1204
+ Get message from conversation.
1205
+
1206
+ :param space_id: str
1207
+ The ID associated with the Genie space where the target conversation is located.
1208
+ :param conversation_id: str
1209
+ The ID associated with the target conversation.
1210
+ :param message_id: str
1211
+ The ID associated with the target message from the identified conversation.
1212
+
1213
+ :returns: :class:`GenieMessage`
1214
+ """
1215
+
1216
+ headers = {'Accept': 'application/json', }
1217
+
1218
+ res = self._api.do(
1219
+ 'GET',
1220
+ f'/api/2.0/genie/spaces/{space_id}/conversations/{conversation_id}/messages/{message_id}',
1221
+ headers=headers)
1222
+ return GenieMessage.from_dict(res)
1223
+
1224
+ def get_message_query_result(self, space_id: str, conversation_id: str,
1225
+ message_id: str) -> GenieGetMessageQueryResultResponse:
1226
+ """Get conversation message SQL query result.
1227
+
1228
+ Get the result of SQL query if the message has a query attachment. This is only available if a message
1229
+ has a query attachment and the message status is `EXECUTING_QUERY`.
1230
+
1231
+ :param space_id: str
1232
+ Genie space ID
1233
+ :param conversation_id: str
1234
+ Conversation ID
1235
+ :param message_id: str
1236
+ Message ID
1237
+
1238
+ :returns: :class:`GenieGetMessageQueryResultResponse`
1239
+ """
1240
+
1241
+ headers = {'Accept': 'application/json', }
1242
+
1243
+ res = self._api.do(
1244
+ 'GET',
1245
+ f'/api/2.0/genie/spaces/{space_id}/conversations/{conversation_id}/messages/{message_id}/query-result',
1246
+ headers=headers)
1247
+ return GenieGetMessageQueryResultResponse.from_dict(res)
1248
+
1249
+ def start_conversation(self, space_id: str, content: str) -> Wait[GenieMessage]:
1250
+ """Start conversation.
1251
+
1252
+ Start a new conversation.
1253
+
1254
+ :param space_id: str
1255
+ The ID associated with the Genie space where you want to start a conversation.
1256
+ :param content: str
1257
+ The text of the message that starts the conversation.
1258
+
1259
+ :returns:
1260
+ Long-running operation waiter for :class:`GenieMessage`.
1261
+ See :method:wait_get_message_genie_completed for more details.
1262
+ """
1263
+ body = {}
1264
+ if content is not None: body['content'] = content
1265
+ headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
1266
+
1267
+ op_response = self._api.do('POST',
1268
+ f'/api/2.0/genie/spaces/{space_id}/start-conversation',
1269
+ body=body,
1270
+ headers=headers)
1271
+ return Wait(self.wait_get_message_genie_completed,
1272
+ response=GenieStartConversationResponse.from_dict(op_response),
1273
+ conversation_id=op_response['conversation_id'],
1274
+ message_id=op_response['message_id'],
1275
+ space_id=space_id)
1276
+
1277
+ def start_conversation_and_wait(self, space_id: str, content: str,
1278
+ timeout=timedelta(minutes=20)) -> GenieMessage:
1279
+ return self.start_conversation(content=content, space_id=space_id).result(timeout=timeout)
1280
+
1281
+
678
1282
  class LakeviewAPI:
679
1283
  """These APIs provide specific management operations for Lakeview dashboards. Generic resource management can
680
1284
  be done with Workspace API (import, export, get-status, list, delete)."""
@@ -911,8 +1515,7 @@ class LakeviewAPI:
911
1515
  The flag to include dashboards located in the trash. If unspecified, only active dashboards will be
912
1516
  returned.
913
1517
  :param view: :class:`DashboardView` (optional)
914
- Indicates whether to include all metadata from the dashboard in the response. If unset, the response
915
- defaults to `DASHBOARD_VIEW_BASIC` which only includes summary metadata from the dashboard.
1518
+ `DASHBOARD_VIEW_BASIC`only includes summary metadata from the dashboard.
916
1519
 
917
1520
  :returns: Iterator over :class:`Dashboard`
918
1521
  """