Skip to content

Commit c2cce49

Browse files
waleedlatif1waleedlatif1
authored
feat(databricks): add Genie agent operations to the Databricks block (#8343)
* feat(databricks): add Genie agent operations to the Databricks block * fix(databricks): fail Genie asks on nested errors and document per-tool outputs
1 parent 65af161 commit c2cce49

24 files changed

Lines changed: 3081 additions & 16 deletions

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

Lines changed: 433 additions & 2 deletions
Large diffs are not rendered by default.

‎apps/sim/blocks/blocks/databricks.ts‎

Lines changed: 379 additions & 6 deletions
Large diffs are not rendered by default.

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

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
import {
2+
type DatabricksGenieAgentAskParams,
3+
type DatabricksGenieAgentAskResponse,
4+
GENIE_AGENT_ITEM_OUTPUT_PROPERTIES,
5+
} from '@/tools/databricks/types'
6+
import {
7+
databricksUrl,
8+
GENIE_SPACE_PARAMS,
9+
genieAgentErrorMessage,
10+
genieAgentPath,
11+
mapGenieAgentItem,
12+
parseGenieAgentQuery,
13+
parseGenieAgentStream,
14+
} from '@/tools/databricks/utils'
15+
import type { ToolConfig } from '@/tools/types'
16+
17+
export const genieAgentAskTool: ToolConfig<
18+
DatabricksGenieAgentAskParams,
19+
DatabricksGenieAgentAskResponse
20+
> = {
21+
id: 'databricks_genie_agent_ask',
22+
name: 'Databricks Genie Agent Ask',
23+
description:
24+
'Ask a Genie agent a question in agent mode, where Genie plans multi-step research, runs several SQL queries, and writes a report with citations. Waits for the final report. Starts a new agent-mode conversation, or continues one when a conversation ID is given. Agent mode is in preview and must be enabled on the workspace.',
25+
version: '1.0.0',
26+
27+
params: {
28+
...GENIE_SPACE_PARAMS,
29+
spaceId: {
30+
...GENIE_SPACE_PARAMS.spaceId,
31+
description: 'The ID of the Genie space (the Genie agent ID)',
32+
},
33+
content: {
34+
type: 'string',
35+
required: true,
36+
visibility: 'user-or-llm',
37+
description: 'The question to ask the Genie agent',
38+
},
39+
conversationId: {
40+
type: 'string',
41+
required: false,
42+
visibility: 'user-or-llm',
43+
description: 'Continue this agent-mode conversation. Omit to start a new conversation.',
44+
},
45+
enableVisualization: {
46+
type: 'boolean',
47+
required: false,
48+
visibility: 'user-or-llm',
49+
description: 'Ask the agent to also generate charts where they fit',
50+
},
51+
},
52+
53+
request: {
54+
modelInput: {
55+
mode: 'project',
56+
select: (params) => ({ content: params.content }),
57+
},
58+
url: (params) => databricksUrl(params.host, `${genieAgentPath(params.spaceId)}/responses`),
59+
method: 'POST',
60+
headers: (params) => ({
61+
'Content-Type': 'application/json',
62+
Accept: 'text/event-stream',
63+
Authorization: `Bearer ${params.apiKey}`,
64+
}),
65+
body: (params) => {
66+
const conversationId = params.conversationId?.trim()
67+
return {
68+
input: [
69+
{
70+
type: 'message',
71+
role: 'user',
72+
content: [{ type: 'input_text', text: params.content }],
73+
},
74+
],
75+
...(conversationId ? { conversation_id: conversationId } : {}),
76+
...(params.enableVisualization ? { enable_viz: true } : {}),
77+
}
78+
},
79+
},
80+
81+
transformResponse: async (response: Response) => {
82+
const { created, final } = parseGenieAgentStream(await response.text())
83+
84+
if (!final) {
85+
throw new Error(
86+
created?.id
87+
? `Genie agent stream ended before response ${created.id} completed`
88+
: 'Genie agent stream ended without a response'
89+
)
90+
}
91+
const responseId = final.id ?? created?.id ?? ''
92+
const conversationId = final.conversation_id ?? created?.conversation_id ?? ''
93+
const context = `Genie agent response ${responseId} in conversation ${conversationId}`
94+
95+
if (final.status === 'failed') {
96+
throw new Error(`${context} failed: ${genieAgentErrorMessage(final.error)}`)
97+
}
98+
99+
const items = (final.output ?? []).map(mapGenieAgentItem)
100+
const joinMessages = (role: string) =>
101+
items
102+
.filter((item) => item.type === 'message' && item.role === role && item.text)
103+
.map((item) => item.text)
104+
.join('\n\n') || null
105+
const report = joinMessages('assistant')
106+
/** A message item can be a system error rather than an assistant report. */
107+
const error = joinMessages('system')
108+
109+
if (!report && error) {
110+
throw new Error(`${context} ended with an error: ${error}`)
111+
}
112+
113+
return {
114+
success: true,
115+
output: {
116+
responseId,
117+
conversationId,
118+
status: final.status ?? 'completed',
119+
report,
120+
queries: items
121+
.filter((item) => item.type === 'function_call' && item.name === 'execute_sql')
122+
.map(parseGenieAgentQuery),
123+
items,
124+
createdAt: final.created_at ?? null,
125+
error,
126+
},
127+
}
128+
},
129+
130+
outputs: {
131+
responseId: { type: 'string', description: 'Agent response ID' },
132+
conversationId: {
133+
type: 'string',
134+
description: 'Agent-mode conversation ID; pass it back to ask a follow-up',
135+
},
136+
status: { type: 'string', description: 'Response status (completed)' },
137+
report: {
138+
type: 'string',
139+
description: "The agent's final report, with citation links",
140+
nullable: true,
141+
},
142+
queries: {
143+
type: 'array',
144+
description: 'SQL queries the agent ran',
145+
items: {
146+
type: 'object',
147+
properties: {
148+
callId: { type: 'string', description: 'Function call ID' },
149+
title: { type: 'string', description: 'Query title', nullable: true },
150+
sql: { type: 'string', description: 'SQL the agent ran', nullable: true },
151+
},
152+
},
153+
},
154+
items: {
155+
type: 'array',
156+
description: 'Every output item: reasoning, SQL calls, query results, and messages',
157+
items: {
158+
type: 'object',
159+
properties: GENIE_AGENT_ITEM_OUTPUT_PROPERTIES,
160+
},
161+
},
162+
createdAt: {
163+
type: 'number',
164+
description: 'When the response was created (Unix epoch seconds)',
165+
nullable: true,
166+
},
167+
error: {
168+
type: 'string',
169+
description: 'System error message the agent reported alongside its report',
170+
nullable: true,
171+
},
172+
},
173+
}

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

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
import {
2+
type DatabricksGenieAgentListItemsParams,
3+
type DatabricksGenieAgentListItemsResponse,
4+
GENIE_AGENT_ITEM_OUTPUT_PROPERTIES,
5+
} from '@/tools/databricks/types'
6+
import {
7+
databricksErrorMessage,
8+
databricksUrl,
9+
GENIE_READ_RETRY,
10+
GENIE_SPACE_PARAMS,
11+
genieAgentPath,
12+
mapGenieAgentItem,
13+
} from '@/tools/databricks/utils'
14+
import type { ToolConfig } from '@/tools/types'
15+
import { safeUrlPathSegment } from '@/tools/url-path'
16+
17+
export const genieAgentListItemsTool: ToolConfig<
18+
DatabricksGenieAgentListItemsParams,
19+
DatabricksGenieAgentListItemsResponse
20+
> = {
21+
id: 'databricks_genie_agent_list_items',
22+
name: 'Databricks Genie Agent List Items',
23+
description:
24+
"List an agent-mode conversation's full history in order: questions, reasoning, SQL calls, query results, and reports. Agent mode is in preview and must be enabled on the workspace.",
25+
version: '1.0.0',
26+
27+
params: {
28+
...GENIE_SPACE_PARAMS,
29+
spaceId: {
30+
...GENIE_SPACE_PARAMS.spaceId,
31+
description: 'The ID of the Genie space (the Genie agent ID)',
32+
},
33+
conversationId: {
34+
type: 'string',
35+
required: true,
36+
visibility: 'user-or-llm',
37+
description: 'The ID of the agent-mode conversation',
38+
},
39+
limit: {
40+
type: 'number',
41+
required: false,
42+
visibility: 'user-or-llm',
43+
description: 'Maximum number of items to return (1-100, default 100)',
44+
},
45+
after: {
46+
type: 'string',
47+
required: false,
48+
visibility: 'user-or-llm',
49+
description: 'Pagination cursor: the lastId from a previous response',
50+
},
51+
order: {
52+
type: 'string',
53+
required: false,
54+
visibility: 'user-or-llm',
55+
description: 'Sort order: asc (oldest first, default) or desc (newest first)',
56+
},
57+
},
58+
59+
request: {
60+
url: (params) => {
61+
const url = new URL(
62+
databricksUrl(
63+
params.host,
64+
`${genieAgentPath(params.spaceId)}/conversations/${safeUrlPathSegment(params.conversationId, 'conversationId')}/items`
65+
)
66+
)
67+
if (params.limit) url.searchParams.set('limit', String(params.limit))
68+
if (params.after) url.searchParams.set('after', params.after)
69+
if (params.order) url.searchParams.set('order', params.order)
70+
return url.toString()
71+
},
72+
method: 'GET',
73+
headers: (params) => ({
74+
Accept: 'application/json',
75+
Authorization: `Bearer ${params.apiKey}`,
76+
}),
77+
retry: GENIE_READ_RETRY,
78+
},
79+
80+
transformResponse: async (response: Response) => {
81+
const data = await response.json()
82+
83+
if (!response.ok) {
84+
throw new Error(databricksErrorMessage(data, 'Failed to list Genie agent conversation items'))
85+
}
86+
87+
return {
88+
success: true,
89+
output: {
90+
items: (data.data ?? []).map(mapGenieAgentItem),
91+
firstId: data.first_id ?? null,
92+
lastId: data.last_id ?? null,
93+
hasMore: data.has_more ?? false,
94+
status: data.status ?? null,
95+
},
96+
}
97+
},
98+
99+
outputs: {
100+
items: {
101+
type: 'array',
102+
description:
103+
'Conversation items in order; user questions are message items whose ID ends in _input',
104+
items: {
105+
type: 'object',
106+
properties: GENIE_AGENT_ITEM_OUTPUT_PROPERTIES,
107+
},
108+
},
109+
firstId: { type: 'string', description: 'ID of the first item in the page', nullable: true },
110+
lastId: {
111+
type: 'string',
112+
description: 'ID of the last item in the page; pass it as After for the next page',
113+
nullable: true,
114+
},
115+
hasMore: { type: 'boolean', description: 'Whether more items follow this page' },
116+
status: {
117+
type: 'string',
118+
description: 'Status of the latest response in the conversation',
119+
nullable: true,
120+
},
121+
},
122+
}

0 commit comments

Comments
 (0)