@databricks/appkit 0.19.0 → 0.19.1
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.
package/dist/appkit/package.js
CHANGED
|
@@ -8,6 +8,12 @@ import * as SDK from "@databricks/sdk-experimental";
|
|
|
8
8
|
const { TimeUnits } = SDK;
|
|
9
9
|
const Time = SDK.Time ?? SDK.default.Time;
|
|
10
10
|
const logger = createLogger("connectors:genie");
|
|
11
|
+
const GenieErrors = {
|
|
12
|
+
SPACE_ACCESS_DENIED: "You don't have access to this Genie Space.",
|
|
13
|
+
TABLE_PERMISSIONS: "You may not have access to the data tables. Please verify your table permissions.",
|
|
14
|
+
REQUEST_FAILED: "Genie request failed",
|
|
15
|
+
QUERY_RESULT_FAILED: "Failed to fetch query result"
|
|
16
|
+
};
|
|
11
17
|
function mapAttachments(message) {
|
|
12
18
|
return message.attachments?.map((att) => ({
|
|
13
19
|
attachmentId: att.attachment_id,
|
|
@@ -32,6 +38,12 @@ function toMessageResponse(message) {
|
|
|
32
38
|
error: message.error?.error
|
|
33
39
|
};
|
|
34
40
|
}
|
|
41
|
+
function classifyGenieError(error) {
|
|
42
|
+
const message = error instanceof Error ? error.message : String(error);
|
|
43
|
+
if (message.includes("RESOURCE_DOES_NOT_EXIST")) return GenieErrors.SPACE_ACCESS_DENIED;
|
|
44
|
+
if (message.includes("failed to reach COMPLETED state") && message.includes("FAILED")) return GenieErrors.TABLE_PERMISSIONS;
|
|
45
|
+
return message || GenieErrors.REQUEST_FAILED;
|
|
46
|
+
}
|
|
35
47
|
var GenieConnector = class {
|
|
36
48
|
config;
|
|
37
49
|
constructor(config = {}) {
|
|
@@ -113,10 +125,10 @@ var GenieConnector = class {
|
|
|
113
125
|
};
|
|
114
126
|
yield* this.emitQueryResults(workspaceClient, spaceId, resultConversationId, messageResponse.messageId, messageResponse);
|
|
115
127
|
} catch (error) {
|
|
116
|
-
logger.error("Genie message error: %O", error);
|
|
128
|
+
logger.error("Genie message error (spaceId=%s, conversationId=%s): %O", spaceId, conversationId ?? "new", error);
|
|
117
129
|
yield {
|
|
118
130
|
type: "error",
|
|
119
|
-
error: error
|
|
131
|
+
error: classifyGenieError(error)
|
|
120
132
|
};
|
|
121
133
|
}
|
|
122
134
|
}
|
|
@@ -136,7 +148,7 @@ var GenieConnector = class {
|
|
|
136
148
|
logger.error("Failed to fetch query result for attachment %s: %O", att.attachmentId, error);
|
|
137
149
|
yield {
|
|
138
150
|
type: "error",
|
|
139
|
-
error:
|
|
151
|
+
error: `${GenieErrors.QUERY_RESULT_FAILED} for attachment ${att.attachmentId}`
|
|
140
152
|
};
|
|
141
153
|
}
|
|
142
154
|
}
|
|
@@ -184,15 +196,15 @@ var GenieConnector = class {
|
|
|
184
196
|
logger.error("Failed to fetch query result: %O", result.reason);
|
|
185
197
|
yield {
|
|
186
198
|
type: "error",
|
|
187
|
-
error: result.reason instanceof Error ? result.reason.message :
|
|
199
|
+
error: result.reason instanceof Error ? result.reason.message : GenieErrors.QUERY_RESULT_FAILED
|
|
188
200
|
};
|
|
189
201
|
}
|
|
190
202
|
}
|
|
191
203
|
} catch (error) {
|
|
192
|
-
logger.error("Genie getConversation error: %O", error);
|
|
204
|
+
logger.error("Genie getConversation error (spaceId=%s, conversationId=%s): %O", spaceId, conversationId, error);
|
|
193
205
|
yield {
|
|
194
206
|
type: "error",
|
|
195
|
-
error: error
|
|
207
|
+
error: classifyGenieError(error)
|
|
196
208
|
};
|
|
197
209
|
}
|
|
198
210
|
}
|
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"client.js","names":[],"sources":["../../../src/connectors/genie/client.ts"],"sourcesContent":["import type { WorkspaceClient } from \"@databricks/sdk-experimental\";\nimport * as SDK from \"@databricks/sdk-experimental\";\nimport type { GenieMessage } from \"@databricks/sdk-experimental/dist/apis/dashboards\";\nimport type { Waiter } from \"@databricks/sdk-experimental/dist/wait\";\nimport { createLogger } from \"../../logging\";\nimport { genieConnectorDefaults } from \"./defaults\";\nimport { pollWaiter } from \"./poll-waiter\";\nimport type {\n GenieAttachmentResponse,\n GenieConversationHistoryResponse,\n GenieMessageResponse,\n GenieStatementResponse,\n GenieStreamEvent,\n} from \"./types\";\n\nconst { TimeUnits } = SDK;\nconst Time = SDK.Time ?? (SDK as any).default.Time;\n\nconst logger = createLogger(\"connectors:genie\");\n\ntype CreateMessageWaiter = Waiter<GenieMessage, GenieMessage>;\n\nexport interface GenieConnectorConfig {\n timeout?: number;\n maxMessages?: number;\n}\n\nfunction mapAttachments(message: GenieMessage): GenieAttachmentResponse[] {\n return (\n message.attachments?.map((att) => ({\n attachmentId: att.attachment_id,\n query: att.query\n ? {\n title: att.query.title,\n description: att.query.description,\n query: att.query.query,\n statementId: att.query.statement_id,\n }\n : undefined,\n text: att.text ? { content: att.text.content } : undefined,\n suggestedQuestions: att.suggested_questions?.questions,\n })) ?? []\n );\n}\n\nfunction toMessageResponse(message: GenieMessage): GenieMessageResponse {\n return {\n messageId: message.message_id,\n conversationId: message.conversation_id,\n spaceId: message.space_id,\n status: message.status ?? \"COMPLETED\",\n content: message.content,\n attachments: mapAttachments(message),\n error: message.error?.error,\n };\n}\n\nexport class GenieConnector {\n private readonly config: Required<GenieConnectorConfig>;\n\n constructor(config: GenieConnectorConfig = {}) {\n this.config = {\n timeout: config.timeout ?? genieConnectorDefaults.timeout,\n maxMessages: config.maxMessages ?? genieConnectorDefaults.maxMessages,\n };\n }\n\n async startMessage(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n content: string,\n conversationId: string | undefined,\n ): Promise<{\n messageWaiter: CreateMessageWaiter;\n conversationId: string;\n messageId: string;\n }> {\n if (conversationId) {\n const waiter = await workspaceClient.genie.createMessage({\n space_id: spaceId,\n conversation_id: conversationId,\n content,\n });\n return {\n messageWaiter: waiter,\n conversationId,\n messageId: waiter.message_id ?? \"\",\n };\n }\n const start = await workspaceClient.genie.startConversation({\n space_id: spaceId,\n content,\n });\n return {\n messageWaiter: start as unknown as CreateMessageWaiter,\n conversationId: start.conversation_id,\n messageId: start.message_id,\n };\n }\n\n async waitForMessage(\n messageWaiter: CreateMessageWaiter,\n options?: { timeout?: number },\n ): Promise<GenieMessage> {\n const timeout = options?.timeout ?? this.config.timeout;\n const waitOptions =\n timeout > 0 ? { timeout: new Time(timeout, TimeUnits.milliseconds) } : {};\n return messageWaiter.wait(waitOptions);\n }\n\n async listConversationMessages(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n conversationId: string,\n options?: { pageSize?: number; pageToken?: string },\n ): Promise<{\n messages: GenieMessageResponse[];\n nextPageToken: string | null;\n }> {\n const pageSize =\n options?.pageSize ?? genieConnectorDefaults.initialPageSize;\n\n const response = await workspaceClient.genie.listConversationMessages({\n space_id: spaceId,\n conversation_id: conversationId,\n page_size: pageSize,\n ...(options?.pageToken ? { page_token: options.pageToken } : {}),\n });\n\n const messages = (response.messages ?? []).reverse().map(toMessageResponse);\n\n return {\n messages,\n nextPageToken: response.next_page_token ?? null,\n };\n }\n\n async getMessageAttachmentQueryResult(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n conversationId: string,\n messageId: string,\n attachmentId: string,\n _signal?: AbortSignal,\n ): Promise<GenieStatementResponse> {\n const response =\n await workspaceClient.genie.getMessageAttachmentQueryResult({\n space_id: spaceId,\n conversation_id: conversationId,\n message_id: messageId,\n attachment_id: attachmentId,\n });\n return response.statement_response as GenieStatementResponse;\n }\n\n async *streamSendMessage(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n content: string,\n conversationId: string | undefined,\n options?: { timeout?: number },\n ): AsyncGenerator<GenieStreamEvent> {\n try {\n const {\n messageWaiter,\n conversationId: resultConversationId,\n messageId: resultMessageId,\n } = await this.startMessage(\n workspaceClient,\n spaceId,\n content,\n conversationId,\n );\n\n yield {\n type: \"message_start\",\n conversationId: resultConversationId,\n messageId: resultMessageId,\n spaceId,\n };\n\n const timeout =\n options?.timeout != null ? options.timeout : this.config.timeout;\n const waitOptions =\n timeout > 0\n ? { timeout: new Time(timeout, TimeUnits.milliseconds) }\n : {};\n\n let completedMessage!: GenieMessage;\n for await (const event of pollWaiter(messageWaiter, waitOptions)) {\n if (event.type === \"progress\" && event.value.status) {\n yield { type: \"status\", status: event.value.status };\n } else if (event.type === \"completed\") {\n completedMessage = event.value;\n }\n }\n\n const messageResponse = toMessageResponse(completedMessage);\n yield { type: \"message_result\", message: messageResponse };\n\n yield* this.emitQueryResults(\n workspaceClient,\n spaceId,\n resultConversationId,\n messageResponse.messageId,\n messageResponse,\n );\n } catch (error) {\n logger.error(\"Genie message error: %O\", error);\n yield {\n type: \"error\",\n error: error instanceof Error ? error.message : \"Genie request failed\",\n };\n }\n }\n\n private async *emitQueryResults(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n conversationId: string,\n messageId: string,\n messageResponse: GenieMessageResponse,\n ): AsyncGenerator<\n Extract<GenieStreamEvent, { type: \"query_result\" } | { type: \"error\" }>\n > {\n const attachments = messageResponse.attachments ?? [];\n for (const att of attachments) {\n if (!att.query?.statementId || !att.attachmentId) continue;\n try {\n const data = await this.getMessageAttachmentQueryResult(\n workspaceClient,\n spaceId,\n conversationId,\n messageId,\n att.attachmentId,\n );\n yield {\n type: \"query_result\",\n attachmentId: att.attachmentId,\n statementId: att.query.statementId,\n data,\n };\n } catch (error) {\n logger.error(\n \"Failed to fetch query result for attachment %s: %O\",\n att.attachmentId,\n error,\n );\n yield {\n type: \"error\",\n error: `Failed to fetch query result for attachment ${att.attachmentId}`,\n };\n }\n }\n }\n\n async *streamConversation(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n conversationId: string,\n options?: {\n includeQueryResults?: boolean;\n pageSize?: number;\n pageToken?: string;\n },\n ): AsyncGenerator<GenieStreamEvent> {\n const includeQueryResults = options?.includeQueryResults !== false;\n\n try {\n const { messages: messageResponses, nextPageToken } =\n await this.listConversationMessages(\n workspaceClient,\n spaceId,\n conversationId,\n { pageSize: options?.pageSize, pageToken: options?.pageToken },\n );\n\n for (const messageResponse of messageResponses) {\n yield { type: \"message_result\", message: messageResponse };\n }\n\n yield {\n type: \"history_info\",\n conversationId,\n spaceId,\n nextPageToken,\n loadedCount: messageResponses.length,\n };\n\n if (includeQueryResults) {\n const queryAttachments: Array<{\n messageId: string;\n attachmentId: string;\n statementId: string;\n }> = [];\n\n for (const msg of messageResponses) {\n for (const att of msg.attachments ?? []) {\n if (att.query?.statementId && att.attachmentId) {\n queryAttachments.push({\n messageId: msg.messageId,\n attachmentId: att.attachmentId,\n statementId: att.query.statementId,\n });\n }\n }\n }\n\n const results = await Promise.allSettled(\n queryAttachments.map(async (att) => {\n const data = await this.getMessageAttachmentQueryResult(\n workspaceClient,\n spaceId,\n conversationId,\n att.messageId,\n att.attachmentId,\n );\n return {\n attachmentId: att.attachmentId,\n statementId: att.statementId,\n data,\n };\n }),\n );\n\n for (const result of results) {\n if (result.status === \"fulfilled\") {\n yield {\n type: \"query_result\",\n attachmentId: result.value.attachmentId,\n statementId: result.value.statementId,\n data: result.value.data,\n };\n } else {\n logger.error(\"Failed to fetch query result: %O\", result.reason);\n yield {\n type: \"error\",\n error:\n result.reason instanceof Error\n ? result.reason.message\n : \"Failed to fetch query result\",\n };\n }\n }\n }\n } catch (error) {\n logger.error(\"Genie getConversation error: %O\", error);\n yield {\n type: \"error\",\n error:\n error instanceof Error\n ? error.message\n : \"Failed to fetch conversation\",\n };\n }\n }\n\n async sendMessage(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n content: string,\n conversationId: string | undefined,\n ): Promise<GenieMessageResponse> {\n const { messageWaiter, conversationId: resultConversationId } =\n await this.startMessage(\n workspaceClient,\n spaceId,\n content,\n conversationId,\n );\n const completedMessage = await this.waitForMessage(messageWaiter);\n const messageResponse = toMessageResponse(completedMessage);\n return {\n ...messageResponse,\n conversationId: resultConversationId,\n };\n }\n\n async getConversation(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n conversationId: string,\n ): Promise<GenieConversationHistoryResponse> {\n const allMessages: GenieMessageResponse[] = [];\n let pageToken: string | undefined;\n\n do {\n const { messages, nextPageToken } = await this.listConversationMessages(\n workspaceClient,\n spaceId,\n conversationId,\n {\n pageSize: genieConnectorDefaults.pageSize,\n pageToken,\n },\n );\n allMessages.push(...messages);\n pageToken = nextPageToken ?? undefined;\n } while (pageToken && allMessages.length < this.config.maxMessages);\n\n return {\n conversationId,\n spaceId,\n messages: allMessages.slice(0, this.config.maxMessages),\n };\n }\n}\n"],"mappings":";;;;;;;AAeA,MAAM,EAAE,cAAc;AACtB,MAAM,OAAO,IAAI,QAAS,IAAY,QAAQ;AAE9C,MAAM,SAAS,aAAa,mBAAmB;AAS/C,SAAS,eAAe,SAAkD;AACxE,QACE,QAAQ,aAAa,KAAK,SAAS;EACjC,cAAc,IAAI;EAClB,OAAO,IAAI,QACP;GACE,OAAO,IAAI,MAAM;GACjB,aAAa,IAAI,MAAM;GACvB,OAAO,IAAI,MAAM;GACjB,aAAa,IAAI,MAAM;GACxB,GACD;EACJ,MAAM,IAAI,OAAO,EAAE,SAAS,IAAI,KAAK,SAAS,GAAG;EACjD,oBAAoB,IAAI,qBAAqB;EAC9C,EAAE,IAAI,EAAE;;AAIb,SAAS,kBAAkB,SAA6C;AACtE,QAAO;EACL,WAAW,QAAQ;EACnB,gBAAgB,QAAQ;EACxB,SAAS,QAAQ;EACjB,QAAQ,QAAQ,UAAU;EAC1B,SAAS,QAAQ;EACjB,aAAa,eAAe,QAAQ;EACpC,OAAO,QAAQ,OAAO;EACvB;;AAGH,IAAa,iBAAb,MAA4B;CAC1B,AAAiB;CAEjB,YAAY,SAA+B,EAAE,EAAE;AAC7C,OAAK,SAAS;GACZ,SAAS,OAAO,WAAW,uBAAuB;GAClD,aAAa,OAAO,eAAe,uBAAuB;GAC3D;;CAGH,MAAM,aACJ,iBACA,SACA,SACA,gBAKC;AACD,MAAI,gBAAgB;GAClB,MAAM,SAAS,MAAM,gBAAgB,MAAM,cAAc;IACvD,UAAU;IACV,iBAAiB;IACjB;IACD,CAAC;AACF,UAAO;IACL,eAAe;IACf;IACA,WAAW,OAAO,cAAc;IACjC;;EAEH,MAAM,QAAQ,MAAM,gBAAgB,MAAM,kBAAkB;GAC1D,UAAU;GACV;GACD,CAAC;AACF,SAAO;GACL,eAAe;GACf,gBAAgB,MAAM;GACtB,WAAW,MAAM;GAClB;;CAGH,MAAM,eACJ,eACA,SACuB;EACvB,MAAM,UAAU,SAAS,WAAW,KAAK,OAAO;EAChD,MAAM,cACJ,UAAU,IAAI,EAAE,SAAS,IAAI,KAAK,SAAS,UAAU,aAAa,EAAE,GAAG,EAAE;AAC3E,SAAO,cAAc,KAAK,YAAY;;CAGxC,MAAM,yBACJ,iBACA,SACA,gBACA,SAIC;EACD,MAAM,WACJ,SAAS,YAAY,uBAAuB;EAE9C,MAAM,WAAW,MAAM,gBAAgB,MAAM,yBAAyB;GACpE,UAAU;GACV,iBAAiB;GACjB,WAAW;GACX,GAAI,SAAS,YAAY,EAAE,YAAY,QAAQ,WAAW,GAAG,EAAE;GAChE,CAAC;AAIF,SAAO;GACL,WAHgB,SAAS,YAAY,EAAE,EAAE,SAAS,CAAC,IAAI,kBAAkB;GAIzE,eAAe,SAAS,mBAAmB;GAC5C;;CAGH,MAAM,gCACJ,iBACA,SACA,gBACA,WACA,cACA,SACiC;AAQjC,UANE,MAAM,gBAAgB,MAAM,gCAAgC;GAC1D,UAAU;GACV,iBAAiB;GACjB,YAAY;GACZ,eAAe;GAChB,CAAC,EACY;;CAGlB,OAAO,kBACL,iBACA,SACA,SACA,gBACA,SACkC;AAClC,MAAI;GACF,MAAM,EACJ,eACA,gBAAgB,sBAChB,WAAW,oBACT,MAAM,KAAK,aACb,iBACA,SACA,SACA,eACD;AAED,SAAM;IACJ,MAAM;IACN,gBAAgB;IAChB,WAAW;IACX;IACD;GAED,MAAM,UACJ,SAAS,WAAW,OAAO,QAAQ,UAAU,KAAK,OAAO;GAC3D,MAAM,cACJ,UAAU,IACN,EAAE,SAAS,IAAI,KAAK,SAAS,UAAU,aAAa,EAAE,GACtD,EAAE;GAER,IAAI;AACJ,cAAW,MAAM,SAAS,WAAW,eAAe,YAAY,CAC9D,KAAI,MAAM,SAAS,cAAc,MAAM,MAAM,OAC3C,OAAM;IAAE,MAAM;IAAU,QAAQ,MAAM,MAAM;IAAQ;YAC3C,MAAM,SAAS,YACxB,oBAAmB,MAAM;GAI7B,MAAM,kBAAkB,kBAAkB,iBAAiB;AAC3D,SAAM;IAAE,MAAM;IAAkB,SAAS;IAAiB;AAE1D,UAAO,KAAK,iBACV,iBACA,SACA,sBACA,gBAAgB,WAChB,gBACD;WACM,OAAO;AACd,UAAO,MAAM,2BAA2B,MAAM;AAC9C,SAAM;IACJ,MAAM;IACN,OAAO,iBAAiB,QAAQ,MAAM,UAAU;IACjD;;;CAIL,OAAe,iBACb,iBACA,SACA,gBACA,WACA,iBAGA;EACA,MAAM,cAAc,gBAAgB,eAAe,EAAE;AACrD,OAAK,MAAM,OAAO,aAAa;AAC7B,OAAI,CAAC,IAAI,OAAO,eAAe,CAAC,IAAI,aAAc;AAClD,OAAI;IACF,MAAM,OAAO,MAAM,KAAK,gCACtB,iBACA,SACA,gBACA,WACA,IAAI,aACL;AACD,UAAM;KACJ,MAAM;KACN,cAAc,IAAI;KAClB,aAAa,IAAI,MAAM;KACvB;KACD;YACM,OAAO;AACd,WAAO,MACL,sDACA,IAAI,cACJ,MACD;AACD,UAAM;KACJ,MAAM;KACN,OAAO,+CAA+C,IAAI;KAC3D;;;;CAKP,OAAO,mBACL,iBACA,SACA,gBACA,SAKkC;EAClC,MAAM,sBAAsB,SAAS,wBAAwB;AAE7D,MAAI;GACF,MAAM,EAAE,UAAU,kBAAkB,kBAClC,MAAM,KAAK,yBACT,iBACA,SACA,gBACA;IAAE,UAAU,SAAS;IAAU,WAAW,SAAS;IAAW,CAC/D;AAEH,QAAK,MAAM,mBAAmB,iBAC5B,OAAM;IAAE,MAAM;IAAkB,SAAS;IAAiB;AAG5D,SAAM;IACJ,MAAM;IACN;IACA;IACA;IACA,aAAa,iBAAiB;IAC/B;AAED,OAAI,qBAAqB;IACvB,MAAM,mBAID,EAAE;AAEP,SAAK,MAAM,OAAO,iBAChB,MAAK,MAAM,OAAO,IAAI,eAAe,EAAE,CACrC,KAAI,IAAI,OAAO,eAAe,IAAI,aAChC,kBAAiB,KAAK;KACpB,WAAW,IAAI;KACf,cAAc,IAAI;KAClB,aAAa,IAAI,MAAM;KACxB,CAAC;IAKR,MAAM,UAAU,MAAM,QAAQ,WAC5B,iBAAiB,IAAI,OAAO,QAAQ;KAClC,MAAM,OAAO,MAAM,KAAK,gCACtB,iBACA,SACA,gBACA,IAAI,WACJ,IAAI,aACL;AACD,YAAO;MACL,cAAc,IAAI;MAClB,aAAa,IAAI;MACjB;MACD;MACD,CACH;AAED,SAAK,MAAM,UAAU,QACnB,KAAI,OAAO,WAAW,YACpB,OAAM;KACJ,MAAM;KACN,cAAc,OAAO,MAAM;KAC3B,aAAa,OAAO,MAAM;KAC1B,MAAM,OAAO,MAAM;KACpB;SACI;AACL,YAAO,MAAM,oCAAoC,OAAO,OAAO;AAC/D,WAAM;MACJ,MAAM;MACN,OACE,OAAO,kBAAkB,QACrB,OAAO,OAAO,UACd;MACP;;;WAIA,OAAO;AACd,UAAO,MAAM,mCAAmC,MAAM;AACtD,SAAM;IACJ,MAAM;IACN,OACE,iBAAiB,QACb,MAAM,UACN;IACP;;;CAIL,MAAM,YACJ,iBACA,SACA,SACA,gBAC+B;EAC/B,MAAM,EAAE,eAAe,gBAAgB,yBACrC,MAAM,KAAK,aACT,iBACA,SACA,SACA,eACD;AAGH,SAAO;GACL,GAFsB,kBADC,MAAM,KAAK,eAAe,cAAc,CACN;GAGzD,gBAAgB;GACjB;;CAGH,MAAM,gBACJ,iBACA,SACA,gBAC2C;EAC3C,MAAM,cAAsC,EAAE;EAC9C,IAAI;AAEJ,KAAG;GACD,MAAM,EAAE,UAAU,kBAAkB,MAAM,KAAK,yBAC7C,iBACA,SACA,gBACA;IACE,UAAU,uBAAuB;IACjC;IACD,CACF;AACD,eAAY,KAAK,GAAG,SAAS;AAC7B,eAAY,iBAAiB;WACtB,aAAa,YAAY,SAAS,KAAK,OAAO;AAEvD,SAAO;GACL;GACA;GACA,UAAU,YAAY,MAAM,GAAG,KAAK,OAAO,YAAY;GACxD"}
|
|
1
|
+
{"version":3,"file":"client.js","names":[],"sources":["../../../src/connectors/genie/client.ts"],"sourcesContent":["import type { WorkspaceClient } from \"@databricks/sdk-experimental\";\nimport * as SDK from \"@databricks/sdk-experimental\";\nimport type { GenieMessage } from \"@databricks/sdk-experimental/dist/apis/dashboards\";\nimport type { Waiter } from \"@databricks/sdk-experimental/dist/wait\";\nimport { createLogger } from \"../../logging\";\nimport { genieConnectorDefaults } from \"./defaults\";\nimport { pollWaiter } from \"./poll-waiter\";\nimport type {\n GenieAttachmentResponse,\n GenieConversationHistoryResponse,\n GenieMessageResponse,\n GenieStatementResponse,\n GenieStreamEvent,\n} from \"./types\";\n\nconst { TimeUnits } = SDK;\nconst Time = SDK.Time ?? (SDK as any).default.Time;\n\nconst logger = createLogger(\"connectors:genie\");\n\nconst GenieErrors = {\n SPACE_ACCESS_DENIED: \"You don't have access to this Genie Space.\",\n TABLE_PERMISSIONS:\n \"You may not have access to the data tables. Please verify your table permissions.\",\n REQUEST_FAILED: \"Genie request failed\",\n QUERY_RESULT_FAILED: \"Failed to fetch query result\",\n} as const;\n\ntype CreateMessageWaiter = Waiter<GenieMessage, GenieMessage>;\n\nexport interface GenieConnectorConfig {\n timeout?: number;\n maxMessages?: number;\n}\n\nfunction mapAttachments(message: GenieMessage): GenieAttachmentResponse[] {\n return (\n message.attachments?.map((att) => ({\n attachmentId: att.attachment_id,\n query: att.query\n ? {\n title: att.query.title,\n description: att.query.description,\n query: att.query.query,\n statementId: att.query.statement_id,\n }\n : undefined,\n text: att.text ? { content: att.text.content } : undefined,\n suggestedQuestions: att.suggested_questions?.questions,\n })) ?? []\n );\n}\n\nfunction toMessageResponse(message: GenieMessage): GenieMessageResponse {\n return {\n messageId: message.message_id,\n conversationId: message.conversation_id,\n spaceId: message.space_id,\n status: message.status ?? \"COMPLETED\",\n content: message.content,\n attachments: mapAttachments(message),\n error: message.error?.error,\n };\n}\n\nfunction classifyGenieError(error: unknown): string {\n const message = error instanceof Error ? error.message : String(error);\n\n if (message.includes(\"RESOURCE_DOES_NOT_EXIST\")) {\n return GenieErrors.SPACE_ACCESS_DENIED;\n }\n\n if (\n message.includes(\"failed to reach COMPLETED state\") &&\n message.includes(\"FAILED\")\n ) {\n return GenieErrors.TABLE_PERMISSIONS;\n }\n\n return message || GenieErrors.REQUEST_FAILED;\n}\n\nexport class GenieConnector {\n private readonly config: Required<GenieConnectorConfig>;\n\n constructor(config: GenieConnectorConfig = {}) {\n this.config = {\n timeout: config.timeout ?? genieConnectorDefaults.timeout,\n maxMessages: config.maxMessages ?? genieConnectorDefaults.maxMessages,\n };\n }\n\n async startMessage(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n content: string,\n conversationId: string | undefined,\n ): Promise<{\n messageWaiter: CreateMessageWaiter;\n conversationId: string;\n messageId: string;\n }> {\n if (conversationId) {\n const waiter = await workspaceClient.genie.createMessage({\n space_id: spaceId,\n conversation_id: conversationId,\n content,\n });\n return {\n messageWaiter: waiter,\n conversationId,\n messageId: waiter.message_id ?? \"\",\n };\n }\n const start = await workspaceClient.genie.startConversation({\n space_id: spaceId,\n content,\n });\n return {\n messageWaiter: start as unknown as CreateMessageWaiter,\n conversationId: start.conversation_id,\n messageId: start.message_id,\n };\n }\n\n async waitForMessage(\n messageWaiter: CreateMessageWaiter,\n options?: { timeout?: number },\n ): Promise<GenieMessage> {\n const timeout = options?.timeout ?? this.config.timeout;\n const waitOptions =\n timeout > 0 ? { timeout: new Time(timeout, TimeUnits.milliseconds) } : {};\n return messageWaiter.wait(waitOptions);\n }\n\n async listConversationMessages(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n conversationId: string,\n options?: { pageSize?: number; pageToken?: string },\n ): Promise<{\n messages: GenieMessageResponse[];\n nextPageToken: string | null;\n }> {\n const pageSize =\n options?.pageSize ?? genieConnectorDefaults.initialPageSize;\n\n const response = await workspaceClient.genie.listConversationMessages({\n space_id: spaceId,\n conversation_id: conversationId,\n page_size: pageSize,\n ...(options?.pageToken ? { page_token: options.pageToken } : {}),\n });\n\n const messages = (response.messages ?? []).reverse().map(toMessageResponse);\n\n return {\n messages,\n nextPageToken: response.next_page_token ?? null,\n };\n }\n\n async getMessageAttachmentQueryResult(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n conversationId: string,\n messageId: string,\n attachmentId: string,\n _signal?: AbortSignal,\n ): Promise<GenieStatementResponse> {\n const response =\n await workspaceClient.genie.getMessageAttachmentQueryResult({\n space_id: spaceId,\n conversation_id: conversationId,\n message_id: messageId,\n attachment_id: attachmentId,\n });\n return response.statement_response as GenieStatementResponse;\n }\n\n async *streamSendMessage(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n content: string,\n conversationId: string | undefined,\n options?: { timeout?: number },\n ): AsyncGenerator<GenieStreamEvent> {\n try {\n const {\n messageWaiter,\n conversationId: resultConversationId,\n messageId: resultMessageId,\n } = await this.startMessage(\n workspaceClient,\n spaceId,\n content,\n conversationId,\n );\n\n yield {\n type: \"message_start\",\n conversationId: resultConversationId,\n messageId: resultMessageId,\n spaceId,\n };\n\n const timeout =\n options?.timeout != null ? options.timeout : this.config.timeout;\n const waitOptions =\n timeout > 0\n ? { timeout: new Time(timeout, TimeUnits.milliseconds) }\n : {};\n\n let completedMessage!: GenieMessage;\n for await (const event of pollWaiter(messageWaiter, waitOptions)) {\n if (event.type === \"progress\" && event.value.status) {\n yield { type: \"status\", status: event.value.status };\n } else if (event.type === \"completed\") {\n completedMessage = event.value;\n }\n }\n\n const messageResponse = toMessageResponse(completedMessage);\n yield { type: \"message_result\", message: messageResponse };\n\n yield* this.emitQueryResults(\n workspaceClient,\n spaceId,\n resultConversationId,\n messageResponse.messageId,\n messageResponse,\n );\n } catch (error) {\n logger.error(\n \"Genie message error (spaceId=%s, conversationId=%s): %O\",\n spaceId,\n conversationId ?? \"new\",\n error,\n );\n yield { type: \"error\", error: classifyGenieError(error) };\n }\n }\n\n private async *emitQueryResults(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n conversationId: string,\n messageId: string,\n messageResponse: GenieMessageResponse,\n ): AsyncGenerator<\n Extract<GenieStreamEvent, { type: \"query_result\" } | { type: \"error\" }>\n > {\n const attachments = messageResponse.attachments ?? [];\n for (const att of attachments) {\n if (!att.query?.statementId || !att.attachmentId) continue;\n try {\n const data = await this.getMessageAttachmentQueryResult(\n workspaceClient,\n spaceId,\n conversationId,\n messageId,\n att.attachmentId,\n );\n yield {\n type: \"query_result\",\n attachmentId: att.attachmentId,\n statementId: att.query.statementId,\n data,\n };\n } catch (error) {\n logger.error(\n \"Failed to fetch query result for attachment %s: %O\",\n att.attachmentId,\n error,\n );\n yield {\n type: \"error\",\n error: `${GenieErrors.QUERY_RESULT_FAILED} for attachment ${att.attachmentId}`,\n };\n }\n }\n }\n\n async *streamConversation(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n conversationId: string,\n options?: {\n includeQueryResults?: boolean;\n pageSize?: number;\n pageToken?: string;\n },\n ): AsyncGenerator<GenieStreamEvent> {\n const includeQueryResults = options?.includeQueryResults !== false;\n\n try {\n const { messages: messageResponses, nextPageToken } =\n await this.listConversationMessages(\n workspaceClient,\n spaceId,\n conversationId,\n { pageSize: options?.pageSize, pageToken: options?.pageToken },\n );\n\n for (const messageResponse of messageResponses) {\n yield { type: \"message_result\", message: messageResponse };\n }\n\n yield {\n type: \"history_info\",\n conversationId,\n spaceId,\n nextPageToken,\n loadedCount: messageResponses.length,\n };\n\n if (includeQueryResults) {\n const queryAttachments: Array<{\n messageId: string;\n attachmentId: string;\n statementId: string;\n }> = [];\n\n for (const msg of messageResponses) {\n for (const att of msg.attachments ?? []) {\n if (att.query?.statementId && att.attachmentId) {\n queryAttachments.push({\n messageId: msg.messageId,\n attachmentId: att.attachmentId,\n statementId: att.query.statementId,\n });\n }\n }\n }\n\n const results = await Promise.allSettled(\n queryAttachments.map(async (att) => {\n const data = await this.getMessageAttachmentQueryResult(\n workspaceClient,\n spaceId,\n conversationId,\n att.messageId,\n att.attachmentId,\n );\n return {\n attachmentId: att.attachmentId,\n statementId: att.statementId,\n data,\n };\n }),\n );\n\n for (const result of results) {\n if (result.status === \"fulfilled\") {\n yield {\n type: \"query_result\",\n attachmentId: result.value.attachmentId,\n statementId: result.value.statementId,\n data: result.value.data,\n };\n } else {\n logger.error(\"Failed to fetch query result: %O\", result.reason);\n yield {\n type: \"error\",\n error:\n result.reason instanceof Error\n ? result.reason.message\n : GenieErrors.QUERY_RESULT_FAILED,\n };\n }\n }\n }\n } catch (error) {\n logger.error(\n \"Genie getConversation error (spaceId=%s, conversationId=%s): %O\",\n spaceId,\n conversationId,\n error,\n );\n yield { type: \"error\", error: classifyGenieError(error) };\n }\n }\n\n async sendMessage(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n content: string,\n conversationId: string | undefined,\n ): Promise<GenieMessageResponse> {\n const { messageWaiter, conversationId: resultConversationId } =\n await this.startMessage(\n workspaceClient,\n spaceId,\n content,\n conversationId,\n );\n const completedMessage = await this.waitForMessage(messageWaiter);\n const messageResponse = toMessageResponse(completedMessage);\n return {\n ...messageResponse,\n conversationId: resultConversationId,\n };\n }\n\n async getConversation(\n workspaceClient: WorkspaceClient,\n spaceId: string,\n conversationId: string,\n ): Promise<GenieConversationHistoryResponse> {\n const allMessages: GenieMessageResponse[] = [];\n let pageToken: string | undefined;\n\n do {\n const { messages, nextPageToken } = await this.listConversationMessages(\n workspaceClient,\n spaceId,\n conversationId,\n {\n pageSize: genieConnectorDefaults.pageSize,\n pageToken,\n },\n );\n allMessages.push(...messages);\n pageToken = nextPageToken ?? undefined;\n } while (pageToken && allMessages.length < this.config.maxMessages);\n\n return {\n conversationId,\n spaceId,\n messages: allMessages.slice(0, this.config.maxMessages),\n };\n }\n}\n"],"mappings":";;;;;;;AAeA,MAAM,EAAE,cAAc;AACtB,MAAM,OAAO,IAAI,QAAS,IAAY,QAAQ;AAE9C,MAAM,SAAS,aAAa,mBAAmB;AAE/C,MAAM,cAAc;CAClB,qBAAqB;CACrB,mBACE;CACF,gBAAgB;CAChB,qBAAqB;CACtB;AASD,SAAS,eAAe,SAAkD;AACxE,QACE,QAAQ,aAAa,KAAK,SAAS;EACjC,cAAc,IAAI;EAClB,OAAO,IAAI,QACP;GACE,OAAO,IAAI,MAAM;GACjB,aAAa,IAAI,MAAM;GACvB,OAAO,IAAI,MAAM;GACjB,aAAa,IAAI,MAAM;GACxB,GACD;EACJ,MAAM,IAAI,OAAO,EAAE,SAAS,IAAI,KAAK,SAAS,GAAG;EACjD,oBAAoB,IAAI,qBAAqB;EAC9C,EAAE,IAAI,EAAE;;AAIb,SAAS,kBAAkB,SAA6C;AACtE,QAAO;EACL,WAAW,QAAQ;EACnB,gBAAgB,QAAQ;EACxB,SAAS,QAAQ;EACjB,QAAQ,QAAQ,UAAU;EAC1B,SAAS,QAAQ;EACjB,aAAa,eAAe,QAAQ;EACpC,OAAO,QAAQ,OAAO;EACvB;;AAGH,SAAS,mBAAmB,OAAwB;CAClD,MAAM,UAAU,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM;AAEtE,KAAI,QAAQ,SAAS,0BAA0B,CAC7C,QAAO,YAAY;AAGrB,KACE,QAAQ,SAAS,kCAAkC,IACnD,QAAQ,SAAS,SAAS,CAE1B,QAAO,YAAY;AAGrB,QAAO,WAAW,YAAY;;AAGhC,IAAa,iBAAb,MAA4B;CAC1B,AAAiB;CAEjB,YAAY,SAA+B,EAAE,EAAE;AAC7C,OAAK,SAAS;GACZ,SAAS,OAAO,WAAW,uBAAuB;GAClD,aAAa,OAAO,eAAe,uBAAuB;GAC3D;;CAGH,MAAM,aACJ,iBACA,SACA,SACA,gBAKC;AACD,MAAI,gBAAgB;GAClB,MAAM,SAAS,MAAM,gBAAgB,MAAM,cAAc;IACvD,UAAU;IACV,iBAAiB;IACjB;IACD,CAAC;AACF,UAAO;IACL,eAAe;IACf;IACA,WAAW,OAAO,cAAc;IACjC;;EAEH,MAAM,QAAQ,MAAM,gBAAgB,MAAM,kBAAkB;GAC1D,UAAU;GACV;GACD,CAAC;AACF,SAAO;GACL,eAAe;GACf,gBAAgB,MAAM;GACtB,WAAW,MAAM;GAClB;;CAGH,MAAM,eACJ,eACA,SACuB;EACvB,MAAM,UAAU,SAAS,WAAW,KAAK,OAAO;EAChD,MAAM,cACJ,UAAU,IAAI,EAAE,SAAS,IAAI,KAAK,SAAS,UAAU,aAAa,EAAE,GAAG,EAAE;AAC3E,SAAO,cAAc,KAAK,YAAY;;CAGxC,MAAM,yBACJ,iBACA,SACA,gBACA,SAIC;EACD,MAAM,WACJ,SAAS,YAAY,uBAAuB;EAE9C,MAAM,WAAW,MAAM,gBAAgB,MAAM,yBAAyB;GACpE,UAAU;GACV,iBAAiB;GACjB,WAAW;GACX,GAAI,SAAS,YAAY,EAAE,YAAY,QAAQ,WAAW,GAAG,EAAE;GAChE,CAAC;AAIF,SAAO;GACL,WAHgB,SAAS,YAAY,EAAE,EAAE,SAAS,CAAC,IAAI,kBAAkB;GAIzE,eAAe,SAAS,mBAAmB;GAC5C;;CAGH,MAAM,gCACJ,iBACA,SACA,gBACA,WACA,cACA,SACiC;AAQjC,UANE,MAAM,gBAAgB,MAAM,gCAAgC;GAC1D,UAAU;GACV,iBAAiB;GACjB,YAAY;GACZ,eAAe;GAChB,CAAC,EACY;;CAGlB,OAAO,kBACL,iBACA,SACA,SACA,gBACA,SACkC;AAClC,MAAI;GACF,MAAM,EACJ,eACA,gBAAgB,sBAChB,WAAW,oBACT,MAAM,KAAK,aACb,iBACA,SACA,SACA,eACD;AAED,SAAM;IACJ,MAAM;IACN,gBAAgB;IAChB,WAAW;IACX;IACD;GAED,MAAM,UACJ,SAAS,WAAW,OAAO,QAAQ,UAAU,KAAK,OAAO;GAC3D,MAAM,cACJ,UAAU,IACN,EAAE,SAAS,IAAI,KAAK,SAAS,UAAU,aAAa,EAAE,GACtD,EAAE;GAER,IAAI;AACJ,cAAW,MAAM,SAAS,WAAW,eAAe,YAAY,CAC9D,KAAI,MAAM,SAAS,cAAc,MAAM,MAAM,OAC3C,OAAM;IAAE,MAAM;IAAU,QAAQ,MAAM,MAAM;IAAQ;YAC3C,MAAM,SAAS,YACxB,oBAAmB,MAAM;GAI7B,MAAM,kBAAkB,kBAAkB,iBAAiB;AAC3D,SAAM;IAAE,MAAM;IAAkB,SAAS;IAAiB;AAE1D,UAAO,KAAK,iBACV,iBACA,SACA,sBACA,gBAAgB,WAChB,gBACD;WACM,OAAO;AACd,UAAO,MACL,2DACA,SACA,kBAAkB,OAClB,MACD;AACD,SAAM;IAAE,MAAM;IAAS,OAAO,mBAAmB,MAAM;IAAE;;;CAI7D,OAAe,iBACb,iBACA,SACA,gBACA,WACA,iBAGA;EACA,MAAM,cAAc,gBAAgB,eAAe,EAAE;AACrD,OAAK,MAAM,OAAO,aAAa;AAC7B,OAAI,CAAC,IAAI,OAAO,eAAe,CAAC,IAAI,aAAc;AAClD,OAAI;IACF,MAAM,OAAO,MAAM,KAAK,gCACtB,iBACA,SACA,gBACA,WACA,IAAI,aACL;AACD,UAAM;KACJ,MAAM;KACN,cAAc,IAAI;KAClB,aAAa,IAAI,MAAM;KACvB;KACD;YACM,OAAO;AACd,WAAO,MACL,sDACA,IAAI,cACJ,MACD;AACD,UAAM;KACJ,MAAM;KACN,OAAO,GAAG,YAAY,oBAAoB,kBAAkB,IAAI;KACjE;;;;CAKP,OAAO,mBACL,iBACA,SACA,gBACA,SAKkC;EAClC,MAAM,sBAAsB,SAAS,wBAAwB;AAE7D,MAAI;GACF,MAAM,EAAE,UAAU,kBAAkB,kBAClC,MAAM,KAAK,yBACT,iBACA,SACA,gBACA;IAAE,UAAU,SAAS;IAAU,WAAW,SAAS;IAAW,CAC/D;AAEH,QAAK,MAAM,mBAAmB,iBAC5B,OAAM;IAAE,MAAM;IAAkB,SAAS;IAAiB;AAG5D,SAAM;IACJ,MAAM;IACN;IACA;IACA;IACA,aAAa,iBAAiB;IAC/B;AAED,OAAI,qBAAqB;IACvB,MAAM,mBAID,EAAE;AAEP,SAAK,MAAM,OAAO,iBAChB,MAAK,MAAM,OAAO,IAAI,eAAe,EAAE,CACrC,KAAI,IAAI,OAAO,eAAe,IAAI,aAChC,kBAAiB,KAAK;KACpB,WAAW,IAAI;KACf,cAAc,IAAI;KAClB,aAAa,IAAI,MAAM;KACxB,CAAC;IAKR,MAAM,UAAU,MAAM,QAAQ,WAC5B,iBAAiB,IAAI,OAAO,QAAQ;KAClC,MAAM,OAAO,MAAM,KAAK,gCACtB,iBACA,SACA,gBACA,IAAI,WACJ,IAAI,aACL;AACD,YAAO;MACL,cAAc,IAAI;MAClB,aAAa,IAAI;MACjB;MACD;MACD,CACH;AAED,SAAK,MAAM,UAAU,QACnB,KAAI,OAAO,WAAW,YACpB,OAAM;KACJ,MAAM;KACN,cAAc,OAAO,MAAM;KAC3B,aAAa,OAAO,MAAM;KAC1B,MAAM,OAAO,MAAM;KACpB;SACI;AACL,YAAO,MAAM,oCAAoC,OAAO,OAAO;AAC/D,WAAM;MACJ,MAAM;MACN,OACE,OAAO,kBAAkB,QACrB,OAAO,OAAO,UACd,YAAY;MACnB;;;WAIA,OAAO;AACd,UAAO,MACL,mEACA,SACA,gBACA,MACD;AACD,SAAM;IAAE,MAAM;IAAS,OAAO,mBAAmB,MAAM;IAAE;;;CAI7D,MAAM,YACJ,iBACA,SACA,SACA,gBAC+B;EAC/B,MAAM,EAAE,eAAe,gBAAgB,yBACrC,MAAM,KAAK,aACT,iBACA,SACA,SACA,eACD;AAGH,SAAO;GACL,GAFsB,kBADC,MAAM,KAAK,eAAe,cAAc,CACN;GAGzD,gBAAgB;GACjB;;CAGH,MAAM,gBACJ,iBACA,SACA,gBAC2C;EAC3C,MAAM,cAAsC,EAAE;EAC9C,IAAI;AAEJ,KAAG;GACD,MAAM,EAAE,UAAU,kBAAkB,MAAM,KAAK,yBAC7C,iBACA,SACA,gBACA;IACE,UAAU,uBAAuB;IACjC;IACD,CACF;AACD,eAAY,KAAK,GAAG,SAAS;AAC7B,eAAY,iBAAiB;WACtB,aAAa,YAAY,SAAS,KAAK,OAAO;AAEvD,SAAO;GACL;GACA;GACA,UAAU,YAAY,MAAM,GAAG,KAAK,OAAO,YAAY;GACxD"}
|