@dbx-tools/genie 0.1.9

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/README.md ADDED
@@ -0,0 +1,162 @@
1
+ # @dbx-tools/node-genie
2
+
3
+ Server-side Databricks Genie chat drivers.
4
+
5
+ Import this package when Node or AppKit backend code needs to run one turn
6
+ against a Genie space and consume either raw Genie message snapshots or a typed
7
+ event stream. It preserves AppKit OBO auth when called inside an AppKit request,
8
+ falls back to the Databricks SDK default auth outside AppKit, and supports
9
+ caller-provided cancellation.
10
+
11
+ Pure Genie schemas and event detector helpers live in
12
+ [`@dbx-tools/shared-genie`](../../shared/genie).
13
+
14
+ Key features:
15
+
16
+ - Starts new Genie conversations or continues an existing `conversationId`.
17
+ - Polls Databricks Genie until terminal status while filtering unchanged
18
+ snapshots.
19
+ - Converts raw Genie messages into semantic events for thinking text, generated
20
+ SQL, row counts, final results, and errors.
21
+ - Preserves AppKit OBO auth when called during an AppKit request, but also works
22
+ from standalone scripts with normal Databricks SDK auth.
23
+ - Accepts SDK `Context` or web `AbortSignal` cancellation for route handlers and
24
+ CLI tools.
25
+ - Fetches Genie space metadata and starter questions for UI suggestions.
26
+
27
+ ## Why Not Just AppKit Genie?
28
+
29
+ Native AppKit's Genie plugin is the right choice for a standalone Genie chat
30
+ experience: it provides named space aliases, SSE status updates, conversation
31
+ history replay, query result fetching, OBO execution, and the AppKit UI
32
+ `GenieChat` component.
33
+
34
+ Use this package when Genie is one capability inside a larger agent or custom
35
+ backend:
36
+
37
+ - You want a low-level async iterator rather than an AppKit HTTP route.
38
+ - You want raw message snapshots or a normalized event stream that can be fed
39
+ into Mastra writer events, logs, tests, or custom SSE endpoints.
40
+ - You need to diff snapshots and emit only newly observed thinking, SQL, rows,
41
+ result, and error events.
42
+ - You want to combine Genie answers with agent-side chart planning, statement
43
+ row fetches, or durable thread storage owned elsewhere.
44
+ - You need the same driver to work inside AppKit with OBO auth and outside
45
+ AppKit from scripts using normal Databricks SDK auth.
46
+
47
+ ## Stream Semantic Events
48
+
49
+ ```ts
50
+ import { chat } from "@dbx-tools/node-genie";
51
+
52
+ for await (const event of chat.genieEventChat(spaceId, "Top stores by revenue?")) {
53
+ switch (event.type) {
54
+ case "thinking":
55
+ console.log(event.thought_type, event.text);
56
+ break;
57
+ case "query":
58
+ console.log(event.sql);
59
+ break;
60
+ case "rows":
61
+ console.log(event.row_count);
62
+ break;
63
+ case "result":
64
+ console.log(event.status);
65
+ break;
66
+ }
67
+ }
68
+ ```
69
+
70
+ `chat.genieEventChat()` wraps the lower-level snapshot stream and yields a
71
+ `GenieChatEvent` union. Use it for SSE streams, log pipelines, and tool writer
72
+ events where consumers care about progress and SQL, not just the terminal
73
+ message.
74
+
75
+ ## Stream Raw Message Snapshots
76
+
77
+ ```ts
78
+ import { chat } from "@dbx-tools/node-genie";
79
+
80
+ for await (const message of chat.genieChat(spaceId, "Top stores by revenue?")) {
81
+ renderSnapshot(message);
82
+ }
83
+ ```
84
+
85
+ `chat.genieChat()` starts a conversation or appends to an existing one, polls
86
+ `client.genie.getMessage`, filters identical consecutive payloads, and stops
87
+ after a terminal status. Use it when you want to run your own diffing or persist
88
+ the raw Genie wire shape.
89
+
90
+ ## Continue A Conversation
91
+
92
+ ```ts
93
+ let conversationId: string | undefined;
94
+
95
+ for (const prompt of prompts) {
96
+ for await (const event of chat.genieEventChat(spaceId, prompt, { conversationId })) {
97
+ if ("conversation_id" in event && event.conversation_id) {
98
+ conversationId = event.conversation_id;
99
+ }
100
+ }
101
+ }
102
+ ```
103
+
104
+ The driver does not own multi-turn state. Callers read the conversation id from
105
+ a yielded message/event and pass it into the next turn. That makes the package
106
+ usable in stateless route handlers, durable thread stores, and one-off scripts.
107
+
108
+ This split is deliberate: the package is a transport/driver layer, not a thread
109
+ store. AppKit-Mastra persists thread state separately and passes the Genie
110
+ conversation id back into this driver when a turn continues.
111
+
112
+ ## Resolve A Workspace Client
113
+
114
+ ```ts
115
+ import { WorkspaceClient } from "@databricks/sdk-experimental";
116
+ import { chat } from "@dbx-tools/node-genie";
117
+
118
+ await chat.genieEventChat(spaceId, content, {
119
+ workspaceClient: new WorkspaceClient({ profile: "dev" }),
120
+ });
121
+ ```
122
+
123
+ Client resolution order:
124
+
125
+ 1. `options.workspaceClient`;
126
+ 2. AppKit execution-context client, when present;
127
+ 3. `new WorkspaceClient({})` using normal Databricks SDK auth.
128
+
129
+ Pass `options.context` as an `AbortSignal` or SDK context to cancel SDK calls and
130
+ the polling sleep.
131
+
132
+ ## Read Space Metadata And Starter Questions
133
+
134
+ ```ts
135
+ import { space } from "@dbx-tools/node-genie";
136
+
137
+ const genieSpace = await space.getGenieSpace(spaceId);
138
+ const questions = space.genieSampleQuestions(genieSpace);
139
+ ```
140
+
141
+ `space.getGenieSpace()` fetches the space definition, including serialized space
142
+ metadata by default. `space.genieSampleQuestions()` extracts curated starter
143
+ questions and returns `[]` when none are configured.
144
+
145
+ ## Options
146
+
147
+ `chat.GenieChatOptions` is shared by both drivers:
148
+
149
+ - `conversationId` - append to an existing Genie conversation.
150
+ - `workspaceClient` - explicit Databricks SDK client.
151
+ - `pollIntervalMs` - polling cadence, default `500`.
152
+ - `context` - SDK `Context` or `AbortSignal` for cancellation.
153
+
154
+ ## Modules
155
+
156
+ - `chat` - `genieChat()` raw snapshot stream and `genieEventChat()` typed event
157
+ stream.
158
+ - `space` - `getGenieSpace()` and `genieSampleQuestions()`.
159
+
160
+ The AppKit-Mastra package builds its Genie tools on top of this driver; see
161
+ [`@dbx-tools/node-appkit-mastra`](../appkit-mastra) for the agent-level
162
+ workflow.
package/index.ts ADDED
@@ -0,0 +1,8 @@
1
+ // GENERATED by projen watch - DO NOT EDIT.
2
+ // Regenerated from the exporting modules in ./src.
3
+ // Hand edits are overwritten on the next watch; this file is read-only.
4
+
5
+ export * as chat from "./src/chat";
6
+ export * as space from "./src/space";
7
+ export type { GenieChatOptions } from "./src/chat";
8
+ export type { GetGenieSpaceOptions } from "./src/space";
package/package.json ADDED
@@ -0,0 +1,54 @@
1
+ {
2
+ "name": "@dbx-tools/genie",
3
+ "repository": {
4
+ "type": "git",
5
+ "url": "git+https://github.com/reggie-db/dbx-tools.git",
6
+ "directory": "workspaces/node/genie"
7
+ },
8
+ "devDependencies": {
9
+ "@databricks/appkit": "^0.43.0",
10
+ "@types/node": "^24.6.0",
11
+ "tsx": "^4.23.0",
12
+ "typescript": "^5.9.3"
13
+ },
14
+ "peerDependencies": {
15
+ "@databricks/appkit": "^0.43.0"
16
+ },
17
+ "dependencies": {
18
+ "@databricks/sdk-experimental": "^0.17.0",
19
+ "@dbx-tools/appkit": "0.1.9",
20
+ "@dbx-tools/shared-core": "0.1.9",
21
+ "@dbx-tools/shared-genie": "0.1.9"
22
+ },
23
+ "main": "index.ts",
24
+ "license": "UNLICENSED",
25
+ "version": "0.1.9",
26
+ "types": "index.ts",
27
+ "type": "module",
28
+ "exports": {
29
+ ".": "./index.ts",
30
+ "./package.json": "./package.json"
31
+ },
32
+ "peerDependenciesMeta": {
33
+ "@databricks/appkit": {
34
+ "optional": true
35
+ }
36
+ },
37
+ "dbxToolsConfig": {
38
+ "tags": [
39
+ "node"
40
+ ]
41
+ },
42
+ "//": "~~ Generated by projen. To modify, edit .projenrc.js and run \"pnpm exec projen\".",
43
+ "scripts": {
44
+ "build": "projen build",
45
+ "compile": "projen compile",
46
+ "default": "projen default",
47
+ "package": "projen package",
48
+ "post-compile": "projen post-compile",
49
+ "pre-compile": "projen pre-compile",
50
+ "test": "projen test",
51
+ "watch": "projen watch",
52
+ "projen": "projen"
53
+ }
54
+ }
package/src/chat.ts ADDED
@@ -0,0 +1,324 @@
1
+ /**
2
+ * Genie chat driver.
3
+ *
4
+ * Drives a single turn against a Genie space from one `content` string;
5
+ * multi-turn conversations are the caller's job (thread the `conversation_id`
6
+ * returned on each `GenieMessage` back into the next turn's
7
+ * `options.conversationId`).
8
+ *
9
+ * Two layers serve two kinds of consumer. The low-level layer yields every
10
+ * poll-observed `GenieMessage` (validated against `GenieMessageSchema`,
11
+ * falling back to the raw snapshot on a schema miss) and owns the messy parts
12
+ * - cancellation, conversation seeding, distinct-filtering, and SDK quirks
13
+ * (Waiter stripping); reach for it when you want the raw stream. The
14
+ * high-level layer wraps it and emits semantic, deduplicated `{ type, payload }`
15
+ * events (see {@link GenieChatEvent}), always closing a successful turn with a
16
+ * terminal `result` event carrying the final `GenieMessage`; errors propagate
17
+ * by throwing, with no `error` variant. Iterating UI / agent code that wants
18
+ * every message verbatim takes the low-level stream; subscribers reacting to
19
+ * "Genie is thinking about X" or "Genie produced text Y" take the event layer.
20
+ */
21
+
22
+ import { async, log, type PollContext } from "@dbx-tools/shared-core";
23
+ import { databricks } from "@dbx-tools/appkit";
24
+ import { event, genieModel, type GenieChatEvent, type GenieMessage } from "@dbx-tools/shared-genie";
25
+ import { WorkspaceClient } from "@databricks/sdk-experimental";
26
+
27
+ const logger = log.logger("genie/chat");
28
+
29
+ /**
30
+ * Validate a polled wire snapshot against {@link GenieMessageSchema}
31
+ * and return the schema-normalized message. Genie's wire occasionally ships a
32
+ * shape the (SDK-derived) schema doesn't model exactly - e.g. an early poll
33
+ * that omits the SDK-required `message_id` - so a miss degrades to the raw
34
+ * snapshot rather than throwing, keeping a single odd poll from aborting the
35
+ * whole turn.
36
+ */
37
+ function validateMessage(raw: GenieMessage): GenieMessage {
38
+ const result = genieModel.GenieMessageSchema.safeParse(raw);
39
+ if (result.success) return result.data;
40
+ logger.debug("wire-message:schema-miss", {
41
+ message_id: raw.message_id ?? raw.id,
42
+ issues: result.error.issues.length,
43
+ });
44
+ return raw;
45
+ }
46
+
47
+ /* -------------------------- shared options -------------------------- */
48
+
49
+ /** Options accepted by both {@link genieChat} and {@link genieEventChat}. */
50
+ export interface GenieChatOptions {
51
+ /**
52
+ * Seed conversation id. When set, this turn appends to the existing
53
+ * conversation (via `createMessage`) instead of opening a new one. Use it to
54
+ * thread a multi-turn conversation: read `conversation_id` off the prior
55
+ * turn's terminal `GenieMessage` (or the `result` event's
56
+ * `payload.conversation_id`) and pass it into the next call.
57
+ */
58
+ conversationId?: string;
59
+ /**
60
+ * Explicit `WorkspaceClient`. Defaults to AppKit's per-request
61
+ * execution-context client when AppKit is installed and we're inside a
62
+ * request; falls back to a fresh `new WorkspaceClient({})` (env-var auth)
63
+ * otherwise.
64
+ */
65
+ workspaceClient?: WorkspaceClient;
66
+ /** Poll cadence in milliseconds between successive `getMessage` calls (default 500). */
67
+ pollIntervalMs?: number;
68
+ /**
69
+ * External cancellation. Accepts a WHATWG `AbortSignal` or a fully-built SDK
70
+ * `Context` (see `databricks.ContextLike`). Aborting it cancels every in-flight
71
+ * SDK call and the next inter-poll sleep.
72
+ */
73
+ context?: databricks.ContextLike;
74
+ }
75
+
76
+ /* ----------------------- low-level: genieChat ----------------------- */
77
+
78
+ /**
79
+ * One turn against a Genie space, yielded as a stream of `GenieMessage`
80
+ * snapshots.
81
+ *
82
+ * Turn lifecycle:
83
+ *
84
+ * - No `options.conversationId`: open a new conversation via
85
+ * `client.genie.startConversation`. The opened conversation id surfaces on
86
+ * every yielded `GenieMessage` (`.conversation_id`) so the caller can
87
+ * thread it into a follow-up call.
88
+ * - With `options.conversationId`: append to that conversation via
89
+ * `client.genie.createMessage`.
90
+ * - In both cases, after the create/start the driver polls
91
+ * `client.genie.getMessage` every `options.pollIntervalMs` (default 500ms)
92
+ * until the message reaches a terminal status, then yields the terminal
93
+ * snapshot and returns.
94
+ *
95
+ * Cancellation: a single internal `AbortController` covers the whole turn.
96
+ * `options.context` is tied into that controller so an external abort tears
97
+ * down every in-flight SDK call AND the inter-poll sleep. Breaking out of the
98
+ * `for await` does the same via the `try / finally`.
99
+ *
100
+ * @example
101
+ * // Single turn.
102
+ * for await (const m of genieChat(spaceId, "Top 5 stores?")) {
103
+ * render(m);
104
+ * }
105
+ *
106
+ * @example
107
+ * // Multi-turn: caller threads the conversation id.
108
+ * let conversationId: string | undefined;
109
+ * for (const question of questions) {
110
+ * for await (const m of genieChat(spaceId, question, { conversationId })) {
111
+ * conversationId = m.conversation_id ?? conversationId;
112
+ * render(m);
113
+ * }
114
+ * }
115
+ */
116
+ export async function* genieChat(
117
+ space_id: string,
118
+ content: string,
119
+ options?: GenieChatOptions,
120
+ ): AsyncGenerator<GenieMessage, void, void> {
121
+ const controller = new AbortController();
122
+ try {
123
+ const client = await getWorkspaceClient(options);
124
+ // Build the SDK Context ONCE. Building it inside the poll producer would
125
+ // re-attach an abort listener to `options.context` on every poll iteration
126
+ // (via `databricks.toContext` -> `async.tieAbortSignal`), eventually tripping
127
+ // Node's `MaxListenersExceededWarning`.
128
+ const ctx = databricks.toContext(controller, options?.context);
129
+ let conversationId = options?.conversationId;
130
+ let messageId: string | undefined;
131
+
132
+ const pollProducer = async (pollCtx: PollContext<GenieMessage>): Promise<GenieMessage> => {
133
+ if (!conversationId) {
134
+ // First poll: open the conversation. Refuse to retry: if
135
+ // `startConversation` returned a response without a `conversation_id`,
136
+ // retrying would just open conversation after conversation.
137
+ if (pollCtx.attempt > 0) {
138
+ throw new Error("Genie did not return a conversation id; refusing to retry");
139
+ }
140
+ const startResponse = await client.genie.startConversation({ space_id, content }, ctx);
141
+ conversationId = startResponse.conversation_id;
142
+ messageId = startResponse.message_id;
143
+ return startResponse.message!;
144
+ }
145
+ if (!messageId) {
146
+ // First poll of a follow-up turn: append to the seeded conversation.
147
+ // `client.genie.createMessage` returns a `Waiter<GenieMessage>`
148
+ // (`{ ...message, wait: async () => ... }`). Strip `wait` here so
149
+ // downstream serializers (e.g. yaml.stringify) don't choke on the
150
+ // AsyncFunction value.
151
+ const { wait: _wait, ...createResponse } = await client.genie.createMessage(
152
+ { space_id, conversation_id: conversationId, content },
153
+ ctx,
154
+ );
155
+ messageId = createResponse.message_id;
156
+ return createResponse;
157
+ }
158
+ // Subsequent polls: re-fetch the current message until its status becomes
159
+ // terminal.
160
+ return await client.genie.getMessage(
161
+ {
162
+ space_id,
163
+ conversation_id: conversationId,
164
+ message_id: messageId,
165
+ },
166
+ ctx,
167
+ );
168
+ };
169
+
170
+ yield* async.poll(
171
+ async (pollCtx: PollContext<GenieMessage>) => validateMessage(await pollProducer(pollCtx)),
172
+ {
173
+ intervalMs: options?.pollIntervalMs ?? 500,
174
+ // Skip yielding identical consecutive snapshots; Genie often returns
175
+ // the exact same payload twice during quiet periods. `poll` does a deep
176
+ // equal on the previous yield.
177
+ filter: "distinct",
178
+ // Stop after the terminal message is yielded. `poll` checks the
179
+ // predicate AFTER yielding, so the terminal message still reaches the
180
+ // consumer.
181
+ predicate: (m) => !genieModel.isTerminalStatus(m.status),
182
+ // Wake the inter-poll sleep on abort so a `for await` break (or
183
+ // external abort) tears down promptly instead of waiting out the
184
+ // interval.
185
+ signal: controller.signal,
186
+ },
187
+ );
188
+ } finally {
189
+ // Cancels any still-pending SDK call and the inter-poll sleep whether we're
190
+ // unwinding from a normal return, a consumer break, or a thrown error.
191
+ // Idempotent.
192
+ controller.abort();
193
+ }
194
+ }
195
+
196
+ /* ---------------------- high-level: genieEventChat ------------------ */
197
+
198
+ /**
199
+ * One turn against a Genie space, yielded as a typed {@link GenieChatEvent}
200
+ * stream. Drives {@link genieChat} underneath and decorates each snapshot with
201
+ * the derived events the field-level diff produced. Stream order:
202
+ *
203
+ * 1. `{ type: "message", message }` - the raw `GenieMessage`, once per poll
204
+ * yield.
205
+ * 2. `{ type: "question", content, message_id, ... }` fires exactly once, on
206
+ * the FIRST `message` yield. We read `content` and `message_id` straight
207
+ * off the snapshot so every downstream event for this turn shares the same
208
+ * `message_id` (the question included) - subscribers can group everything
209
+ * for one Genie call under that one key.
210
+ * 3. Any of `status` / `attachment` / `thinking` / `text` / `query` /
211
+ * `statement` / `rows` / `suggested_questions` the diff against the prior
212
+ * snapshot produced.
213
+ * 4. On the terminal snapshot, `{ type: "result", ... }` as the final yield.
214
+ *
215
+ * Errors propagate by the generator throwing - there's no `"error"` variant.
216
+ * Wrap the `for await` in `try/catch` if you need to handle failures.
217
+ *
218
+ * @example
219
+ * for await (const evt of genieEventChat(spaceId, "Top stores?")) {
220
+ * switch (evt.type) {
221
+ * case "thinking":
222
+ * console.log("[thinking]", evt.thought_type, evt.text);
223
+ * break;
224
+ * case "text":
225
+ * console.log("[text]", evt.text);
226
+ * break;
227
+ * case "result":
228
+ * console.log("[done]", evt.status);
229
+ * break;
230
+ * }
231
+ * }
232
+ */
233
+ export async function* genieEventChat(
234
+ space_id: string,
235
+ content: string,
236
+ options?: GenieChatOptions,
237
+ ): AsyncGenerator<GenieChatEvent, void, void> {
238
+ // Diff source for the current turn. Always `undefined` on the first snapshot
239
+ // so the initial status / attachments emit fresh; updated to the most recent
240
+ // snapshot after each yield.
241
+ let previous: GenieMessage | undefined;
242
+ // The `question` event is deferred to the first `message` yield so it can
243
+ // carry the assigned `message_id` (subscribers use it as the grouping key for
244
+ // every event in this turn). The first snapshot is the earliest point that id
245
+ // exists.
246
+ let questionEmitted = false;
247
+ for await (const rawMessage of genieChat(space_id, content, options)) {
248
+ // Normalize `message_id` from the legacy `id` field when Genie's wire
249
+ // response only populates one of them. The SDK schema marks both as
250
+ // required, but in practice the `startConversation` / `createMessage` inner
251
+ // `message` payload sometimes ships only `id` while the new `message_id`
252
+ // field lands undefined. Every downstream event detector keys grouping off
253
+ // `message_id`; the fallback keeps one Genie turn's events from splitting
254
+ // across an anon group + the real-id group when subscribers bucket by
255
+ // `message_id`.
256
+ const message: GenieMessage = rawMessage.message_id
257
+ ? rawMessage
258
+ : { ...rawMessage, message_id: rawMessage.id };
259
+ yield {
260
+ type: "message",
261
+ space_id: message.space_id,
262
+ message_id: message.message_id,
263
+ message,
264
+ };
265
+ if (!questionEmitted) {
266
+ yield {
267
+ type: "question",
268
+ space_id: message.space_id,
269
+ ...(message.conversation_id ? { conversation_id: message.conversation_id } : {}),
270
+ ...(message.message_id ? { message_id: message.message_id } : {}),
271
+ content: message.content,
272
+ };
273
+ questionEmitted = true;
274
+ }
275
+ yield* event.eventsFromMessage(message, previous, message.space_id);
276
+ const status = message.status;
277
+ if (genieModel.isTerminalStatus(status)) {
278
+ yield {
279
+ type: "result",
280
+ space_id: message.space_id,
281
+ conversation_id: message.conversation_id,
282
+ message_id: message.message_id,
283
+ status,
284
+ message,
285
+ };
286
+ }
287
+ previous = message;
288
+ }
289
+ }
290
+
291
+ /* ---------------------- workspace client helper --------------------- */
292
+
293
+ /**
294
+ * Resolve a `WorkspaceClient` in this preference order:
295
+ *
296
+ * 1. Caller-supplied `options.workspaceClient`.
297
+ * 2. AppKit's per-request execution-context client, when AppKit is installed
298
+ * AND we're inside a request scope.
299
+ * 3. Fresh `new WorkspaceClient({})` (env-var auth via
300
+ * `DATABRICKS_CONFIG_PROFILE` / `DATABRICKS_HOST` / `DATABRICKS_TOKEN`).
301
+ *
302
+ * AppKit is loaded lazily so this package stays usable in non-AppKit
303
+ * environments.
304
+ */
305
+ async function getWorkspaceClient(options?: GenieChatOptions): Promise<WorkspaceClient> {
306
+ if (options?.workspaceClient) return options.workspaceClient;
307
+ const appkit = await getAppKit();
308
+ if (appkit) {
309
+ try {
310
+ return appkit.getExecutionContext().client;
311
+ } catch {
312
+ // Not inside an AppKit request context; fall through to env.
313
+ }
314
+ }
315
+ return new WorkspaceClient({});
316
+ }
317
+
318
+ async function getAppKit() {
319
+ try {
320
+ return await import("@databricks/appkit");
321
+ } catch {
322
+ return undefined;
323
+ }
324
+ }
package/src/space.ts ADDED
@@ -0,0 +1,113 @@
1
+ /**
2
+ * Genie space metadata helpers.
3
+ *
4
+ * Fetches a Genie space's definition (including the opt-in `serialized_space`
5
+ * blob) and extracts the curated starter questions an author configured on the
6
+ * space. The typed SDK `client.genie.getSpace` only returns the
7
+ * directory-listing surface (`title` / `description` / `warehouse_id`); the
8
+ * sample questions live inside `serialized_space`, which the REST API returns
9
+ * only when `include_serialized_space=true`. We hit that endpoint through the
10
+ * workspace client's raw `apiClient` since the typed request shape has no flag
11
+ * for it.
12
+ */
13
+
14
+ import { error, log, string } from "@dbx-tools/shared-core";
15
+ import { databricks } from "@dbx-tools/appkit";
16
+ import { genieModel, type GenieSpace } from "@dbx-tools/shared-genie";
17
+ import { WorkspaceClient } from "@databricks/sdk-experimental";
18
+
19
+ const logger = log.logger("genie/space");
20
+
21
+ /** Options for {@link getGenieSpace}. */
22
+ export interface GetGenieSpaceOptions {
23
+ /**
24
+ * Explicit `WorkspaceClient`. Defaults to a fresh `new WorkspaceClient({})`
25
+ * (env-var auth). Server callers should pass their OBO-scoped client so the
26
+ * lookup runs as the user.
27
+ */
28
+ workspaceClient?: WorkspaceClient;
29
+ /**
30
+ * Request the `serialized_space` blob (catalogs, tables, sample questions,
31
+ * prompts). Defaults to `true` - the only reason to skip it is when the
32
+ * caller just needs title / description and wants the smaller payload.
33
+ */
34
+ serialized?: boolean;
35
+ /**
36
+ * External cancellation. Accepts a WHATWG `AbortSignal` or a fully-built SDK
37
+ * `Context` (see `databricks.ContextLike`).
38
+ */
39
+ context?: databricks.ContextLike;
40
+ }
41
+
42
+ /**
43
+ * Fetch a Genie space by id, optionally including its serialized definition.
44
+ * Hits `GET /api/2.0/genie/spaces/<id>` with `include_serialized_space=true`
45
+ * through the raw `apiClient`, then validates the response against
46
+ * {@link GenieSpaceSchema} (unknown fields like `etag` /
47
+ * `parent_path` are stripped).
48
+ */
49
+ export async function getGenieSpace(
50
+ spaceId: string,
51
+ options?: GetGenieSpaceOptions,
52
+ ): Promise<GenieSpace> {
53
+ const client = options?.workspaceClient ?? new WorkspaceClient({});
54
+ const serialized = options?.serialized !== false;
55
+ const ctx = options?.context ? databricks.toContext(options.context) : undefined;
56
+ const raw = await client.apiClient.request(
57
+ {
58
+ path: `/api/2.0/genie/spaces/${encodeURIComponent(spaceId)}`,
59
+ method: "GET",
60
+ query: serialized ? { include_serialized_space: true } : {},
61
+ headers: new Headers(),
62
+ raw: false,
63
+ },
64
+ ctx,
65
+ );
66
+ return genieModel.GenieSpaceSchema.parse(raw);
67
+ }
68
+
69
+ /**
70
+ * One entry in a serialized space's `config.sample_questions`. The
71
+ * author-facing field is `question`, which the wire format models as a string
72
+ * array (a single multi-line question is split across entries); we treat the
73
+ * first non-empty entry as the displayable question text.
74
+ */
75
+ interface SerializedSampleQuestion {
76
+ question?: unknown;
77
+ }
78
+
79
+ /**
80
+ * Extract the curated starter questions an author configured on a Genie space.
81
+ * Reads `serialized_space -> config.sample_questions[*].question`. Returns `[]`
82
+ * when the space carries no serialized blob, the blob is unparseable, or no
83
+ * sample questions are configured - so a missing or misconfigured space
84
+ * degrades to "no suggestions" rather than throwing. Order is preserved (the
85
+ * author's ordering) and duplicates are dropped.
86
+ */
87
+ export function genieSampleQuestions(space: GenieSpace): string[] {
88
+ const serialized = space.serialized_space;
89
+ if (!serialized) return [];
90
+ let parsed: unknown;
91
+ try {
92
+ parsed = JSON.parse(serialized);
93
+ } catch (err) {
94
+ logger.warn("serialized-space:parse-error", {
95
+ spaceId: space.space_id,
96
+ error: error.errorMessage(err),
97
+ });
98
+ return [];
99
+ }
100
+ const sampleQuestions = (parsed as { config?: { sample_questions?: unknown } } | null)?.config
101
+ ?.sample_questions;
102
+ if (!Array.isArray(sampleQuestions)) return [];
103
+
104
+ const seen = new Set<string>();
105
+ const out: string[] = [];
106
+ for (const entry of sampleQuestions as SerializedSampleQuestion[]) {
107
+ const text = string.firstNonEmpty(entry?.question);
108
+ if (!text || seen.has(text)) continue;
109
+ seen.add(text);
110
+ out.push(text);
111
+ }
112
+ return out;
113
+ }
package/tsconfig.json ADDED
@@ -0,0 +1,41 @@
1
+ // ~~ Generated by projen. To modify, edit .projenrc.js and run "pnpm exec projen".
2
+ {
3
+ "compilerOptions": {
4
+ "rootDir": "src",
5
+ "outDir": "lib",
6
+ "alwaysStrict": true,
7
+ "declaration": true,
8
+ "esModuleInterop": true,
9
+ "experimentalDecorators": true,
10
+ "inlineSourceMap": true,
11
+ "inlineSources": true,
12
+ "lib": [
13
+ "ES2022"
14
+ ],
15
+ "module": "ESNext",
16
+ "noEmitOnError": false,
17
+ "noFallthroughCasesInSwitch": true,
18
+ "noImplicitAny": true,
19
+ "noImplicitReturns": true,
20
+ "noImplicitThis": true,
21
+ "noUnusedLocals": true,
22
+ "noUnusedParameters": true,
23
+ "resolveJsonModule": true,
24
+ "strict": true,
25
+ "strictNullChecks": true,
26
+ "strictPropertyInitialization": true,
27
+ "stripInternal": true,
28
+ "target": "ES2022",
29
+ "types": [
30
+ "node"
31
+ ],
32
+ "moduleResolution": "bundler",
33
+ "skipLibCheck": true
34
+ },
35
+ "include": [
36
+ "src/**/*.ts"
37
+ ],
38
+ "exclude": [
39
+ "node_modules"
40
+ ]
41
+ }