@dbx-tools/genie 0.1.18

package/index.ts

@@ -0,0 +1,34 @@
+/**
+ * `@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/package.json

@@ -0,0 +1,35 @@
+{
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    ".": {
+      "source": "./index.ts",
+      "types": "./dist/index.d.ts",
+      "default": "./dist/index.js"
+    }
+  },
+  "name": "@dbx-tools/genie",
+  "version": "0.1.18",
+  "dependencies": {
+    "@databricks/sdk-experimental": "^0.17",
+    "@dbx-tools/genie-shared": "0.1.18",
+    "@dbx-tools/shared": "0.1.18"
+  },
+  "module": "index.ts",
+  "type": "module",
+  "files": [
+    "dist",
+    "index*.ts",
+    "src"
+  ],
+  "license": "Apache-2.0",
+  "homepage": "https://github.com/reggie-db/dbx-tools-appkit#readme",
+  "bugs": {
+    "url": "https://github.com/reggie-db/dbx-tools-appkit/issues"
+  },
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/reggie-db/dbx-tools-appkit.git",
+    "directory": "packages/genie"
+  }
+}

package/src/chat.ts

@@ -0,0 +1,317 @@
+/**
+ * `@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,
+  type GenieChatEvent,
+  type GenieMessage,
+} from "@dbx-tools/genie-shared";
+import { apiUtils, commonUtils } from "@dbx-tools/shared";
+
+/* -------------------------- shared options -------------------------- */
+
+/** 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;
+}
+
+/* ----------------------- 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: string,
+  content: string,
+  options?: GenieChatOptions,
+): AsyncGenerator<GenieMessage, void, void> {
+  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: string | undefined;
+
+    const pollProducer = async (
+      ctx: commonUtils.PollContext<GenieMessage>,
+    ): Promise<GenieMessage> => {
+      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: string,
+  content: string,
+  options?: GenieChatOptions,
+): AsyncGenerator<GenieChatEvent, void, void> {
+  // 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: GenieMessage | undefined;
+  // 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?: GenieChatOptions,
+): Promise<WorkspaceClient> {
+  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;
+  }
+}