@dbx-tools/genie 0.1.111 → 0.3.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.
Files changed (8) hide show
  1. package/README.md +115 -155
  2. package/index.ts +8 -0
  3. package/package.json +48 -16
  4. package/src/chat.ts +324 -0
  5. package/src/space.ts +111 -0
  6. package/tsconfig.json +41 -0
  7. package/dist/index.d.ts +0 -157
  8. package/dist/index.js +0 -307
package/README.md CHANGED
@@ -1,202 +1,162 @@
1
1
  # @dbx-tools/genie
2
2
 
3
- Server-side Genie chat drivers. Two async generators that take a
4
- `space_id` + a single `content` string for one turn against a Genie
5
- space and yield either the raw `GenieMessage` snapshots
6
- (`genieChat`) or a typed `GenieChatEvent` stream of flat
7
- `{ type, ...fields }` records (`genieEventChat`).
8
-
9
- Multi-turn conversations are caller-driven: read `conversation_id`
10
- off the prior turn's terminal `GenieMessage` (or the `result`
11
- event's `conversation_id`) and thread it into the next call's
12
- `options.conversationId`.
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
13
48
 
14
49
  ```ts
15
- import { genieEventChat } from "@dbx-tools/genie";
50
+ import { chat } from "@dbx-tools/genie";
16
51
 
17
- for await (const event of genieEventChat(spaceId, "Top 5 stores?")) {
52
+ for await (const event of chat.genieEventChat(spaceId, "Top stores by revenue?")) {
18
53
  switch (event.type) {
19
54
  case "thinking":
20
- console.log("[think]", event.thought_type, event.text);
55
+ console.log(event.thought_type, event.text);
21
56
  break;
22
57
  case "query":
23
- console.log("[sql]", event.title, "\n", event.sql);
58
+ console.log(event.sql);
59
+ break;
60
+ case "rows":
61
+ console.log(event.row_count);
24
62
  break;
25
63
  case "result":
26
- console.log("[done]", event.status);
64
+ console.log(event.status);
27
65
  break;
28
66
  }
29
67
  }
30
68
  ```
31
69
 
32
- Browser safety: `chat.ts` pulls in `WorkspaceClient` from
33
- `@databricks/sdk-experimental` and is Node-only. Browser bundles
34
- should import from [`@dbx-tools/genie-shared`](../genie-shared)
35
- directly - the protocol types, the `GenieChatEvent` union, the pure
36
- detectors, and `eventsFromMessage` all live there and are also
37
- re-exported from `@dbx-tools/genie` for server-side convenience.
38
-
39
- ## `genieChat` - low-level snapshot stream
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.
40
74
 
41
- Yields every poll-observed `GenieMessage` for a single turn. Use it
42
- when you want the raw stream (e.g. to drive a custom UI off the wire
43
- shape, or to derive events yourself with the detectors in
44
- `@dbx-tools/genie-shared`).
75
+ ## Stream Raw Message Snapshots
45
76
 
46
77
  ```ts
47
- import { genieChat } from "@dbx-tools/genie";
78
+ import { chat } from "@dbx-tools/genie";
48
79
 
49
- // Single turn.
50
- for await (const m of genieChat(spaceId, "Top 5 stores?")) {
51
- render(m);
52
- }
53
-
54
- // Multi-turn: caller threads the conversation id.
55
- let conversationId: string | undefined;
56
- for (const question of questions) {
57
- for await (const m of genieChat(spaceId, question, { conversationId })) {
58
- conversationId = m.conversation_id ?? conversationId;
59
- render(m);
60
- }
80
+ for await (const message of chat.genieChat(spaceId, "Top stores by revenue?")) {
81
+ renderSnapshot(message);
61
82
  }
62
83
  ```
63
84
 
64
- Turn lifecycle:
65
-
66
- - **No `options.conversationId`**: opens a new conversation via
67
- `client.genie.startConversation`. The assigned id surfaces on
68
- every yielded `GenieMessage.conversation_id`.
69
- - **With `options.conversationId`**: appends to that conversation
70
- via `client.genie.createMessage`.
71
- - In both cases, after the create/start the driver polls
72
- `client.genie.getMessage` every `options.pollIntervalMs` (default
73
- 500ms) until the message hits a terminal status
74
- (`COMPLETED` / `FAILED` / `CANCELLED`), then yields the terminal
75
- snapshot and returns.
76
-
77
- Identical consecutive snapshots are filtered out (deep equal)
78
- because Genie often returns the exact same payload twice during
79
- quiet periods.
80
-
81
- ## `genieEventChat` - high-level typed events
82
-
83
- Wraps `genieChat` and yields a `GenieChatEvent` discriminated union.
84
- Stream order per turn:
85
-
86
- 1. `{ type: "message", message }` - the raw `GenieMessage`, once
87
- per poll yield.
88
- 2. `{ type: "question", content, message_id, ... }` - fires
89
- exactly once, on the first `message` yield. Carries the prompt
90
- text Genie echoed back and the assigned `message_id` so
91
- subscribers can group everything for one Genie call under that
92
- one key.
93
- 3. Any of `status` / `attachment` / `thinking` / `text` / `query` /
94
- `statement` / `rows` / `suggested_questions` the diff against
95
- the prior snapshot produced.
96
- 4. On the terminal snapshot, `{ type: "result", status, message }`
97
- as the final yield.
98
-
99
- See [`@dbx-tools/genie-shared`](../genie-shared) for the full event
100
- catalogue, field shapes, and the pure detectors that derive each
101
- event from a snapshot diff.
102
-
103
- Errors propagate by the generator throwing - there is no `error`
104
- variant. Wrap the `for await` in `try / catch` if you need to handle
105
- failures.
106
-
107
- ## Space metadata
108
-
109
- Beyond the chat drivers, the package exposes two helpers for reading a
110
- space's definition (both re-exported from the package root):
111
-
112
- - `getGenieSpace(spaceId, options?)` - fetch a `GenieSpace` via
113
- `GET /api/2.0/genie/spaces/<id>`. Includes the `serialized_space`
114
- blob (catalogs, tables, sample questions, prompts) by default; pass
115
- `{ serialized: false }` for the lighter title/description-only
116
- payload. Takes the same `workspaceClient` / `context` options as the
117
- drivers.
118
- - `genieSampleQuestions(space)` - pull the author-curated starter
119
- questions out of a space's serialized blob
120
- (`config.sample_questions[*].question`), order preserved and
121
- duplicates dropped. Returns `[]` when the space carries no serialized
122
- blob or no configured questions, so a missing/misconfigured space
123
- degrades to "no suggestions" rather than throwing.
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
124
91
 
125
92
  ```ts
126
- import { getGenieSpace, genieSampleQuestions } from "@dbx-tools/genie";
93
+ let conversationId: string | undefined;
127
94
 
128
- const space = await getGenieSpace(spaceId);
129
- const prompts = genieSampleQuestions(space); // string[]
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
+ }
130
102
  ```
131
103
 
132
- ## Options
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.
133
107
 
134
- `GenieChatOptions` is the same shape for both drivers:
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.
135
111
 
136
- | Option | Default | Description |
137
- | ----------------- | --------------------------------------- | ----------------------------------------------------------------------------------------------------------- |
138
- | `conversationId` | `undefined` | Seed conversation id. When set, this turn appends to the existing conversation via `createMessage`. |
139
- | `workspaceClient` | resolved (see below) | Explicit `WorkspaceClient`. Defaults to AppKit's per-request client when available; otherwise env-var auth. |
140
- | `pollIntervalMs` | `500` | Cadence between successive `getMessage` calls. |
141
- | `context` | new internal `AbortController` per call | External cancellation. Accepts an `AbortSignal` or a fully-built SDK `Context` (`apiUtils.ContextLike`). |
112
+ ## Resolve A Workspace Client
142
113
 
143
- ### Workspace client resolution
114
+ ```ts
115
+ import { WorkspaceClient } from "@databricks/sdk-experimental";
116
+ import { chat } from "@dbx-tools/genie";
144
117
 
145
- Resolved in priority order:
118
+ await chat.genieEventChat(spaceId, content, {
119
+ workspaceClient: new WorkspaceClient({ profile: "dev" }),
120
+ });
121
+ ```
146
122
 
147
- 1. Caller-supplied `options.workspaceClient`.
148
- 2. AppKit's per-request execution-context client, when
149
- `@databricks/appkit` is installed AND we're inside an active
150
- request scope. OBO auth is preserved automatically.
151
- 3. Fresh `new WorkspaceClient({})` (env-var auth via
152
- `DATABRICKS_CONFIG_PROFILE` / `DATABRICKS_HOST` /
153
- `DATABRICKS_TOKEN`).
123
+ Client resolution order:
154
124
 
155
- AppKit is loaded lazily, so this package is usable from non-AppKit
156
- environments (smoke scripts, batch jobs, tests, ...).
125
+ 1. `options.workspaceClient`;
126
+ 2. AppKit execution-context client, when present;
127
+ 3. `new WorkspaceClient({})` using normal Databricks SDK auth.
157
128
 
158
- ### Cancellation
129
+ Pass `options.context` as an `AbortSignal` or SDK context to cancel SDK calls and
130
+ the polling sleep.
159
131
 
160
- A single internal `AbortController` covers each call. `options.context`
161
- ties into that controller, so an external abort tears down every
162
- in-flight SDK call AND the inter-poll sleep. Breaking out of the
163
- `for await` does the same via the `try / finally`.
132
+ ## Read Space Metadata And Starter Questions
164
133
 
165
134
  ```ts
166
- const ac = new AbortController();
167
- setTimeout(() => ac.abort(), 30_000); // hard 30s ceiling
135
+ import { space } from "@dbx-tools/genie";
168
136
 
169
- try {
170
- for await (const event of genieEventChat(spaceId, content, { context: ac.signal })) {
171
- handle(event);
172
- }
173
- } catch (err) {
174
- if (ac.signal.aborted) console.log("timed out");
175
- else throw err;
176
- }
137
+ const genieSpace = await space.getGenieSpace(spaceId);
138
+ const questions = space.genieSampleQuestions(genieSpace);
177
139
  ```
178
140
 
179
- ## Smoke test
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.
180
144
 
181
- [`test/poll-chat.ts`](test/poll-chat.ts) drives `genieEventChat` from
182
- argv / stdin / a REPL and writes every emitted event to a per-run tmp
183
- directory so you can inspect deltas after the fact. One subdirectory
184
- per `GenieChatEvent` variant, plus a `rows-data/` directory for paired
185
- SQL fetches on terminals that carry a `query_result.statement_id`.
145
+ ## Options
186
146
 
187
- ```bash
188
- # Required env (already in repo .env):
189
- # DATABRICKS_GENIE_SPACE_ID
190
- # DATABRICKS_CONFIG_PROFILE (or any SDK auth env)
147
+ `chat.GenieChatOptions` is shared by both drivers:
191
148
 
192
- bun packages/genie/test/poll-chat.ts # REPL
193
- bun packages/genie/test/poll-chat.ts "Top 5 stores by revenue?"
194
- echo "Top 5 stores by revenue?" | bun packages/genie/test/poll-chat.ts
195
- ```
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.
196
153
 
197
- Stdout carries a one-line summary per event with the relative path;
198
- the per-run directory is printed at the top.
154
+ ## Modules
199
155
 
200
- ## License
156
+ - `chat` - `genieChat()` raw snapshot stream and `genieEventChat()` typed event
157
+ stream.
158
+ - `space` - `getGenieSpace()` and `genieSampleQuestions()`.
201
159
 
202
- Apache-2.0
160
+ The AppKit-Mastra package builds its Genie tools on top of this driver; see
161
+ [`@dbx-tools/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 CHANGED
@@ -1,25 +1,57 @@
1
1
  {
2
2
  "name": "@dbx-tools/genie",
3
- "version": "0.1.111",
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
+ },
4
17
  "dependencies": {
5
18
  "@databricks/sdk-experimental": "^0.17.0",
6
- "@dbx-tools/genie-shared": "0.1.111",
7
- "@dbx-tools/shared": "0.1.111"
19
+ "@dbx-tools/appkit": "0.3.1",
20
+ "@dbx-tools/shared-genie": "0.3.1",
21
+ "@dbx-tools/shared-core": "0.3.1"
22
+ },
23
+ "main": "index.ts",
24
+ "license": "UNLICENSED",
25
+ "publishConfig": {
26
+ "access": "public"
8
27
  },
28
+ "version": "0.3.1",
29
+ "types": "index.ts",
30
+ "type": "module",
9
31
  "exports": {
10
- ".": {
11
- "types": "./dist/index.d.ts",
12
- "default": "./dist/index.js"
32
+ ".": "./index.ts",
33
+ "./package.json": "./package.json"
34
+ },
35
+ "peerDependenciesMeta": {
36
+ "@databricks/appkit": {
37
+ "optional": true
13
38
  }
14
39
  },
15
- "type": "module",
16
- "main": "./dist/index.js",
17
- "types": "./dist/index.d.ts",
18
- "files": [
19
- "dist"
20
- ],
21
- "license": "Apache-2.0",
22
- "publishConfig": {
23
- "access": "public"
40
+ "dbxToolsConfig": {
41
+ "tags": [
42
+ "node"
43
+ ]
44
+ },
45
+ "//": "~~ Generated by projen. To modify, edit .projenrc.js and run \"pnpm exec projen\".",
46
+ "scripts": {
47
+ "build": "projen build",
48
+ "compile": "projen compile",
49
+ "default": "projen default",
50
+ "package": "projen package",
51
+ "post-compile": "projen post-compile",
52
+ "pre-compile": "projen pre-compile",
53
+ "test": "projen test",
54
+ "watch": "projen watch",
55
+ "projen": "projen"
24
56
  }
25
- }
57
+ }
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,111 @@
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, object, 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
+ return [
105
+ ...object
106
+ .sequence(sampleQuestions as SerializedSampleQuestion[])
107
+ .map((entry) => string.firstNonEmpty(entry?.question))
108
+ .nonNull()
109
+ .distinct(),
110
+ ];
111
+ }
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
+ }
package/dist/index.d.ts DELETED
@@ -1,157 +0,0 @@
1
- import { GenieChatEvent, GenieMessage, GenieSpace } from "@dbx-tools/genie-shared";
2
- import { WorkspaceClient } from "@databricks/sdk-experimental";
3
- import { apiUtils } from "@dbx-tools/shared";
4
- export * from "@dbx-tools/genie-shared";
5
-
6
- //#region packages/genie/src/chat.d.ts
7
- /** Options accepted by both {@link genieChat} and {@link genieEventChat}. */
8
- interface GenieChatOptions {
9
- /**
10
- * Seed conversation id. When set, this turn appends to the
11
- * existing conversation (via `createMessage`) instead of opening
12
- * a new one. Use it to thread a multi-turn conversation: read
13
- * `conversation_id` off the prior turn's terminal `GenieMessage`
14
- * (or the `result` event's `payload.conversation_id`) and pass
15
- * it into the next call.
16
- */
17
- conversationId?: string;
18
- /**
19
- * Explicit `WorkspaceClient`. Defaults to AppKit's per-request
20
- * execution-context client when AppKit is installed and we're
21
- * inside a request; falls back to a fresh `new WorkspaceClient({})`
22
- * (env-var auth) otherwise.
23
- */
24
- workspaceClient?: WorkspaceClient;
25
- /** Poll cadence in milliseconds between successive `getMessage` calls (default 500). */
26
- pollIntervalMs?: number;
27
- /**
28
- * External cancellation. Accepts a WHATWG `AbortSignal` or a
29
- * fully-built SDK `Context` (see `apiUtils.ContextLike`).
30
- * Aborting it cancels every in-flight SDK call and the next
31
- * inter-poll sleep.
32
- */
33
- context?: apiUtils.ContextLike;
34
- }
35
- /**
36
- * One turn against a Genie space, yielded as a stream of
37
- * `GenieMessage` snapshots.
38
- *
39
- * Turn lifecycle:
40
- *
41
- * - No `options.conversationId`: open a new conversation via
42
- * `client.genie.startConversation`. The opened conversation id
43
- * surfaces on every yielded `GenieMessage` (`.conversation_id`)
44
- * so the caller can thread it into a follow-up call.
45
- * - With `options.conversationId`: append to that conversation
46
- * via `client.genie.createMessage`.
47
- * - In both cases, after the create/start the driver polls
48
- * `client.genie.getMessage` every `options.pollIntervalMs`
49
- * (default 500ms) until the message reaches a terminal
50
- * status, then yields the terminal snapshot and returns.
51
- *
52
- * Cancellation: a single internal `AbortController` covers the
53
- * whole turn. `options.context` is tied into that controller so an
54
- * external abort tears down every in-flight SDK call AND the
55
- * inter-poll sleep. Breaking out of the `for await` does the same
56
- * via the `try / finally`.
57
- *
58
- * @example
59
- * // Single turn.
60
- * for await (const m of genieChat(spaceId, "Top 5 stores?")) {
61
- * render(m);
62
- * }
63
- *
64
- * @example
65
- * // Multi-turn: caller threads the conversation id.
66
- * let conversationId: string | undefined;
67
- * for (const question of questions) {
68
- * for await (const m of genieChat(spaceId, question, { conversationId })) {
69
- * conversationId = m.conversation_id ?? conversationId;
70
- * render(m);
71
- * }
72
- * }
73
- */
74
- declare function genieChat(space_id: string, content: string, options?: GenieChatOptions): AsyncGenerator<GenieMessage, void, void>;
75
- /**
76
- * One turn against a Genie space, yielded as a typed
77
- * {@link GenieChatEvent} stream. Drives {@link genieChat}
78
- * underneath and decorates each snapshot with the derived events
79
- * the field-level diff produced. Stream order:
80
- *
81
- * 1. `{ type: "message", message }` - the raw `GenieMessage`,
82
- * once per poll yield.
83
- * 2. `{ type: "question", content, message_id, ... }` fires
84
- * exactly once, on the FIRST `message` yield. We read
85
- * `content` and `message_id` straight off the snapshot so
86
- * every downstream event for this turn shares the same
87
- * `message_id` (the question included) - subscribers can
88
- * group everything for one Genie call under that one key.
89
- * 3. Any of `status` / `attachment` / `thinking` / `text` /
90
- * `query` / `statement` / `rows` / `suggested_questions` the
91
- * diff against the prior snapshot produced.
92
- * 4. On the terminal snapshot, `{ type: "result", ... }` as
93
- * the final yield.
94
- *
95
- * Errors propagate by the generator throwing - there's no
96
- * `"error"` variant. Wrap the `for await` in `try/catch` if you
97
- * need to handle failures.
98
- *
99
- * @example
100
- * for await (const event of genieEventChat(spaceId, "Top stores?")) {
101
- * switch (event.type) {
102
- * case "thinking":
103
- * console.log("[thinking]", event.thought_type, event.text);
104
- * break;
105
- * case "text":
106
- * console.log("[text]", event.text);
107
- * break;
108
- * case "result":
109
- * console.log("[done]", event.status);
110
- * break;
111
- * }
112
- * }
113
- */
114
- declare function genieEventChat(space_id: string, content: string, options?: GenieChatOptions): AsyncGenerator<GenieChatEvent, void, void>;
115
- //#endregion
116
- //#region packages/genie/src/space.d.ts
117
- /** Options for {@link getGenieSpace}. */
118
- interface GetGenieSpaceOptions {
119
- /**
120
- * Explicit `WorkspaceClient`. Defaults to a fresh
121
- * `new WorkspaceClient({})` (env-var auth). Server callers should
122
- * pass their OBO-scoped client so the lookup runs as the user.
123
- */
124
- workspaceClient?: WorkspaceClient;
125
- /**
126
- * Request the `serialized_space` blob (catalogs, tables, sample
127
- * questions, prompts). Defaults to `true` - the only reason to
128
- * skip it is when the caller just needs title / description and
129
- * wants the smaller payload.
130
- */
131
- serialized?: boolean;
132
- /**
133
- * External cancellation. Accepts a WHATWG `AbortSignal` or a
134
- * fully-built SDK `Context` (see `apiUtils.ContextLike`).
135
- */
136
- context?: apiUtils.ContextLike;
137
- }
138
- /**
139
- * Fetch a Genie space by id, optionally including its serialized
140
- * definition. Hits `GET /api/2.0/genie/spaces/<id>` with
141
- * `include_serialized_space=true` through the raw `apiClient`, then
142
- * validates the response against {@link GenieSpaceSchema} (unknown
143
- * fields like `etag` / `parent_path` are stripped).
144
- */
145
- declare function getGenieSpace(spaceId: string, options?: GetGenieSpaceOptions): Promise<GenieSpace>;
146
- /**
147
- * Extract the curated starter questions an author configured on a
148
- * Genie space. Reads `serialized_space -> config.sample_questions[*]
149
- * .question`. Returns `[]` when the space carries no serialized blob,
150
- * the blob is unparseable, or no sample questions are configured -
151
- * so a missing or misconfigured space degrades to "no suggestions"
152
- * rather than throwing. Order is preserved (the author's ordering)
153
- * and duplicates are dropped.
154
- */
155
- declare function genieSampleQuestions(space: GenieSpace): string[];
156
- //#endregion
157
- export { GenieChatOptions, GetGenieSpaceOptions, genieChat, genieEventChat, genieSampleQuestions, getGenieSpace };
package/dist/index.js DELETED
@@ -1,307 +0,0 @@
1
- import { GenieMessageSchema, GenieSpaceSchema, eventsFromMessage, isTerminalStatus } from "@dbx-tools/genie-shared";
2
- import { WorkspaceClient } from "@databricks/sdk-experimental";
3
- import { apiUtils, commonUtils, logUtils, stringUtils } from "@dbx-tools/shared";
4
-
5
- export * from "@dbx-tools/genie-shared"
6
-
7
- //#region packages/genie/src/chat.ts
8
- /**
9
- * `@dbx-tools/genie` chat driver.
10
- *
11
- * Drives a single turn against a Genie space from one `content`
12
- * string; multi-turn conversations are the caller's job (thread
13
- * the `conversation_id` returned on each `GenieMessage` back into
14
- * the next turn's `options.conversationId`).
15
- *
16
- * Two layers serve two kinds of consumer. The low-level layer
17
- * yields every poll-observed `GenieMessage` (validated against
18
- * `GenieMessageSchema`, falling back to the raw snapshot on a schema
19
- * miss) and owns the messy parts - cancellation, conversation
20
- * seeding, distinct-filtering, and SDK quirks (Waiter stripping);
21
- * reach for it when you want the raw stream. The high-level layer wraps it
22
- * and emits semantic, deduplicated `{ type, payload }` events
23
- * (see {@link GenieChatEvent}), always closing a successful turn
24
- * with a terminal `result` event carrying the final
25
- * `GenieMessage`; errors propagate by throwing, with no `error`
26
- * variant. Iterating UI / agent code that wants every message
27
- * verbatim takes the low-level stream; subscribers reacting to
28
- * "Genie is thinking about X" or "Genie produced text Y" take the
29
- * event layer.
30
- */
31
- const log$1 = logUtils.logger("genie/chat");
32
- /**
33
- * Validate a polled wire snapshot against {@link GenieMessageSchema}
34
- * and return the schema-normalized message. Genie's wire occasionally
35
- * ships a shape the (SDK-derived) schema doesn't model exactly - e.g.
36
- * an early poll that omits the SDK-required `message_id` - so a miss
37
- * degrades to the raw snapshot rather than throwing, keeping a single
38
- * odd poll from aborting the whole turn.
39
- */
40
- function validateMessage(raw) {
41
- const result = GenieMessageSchema.safeParse(raw);
42
- if (result.success) return result.data;
43
- log$1.debug("wire-message:schema-miss", {
44
- message_id: raw.message_id ?? raw.id,
45
- issues: result.error.issues.length
46
- });
47
- return raw;
48
- }
49
- /**
50
- * One turn against a Genie space, yielded as a stream of
51
- * `GenieMessage` snapshots.
52
- *
53
- * Turn lifecycle:
54
- *
55
- * - No `options.conversationId`: open a new conversation via
56
- * `client.genie.startConversation`. The opened conversation id
57
- * surfaces on every yielded `GenieMessage` (`.conversation_id`)
58
- * so the caller can thread it into a follow-up call.
59
- * - With `options.conversationId`: append to that conversation
60
- * via `client.genie.createMessage`.
61
- * - In both cases, after the create/start the driver polls
62
- * `client.genie.getMessage` every `options.pollIntervalMs`
63
- * (default 500ms) until the message reaches a terminal
64
- * status, then yields the terminal snapshot and returns.
65
- *
66
- * Cancellation: a single internal `AbortController` covers the
67
- * whole turn. `options.context` is tied into that controller so an
68
- * external abort tears down every in-flight SDK call AND the
69
- * inter-poll sleep. Breaking out of the `for await` does the same
70
- * via the `try / finally`.
71
- *
72
- * @example
73
- * // Single turn.
74
- * for await (const m of genieChat(spaceId, "Top 5 stores?")) {
75
- * render(m);
76
- * }
77
- *
78
- * @example
79
- * // Multi-turn: caller threads the conversation id.
80
- * let conversationId: string | undefined;
81
- * for (const question of questions) {
82
- * for await (const m of genieChat(spaceId, question, { conversationId })) {
83
- * conversationId = m.conversation_id ?? conversationId;
84
- * render(m);
85
- * }
86
- * }
87
- */
88
- async function* genieChat(space_id, content, options) {
89
- const controller = new AbortController();
90
- try {
91
- const client = await getWorkspaceClient(options);
92
- const context = apiUtils.toContext(controller, options?.context);
93
- let conversationId = options?.conversationId;
94
- let messageId;
95
- const pollProducer = async (ctx) => {
96
- if (!conversationId) {
97
- if (ctx.attempt > 0) throw new Error("Genie did not return a conversation id; refusing to retry");
98
- const startResponse = await client.genie.startConversation({
99
- space_id,
100
- content
101
- }, context);
102
- conversationId = startResponse.conversation_id;
103
- messageId = startResponse.message_id;
104
- return startResponse.message;
105
- }
106
- if (!messageId) {
107
- const { wait: _wait, ...createResponse } = await client.genie.createMessage({
108
- space_id,
109
- conversation_id: conversationId,
110
- content
111
- }, context);
112
- messageId = createResponse.message_id;
113
- return createResponse;
114
- }
115
- return await client.genie.getMessage({
116
- space_id,
117
- conversation_id: conversationId,
118
- message_id: messageId
119
- }, context);
120
- };
121
- yield* commonUtils.poll(async (ctx) => validateMessage(await pollProducer(ctx)), {
122
- intervalMs: options?.pollIntervalMs ?? 500,
123
- filter: "distinct",
124
- predicate: (m) => !isTerminalStatus(m.status),
125
- signal: controller.signal
126
- });
127
- } finally {
128
- controller.abort();
129
- }
130
- }
131
- /**
132
- * One turn against a Genie space, yielded as a typed
133
- * {@link GenieChatEvent} stream. Drives {@link genieChat}
134
- * underneath and decorates each snapshot with the derived events
135
- * the field-level diff produced. Stream order:
136
- *
137
- * 1. `{ type: "message", message }` - the raw `GenieMessage`,
138
- * once per poll yield.
139
- * 2. `{ type: "question", content, message_id, ... }` fires
140
- * exactly once, on the FIRST `message` yield. We read
141
- * `content` and `message_id` straight off the snapshot so
142
- * every downstream event for this turn shares the same
143
- * `message_id` (the question included) - subscribers can
144
- * group everything for one Genie call under that one key.
145
- * 3. Any of `status` / `attachment` / `thinking` / `text` /
146
- * `query` / `statement` / `rows` / `suggested_questions` the
147
- * diff against the prior snapshot produced.
148
- * 4. On the terminal snapshot, `{ type: "result", ... }` as
149
- * the final yield.
150
- *
151
- * Errors propagate by the generator throwing - there's no
152
- * `"error"` variant. Wrap the `for await` in `try/catch` if you
153
- * need to handle failures.
154
- *
155
- * @example
156
- * for await (const event of genieEventChat(spaceId, "Top stores?")) {
157
- * switch (event.type) {
158
- * case "thinking":
159
- * console.log("[thinking]", event.thought_type, event.text);
160
- * break;
161
- * case "text":
162
- * console.log("[text]", event.text);
163
- * break;
164
- * case "result":
165
- * console.log("[done]", event.status);
166
- * break;
167
- * }
168
- * }
169
- */
170
- async function* genieEventChat(space_id, content, options) {
171
- let previous;
172
- let questionEmitted = false;
173
- for await (const rawMessage of genieChat(space_id, content, options)) {
174
- const message = rawMessage.message_id ? rawMessage : {
175
- ...rawMessage,
176
- message_id: rawMessage.id
177
- };
178
- yield {
179
- type: "message",
180
- space_id: message.space_id,
181
- message_id: message.message_id,
182
- message
183
- };
184
- if (!questionEmitted) {
185
- yield {
186
- type: "question",
187
- space_id: message.space_id,
188
- ...message.conversation_id ? { conversation_id: message.conversation_id } : {},
189
- ...message.message_id ? { message_id: message.message_id } : {},
190
- content: message.content
191
- };
192
- questionEmitted = true;
193
- }
194
- yield* eventsFromMessage(message, previous, message.space_id);
195
- if (isTerminalStatus(message.status)) yield {
196
- type: "result",
197
- space_id: message.space_id,
198
- conversation_id: message.conversation_id,
199
- message_id: message.message_id,
200
- status: message.status,
201
- message
202
- };
203
- previous = message;
204
- }
205
- }
206
- /**
207
- * Resolve a `WorkspaceClient` in this preference order:
208
- *
209
- * 1. Caller-supplied `options.workspaceClient`.
210
- * 2. AppKit's per-request execution-context client, when AppKit
211
- * is installed AND we're inside a request scope.
212
- * 3. Fresh `new WorkspaceClient({})` (env-var auth via
213
- * `DATABRICKS_CONFIG_PROFILE` / `DATABRICKS_HOST` /
214
- * `DATABRICKS_TOKEN`).
215
- *
216
- * AppKit is loaded lazily so this package stays usable in
217
- * non-AppKit environments (e.g. the `poll-chat` smoke test).
218
- */
219
- async function getWorkspaceClient(options) {
220
- if (options?.workspaceClient) return options.workspaceClient;
221
- const appkit = await getAppKit();
222
- if (appkit) try {
223
- return appkit.getExecutionContext().client;
224
- } catch {}
225
- return new WorkspaceClient({});
226
- }
227
- async function getAppKit() {
228
- try {
229
- return await import("@databricks/appkit");
230
- } catch {
231
- return;
232
- }
233
- }
234
-
235
- //#endregion
236
- //#region packages/genie/src/space.ts
237
- /**
238
- * `@dbx-tools/genie` space metadata helpers.
239
- *
240
- * Fetches a Genie space's definition (including the opt-in
241
- * `serialized_space` blob) and extracts the curated starter
242
- * questions an author configured on the space. The typed SDK
243
- * `client.genie.getSpace` only returns the directory-listing surface
244
- * (`title` / `description` / `warehouse_id`); the sample questions
245
- * live inside `serialized_space`, which the REST API returns only
246
- * when `include_serialized_space=true`. We hit that endpoint through
247
- * the workspace client's raw `apiClient` since the typed request
248
- * shape has no flag for it.
249
- */
250
- const log = logUtils.logger("genie/space");
251
- /**
252
- * Fetch a Genie space by id, optionally including its serialized
253
- * definition. Hits `GET /api/2.0/genie/spaces/<id>` with
254
- * `include_serialized_space=true` through the raw `apiClient`, then
255
- * validates the response against {@link GenieSpaceSchema} (unknown
256
- * fields like `etag` / `parent_path` are stripped).
257
- */
258
- async function getGenieSpace(spaceId, options) {
259
- const client = options?.workspaceClient ?? new WorkspaceClient({});
260
- const serialized = options?.serialized !== false;
261
- const context = options?.context ? apiUtils.toContext(options.context) : void 0;
262
- const raw = await client.apiClient.request({
263
- path: `/api/2.0/genie/spaces/${encodeURIComponent(spaceId)}`,
264
- method: "GET",
265
- query: serialized ? { include_serialized_space: true } : {},
266
- headers: new Headers(),
267
- raw: false
268
- }, context);
269
- return GenieSpaceSchema.parse(raw);
270
- }
271
- /**
272
- * Extract the curated starter questions an author configured on a
273
- * Genie space. Reads `serialized_space -> config.sample_questions[*]
274
- * .question`. Returns `[]` when the space carries no serialized blob,
275
- * the blob is unparseable, or no sample questions are configured -
276
- * so a missing or misconfigured space degrades to "no suggestions"
277
- * rather than throwing. Order is preserved (the author's ordering)
278
- * and duplicates are dropped.
279
- */
280
- function genieSampleQuestions(space) {
281
- const serialized = space.serialized_space;
282
- if (!serialized) return [];
283
- let parsed;
284
- try {
285
- parsed = JSON.parse(serialized);
286
- } catch (err) {
287
- log.warn("serialized-space:parse-error", {
288
- spaceId: space.space_id,
289
- error: commonUtils.errorMessage(err)
290
- });
291
- return [];
292
- }
293
- const sampleQuestions = parsed?.config?.sample_questions;
294
- if (!Array.isArray(sampleQuestions)) return [];
295
- const seen = /* @__PURE__ */ new Set();
296
- const out = [];
297
- for (const entry of sampleQuestions) {
298
- const text = stringUtils.firstNonEmpty(entry?.question);
299
- if (!text || seen.has(text)) continue;
300
- seen.add(text);
301
- out.push(text);
302
- }
303
- return out;
304
- }
305
-
306
- //#endregion
307
- export { genieChat, genieEventChat, genieSampleQuestions, getGenieSpace };