@dbx-tools/genie 0.1.67 → 0.1.68

package/package.json

@@ -1,14 +1,13 @@
 {
   "name": "@dbx-tools/genie",
-  "version": "0.1.67",
+  "version": "0.1.68",
   "dependencies": {
     "@databricks/sdk-experimental": "^0.17",
-    "@dbx-tools/genie-shared": "0.1.67",
-    "@dbx-tools/shared": "0.1.67"
+    "@dbx-tools/genie-shared": "0.1.68",
+    "@dbx-tools/shared": "0.1.68"
   },
   "exports": {
     ".": {
-      "source": "./src/index.ts",
       "types": "./dist/index.d.ts",
       "default": "./dist/index.js"
     }
@@ -17,8 +16,7 @@
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "files": [
-    "dist",
-    "src"
+    "dist"
   ],
   "license": "Apache-2.0"
 }

package/src/chat.ts

@@ -1,353 +0,0 @@
-/**
- * `@dbx-tools/genie` chat driver.
- *
- * Drives a single turn against a Genie space from one `content`
- * string; multi-turn conversations are the caller's job (thread
- * the `conversation_id` returned on each `GenieMessage` back into
- * the next turn's `options.conversationId`).
- *
- * Two layers serve two kinds of consumer. The low-level layer
- * yields every poll-observed `GenieMessage` (validated against
- * `GenieMessageSchema`, falling back to the raw snapshot on a schema
- * miss) and owns the messy parts - cancellation, conversation
- * seeding, distinct-filtering, and SDK quirks (Waiter stripping);
- * reach for it when you want the raw stream. The high-level layer wraps it
- * and emits semantic, deduplicated `{ type, payload }` events
- * (see {@link GenieChatEvent}), always closing a successful turn
- * with a terminal `result` event carrying the final
- * `GenieMessage`; errors propagate by throwing, with no `error`
- * variant. Iterating UI / agent code that wants every message
- * verbatim takes the low-level stream; subscribers reacting to
- * "Genie is thinking about X" or "Genie produced text Y" take the
- * event layer.
- */
-
-import { WorkspaceClient } from "@databricks/sdk-experimental";
-import {
-  eventsFromMessage,
-  GenieMessageSchema,
-  isTerminalStatus,
-  type GenieChatEvent,
-  type GenieMessage,
-} from "@dbx-tools/genie-shared";
-import { apiUtils, commonUtils, logUtils } from "@dbx-tools/shared";
-
-const log = logUtils.logger("genie/chat");
-
-/**
- * Validate a polled wire snapshot against {@link GenieMessageSchema}
- * and return the schema-normalized message. Genie's wire occasionally
- * ships a shape the (SDK-derived) schema doesn't model exactly - e.g.
- * an early poll that omits the SDK-required `message_id` - so a miss
- * degrades to the raw snapshot rather than throwing, keeping a single
- * odd poll from aborting the whole turn.
- */
-function validateMessage(raw: GenieMessage): GenieMessage {
-  const result = GenieMessageSchema.safeParse(raw);
-  if (result.success) return result.data;
-  log.debug("wire-message:schema-miss", {
-    message_id: raw.message_id ?? raw.id,
-    issues: result.error.issues.length,
-  });
-  return raw;
-}
-
-/* -------------------------- 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(
-      async (ctx: commonUtils.PollContext<GenieMessage>) =>
-        validateMessage(await pollProducer(ctx)),
-      {
-        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 rawMessage of genieChat(space_id, content, options)) {
-    // Normalize `message_id` from the legacy `id` field when Genie's
-    // wire response only populates one of them. The SDK schema marks
-    // both as required, but in practice the `startConversation` /
-    // `createMessage` inner `message` payload sometimes ships only
-    // `id` while the new `message_id` field lands undefined. Every
-    // downstream event detector keys grouping off `message_id`; the
-    // fallback keeps one Genie turn's events from splitting across
-    // an anon group + the real-id group when subscribers bucket by
-    // `message_id` (see `summarizeProgress` in the demo UI).
-    const message: GenieMessage = rawMessage.message_id
-      ? rawMessage
-      : { ...rawMessage, message_id: rawMessage.id };
-    yield {
-      type: "message",
-      space_id: message.space_id,
-      message_id: message.message_id,
-      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;
-  }
-}

package/src/index.ts

@@ -1,18 +0,0 @@
-/**
- * `@dbx-tools/genie` public surface.
- *
- * Bundles the package's Node-side drivers (live chat plus space
- * metadata / curated-question lookup) with a re-export of the pure
- * `@dbx-tools/genie-shared` wire vocabulary so a single
- * `from "@dbx-tools/genie"` import serves server-side consumers
- * that need both the live driver and the protocol types.
- *
- * Browser-side consumers should import `@dbx-tools/genie-shared`
- * directly: the chat driver pulls in `WorkspaceClient` and is
- * Node-only, whereas the re-exported protocol surface is pure
- * (types + sync functions) and safe for any runtime.
- */
-
-export * from "@dbx-tools/genie-shared";
-export * from "./chat.js";
-export * from "./space.js";

package/src/space.ts

@@ -1,130 +0,0 @@
-/**
- * `@dbx-tools/genie` space metadata helpers.
- *
- * Fetches a Genie space's definition (including the opt-in
- * `serialized_space` blob) and extracts the curated starter
- * questions an author configured on the space. The typed SDK
- * `client.genie.getSpace` only returns the directory-listing surface
- * (`title` / `description` / `warehouse_id`); the sample questions
- * live inside `serialized_space`, which the REST API returns only
- * when `include_serialized_space=true`. We hit that endpoint through
- * the workspace client's raw `apiClient` since the typed request
- * shape has no flag for it.
- */
-
-import { WorkspaceClient } from "@databricks/sdk-experimental";
-import { GenieSpaceSchema, type GenieSpace } from "@dbx-tools/genie-shared";
-import { apiUtils, commonUtils, logUtils } from "@dbx-tools/shared";
-
-const log = logUtils.logger("genie/space");
-
-/** Options for {@link getGenieSpace}. */
-export interface GetGenieSpaceOptions {
-  /**
-   * Explicit `WorkspaceClient`. Defaults to a fresh
-   * `new WorkspaceClient({})` (env-var auth). Server callers should
-   * pass their OBO-scoped client so the lookup runs as the user.
-   */
-  workspaceClient?: WorkspaceClient;
-  /**
-   * Request the `serialized_space` blob (catalogs, tables, sample
-   * questions, prompts). Defaults to `true` - the only reason to
-   * skip it is when the caller just needs title / description and
-   * wants the smaller payload.
-   */
-  serialized?: boolean;
-  /**
-   * External cancellation. Accepts a WHATWG `AbortSignal` or a
-   * fully-built SDK `Context` (see `apiUtils.ContextLike`).
-   */
-  context?: apiUtils.ContextLike;
-}
-
-/**
- * Fetch a Genie space by id, optionally including its serialized
- * definition. Hits `GET /api/2.0/genie/spaces/<id>` with
- * `include_serialized_space=true` through the raw `apiClient`, then
- * validates the response against {@link GenieSpaceSchema} (unknown
- * fields like `etag` / `parent_path` are stripped).
- */
-export async function getGenieSpace(
-  spaceId: string,
-  options?: GetGenieSpaceOptions,
-): Promise<GenieSpace> {
-  const client = options?.workspaceClient ?? new WorkspaceClient({});
-  const serialized = options?.serialized !== false;
-  const context = options?.context ? apiUtils.toContext(options.context) : undefined;
-  const raw = await client.apiClient.request(
-    {
-      path: `/api/2.0/genie/spaces/${encodeURIComponent(spaceId)}`,
-      method: "GET",
-      query: serialized ? { include_serialized_space: true } : {},
-      headers: new Headers(),
-      raw: false,
-    },
-    context,
-  );
-  return GenieSpaceSchema.parse(raw);
-}
-
-/**
- * One entry in a serialized space's `config.sample_questions`. The
- * author-facing field is `question`, which the wire format models as
- * a string array (a single multi-line question is split across
- * entries); we treat the first non-empty entry as the displayable
- * question text.
- */
-interface SerializedSampleQuestion {
-  question?: unknown;
-}
-
-/** Pull the first non-empty string out of a `question` field (string | string[]). */
-function questionText(question: unknown): string | undefined {
-  if (typeof question === "string") {
-    const trimmed = question.trim();
-    return trimmed.length > 0 ? trimmed : undefined;
-  }
-  if (Array.isArray(question)) {
-    for (const part of question) {
-      if (typeof part === "string" && part.trim().length > 0) return part.trim();
-    }
-  }
-  return undefined;
-}
-
-/**
- * Extract the curated starter questions an author configured on a
- * Genie space. Reads `serialized_space -> config.sample_questions[*]
- * .question`. Returns `[]` when the space carries no serialized blob,
- * the blob is unparseable, or no sample questions are configured -
- * so a missing or misconfigured space degrades to "no suggestions"
- * rather than throwing. Order is preserved (the author's ordering)
- * and duplicates are dropped.
- */
-export function genieSampleQuestions(space: GenieSpace): string[] {
-  const serialized = space.serialized_space;
-  if (!serialized) return [];
-  let parsed: unknown;
-  try {
-    parsed = JSON.parse(serialized);
-  } catch (err) {
-    log.warn("serialized-space:parse-error", {
-      spaceId: space.space_id,
-      error: commonUtils.errorMessage(err),
-    });
-    return [];
-  }
-  const sampleQuestions = (parsed as { config?: { sample_questions?: unknown } } | null)
-    ?.config?.sample_questions;
-  if (!Array.isArray(sampleQuestions)) return [];
-
-  const seen = new Set<string>();
-  const out: string[] = [];
-  for (const entry of sampleQuestions as SerializedSampleQuestion[]) {
-    const text = questionText(entry?.question);
-    if (!text || seen.has(text)) continue;
-    seen.add(text);
-    out.push(text);
-  }
-  return out;
-}