@dbx-tools/genie 0.1.18

package/dist/index.d.ts

@@ -0,0 +1,33 @@
+/**
+ * `@dbx-tools/genie` public surface.
+ *
+ *   - {@link genieChat}: low-level async generator. Yields every
+ *     poll-observed `GenieMessage` for a single turn against a
+ *     Genie space. Multi-turn conversations are driven by the
+ *     caller (thread the `conversation_id` off each yielded
+ *     message into the next call's `options.conversationId`).
+ *   - {@link genieEventChat}: high-level async generator. Drives
+ *     `genieChat` and yields a strongly-typed
+ *     {@link GenieChatEvent} stream of flat `{type, ...fields}`
+ *     records (`message`, `status`, `attachment`, `thinking`,
+ *     `text`, `query`, `statement`, `rows`,
+ *     `suggested_questions`, `result`). The final yield for any
+ *     terminal turn is always `{ type: "result", ... }`.
+ *
+ * Wire-format types, the {@link GenieChatEvent} discriminated
+ * union, per-event detectors (`detectStatus`, `detectThinking`,
+ * `detectAttachmentAdded`, `detectText`, `detectQuery`,
+ * `detectStatement`, `detectRows`, `detectSuggestedQuestions`),
+ * the `eventsFromMessage` sync generator, and terminal-status
+ * helpers all live in `@dbx-tools/genie-shared` and are
+ * re-exported here so a single `from "@dbx-tools/genie"` import
+ * works for server-side consumers. Browser-side consumers should
+ * import from `@dbx-tools/genie-shared` directly to avoid pulling
+ * in the Node-only `chat.ts` runtime.
+ *
+ * Browser safety: `chat.ts` pulls in `WorkspaceClient` and is
+ * Node-only; the re-exports from `@dbx-tools/genie-shared` are
+ * pure (types + sync functions) and safe for any runtime.
+ */
+export * from "./src/chat.js";
+export * from "@dbx-tools/genie-shared";

package/dist/index.js

@@ -0,0 +1,33 @@
+/**
+ * `@dbx-tools/genie` public surface.
+ *
+ *   - {@link genieChat}: low-level async generator. Yields every
+ *     poll-observed `GenieMessage` for a single turn against a
+ *     Genie space. Multi-turn conversations are driven by the
+ *     caller (thread the `conversation_id` off each yielded
+ *     message into the next call's `options.conversationId`).
+ *   - {@link genieEventChat}: high-level async generator. Drives
+ *     `genieChat` and yields a strongly-typed
+ *     {@link GenieChatEvent} stream of flat `{type, ...fields}`
+ *     records (`message`, `status`, `attachment`, `thinking`,
+ *     `text`, `query`, `statement`, `rows`,
+ *     `suggested_questions`, `result`). The final yield for any
+ *     terminal turn is always `{ type: "result", ... }`.
+ *
+ * Wire-format types, the {@link GenieChatEvent} discriminated
+ * union, per-event detectors (`detectStatus`, `detectThinking`,
+ * `detectAttachmentAdded`, `detectText`, `detectQuery`,
+ * `detectStatement`, `detectRows`, `detectSuggestedQuestions`),
+ * the `eventsFromMessage` sync generator, and terminal-status
+ * helpers all live in `@dbx-tools/genie-shared` and are
+ * re-exported here so a single `from "@dbx-tools/genie"` import
+ * works for server-side consumers. Browser-side consumers should
+ * import from `@dbx-tools/genie-shared` directly to avoid pulling
+ * in the Node-only `chat.ts` runtime.
+ *
+ * Browser safety: `chat.ts` pulls in `WorkspaceClient` and is
+ * Node-only; the re-exports from `@dbx-tools/genie-shared` are
+ * pure (types + sync functions) and safe for any runtime.
+ */
+export * from "./src/chat.js";
+export * from "@dbx-tools/genie-shared";

package/dist/src/chat.d.ts

@@ -0,0 +1,137 @@
+/**
+ * `@dbx-tools/genie` chat driver.
+ *
+ * Two async generators ship from this file. Both take a single
+ * `content` string for one turn against a Genie space; for
+ * multi-turn conversations, the caller drives the loop and
+ * threads the `conversation_id` returned on each
+ * `GenieMessage` into the next call's `options.conversationId`.
+ *
+ *   - {@link genieChat}: low-level. Yields every poll-observed
+ *     `GenieMessage` verbatim. Cancellation, conversation seeding,
+ *     distinct-filtering, and SDK quirks (Waiter stripping) live
+ *     here. Use directly when you want the raw stream.
+ *   - {@link genieEventChat}: high-level. Wraps `genieChat` and
+ *     emits semantic, deduplicated events as a typed
+ *     `{ type, payload }` stream (see {@link GenieChatEvent}). The
+ *     final yield for a successful turn is always
+ *     `{ type: "result", payload }` carrying the terminal
+ *     `GenieMessage`. Errors propagate by the generator throwing -
+ *     there is no `"error"` variant.
+ *
+ * Pick the layer that matches your consumer. Iterating UI / agent
+ * code that wants every message verbatim should use `genieChat`.
+ * Subscribers that want to react to "Genie is thinking about X"
+ * or "Genie produced text Y" should use `genieEventChat`.
+ */
+import { WorkspaceClient } from "@databricks/sdk-experimental";
+import { type GenieChatEvent, type GenieMessage } from "@dbx-tools/genie-shared";
+import { apiUtils } from "@dbx-tools/shared";
+/** Options accepted by both {@link genieChat} and {@link genieEventChat}. */
+export interface GenieChatOptions {
+    /**
+     * Seed conversation id. When set, this turn appends to the
+     * existing conversation (via `createMessage`) instead of opening
+     * a new one. Use it to thread a multi-turn conversation: read
+     * `conversation_id` off the prior turn's terminal `GenieMessage`
+     * (or the `result` event's `payload.conversation_id`) and pass
+     * it into the next call.
+     */
+    conversationId?: string;
+    /**
+     * Explicit `WorkspaceClient`. Defaults to AppKit's per-request
+     * execution-context client when AppKit is installed and we're
+     * inside a request; falls back to a fresh `new WorkspaceClient({})`
+     * (env-var auth) otherwise.
+     */
+    workspaceClient?: WorkspaceClient;
+    /** Poll cadence in milliseconds between successive `getMessage` calls (default 500). */
+    pollIntervalMs?: number;
+    /**
+     * External cancellation. Accepts a WHATWG `AbortSignal` or a
+     * fully-built SDK `Context` (see `apiUtils.ContextLike`).
+     * Aborting it cancels every in-flight SDK call and the next
+     * inter-poll sleep.
+     */
+    context?: apiUtils.ContextLike;
+}
+/**
+ * One turn against a Genie space, yielded as a stream of
+ * `GenieMessage` snapshots.
+ *
+ * Turn lifecycle:
+ *
+ *   - No `options.conversationId`: open a new conversation via
+ *     `client.genie.startConversation`. The opened conversation id
+ *     surfaces on every yielded `GenieMessage` (`.conversation_id`)
+ *     so the caller can thread it into a follow-up call.
+ *   - With `options.conversationId`: append to that conversation
+ *     via `client.genie.createMessage`.
+ *   - In both cases, after the create/start the driver polls
+ *     `client.genie.getMessage` every `options.pollIntervalMs`
+ *     (default 500ms) until the message reaches a terminal
+ *     status, then yields the terminal snapshot and returns.
+ *
+ * Cancellation: a single internal `AbortController` covers the
+ * whole turn. `options.context` is tied into that controller so an
+ * external abort tears down every in-flight SDK call AND the
+ * inter-poll sleep. Breaking out of the `for await` does the same
+ * via the `try / finally`.
+ *
+ * @example
+ * // Single turn.
+ * for await (const m of genieChat(spaceId, "Top 5 stores?")) {
+ *   render(m);
+ * }
+ *
+ * @example
+ * // Multi-turn: caller threads the conversation id.
+ * let conversationId: string | undefined;
+ * for (const question of questions) {
+ *   for await (const m of genieChat(spaceId, question, { conversationId })) {
+ *     conversationId = m.conversation_id ?? conversationId;
+ *     render(m);
+ *   }
+ * }
+ */
+export declare function genieChat(space_id: string, content: string, options?: GenieChatOptions): AsyncGenerator<GenieMessage, void, void>;
+/**
+ * One turn against a Genie space, yielded as a typed
+ * {@link GenieChatEvent} stream. Drives {@link genieChat}
+ * underneath and decorates each snapshot with the derived events
+ * the field-level diff produced. Stream order:
+ *
+ *   1. `{ type: "message", message }` - the raw `GenieMessage`,
+ *      once per poll yield.
+ *   2. `{ type: "question", content, message_id, ... }` fires
+ *      exactly once, on the FIRST `message` yield. We read
+ *      `content` and `message_id` straight off the snapshot so
+ *      every downstream event for this turn shares the same
+ *      `message_id` (the question included) - subscribers can
+ *      group everything for one Genie call under that one key.
+ *   3. Any of `status` / `attachment` / `thinking` / `text` /
+ *      `query` / `statement` / `rows` / `suggested_questions` the
+ *      diff against the prior snapshot produced.
+ *   4. On the terminal snapshot, `{ type: "result", ... }` as
+ *      the final yield.
+ *
+ * Errors propagate by the generator throwing - there's no
+ * `"error"` variant. Wrap the `for await` in `try/catch` if you
+ * need to handle failures.
+ *
+ * @example
+ * for await (const event of genieEventChat(spaceId, "Top stores?")) {
+ *   switch (event.type) {
+ *     case "thinking":
+ *       console.log("[thinking]", event.thought_type, event.text);
+ *       break;
+ *     case "text":
+ *       console.log("[text]", event.text);
+ *       break;
+ *     case "result":
+ *       console.log("[done]", event.status);
+ *       break;
+ *   }
+ * }
+ */
+export declare function genieEventChat(space_id: string, content: string, options?: GenieChatOptions): AsyncGenerator<GenieChatEvent, void, void>;

package/dist/src/chat.js

@@ -0,0 +1,251 @@
+/**
+ * `@dbx-tools/genie` chat driver.
+ *
+ * Two async generators ship from this file. Both take a single
+ * `content` string for one turn against a Genie space; for
+ * multi-turn conversations, the caller drives the loop and
+ * threads the `conversation_id` returned on each
+ * `GenieMessage` into the next call's `options.conversationId`.
+ *
+ *   - {@link genieChat}: low-level. Yields every poll-observed
+ *     `GenieMessage` verbatim. Cancellation, conversation seeding,
+ *     distinct-filtering, and SDK quirks (Waiter stripping) live
+ *     here. Use directly when you want the raw stream.
+ *   - {@link genieEventChat}: high-level. Wraps `genieChat` and
+ *     emits semantic, deduplicated events as a typed
+ *     `{ type, payload }` stream (see {@link GenieChatEvent}). The
+ *     final yield for a successful turn is always
+ *     `{ type: "result", payload }` carrying the terminal
+ *     `GenieMessage`. Errors propagate by the generator throwing -
+ *     there is no `"error"` variant.
+ *
+ * Pick the layer that matches your consumer. Iterating UI / agent
+ * code that wants every message verbatim should use `genieChat`.
+ * Subscribers that want to react to "Genie is thinking about X"
+ * or "Genie produced text Y" should use `genieEventChat`.
+ */
+import { WorkspaceClient } from "@databricks/sdk-experimental";
+import { eventsFromMessage, isTerminalStatus, } from "@dbx-tools/genie-shared";
+import { apiUtils, commonUtils } from "@dbx-tools/shared";
+/* ----------------------- low-level: genieChat ----------------------- */
+/**
+ * One turn against a Genie space, yielded as a stream of
+ * `GenieMessage` snapshots.
+ *
+ * Turn lifecycle:
+ *
+ *   - No `options.conversationId`: open a new conversation via
+ *     `client.genie.startConversation`. The opened conversation id
+ *     surfaces on every yielded `GenieMessage` (`.conversation_id`)
+ *     so the caller can thread it into a follow-up call.
+ *   - With `options.conversationId`: append to that conversation
+ *     via `client.genie.createMessage`.
+ *   - In both cases, after the create/start the driver polls
+ *     `client.genie.getMessage` every `options.pollIntervalMs`
+ *     (default 500ms) until the message reaches a terminal
+ *     status, then yields the terminal snapshot and returns.
+ *
+ * Cancellation: a single internal `AbortController` covers the
+ * whole turn. `options.context` is tied into that controller so an
+ * external abort tears down every in-flight SDK call AND the
+ * inter-poll sleep. Breaking out of the `for await` does the same
+ * via the `try / finally`.
+ *
+ * @example
+ * // Single turn.
+ * for await (const m of genieChat(spaceId, "Top 5 stores?")) {
+ *   render(m);
+ * }
+ *
+ * @example
+ * // Multi-turn: caller threads the conversation id.
+ * let conversationId: string | undefined;
+ * for (const question of questions) {
+ *   for await (const m of genieChat(spaceId, question, { conversationId })) {
+ *     conversationId = m.conversation_id ?? conversationId;
+ *     render(m);
+ *   }
+ * }
+ */
+export async function* genieChat(space_id, content, options) {
+    const controller = new AbortController();
+    try {
+        const client = await getWorkspaceClient(options);
+        // Build the SDK Context ONCE. Building it inside the poll
+        // producer would re-attach an abort listener to
+        // `options.context` on every poll iteration (via
+        // `apiUtils.toContext` -> `tieAbortSignal`), eventually
+        // tripping Node's `MaxListenersExceededWarning`.
+        const context = apiUtils.toContext(controller, options?.context);
+        let conversationId = options?.conversationId;
+        let messageId;
+        const pollProducer = async (ctx) => {
+            if (!conversationId) {
+                // First poll: open the conversation. Refuse to retry: if
+                // `startConversation` returned a response without a
+                // `conversation_id`, retrying would just open conversation
+                // after conversation.
+                if (ctx.attempt > 0) {
+                    throw new Error("Genie did not return a conversation id; refusing to retry");
+                }
+                const startResponse = await client.genie.startConversation({ space_id, content }, context);
+                conversationId = startResponse.conversation_id;
+                messageId = startResponse.message_id;
+                return startResponse.message;
+            }
+            if (!messageId) {
+                // First poll of a follow-up turn: append to the seeded
+                // conversation. `client.genie.createMessage` returns a
+                // `Waiter<GenieMessage>` (`{ ...message, wait: async () =>
+                // ... }`). Strip `wait` here so downstream serializers
+                // (e.g. yaml.stringify in poll-chat) don't choke on the
+                // AsyncFunction value.
+                const { wait: _wait, ...createResponse } = await client.genie.createMessage({ space_id, conversation_id: conversationId, content }, context);
+                messageId = createResponse.message_id;
+                return createResponse;
+            }
+            // Subsequent polls: re-fetch the current message until its
+            // status becomes terminal.
+            return await client.genie.getMessage({
+                space_id,
+                conversation_id: conversationId,
+                message_id: messageId,
+            }, context);
+        };
+        yield* commonUtils.poll(pollProducer, {
+            intervalMs: options?.pollIntervalMs ?? 500,
+            // Skip yielding identical consecutive snapshots; Genie
+            // often returns the exact same payload twice during quiet
+            // periods. `poll` does a deep equal on the previous yield.
+            filter: "distinct",
+            // Stop after the terminal message is yielded. `poll` checks
+            // the predicate AFTER yielding, so the terminal message
+            // still reaches the consumer.
+            predicate: (m) => !isTerminalStatus(m.status),
+            // Wake the inter-poll sleep on abort so a `for await` break
+            // (or external abort) tears down promptly instead of waiting
+            // out the interval.
+            signal: controller.signal,
+        });
+    }
+    finally {
+        // Cancels any still-pending SDK call and the inter-poll sleep
+        // whether we're unwinding from a normal return, a consumer
+        // break, or a thrown error. Idempotent.
+        controller.abort();
+    }
+}
+/* ---------------------- high-level: genieEventChat ------------------ */
+/**
+ * One turn against a Genie space, yielded as a typed
+ * {@link GenieChatEvent} stream. Drives {@link genieChat}
+ * underneath and decorates each snapshot with the derived events
+ * the field-level diff produced. Stream order:
+ *
+ *   1. `{ type: "message", message }` - the raw `GenieMessage`,
+ *      once per poll yield.
+ *   2. `{ type: "question", content, message_id, ... }` fires
+ *      exactly once, on the FIRST `message` yield. We read
+ *      `content` and `message_id` straight off the snapshot so
+ *      every downstream event for this turn shares the same
+ *      `message_id` (the question included) - subscribers can
+ *      group everything for one Genie call under that one key.
+ *   3. Any of `status` / `attachment` / `thinking` / `text` /
+ *      `query` / `statement` / `rows` / `suggested_questions` the
+ *      diff against the prior snapshot produced.
+ *   4. On the terminal snapshot, `{ type: "result", ... }` as
+ *      the final yield.
+ *
+ * Errors propagate by the generator throwing - there's no
+ * `"error"` variant. Wrap the `for await` in `try/catch` if you
+ * need to handle failures.
+ *
+ * @example
+ * for await (const event of genieEventChat(spaceId, "Top stores?")) {
+ *   switch (event.type) {
+ *     case "thinking":
+ *       console.log("[thinking]", event.thought_type, event.text);
+ *       break;
+ *     case "text":
+ *       console.log("[text]", event.text);
+ *       break;
+ *     case "result":
+ *       console.log("[done]", event.status);
+ *       break;
+ *   }
+ * }
+ */
+export async function* genieEventChat(space_id, content, options) {
+    // Diff source for the current turn. Always `undefined` on the
+    // first snapshot so the initial status / attachments emit
+    // fresh; updated to the most recent snapshot after each yield.
+    let previous;
+    // The `question` event is deferred to the first `message` yield
+    // so it can carry the assigned `message_id` (subscribers use it
+    // as the grouping key for every event in this turn). The first
+    // snapshot is the earliest point that id exists.
+    let questionEmitted = false;
+    for await (const message of genieChat(space_id, content, options)) {
+        yield { type: "message", message };
+        if (!questionEmitted) {
+            yield {
+                type: "question",
+                space_id: message.space_id,
+                ...(message.conversation_id
+                    ? { conversation_id: message.conversation_id }
+                    : {}),
+                ...(message.message_id ? { message_id: message.message_id } : {}),
+                content: message.content,
+            };
+            questionEmitted = true;
+        }
+        yield* eventsFromMessage(message, previous, message.space_id);
+        if (isTerminalStatus(message.status)) {
+            yield {
+                type: "result",
+                space_id: message.space_id,
+                conversation_id: message.conversation_id,
+                message_id: message.message_id,
+                status: message.status,
+                message,
+            };
+        }
+        previous = message;
+    }
+}
+/* ---------------------- workspace client helper --------------------- */
+/**
+ * Resolve a `WorkspaceClient` in this preference order:
+ *
+ *   1. Caller-supplied `options.workspaceClient`.
+ *   2. AppKit's per-request execution-context client, when AppKit
+ *      is installed AND we're inside a request scope.
+ *   3. Fresh `new WorkspaceClient({})` (env-var auth via
+ *      `DATABRICKS_CONFIG_PROFILE` / `DATABRICKS_HOST` /
+ *      `DATABRICKS_TOKEN`).
+ *
+ * AppKit is loaded lazily so this package stays usable in
+ * non-AppKit environments (e.g. the `poll-chat` smoke test).
+ */
+async function getWorkspaceClient(options) {
+    if (options?.workspaceClient)
+        return options.workspaceClient;
+    const appkit = await getAppKit();
+    if (appkit) {
+        try {
+            return appkit.getExecutionContext().client;
+        }
+        catch {
+            // Not inside an AppKit request context; fall through to env.
+        }
+    }
+    return new WorkspaceClient({});
+}
+async function getAppKit() {
+    try {
+        return await import("@databricks/appkit");
+    }
+    catch {
+        return undefined;
+    }
+}