@dbx-tools/shared-genie 0.1.9

package/src/genie-model.ts

@@ -0,0 +1,554 @@
+/**
+ * Wire-format zod schemas + types and high-level event vocabulary
+ * for the Genie chat driver.
+ *
+ * Two related layers live here:
+ *
+ *   1. Genie wire shapes derived from `@dbx-tools/shared-sdk-model`'s
+ *      generated `dashboards` schemas (regenerated from the upstream
+ *      `@databricks/sdk-experimental` `apis/dashboards/model.d.ts`
+ *      by `dbxtools codegen`). We extend the SDK schemas where Genie
+ *      ships fields on the wire that the SDK doesn't currently type:
+ *
+ *        - `GenieMessage.auto_regenerate_count: number` (stamped
+ *          on every wire message, omitted from the SDK shape).
+ *        - `GenieQueryAttachment.thoughts: GenieThought[]` (the
+ *          streamed reasoning payload with `DESCRIPTION`,
+ *          `DATA_SOURCING`, `STEPS`, and `UNDERSTANDING` thought
+ *          kinds).
+ *        - `GenieAttachment.attachment_type: AttachmentType`
+ *          (an optional discriminator literal so callers can
+ *          `switch (att.attachment_type)` instead of probing
+ *          which sub-key is populated; compute it with
+ *          {@link detectAttachmentType}).
+ *
+ *   2. Event vocabulary for the high-level `genieEventChat`
+ *      driver. Each event is a flat `z.object` with a `type`
+ *      literal discriminator and snake_case payload fields hoisted
+ *      to the top level (no `payload` wrapper). Events that share
+ *      the attachment-scoped location use
+ *      {@link GenieChatLocationSchema} as a base; events that
+ *      don't carry an `attachment_id` (status, result) omit it.
+ *      {@link GenieChatEventSchema} bundles the variants into one
+ *      discriminated union.
+ *
+ * Pure types: no runtime imports beyond zod + the generated
+ * sdk-type schemas, no Node-only code, safe for browser bundles.
+ */
+
+import { string } from "@dbx-tools/shared-core";
+import { dashboards } from "@dbx-tools/shared-sdk-model";
+import { z } from "zod";
+
+// The generated SDK schemas live under the `dashboards` namespace of
+// @dbx-tools/shared-sdk-model. Alias the handful Genie extends so the
+// widened schemas below read against short, local names.
+const {
+  messageStatusSchema: MessageStatusSchema,
+  genieAttachmentSchema: SDKGenieAttachmentSchema,
+  genieMessageSchema: SDKGenieMessageSchema,
+  genieQueryAttachmentSchema: SDKGenieQueryAttachmentSchema,
+  genieSpaceSchema: SDKGenieSpaceSchema,
+} = dashboards;
+
+/**
+ * The SDK's `MessageStatus` type, re-exported so consumers of the
+ * Genie contract get the full wire vocabulary from a single import
+ * path - mirrors the re-exports of the other SDK-derived schemas
+ * below.
+ */
+export type MessageStatus = dashboards.MessageStatus;
+
+// Re-export the SDK's `MessageStatus` schema so callers can parse a
+// raw status without reaching back into the generated namespace.
+export { MessageStatusSchema };
+
+/* ----------------------------- thoughts ---------------------------- */
+
+/**
+ * Genie's per-query "thoughts" surface. These appear on the
+ * `/messages/{id}` wire under `attachments[i].query.thoughts[]` but
+ * are not typed on the SDK's `GenieQueryAttachment` at v0.17.
+ *
+ * Known thought types observed in production polls:
+ *
+ *   - `THOUGHT_TYPE_DESCRIPTION`: a one-paragraph restatement of
+ *     what the user asked. The final `query.description` field on
+ *     the attachment carries the same text.
+ *   - `THOUGHT_TYPE_DATA_SOURCING`: markdown bullets of the
+ *     fully-qualified `catalog.schema.table` sources Genie chose.
+ *   - `THOUGHT_TYPE_STEPS`: the high-level plan Genie wrote before
+ *     running SQL (one bullet per step).
+ *   - `THOUGHT_TYPE_UNDERSTANDING`: ambiguity / interpretation
+ *     notes (e.g. "'revenue' could be interpreted as gross,
+ *     net, or recognized revenue...").
+ *
+ * Open at the type level (`| (string & {})`) so a new server-side
+ * thought type doesn't break compilation; the four known types
+ * still narrow correctly under `switch`.
+ */
+export type GenieThoughtType =
+  | "THOUGHT_TYPE_DESCRIPTION"
+  | "THOUGHT_TYPE_DATA_SOURCING"
+  | "THOUGHT_TYPE_STEPS"
+  | "THOUGHT_TYPE_UNDERSTANDING"
+  | (string & {});
+
+/**
+ * One reasoning step on a query attachment. See
+ * {@link GenieThoughtType} for the known `thought_type` values.
+ */
+export const GenieThoughtSchema = z.object({
+  thought_type: z.custom<GenieThoughtType>((v) => typeof v === "string"),
+  content: z.string(),
+});
+export type GenieThought = z.infer<typeof GenieThoughtSchema>;
+
+/* ----------------------- attachment discriminator ---------------------- */
+
+/**
+ * Discriminator for what's inside a `GenieAttachment`. Genie only
+ * ever populates one of `query` / `text` / `suggested_questions`
+ * per attachment. Open via `(string & {})` so a new server-side
+ * type doesn't break compilation; the three known types still
+ * narrow correctly under `switch`.
+ *
+ * Lifted into the schema as the optional
+ * {@link GenieAttachmentSchema}.`attachment_type` field (compute
+ * it with {@link detectAttachmentType}) so consumers can branch
+ * on a literal instead of probing which sub-object key is
+ * present. Also surfaced on {@link AttachmentEvent}'s payload as
+ * `attachment_type` (vs a bare `type`) to keep the field clear of
+ * the event-union discriminator key.
+ */
+export const ATTACHMENT_TYPES = [
+  "query",
+  "text",
+  "suggested_questions",
+] as const satisfies readonly string[];
+export type KnownAttachmentType = (typeof ATTACHMENT_TYPES)[number];
+export type AttachmentType = KnownAttachmentType | (string & {});
+
+/* ------------------- widened wire schemas + types ------------------ */
+
+/**
+ * `GenieQueryAttachment` widened with the `thoughts[]` field the
+ * Genie wire exposes but the SDK doesn't currently type. Every
+ * other field (`description`, `query`, `statement_id`,
+ * `query_result_metadata`, `title`, `parameters`,
+ * `last_updated_timestamp`, `id`) flows through verbatim from
+ * the SDK schema.
+ */
+export const GenieQueryAttachmentSchema = SDKGenieQueryAttachmentSchema.extend({
+  thoughts: z.array(GenieThoughtSchema).optional(),
+});
+export type GenieQueryAttachment = z.infer<typeof GenieQueryAttachmentSchema>;
+
+/**
+ * `GenieAttachment` with:
+ *
+ *   - `query` re-typed to the thoughts-aware
+ *     {@link GenieQueryAttachmentSchema}.
+ *   - `attachment_type` discriminator literal
+ *     ({@link AttachmentType}) so consumers can `switch
+ *     (att.attachment_type)` to narrow which sub-object is
+ *     populated. Optional on the wire (Genie doesn't send it);
+ *     callers compute it with {@link detectAttachmentType} when
+ *     they want the literal up-front.
+ *
+ * `attachment_id`, `text`, and `suggested_questions` pass through
+ * unchanged. `attachment_id` is genuinely optional on the wire:
+ * the first text attachment Genie emits per turn (the "main
+ * answer text") arrives with no id, only the follow-up text
+ * attachment gets one.
+ */
+export const GenieAttachmentSchema = SDKGenieAttachmentSchema.extend({
+  query: GenieQueryAttachmentSchema.optional(),
+  attachment_type: z.custom<AttachmentType>((v) => typeof v === "string").optional(),
+});
+export type GenieAttachment = z.infer<typeof GenieAttachmentSchema>;
+
+/**
+ * `GenieMessage` widened with:
+ *
+ *   - `auto_regenerate_count`: number stamped on every wire
+ *     message that the SDK type omits at v0.17.
+ *   - `attachments`: re-typed to the local
+ *     {@link GenieAttachmentSchema} so
+ *     `attachment.query?.thoughts` (and the
+ *     `attachment_type` discriminator) is reachable without a
+ *     cast.
+ *
+ * Every other field (`id`, `space_id`, `conversation_id`,
+ * `user_id`, `created_timestamp`, `last_updated_timestamp`,
+ * `status`, `content`, `message_id`, `query_result`, `error`,
+ * `feedback`) passes through from the SDK schema.
+ */
+export const GenieMessageSchema = SDKGenieMessageSchema.extend({
+  attachments: z.array(GenieAttachmentSchema).optional(),
+  auto_regenerate_count: z.number().optional(),
+});
+export type GenieMessage = z.infer<typeof GenieMessageSchema>;
+
+/**
+ * Domain projection of {@link GenieMessageSchema} that drops every
+ * turn-level field (`id`, `status`, `created_timestamp`,
+ * `query_result`, ...) and yields just the response payload as a
+ * flat `GenieAttachment[]`. Parses the same wire shape as
+ * {@link GenieMessageSchema}; missing `attachments` defaults to
+ * `[]` so the output is always an array.
+ *
+ * Use when consumers only care about the model's response
+ * attachments and not the surrounding message envelope.
+ */
+export const GenieResponseSchema = GenieMessageSchema.transform((m) => m.attachments ?? []);
+export type GenieResponse = z.infer<typeof GenieResponseSchema>;
+
+/**
+ * `GenieSpace` widened with an optional `serialized_space` field
+ * carrying a JSON-string representation of the space's full
+ * definition (catalogs, schemas, tables, sample questions, prompts,
+ * etc.) as exported by the workspace export endpoint.
+ *
+ * The base SDK shape (`description`, `space_id`, `title`,
+ * `warehouse_id`) covers the directory-listing surface - enough to
+ * route a chat. `serialized_space` is opt-in so callers that need
+ * the full space schema (e.g. to seed a system prompt or to ship
+ * the definition to an agent) can attach it without a second type.
+ */
+export const GenieSpaceSchema = SDKGenieSpaceSchema.extend({
+  serialized_space: z
+    .string()
+    .optional()
+    .describe(
+      "The contents of the Genie Space in serialized string form. This field is excluded in List Genie spaces responses. This field provides the structure of the JSON string that represents the space's layout and components.",
+    ),
+});
+export type GenieSpace = z.infer<typeof GenieSpaceSchema>;
+
+/* ------------------------ terminal helpers ------------------------- */
+
+/**
+ * Terminal Genie message statuses. The polling loop in the chat
+ * driver stops as soon as the latest message has one of these
+ * statuses.
+ */
+export const TERMINAL_STATUSES = ["COMPLETED", "FAILED", "CANCELLED"] as const;
+export type TerminalStatus = (typeof TERMINAL_STATUSES)[number];
+
+/** Narrow `MessageStatus | undefined` to a {@link TerminalStatus}. */
+export function isTerminalStatus(s: MessageStatus | undefined): s is TerminalStatus {
+  return s !== undefined && (TERMINAL_STATUSES as readonly string[]).includes(s);
+}
+
+/**
+ * Convert a raw Genie wire status (`FETCHING_METADATA`,
+ * `ASKING_AI`, `EXECUTING_QUERY`, ...) into a short, sentence-cased
+ * label safe to drop into a UI pill. Known states get a curated
+ * label; unknown states fall back to `string.tokenizeWithOptions`
+ * so new states still render cleanly without code changes.
+ *
+ * Pure (no Node-only deps), safe for browser bundles. Both the
+ * Genie agent (server) and any UI that subscribes to status events
+ * call this so labels stay in lock-step across the wire.
+ */
+export function humanizeStatus(status: MessageStatus): string {
+  switch (status) {
+    case "FETCHING_METADATA":
+      return "Fetching metadata";
+    case "ASKING_AI":
+      return "Asking Genie";
+    case "EXECUTING_QUERY":
+      return "Running SQL query";
+    case "FILTERING_CONTEXT":
+      return "Filtering context";
+    case "PENDING_WAREHOUSE":
+      return "Waiting for warehouse";
+    case "COMPLETED":
+      return "Completed";
+    case "FAILED":
+      return "Failed";
+    case "CANCELLED":
+      return "Cancelled";
+    default:
+      return [...string.tokenizeWithOptions({ capitalize: true, lowerCase: true }, status)].join(
+        " ",
+      );
+  }
+}
+
+/* ----------------------- attachment type helper ---------------------- */
+
+/**
+ * Inspect a {@link GenieAttachment} and return the
+ * {@link AttachmentType} of payload it carries. Returns the first
+ * known sub-object that's present; if none of the known ones are
+ * populated, falls back to the first non-bookkeeping key
+ * (forward-compat for types we don't model yet), else `"unknown"`.
+ *
+ * Honors a pre-set `attachment_type` if one is already on the
+ * value, so this is idempotent across re-detections.
+ */
+export function detectAttachmentType(att: GenieAttachment): AttachmentType {
+  if (att.attachment_type) return att.attachment_type;
+  if (att.query) return "query";
+  if (att.text) return "text";
+  if (att.suggested_questions) return "suggested_questions";
+  for (const k of Object.keys(att)) {
+    if (k !== "attachment_id" && k !== "attachment_type") return k;
+  }
+  return "unknown";
+}
+
+/* ------------------------- GenieChat events ------------------------ */
+
+/**
+ * Where on the wire an event was observed. Spread into every
+ * attachment-scoped event payload so subscribers can route, log,
+ * or correlate without re-walking the message.
+ *
+ * Fields:
+ *   - `space_id`: the Genie space the conversation lives in.
+ *   - `conversation_id`: conversation id once Genie has assigned one.
+ *   - `message_id`: Genie message id for the active turn.
+ *   - `attachment_id`: attachment the event came from. Optional
+ *     because Genie sometimes emits an anonymous main-answer
+ *     attachment without an id; those still get events, just with
+ *     `attachment_id` undefined.
+ */
+export const GenieChatLocationSchema = z.object({
+  space_id: z.string(),
+  conversation_id: z.string().optional(),
+  message_id: z.string().optional(),
+  attachment_id: z.string().optional(),
+});
+export type GenieChatLocation = z.infer<typeof GenieChatLocationSchema>;
+
+/**
+ * Lifecycle event: the question this turn is asking Genie. Fires
+ * once per `genieEventChat` call, the first time the underlying
+ * `genieChat` loop yields a `GenieMessage`. Carries the prompt
+ * text Genie echoed back on `message.content` and the assigned
+ * `message_id` so subscribers can group every subsequent event
+ * for this turn under one stable key. `conversation_id` is
+ * populated for both opening and follow-up turns (Genie assigns
+ * it on `startConversation`).
+ *
+ * Deferred (instead of fired synchronously on entry) so the
+ * `message_id` is available when the event lands - that one round
+ * trip costs ~200ms but lets UIs render question / thinking /
+ * query / text as a single grouped block per Genie call instead
+ * of a flat stream.
+ *
+ * `attachment_id` is intentionally absent - questions are turn
+ * scoped, not attachment scoped.
+ */
+export const QuestionEventSchema = GenieChatLocationSchema.omit({
+  attachment_id: true,
+}).extend({
+  type: z.literal("question"),
+  content: z.string(),
+});
+export type QuestionEvent = z.infer<typeof QuestionEventSchema>;
+
+/**
+ * Lifecycle event: raw `GenieMessage` snapshot for the active
+ * turn. Fires once per poll yield. The full message shape
+ * (including the widened thought / regen / attachment-type
+ * fields) is exposed inline as `message`; consumers narrow on
+ * `event.type === "message"` and reach for `event.message`.
+ */
+export const MessageEventSchema = GenieChatLocationSchema.omit({
+  attachment_id: true,
+}).extend({
+  type: z.literal("message"),
+  message: GenieMessageSchema,
+});
+export type MessageEvent = z.infer<typeof MessageEventSchema>;
+
+/**
+ * Top-level `message.status` transitioned. Fires for every
+ * distinct status seen on the wire (e.g. `SUBMITTED` ->
+ * `FILTERING_CONTEXT` -> `ASKING_AI` -> `PENDING_WAREHOUSE` ->
+ * `ASKING_AI` -> `COMPLETED`).
+ */
+export const StatusEventSchema = GenieChatLocationSchema.omit({
+  attachment_id: true,
+}).extend({
+  type: z.literal("status"),
+  status: MessageStatusSchema,
+  previous_status: MessageStatusSchema.optional(),
+});
+export type StatusEvent = z.infer<typeof StatusEventSchema>;
+
+/**
+ * A new attachment slot appeared in `message.attachments[]`.
+ * Fires exactly once per attachment (matched by `attachment_id`,
+ * positionally for anonymous attachments) the first time we see
+ * it. The slot's payload kind lands in `attachment_type` so the
+ * outer `type` discriminator stays unambiguous.
+ */
+export const AttachmentEventSchema = GenieChatLocationSchema.extend({
+  type: z.literal("attachment"),
+  index: z.number(),
+  attachment_type: z.custom<AttachmentType>((v) => typeof v === "string"),
+});
+export type AttachmentEvent = z.infer<typeof AttachmentEventSchema>;
+
+/**
+ * A new reasoning step appeared on a query attachment
+ * (`attachments[i].query.thoughts[i]`). Deduplicated per
+ * `(thought_type, content)` tuple within an attachment so
+ * subscribers don't see the same thought multiple times - Genie
+ * sometimes mutates existing thought slots in place, so the diff
+ * is value-based rather than positional.
+ */
+export const ThinkingEventSchema = GenieChatLocationSchema.extend({
+  type: z.literal("thinking"),
+  text: z.string(),
+  thought_type: z.custom<GenieThoughtType>((v) => typeof v === "string"),
+});
+export type ThinkingEvent = z.infer<typeof ThinkingEventSchema>;
+
+/**
+ * A text-attachment `content` field appeared or changed
+ * (`attachments[i].text.content`). Fires whenever the snapshot
+ * value differs from the previous one for the same attachment.
+ */
+export const TextEventSchema = GenieChatLocationSchema.extend({
+  type: z.literal("text"),
+  text: z.string(),
+});
+export type TextEvent = z.infer<typeof TextEventSchema>;
+
+/**
+ * SQL was finalized on a query attachment
+ * (`attachments[i].query.query`). Fires once when the SQL string
+ * transitions from undefined to defined, and again if Genie ever
+ * rewrites it. `title` and `description` are denormalised off the
+ * attachment's `query.title` / `query.description` so consumers
+ * can label the SQL pill without re-walking the message.
+ */
+export const QueryEventSchema = GenieChatLocationSchema.extend({
+  type: z.literal("query"),
+  sql: z.string(),
+  title: z.string().optional(),
+  description: z.string().optional(),
+});
+export type QueryEvent = z.infer<typeof QueryEventSchema>;
+
+/**
+ * SQL was submitted to a SQL warehouse and a statement id was
+ * assigned (`attachments[i].query.statement_id`). Fires when the
+ * statement id transitions from undefined to defined - this is the
+ * point at which `client.statementExecution.getStatement({ statement_id })`
+ * becomes a valid call.
+ */
+export const StatementEventSchema = GenieChatLocationSchema.extend({
+  type: z.literal("statement"),
+  statement_id: z.string(),
+});
+export type StatementEvent = z.infer<typeof StatementEventSchema>;
+
+/**
+ * Row count for a query attachment's result changed
+ * (`attachments[i].query.query_result_metadata.row_count`). Fires
+ * on every change, including the initial `undefined -> 0` and the
+ * later `0 -> N` once the warehouse finishes execution.
+ */
+export const RowsEventSchema = GenieChatLocationSchema.extend({
+  type: z.literal("rows"),
+  row_count: z.number(),
+  previous_row_count: z.number().optional(),
+  statement_id: z.string().optional(),
+});
+export type RowsEvent = z.infer<typeof RowsEventSchema>;
+
+/**
+ * Genie produced a follow-up suggested-questions list
+ * (`attachments[i].suggested_questions.questions[]`). Fires once
+ * when the list appears, and again if Genie rewrites it.
+ */
+export const SuggestedQuestionsEventSchema = GenieChatLocationSchema.extend({
+  type: z.literal("suggested_questions"),
+  questions: z.array(z.string()),
+});
+export type SuggestedQuestionsEvent = z.infer<typeof SuggestedQuestionsEventSchema>;
+
+/**
+ * The active turn reached a terminal status. Always fires once
+ * per turn (immediately after the terminal `message` event).
+ * Carries the final `GenieMessage` snapshot inline so subscribers
+ * don't need to keep their own copy of the last message.
+ */
+export const ResultEventSchema = GenieChatLocationSchema.omit({
+  attachment_id: true,
+}).extend({
+  type: z.literal("result"),
+  status: z.enum(TERMINAL_STATUSES),
+  message: GenieMessageSchema,
+});
+export type ResultEvent = z.infer<typeof ResultEventSchema>;
+
+/**
+ * Discriminated union yielded by `genieEventChat`. Each variant
+ * is a single flat object with `type` as the discriminator and
+ * the payload fields hoisted directly to the top level - no
+ * `payload` wrapper. Consumers narrow on `type` and read fields
+ * inline:
+ *
+ * @example
+ * for await (const event of genieEventChat(spaceId, "Top stores?")) {
+ *   switch (event.type) {
+ *     case "thinking":
+ *       console.log(event.thought_type, event.text);
+ *       break;
+ *     case "result":
+ *       console.log("done:", event.status);
+ *       break;
+ *   }
+ * }
+ *
+ * Stream order per turn:
+ *
+ *   1. `question`, emitted with the first `message` yield (deferred so
+ *      the assigned `message_id` is available), carrying the prompt
+ *      this turn sent to Genie.
+ *   2. `message` for every poll yield (raw `GenieMessage` on
+ *      `event.message`).
+ *   3. Any derived events the snapshot diff produced (`status`,
+ *      `attachment`, `thinking`, `text`, `query`, `statement`,
+ *      `rows`, `suggested_questions`) in that fixed order.
+ *   4. On the terminal snapshot, a final `result` event.
+ *
+ * Errors propagate via the generator throwing (`try`/`catch` the
+ * `for await`), not via an `error` variant on this union.
+ */
+export const GenieChatEventSchema = z.discriminatedUnion("type", [
+  QuestionEventSchema,
+  MessageEventSchema,
+  StatusEventSchema,
+  AttachmentEventSchema,
+  ThinkingEventSchema,
+  TextEventSchema,
+  QueryEventSchema,
+  StatementEventSchema,
+  RowsEventSchema,
+  SuggestedQuestionsEventSchema,
+  ResultEventSchema,
+]);
+export type GenieChatEvent = z.infer<typeof GenieChatEventSchema>;
+
+/** Discriminator type for {@link GenieChatEvent}. */
+export type GenieChatEventType = GenieChatEvent["type"];
+
+/**
+ * Field set for a given {@link GenieChatEventType} - the variant
+ * with the `type` discriminator stripped. Used by detectors in
+ * `event.ts` so each detector returns just the payload fields and
+ * the orchestrator stamps `type` at yield time.
+ */
+export type GenieChatEventFields<T extends GenieChatEventType> = Omit<
+  Extract<GenieChatEvent, { type: T }>,
+  "type"
+>;