com.databricks.sdk.service.dashboards.GenieAPI Maven / Gradle / Ivy
// Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
package com.databricks.sdk.service.dashboards;
import com.databricks.sdk.core.ApiClient;
import com.databricks.sdk.core.DatabricksException;
import com.databricks.sdk.support.Generated;
import com.databricks.sdk.support.Wait;
import java.time.Duration;
import java.util.Arrays;
import java.util.concurrent.TimeoutException;
import java.util.function.Consumer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Genie provides a no-code experience for business users, powered by AI/BI. Analysts set up spaces
* that business users can use to ask questions using natural language. Genie uses data registered
* to Unity Catalog and requires at least CAN USE permission on a Pro or Serverless SQL warehouse.
* Also, Databricks Assistant must be enabled.
*/
@Generated
public class GenieAPI {
private static final Logger LOG = LoggerFactory.getLogger(GenieAPI.class);
private final GenieService impl;
/** Regular-use constructor */
public GenieAPI(ApiClient apiClient) {
impl = new GenieImpl(apiClient);
}
/** Constructor for mocks */
public GenieAPI(GenieService mock) {
impl = mock;
}
public GenieMessage waitGetMessageGenieCompleted(
String conversationId, String messageId, String spaceId) throws TimeoutException {
return waitGetMessageGenieCompleted(
conversationId, messageId, spaceId, Duration.ofMinutes(20), null);
}
public GenieMessage waitGetMessageGenieCompleted(
String conversationId,
String messageId,
String spaceId,
Duration timeout,
Consumer callback)
throws TimeoutException {
long deadline = System.currentTimeMillis() + timeout.toMillis();
java.util.List targetStates = Arrays.asList(MessageStatus.COMPLETED);
java.util.List failureStates = Arrays.asList(MessageStatus.FAILED);
String statusMessage = "polling...";
int attempt = 1;
while (System.currentTimeMillis() < deadline) {
GenieMessage poll =
getMessage(
new GenieGetConversationMessageRequest()
.setConversationId(conversationId)
.setMessageId(messageId)
.setSpaceId(spaceId));
MessageStatus status = poll.getStatus();
statusMessage = String.format("current status: %s", status);
if (targetStates.contains(status)) {
return poll;
}
if (callback != null) {
callback.accept(poll);
}
if (failureStates.contains(status)) {
String msg = String.format("failed to reach COMPLETED, got %s: %s", status, statusMessage);
throw new IllegalStateException(msg);
}
String prefix =
String.format(
"conversationId=%s, messageId=%s, spaceId=%s", conversationId, messageId, spaceId);
int sleep = attempt;
if (sleep > 10) {
// sleep 10s max per attempt
sleep = 10;
}
LOG.info("{}: ({}) {} (sleeping ~{}s)", prefix, status, statusMessage, sleep);
try {
Thread.sleep((long) (sleep * 1000L + Math.random() * 1000));
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new DatabricksException("Current thread was interrupted", e);
}
attempt++;
}
throw new TimeoutException(String.format("timed out after %s: %s", timeout, statusMessage));
}
public Wait createMessage(
String spaceId, String conversationId, String content) {
return createMessage(
new GenieCreateConversationMessageRequest()
.setSpaceId(spaceId)
.setConversationId(conversationId)
.setContent(content));
}
/**
* Create conversation message.
*
* Create new message in [conversation](:method:genie/startconversation). The AI response uses
* all previously created messages in the conversation to respond.
*/
public Wait createMessage(
GenieCreateConversationMessageRequest request) {
GenieMessage response = impl.createMessage(request);
return new Wait<>(
(timeout, callback) ->
waitGetMessageGenieCompleted(
request.getConversationId(),
response.getId(),
request.getSpaceId(),
timeout,
callback),
response);
}
public GenieGetMessageQueryResultResponse executeMessageQuery(
String spaceId, String conversationId, String messageId) {
return executeMessageQuery(
new ExecuteMessageQueryRequest()
.setSpaceId(spaceId)
.setConversationId(conversationId)
.setMessageId(messageId));
}
/**
* Execute SQL query in a conversation message.
*
* Execute the SQL query in the message.
*/
public GenieGetMessageQueryResultResponse executeMessageQuery(
ExecuteMessageQueryRequest request) {
return impl.executeMessageQuery(request);
}
public GenieMessage getMessage(String spaceId, String conversationId, String messageId) {
return getMessage(
new GenieGetConversationMessageRequest()
.setSpaceId(spaceId)
.setConversationId(conversationId)
.setMessageId(messageId));
}
/**
* Get conversation message.
*
*
Get message from conversation.
*/
public GenieMessage getMessage(GenieGetConversationMessageRequest request) {
return impl.getMessage(request);
}
public GenieGetMessageQueryResultResponse getMessageQueryResult(
String spaceId, String conversationId, String messageId) {
return getMessageQueryResult(
new GenieGetMessageQueryResultRequest()
.setSpaceId(spaceId)
.setConversationId(conversationId)
.setMessageId(messageId));
}
/**
* Get conversation message SQL query result.
*
*
Get the result of SQL query if the message has a query attachment. This is only available if
* a message has a query attachment and the message status is `EXECUTING_QUERY`.
*/
public GenieGetMessageQueryResultResponse getMessageQueryResult(
GenieGetMessageQueryResultRequest request) {
return impl.getMessageQueryResult(request);
}
public Wait startConversation(
String spaceId, String content) {
return startConversation(
new GenieStartConversationMessageRequest().setSpaceId(spaceId).setContent(content));
}
/**
* Start conversation.
*
* Start a new conversation.
*/
public Wait startConversation(
GenieStartConversationMessageRequest request) {
GenieStartConversationResponse response = impl.startConversation(request);
return new Wait<>(
(timeout, callback) ->
waitGetMessageGenieCompleted(
response.getConversationId(),
response.getMessageId(),
request.getSpaceId(),
timeout,
callback),
response);
}
public GenieService impl() {
return impl;
}
}