@dbx-tools/appkit-mastra 0.1.4 → 0.1.12

package/dist/src/genie.js

@@ -11,32 +11,226 @@
  * upstream change in `@databricks/appkit` flows in automatically.
  *
  * As Genie streams its long-running events (`FETCHING_METADATA` →
- * `ASKING_AI` → `EXECUTING_QUERY` → `COMPLETED`, plus SQL queries and
- * row data in `message_result.attachments` / `query_result`), the tool
- * forwards a normalised {@link GenieProgress} discriminated union out
- * through `ctx.writer` so the client can render incremental feedback
- * (status pill, SQL code block, row count) while the LLM still sees a
- * single clean final payload.
+ * `ASKING_AI` → `EXECUTING_QUERY` → `COMPLETED`, plus SQL text and
+ * follow-ups in `message_result.attachments`), the tool forwards a
+ * normalised {@link GenieProgress} discriminated union out through
+ * `ctx.writer` so the client can render an incremental loading pill.
+ * Row payloads from `query_result` are intentionally discarded - the
+ * LLM never sees rows, and charts come from the separate
+ * `render_data` tool when the model decides one is useful.
  */
 import { genie } from "@databricks/appkit";
-import { stringUtils } from "@dbx-tools/appkit-shared";
+import { logUtils, stringUtils } from "@dbx-tools/appkit-shared";
 import { createTool } from "@mastra/core/tools";
 import { z } from "zod";
+import { emitChartWithPlanning } from "./chart.js";
+/**
+ * Module-level logger tagged `[mastra/genie]`. Uses the shared
+ * {@link logUtils.logger} so calls below `LOG_LEVEL` are
+ * discarded for free. Default `LOG_LEVEL` is `info`; flip to
+ * `debug` to see per-turn timing (`query_result` → planner
+ * waits → `drain:return`).
+ */
+const log = logUtils.logger("mastra/genie");
+/**
+ * Per-dataset metadata surfaced to the LLM. The actual rows are
+ * dispatched separately as a `kind: "chart"` writer event so the
+ * model never has the rows in its context (token cost stays flat
+ * regardless of dataset size). The model uses `chartId` to
+ * reference the chart inline via the `[[chart:<chartId>]]` marker.
+ */
+const datasetSchema = z.object({
+    chartId: z.string().describe(stringUtils.toDescription `
+    Short id (8 hex chars) for the chart-render slot the host UI
+    has staged for this dataset. Embed
+    \`[[chart:<chartId>]]\` on its own line in your reply at the
+    position you want the chart to appear; the client renders it
+    inline. Do not paraphrase the dataset's rows in prose - the
+    chart is the rendering.
+  `),
+    title: z.string().optional().describe(stringUtils.toDescription `
+    Genie's own title for the SQL that produced this dataset.
+    Useful as a label when you reference the chart in prose.
+  `),
+    description: z.string().optional().describe(stringUtils.toDescription `
+    Genie's prose description of the SQL, if any.
+  `),
+    columns: z.array(z.string()).describe(stringUtils.toDescription `
+    Column names in display order. Use these when describing what
+    is being charted (e.g. "trend of fill_rate over date").
+  `),
+    rowCount: z.number().describe(stringUtils.toDescription `
+    Total rows in this dataset. Mention only if it adds context
+    (e.g. "across the last 90 days").
+  `),
+    sql: z
+        .string()
+        .optional()
+        .describe(stringUtils.toDescription `
+      SQL Genie generated and executed. The host UI shows this on
+      demand; you do not need to repeat it.
+    `),
+});
+/**
+ * Top-level output schema returned to the LLM from a Genie tool
+ * call. The `datasets` array is intentionally metadata-only - row
+ * data rides a writer event the host UI consumes directly and is
+ * not in the model's context.
+ */
+const genieToolOutputSchema = z.object({
+    conversationId: z
+        .string()
+        .optional()
+        .describe(stringUtils.toDescription `
+      Pass back on the next call to continue the same Genie thread.
+    `),
+    genieAnswer: z
+        .string()
+        .optional()
+        .describe(stringUtils.toDescription `
+      Genie's natural-language answer to the question. Pass this
+      through to the user (verbatim, or as the basis of your
+      reply). Genie may have run multiple SQL queries and tools to
+      produce this; the full text is the answer.
+    `),
+    datasets: z
+        .array(datasetSchema)
+        .optional()
+        .describe(stringUtils.toDescription `
+      Datasets Genie produced for this turn (one per executed SQL
+      statement). Each entry is metadata only; the rows are
+      streamed to the host UI out-of-band. To render any of these
+      as a chart inline in your reply, embed
+      \`[[chart:<chartId>]]\` where you want the chart to appear.
+      Do not paraphrase the rows - the chart is what the user
+      should see; your prose should add interpretation
+      (highlights, deltas, anomalies) around the chart.
+    `),
+    suggestedFollowUps: z
+        .array(z.string())
+        .optional()
+        .describe(stringUtils.toDescription `
+      Follow-up question suggestions Genie produced. The host UI
+      renders these as clickable buttons; you do not need to list
+      them in your reply.
+    `),
+    error: z
+        .string()
+        .optional()
+        .describe(stringUtils.toDescription `
+      Genie-side error message if the request failed.
+    `),
+});
 const sendMessageSchema = z.object({
-    content: z.string().describe("Natural-language question to send to the Genie space."),
+    content: z.string().describe(stringUtils.toDescription `
+    Natural-language question to send to the Genie space.
+  `),
     conversationId: z
         .string()
         .optional()
-        .describe("Optional Genie conversation id to continue an earlier thread. " +
-        "Omit on the first call; pass the id returned in the previous " +
-        "result's `conversationId` to follow up."),
+        .describe(stringUtils.toDescription `
+      Optional Genie conversation id to continue an earlier thread.
+      Omit on the first call; pass the id returned in the previous
+      result's \`conversationId\` to follow up.
+    `),
 });
 const getConversationSchema = z.object({
-    alias: z
-        .string()
-        .describe("Alias of the Genie space the conversation belongs to (matches the " +
-        "key in the genie plugin's `spaces` config)."),
-    conversationId: z.string().describe("Genie conversation id whose history to fetch."),
+    alias: z.string().describe(stringUtils.toDescription `
+    Alias of the Genie space the conversation belongs to (matches
+    the key in the genie plugin's \`spaces\` config).
+  `),
+    conversationId: z.string().describe(stringUtils.toDescription `
+    Genie conversation id whose history to fetch.
+  `),
+});
+/** Per-attachment shape returned inside a stored Genie message. */
+const genieAttachmentSchema = z.object({
+    attachmentId: z.string().optional().describe(stringUtils.toDescription `
+    Genie attachment id; internal bookkeeping.
+  `),
+    query: z
+        .object({
+        title: z.string().optional().describe(stringUtils.toDescription `
+        Genie's title for the SQL, if any.
+      `),
+        description: z.string().optional().describe(stringUtils.toDescription `
+        Genie's prose description of the SQL, if any.
+      `),
+        query: z.string().optional().describe(stringUtils.toDescription `
+        SQL Genie generated and executed.
+      `),
+        statementId: z.string().optional().describe(stringUtils.toDescription `
+        Statement-execution id; internal bookkeeping.
+      `),
+    })
+        .optional()
+        .describe(stringUtils.toDescription `
+      SQL Genie attached to this message, if it ran any.
+    `),
+    text: z
+        .object({
+        content: z.string().optional().describe(stringUtils.toDescription `
+        Genie's natural-language answer text for this attachment.
+      `),
+    })
+        .optional()
+        .describe(stringUtils.toDescription `
+      Per-attachment text content (independent of the message-level
+      \`content\` field).
+    `),
+    suggestedQuestions: z
+        .array(z.string())
+        .optional()
+        .describe(stringUtils.toDescription `
+      Follow-up question suggestions Genie generated for this turn.
+    `),
+});
+/** Single message inside a Genie conversation history page. */
+const genieMessageSchema = z.object({
+    messageId: z.string().describe(stringUtils.toDescription `
+    Genie message id; internal bookkeeping.
+  `),
+    conversationId: z.string().describe(stringUtils.toDescription `
+    Conversation id this message belongs to.
+  `),
+    spaceId: z.string().describe(stringUtils.toDescription `
+    Genie space id this message belongs to.
+  `),
+    status: z.string().describe(stringUtils.toDescription `
+    Genie message status (\`COMPLETED\`, \`FAILED\`, etc.).
+  `),
+    content: z.string().describe(stringUtils.toDescription `
+    Outer message-level natural-language content Genie wrote.
+  `),
+    attachments: z
+        .array(genieAttachmentSchema)
+        .optional()
+        .describe(stringUtils.toDescription `
+      Attachments (SQL queries, text blocks, suggested follow-ups)
+      Genie produced for this message.
+    `),
+    error: z.string().optional().describe(stringUtils.toDescription `
+    Genie-side error attached to this message, if any.
+  `),
+});
+/**
+ * Output schema for the \`genie_get_conversation\` tool. Mirrors
+ * AppKit's \`GenieConversationHistoryResponse\` so the model gets a
+ * clear, typed view of prior messages instead of an opaque blob.
+ */
+const genieGetConversationOutputSchema = z.object({
+    conversationId: z.string().describe(stringUtils.toDescription `
+    Conversation id you fetched.
+  `),
+    spaceId: z.string().describe(stringUtils.toDescription `
+    Genie space the conversation belongs to.
+  `),
+    messages: z.array(genieMessageSchema).describe(stringUtils.toDescription `
+    Messages in the conversation, oldest to newest. Each
+    \`message.content\` is Genie's natural-language answer for
+    that turn; attachments carry the SQL and follow-ups Genie
+    produced.
+  `),
 });
 /**
  * Default tool name for a wired Genie alias. The well-known `default`
@@ -54,6 +248,11 @@ export function defaultGenieToolName(alias) {
  * Build one `sendMessage` tool per configured Genie alias plus a single
  * `getConversation` tool. Returns a record keyed by tool id, ready to
  * spread into an `Agent`'s `tools` map.
+ *
+ * `config` must be the active plugin config; Genie's
+ * `query_result` events are routed through
+ * {@link emitChartWithPlanning} which uses it to resolve the
+ * chart-planner's model.
  */
 export function buildGenieTools(opts) {
     const tools = {};
@@ -61,27 +260,45 @@ export function buildGenieTools(opts) {
         const id = defaultGenieToolName(alias);
         tools[id] = createTool({
             id,
-            description: `Ask the Databricks Genie space "${alias}" a natural-language ` +
-                "question. Genie translates the question to SQL, runs it against " +
-                "the configured datasets, and returns a written answer plus any " +
-                "SQL statements it executed. Returns `{ conversationId, content, " +
-                "queries, ... }`; pass `conversationId` back in to follow up in " +
-                "the same Genie thread.",
+            description: stringUtils.toDescription `
+        Ask the Databricks Genie space "${alias}" a single
+        natural-language question. Genie translates it to SQL,
+        runs it, and returns \`genieAnswer\` (prose) plus
+        \`datasets[]\` (one entry per executed query, each with
+        a short \`chartId\`). Embed \`[[chart:<chartId>]]\` on
+        its own line at the position you want that data rendered
+        as an inline chart. Add interpretation around the chart
+        (deltas, anomalies, takeaways); do not paraphrase row
+        values.
+
+        Issue ONE focused question per user turn. Prefer
+        aggregated queries over raw-row queries for time-series
+        and distributions.
+      `,
             inputSchema: sendMessageSchema,
+            outputSchema: genieToolOutputSchema,
             execute: async ({ content, conversationId }, ctx) => {
                 const stream = opts.exports.sendMessage(alias, content, conversationId, {
                     signal: opts.signal,
                 });
-                return drainGenieStream(stream, ctx.writer);
+                const requestContext = ctx
+                    ?.requestContext;
+                return drainGenieStream(stream, ctx.writer, {
+                    config: opts.config,
+                    ...(requestContext ? { requestContext } : {}),
+                });
             },
         });
     }
     tools.genie_get_conversation = createTool({
         id: "genie_get_conversation",
-        description: "Fetch the full message history of a prior Genie conversation by id. " +
-            "Use when the user references an earlier Genie thread by id, or to " +
-            "inspect attachments / SQL from previous turns.",
+        description: stringUtils.toDescription `
+      Fetch the full message history of a prior Genie conversation
+      by id. Use when the user references an earlier Genie thread
+      by id, or to inspect attachments / SQL from previous turns.
+    `,
         inputSchema: getConversationSchema,
+        outputSchema: genieGetConversationOutputSchema,
         execute: async ({ alias, conversationId }) => {
             return opts.exports.getConversation(alias, conversationId, opts.signal);
         },
@@ -89,29 +306,66 @@ export function buildGenieTools(opts) {
     return tools;
 }
 /**
- * Drain the genie `sendMessage` AsyncGenerator into a flat result the
- * agent's calling LLM can reason about. Final assistant text is pulled
- * from the last `message_result`; SQL statements are extracted from
- * `query_result` events; conversation / message ids are surfaced so the
- * caller can pass `conversationId` back into a follow-up tool call.
+ * Drain the genie `sendMessage` AsyncGenerator into a flat result
+ * the agent's calling LLM can reason about, while forwarding
+ * progress and chart events to the host UI.
+ *
+ * Three streams of output happen in parallel:
+ *
+ * 1. {@link GenieProgress} pill events on the writer (`started`,
+ *    `status`, `sql`, `suggested`, `error`) drive the loading
+ *    pill in the chat bubble.
+ * 2. `kind: "chart"` events on the writer (emitted via
+ *    {@link emitChartWithPlanning}) carry the row payload from
+ *    each Genie SQL statement and, on planner success, a
+ *    follow-up event with the rendered Echarts spec. The host
+ *    UI's `<ChartSlot>` merges the two by `chartId` and
+ *    renders inline at the marker position the model picked.
+ *    The data never reaches the LLM.
+ * 3. The `DrainResult` returned to the LLM contains Genie's
+ *    prose answer plus a `datasets[]` array of metadata
+ *    (chartId, title, columns, rowCount, sql) the model uses
+ *    to cite charts via `[[chart:<chartId>]]` markers.
  *
- * When a Mastra `writer` is passed (i.e. the tool runs inside an agent
- * stream), normalised {@link GenieProgress} events are pushed mid-flight
- * so the UI can show status changes, SQL, and row counts as they
- * happen instead of staring at a spinner for the full Genie round-trip.
+ * `query_result` and `message_result` events arrive in either
+ * order; we buffer per-statement scratch keyed by `statementId`
+ * so each half can fill in what it knows. The chart event
+ * fires the moment `query_result` lands; the planner runs in
+ * the background. We `Promise.allSettled` every planner promise
+ * before returning so all chart work is attributed to the tool's
+ * trace span and so the LLM's `datasets[]` includes every
+ * chartId that has actually been queued.
  */
-async function drainGenieStream(stream, writer) {
+async function drainGenieStream(stream, writer, opts) {
+    const { config, requestContext } = opts;
     let conversationId;
-    let messageId;
-    let spaceId;
-    let status;
-    let content;
-    let attachments;
+    let genieAnswer;
+    let suggestedFollowUps;
     let error;
-    const queries = [];
-    // Best-effort progress emission. Awaited so the underlying agent
-    // stream sees events in order; write failures are swallowed so a
-    // dead writer (e.g. closed downstream) can't take the tool down.
+    // AppKit's `streamSendMessage` forwards every SDK `onProgress`
+    // callback verbatim - the same `EXECUTING_QUERY` can fire several
+    // times during a single poll loop. AppKit's other path,
+    // `streamGetMessage`, dedupes on the connector side; we mirror that
+    // behaviour here so the UI status pill doesn't flicker and we don't
+    // burn writer bytes on no-op events.
+    let lastStatus;
+    const scratchByStatementId = new Map();
+    const getScratch = (statementId) => {
+        let s = scratchByStatementId.get(statementId);
+        if (!s) {
+            s = { statementId, columns: [], rowCount: 0 };
+            scratchByStatementId.set(statementId, s);
+        }
+        return s;
+    };
+    /**
+     * Planner promises kicked off per `query_result`. Awaited
+     * (Promise.allSettled) before drainGenieStream returns so the
+     * Genie tool's trace span covers the chart work and the LLM's
+     * `datasets[]` accurately reflects every chartId that's been
+     * queued for rendering.
+     */
+    const plannerPromises = [];
     const emit = async (event) => {
         if (!writer)
             return;
@@ -123,20 +377,26 @@ async function drainGenieStream(stream, writer) {
         }
     };
     for await (const event of stream) {
+        // Per-event raw payload for tuning the pill / answer pipeline
+        // against real Genie traffic. At `info` (the default) this is
+        // discarded for free; flip `LOG_LEVEL=debug` to see every
+        // raw wire event before the switch routes it through writer
+        // and DrainResult.
+        log.debug("event", { type: event.type, payload: event });
         switch (event.type) {
             case "message_start":
                 conversationId = event.conversationId;
-                messageId = event.messageId;
-                spaceId = event.spaceId;
                 await emit({
                     kind: "started",
-                    conversationId,
-                    messageId,
-                    spaceId,
+                    conversationId: event.conversationId,
+                    messageId: event.messageId,
+                    spaceId: event.spaceId,
                 });
                 break;
             case "status":
-                status = event.status;
+                if (event.status === lastStatus)
+                    break;
+                lastStatus = event.status;
                 await emit({
                     kind: "status",
                     status: event.status,
@@ -144,34 +404,67 @@ async function drainGenieStream(stream, writer) {
                 });
                 break;
             case "query_result": {
-                queries.push({
-                    attachmentId: event.attachmentId,
+                const columns = (event.data?.manifest?.schema?.columns ?? []).map((c) => c.name);
+                const dataArray = (event.data?.result?.data_array ?? []);
+                const rows = genieRowsToObjects(columns, dataArray);
+                const scratch = getScratch(event.statementId);
+                // emitChartWithPlanning emits the dataset event immediately
+                // and kicks off the chart-planner agent in the background.
+                // It returns the chartId synchronously; the plannerPromise
+                // is awaited at end-of-stream so chart work shows up under
+                // this tool's trace span.
+                const { chartId, plannerPromise } = await emitChartWithPlanning({
+                    ...(writer ? { writer } : {}),
+                    config,
+                    ...(requestContext ? { requestContext } : {}),
+                    title: scratch.title ?? `Genie query`,
+                    ...(scratch.description ? { description: scratch.description } : {}),
+                    data: rows,
+                });
+                scratch.chartId = chartId;
+                scratch.columns = columns;
+                scratch.rowCount = rows.length;
+                plannerPromises.push(plannerPromise);
+                log.debug("query_result", {
                     statementId: event.statementId,
-                    data: event.data,
+                    chartId,
+                    rows: rows.length,
+                    columns,
                 });
-                const rowCount = event.data?.result?.data_array?.length ?? 0;
-                const columns = (event.data?.manifest?.schema?.columns ?? []).map((c) => c.name);
-                await emit({ kind: "data", rowCount, columns });
                 break;
             }
             case "message_result":
-                content = event.message.content;
-                attachments = event.message.attachments;
-                status = event.message.status;
-                for (const attachment of attachments ?? []) {
-                    if (attachment.query?.query) {
+                genieAnswer = event.message.content;
+                for (const attachment of event.message.attachments ?? []) {
+                    const sqlText = attachment.query?.query;
+                    const stmtId = attachment.query?.statementId;
+                    if (stmtId) {
+                        const scratch = getScratch(stmtId);
+                        if (sqlText)
+                            scratch.sql = sqlText;
+                        if (attachment.query?.title)
+                            scratch.title = attachment.query.title;
+                        if (attachment.query?.description) {
+                            scratch.description = attachment.query.description;
+                        }
+                    }
+                    if (sqlText) {
                         await emit({
                             kind: "sql",
-                            sql: attachment.query.query,
-                            title: attachment.query.title,
-                            description: attachment.query.description,
-                            statementId: attachment.query.statementId,
+                            sql: sqlText,
+                            title: attachment.query?.title,
+                            description: attachment.query?.description,
+                            statementId: stmtId,
                         });
                     }
                     if (attachment.text?.content) {
                         await emit({ kind: "text", content: attachment.text.content });
                     }
                     if (attachment.suggestedQuestions?.length) {
+                        // Last attachment with suggestions wins (same merge rule
+                        // the UI uses via `collectSuggestions`); keeping just one
+                        // copy per turn caps token usage.
+                        suggestedFollowUps = attachment.suggestedQuestions;
                         await emit({
                             kind: "suggested",
                             questions: attachment.suggestedQuestions,
@@ -187,17 +480,80 @@ async function drainGenieStream(stream, writer) {
                 break;
         }
     }
-    return {
+    // Wait for all chart planners to settle before returning so the
+    // tool's trace span covers chart work and the LLM's
+    // `datasets[]` reflects only chartIds the client has actually
+    // received writer events for. Failures in `emitChartWithPlanning`
+    // are already swallowed inside the helper, so this never
+    // throws.
+    log.debug("planners:awaiting", { count: plannerPromises.length });
+    await Promise.allSettled(plannerPromises);
+    log.debug("planners:settled", { count: plannerPromises.length });
+    // Build the LLM-bound `datasets[]` from scratch entries that
+    // actually ran a query (chartId is assigned at `query_result`
+    // time). Entries that only saw `message_result` metadata
+    // without a row payload are skipped.
+    const datasets = [];
+    for (const scratch of scratchByStatementId.values()) {
+        if (!scratch.chartId)
+            continue;
+        datasets.push({
+            chartId: scratch.chartId,
+            ...(scratch.title ? { title: scratch.title } : {}),
+            ...(scratch.description ? { description: scratch.description } : {}),
+            columns: scratch.columns,
+            rowCount: scratch.rowCount,
+            ...(scratch.sql ? { sql: scratch.sql } : {}),
+        });
+    }
+    log.debug("drain:return", {
         conversationId,
-        messageId,
-        spaceId,
-        status,
-        content,
-        attachments,
-        queries,
+        hasAnswer: typeof genieAnswer === "string",
+        answerLength: genieAnswer?.length ?? 0,
+        chartIds: datasets.map((d) => d.chartId),
+        suggestedCount: suggestedFollowUps?.length ?? 0,
         error,
+    });
+    return {
+        ...(conversationId ? { conversationId } : {}),
+        ...(genieAnswer ? { genieAnswer } : {}),
+        ...(datasets.length > 0 ? { datasets } : {}),
+        ...(suggestedFollowUps ? { suggestedFollowUps } : {}),
+        ...(error ? { error } : {}),
     };
 }
+/**
+ * Convert Genie's `data_array` (column-positional `string | null`
+ * tuples) into plain JS row objects keyed by column name. Numeric
+ * strings are coerced to numbers so the chart-planner picks
+ * `value` axes instead of `category` axes; everything else passes
+ * through verbatim. `null` becomes `null`.
+ */
+function genieRowsToObjects(columns, dataArray) {
+    const out = [];
+    for (const row of dataArray) {
+        const obj = {};
+        columns.forEach((col, i) => {
+            const cell = row[i] ?? null;
+            obj[col] = coerceCell(cell);
+        });
+        out.push(obj);
+    }
+    return out;
+}
+/** Best-effort numeric coercion for Genie's all-strings cells. */
+function coerceCell(cell) {
+    if (cell === null)
+        return null;
+    // Anchored to keep `12.5px` / `123abc` as strings; only fully
+    // numeric values become JS numbers.
+    if (/^-?\d+(\.\d+)?$/.test(cell)) {
+        const n = Number(cell);
+        if (Number.isFinite(n))
+            return n;
+    }
+    return cell;
+}
 /**
  * Toolkit provider built from a live AppKit `GeniePlugin` instance.
  * Returned by {@link buildGenieProvider} so that
@@ -217,7 +573,7 @@ async function drainGenieStream(stream, writer) {
  * all-or-nothing bundle. Wire `only` / `except` / `prefix` / `rename`
  * later if a caller needs them.
  */
-export function buildGenieProvider(plugin) {
+export function buildGenieProvider(plugin, opts) {
     return {
         toolkit(_opts) {
             const aliases = extractGenieAliases(plugin);
@@ -227,6 +583,7 @@ export function buildGenieProvider(plugin) {
                     sendMessage: plugin.sendMessage.bind(plugin),
                     getConversation: plugin.getConversation.bind(plugin),
                 },
+                config: opts.config,
             });
         },
     };
@@ -266,6 +623,8 @@ function humanizeGenieStatus(status) {
         case "FAILED":
             return "Failed";
         default:
-            return status.toLowerCase().replace(/_/g, " ");
+            return [
+                ...stringUtils.tokenizeWithOptions({ capitalize: true, lowerCase: true }, status),
+            ].join(" ");
     }
 }