Skip to content

Commit efcd0a8

Browse files
waleedlatif1waleedlatif1
committed
fix(databricks): fail Genie asks on nested errors and document per-tool outputs
1 parent 3c233a7 commit efcd0a8

15 files changed

Lines changed: 434 additions & 591 deletions

‎apps/docs/content/docs/integrations/databricks.mdx‎

Lines changed: 114 additions & 296 deletions
Large diffs are not rendered by default.

‎apps/sim/tools/databricks/genie_agent_ask.ts‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
1-
import type {
2-
DatabricksGenieAgentAskParams,
3-
DatabricksGenieAgentAskResponse,
1+
import {
2+
type DatabricksGenieAgentAskParams,
3+
type DatabricksGenieAgentAskResponse,
4+
GENIE_AGENT_ITEM_OUTPUT_PROPERTIES,
45
} from '@/tools/databricks/types'
56
import {
67
databricksUrl,
7-
GENIE_AGENT_ITEM_OUTPUT_PROPERTIES,
88
GENIE_SPACE_PARAMS,
99
genieAgentErrorMessage,
1010
genieAgentPath,
@@ -137,7 +137,7 @@ export const genieAgentAskTool: ToolConfig<
137137
report: {
138138
type: 'string',
139139
description: "The agent's final report, with citation links",
140-
optional: true,
140+
nullable: true,
141141
},
142142
queries: {
143143
type: 'array',
@@ -146,8 +146,8 @@ export const genieAgentAskTool: ToolConfig<
146146
type: 'object',
147147
properties: {
148148
callId: { type: 'string', description: 'Function call ID' },
149-
title: { type: 'string', description: 'Query title', optional: true },
150-
sql: { type: 'string', description: 'SQL the agent ran', optional: true },
149+
title: { type: 'string', description: 'Query title', nullable: true },
150+
sql: { type: 'string', description: 'SQL the agent ran', nullable: true },
151151
},
152152
},
153153
},
@@ -162,12 +162,12 @@ export const genieAgentAskTool: ToolConfig<
162162
createdAt: {
163163
type: 'number',
164164
description: 'When the response was created (Unix epoch seconds)',
165-
optional: true,
165+
nullable: true,
166166
},
167167
error: {
168168
type: 'string',
169169
description: 'System error message the agent reported alongside its report',
170-
optional: true,
170+
nullable: true,
171171
},
172172
},
173173
}

‎apps/sim/tools/databricks/genie_agent_list_items.ts‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
1-
import type {
2-
DatabricksGenieAgentListItemsParams,
3-
DatabricksGenieAgentListItemsResponse,
1+
import {
2+
type DatabricksGenieAgentListItemsParams,
3+
type DatabricksGenieAgentListItemsResponse,
4+
GENIE_AGENT_ITEM_OUTPUT_PROPERTIES,
45
} from '@/tools/databricks/types'
56
import {
67
databricksErrorMessage,
78
databricksUrl,
8-
GENIE_AGENT_ITEM_OUTPUT_PROPERTIES,
99
GENIE_READ_RETRY,
1010
GENIE_SPACE_PARAMS,
1111
genieAgentPath,
@@ -106,17 +106,17 @@ export const genieAgentListItemsTool: ToolConfig<
106106
properties: GENIE_AGENT_ITEM_OUTPUT_PROPERTIES,
107107
},
108108
},
109-
firstId: { type: 'string', description: 'ID of the first item in the page', optional: true },
109+
firstId: { type: 'string', description: 'ID of the first item in the page', nullable: true },
110110
lastId: {
111111
type: 'string',
112112
description: 'ID of the last item in the page; pass it as After for the next page',
113-
optional: true,
113+
nullable: true,
114114
},
115115
hasMore: { type: 'boolean', description: 'Whether more items follow this page' },
116116
status: {
117117
type: 'string',
118118
description: 'Status of the latest response in the conversation',
119-
optional: true,
119+
nullable: true,
120120
},
121121
},
122122
}

‎apps/sim/tools/databricks/genie_ask.ts‎

Lines changed: 70 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,23 @@
11
import { createLogger } from '@sim/logger'
2+
import { getErrorMessage } from '@sim/utils/errors'
23
import { sleep } from '@sim/utils/helpers'
34
import { DEFAULT_EXECUTION_TIMEOUT_MS } from '@/lib/core/execution-limits'
4-
import type {
5-
DatabricksGenieAskParams,
6-
DatabricksGenieAskResponse,
7-
DatabricksGenieGetMessageResponse,
8-
DatabricksGenieQueryResultResponse,
5+
import {
6+
type DatabricksGenieAskParams,
7+
type DatabricksGenieAskResponse,
8+
type DatabricksGenieGetMessageResponse,
9+
type DatabricksGenieQueryResultResponse,
10+
GENIE_MESSAGE_OUTPUT_PROPERTIES,
11+
STATEMENT_RESULT_OUTPUT_PROPERTIES,
912
} from '@/tools/databricks/types'
1013
import {
1114
databricksErrorMessage,
1215
databricksUrl,
13-
GENIE_MESSAGE_OUTPUT_PROPERTIES,
1416
GENIE_SPACE_PARAMS,
1517
GENIE_TERMINAL_STATUSES,
1618
genieConversationPath,
1719
genieSpacePath,
1820
mapGenieMessage,
19-
STATEMENT_RESULT_OUTPUT_PROPERTIES,
2021
} from '@/tools/databricks/utils'
2122
import type { ToolConfig } from '@/tools/types'
2223

@@ -129,75 +130,93 @@ export const genieAskTool: ToolConfig<DatabricksGenieAskParams, DatabricksGenieA
129130
}
130131

131132
let message = result.output
132-
const startedAt = Date.now()
133-
let attempt = 0
134133

135-
while (!GENIE_TERMINAL_STATUSES.has(message.status)) {
136-
if (Date.now() - startedAt >= MAX_POLL_TIME_MS) {
137-
logger.warn(`Genie message ${messageId} did not complete within ${MAX_POLL_TIME_MS} ms`)
134+
/**
135+
* A nested call can throw instead of returning a failed result, and the executor falls back
136+
* to the pending first response when post-processing throws. Convert every failure here so
137+
* the ask fails with its conversation and message IDs intact.
138+
*/
139+
try {
140+
const startedAt = Date.now()
141+
let attempt = 0
142+
143+
while (!GENIE_TERMINAL_STATUSES.has(message.status)) {
144+
if (Date.now() - startedAt >= MAX_POLL_TIME_MS) {
145+
logger.warn(`Genie message ${messageId} did not complete within ${MAX_POLL_TIME_MS} ms`)
146+
return {
147+
...result,
148+
output: message,
149+
success: false,
150+
error: `Genie did not answer within ${MAX_POLL_TIME_MS / 1000}s (last status: ${message.status}). Use Get Genie Message with this conversation and message ID to check on it.`,
151+
}
152+
}
153+
154+
attempt += 1
155+
const intervalMs = Math.min(1000 * attempt, MAX_POLL_INTERVAL_MS)
156+
await sleep(intervalMs)
157+
158+
const polled = (await executeTool(
159+
'databricks_genie_get_message',
160+
messageParams
161+
)) as DatabricksGenieGetMessageResponse
162+
if (!polled.success) {
163+
return {
164+
...result,
165+
output: message,
166+
success: false,
167+
error: polled.error ?? 'Failed to poll Genie message',
168+
}
169+
}
170+
message = { ...message, ...polled.output }
171+
}
172+
173+
if (message.status === 'FAILED' || message.status === 'CANCELLED') {
138174
return {
139175
...result,
140176
output: message,
141177
success: false,
142-
error: `Genie did not answer within ${MAX_POLL_TIME_MS / 1000}s (last status: ${message.status}). Use Get Genie Message with this conversation and message ID to check on it.`,
178+
error: message.error ?? `Genie message ${message.status.toLowerCase()}`,
143179
}
144180
}
145181

146-
attempt += 1
147-
const intervalMs = Math.min(1000 * attempt, MAX_POLL_INTERVAL_MS)
148-
await sleep(intervalMs)
182+
if (message.status !== 'COMPLETED' || !message.queryAttachmentId) {
183+
return { ...result, output: message }
184+
}
149185

150-
const polled = (await executeTool(
151-
'databricks_genie_get_message',
152-
messageParams
153-
)) as DatabricksGenieGetMessageResponse
154-
if (!polled.success) {
186+
const queryResult = (await executeTool('databricks_genie_get_query_result', {
187+
...messageParams,
188+
attachmentId: message.queryAttachmentId,
189+
})) as DatabricksGenieQueryResultResponse
190+
if (!queryResult.success) {
155191
return {
156192
...result,
157193
output: message,
158194
success: false,
159-
error: polled.error ?? 'Failed to poll Genie message',
195+
error: queryResult.error ?? 'Failed to get Genie query result',
160196
}
161197
}
162-
message = { ...message, ...polled.output }
163-
}
164198

165-
if (message.status === 'FAILED' || message.status === 'CANCELLED') {
166199
return {
167200
...result,
168-
output: message,
169-
success: false,
170-
error: message.error ?? `Genie message ${message.status.toLowerCase()}`,
201+
output: {
202+
...message,
203+
columns: queryResult.output.columns,
204+
data: queryResult.output.data,
205+
totalRows: queryResult.output.totalRows,
206+
truncated: queryResult.output.truncated,
207+
},
171208
}
172-
}
173-
174-
if (message.status !== 'COMPLETED' || !message.queryAttachmentId) {
175-
return { ...result, output: message }
176-
}
177-
178-
const queryResult = (await executeTool('databricks_genie_get_query_result', {
179-
...messageParams,
180-
attachmentId: message.queryAttachmentId,
181-
})) as DatabricksGenieQueryResultResponse
182-
if (!queryResult.success) {
209+
} catch (error) {
210+
logger.error(`Error waiting for Genie message ${messageId}`, {
211+
message: getErrorMessage(error, 'Unknown error'),
212+
})
183213
return {
184214
...result,
185215
output: message,
186216
success: false,
187-
error: queryResult.error ?? 'Failed to get Genie query result',
217+
error: `Error waiting for Genie to answer: ${getErrorMessage(error, 'Unknown error')}`,
188218
}
189219
}
190-
191-
return {
192-
...result,
193-
output: {
194-
...message,
195-
columns: queryResult.output.columns,
196-
data: queryResult.output.data,
197-
totalRows: queryResult.output.totalRows,
198-
truncated: queryResult.output.truncated,
199-
},
200-
}
201220
},
202221

203222
outputs: {

‎apps/sim/tools/databricks/genie_download_visualization.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,12 @@ import type {
22
DatabricksGenieAttachmentParams,
33
DatabricksGenieDownloadVisualizationResponse,
44
} from '@/tools/databricks/types'
5-
import { databricksUrl, GENIE_MESSAGE_PARAMS, genieAttachmentPath } from '@/tools/databricks/utils'
5+
import {
6+
databricksUrl,
7+
GENIE_MESSAGE_PARAMS,
8+
GENIE_READ_RETRY,
9+
genieAttachmentPath,
10+
} from '@/tools/databricks/utils'
611
import type { ToolConfig } from '@/tools/types'
712

813
export const genieDownloadVisualizationTool: ToolConfig<
@@ -34,6 +39,7 @@ export const genieDownloadVisualizationTool: ToolConfig<
3439
Authorization: `Bearer ${params.apiKey}`,
3540
}),
3641
responseType: 'binary',
42+
retry: GENIE_READ_RETRY,
3743
},
3844

3945
transformResponse: async (response: Response, params) => {

‎apps/sim/tools/databricks/genie_execute_query.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
1-
import type {
2-
DatabricksGenieAttachmentParams,
3-
DatabricksGenieQueryResultResponse,
1+
import {
2+
type DatabricksGenieAttachmentParams,
3+
type DatabricksGenieQueryResultResponse,
4+
GENIE_QUERY_RESULT_OUTPUTS,
45
} from '@/tools/databricks/types'
56
import {
67
databricksErrorMessage,
78
databricksUrl,
89
GENIE_QUERY_ATTACHMENT_PARAMS,
9-
GENIE_QUERY_RESULT_OUTPUTS,
1010
genieAttachmentPath,
1111
mapStatementResult,
1212
} from '@/tools/databricks/utils'

‎apps/sim/tools/databricks/genie_get_message.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
1-
import type {
2-
DatabricksGenieGetMessageResponse,
3-
DatabricksGenieMessageParams,
1+
import {
2+
type DatabricksGenieGetMessageResponse,
3+
type DatabricksGenieMessageParams,
4+
GENIE_MESSAGE_OUTPUT_PROPERTIES,
45
} from '@/tools/databricks/types'
56
import {
67
databricksErrorMessage,
78
databricksUrl,
8-
GENIE_MESSAGE_OUTPUT_PROPERTIES,
99
GENIE_MESSAGE_PARAMS,
1010
GENIE_READ_RETRY,
1111
genieMessagePath,

‎apps/sim/tools/databricks/genie_get_query_result.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
1-
import type {
2-
DatabricksGenieAttachmentParams,
3-
DatabricksGenieQueryResultResponse,
1+
import {
2+
type DatabricksGenieAttachmentParams,
3+
type DatabricksGenieQueryResultResponse,
4+
GENIE_QUERY_RESULT_OUTPUTS,
45
} from '@/tools/databricks/types'
56
import {
67
databricksErrorMessage,
78
databricksUrl,
89
GENIE_QUERY_ATTACHMENT_PARAMS,
9-
GENIE_QUERY_RESULT_OUTPUTS,
1010
GENIE_READ_RETRY,
1111
genieAttachmentPath,
1212
mapStatementResult,

‎apps/sim/tools/databricks/genie_get_space.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
1-
import type {
2-
DatabricksGenieGetSpaceParams,
3-
DatabricksGenieGetSpaceResponse,
1+
import {
2+
type DatabricksGenieGetSpaceParams,
3+
type DatabricksGenieGetSpaceResponse,
4+
GENIE_SPACE_OUTPUT_PROPERTIES,
45
} from '@/tools/databricks/types'
56
import {
67
databricksErrorMessage,
78
databricksUrl,
89
GENIE_READ_RETRY,
9-
GENIE_SPACE_OUTPUT_PROPERTIES,
1010
GENIE_SPACE_PARAMS,
1111
genieSpacePath,
1212
mapGenieSpace,
@@ -70,7 +70,7 @@ export const genieGetSpaceTool: ToolConfig<
7070
type: 'string',
7171
description:
7272
'Serialized space configuration as a JSON string (data sources, instructions, sample questions)',
73-
optional: true,
73+
nullable: true,
7474
},
7575
},
7676
}

‎apps/sim/tools/databricks/genie_list_conversations.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -100,25 +100,25 @@ export const genieListConversationsTool: ToolConfig<
100100
type: 'object',
101101
properties: {
102102
conversationId: { type: 'string', description: 'Conversation ID' },
103-
title: { type: 'string', description: 'Conversation title', optional: true },
103+
title: { type: 'string', description: 'Conversation title', nullable: true },
104104
createdTimestamp: {
105105
type: 'number',
106106
description: 'When the conversation was created',
107-
optional: true,
107+
nullable: true,
108108
},
109109
agentType: {
110110
type: 'string',
111111
description:
112112
'Conversation mode (GENIE_CONVERSATION_TYPE_CHAT or GENIE_CONVERSATION_TYPE_AGENT)',
113-
optional: true,
113+
nullable: true,
114114
},
115115
},
116116
},
117117
},
118118
nextPageToken: {
119119
type: 'string',
120120
description: 'Token for the next page of results',
121-
optional: true,
121+
nullable: true,
122122
},
123123
},
124124
}

0 commit comments

Comments
 (0)