@dbx-tools/appkit-mastra 0.3.44 → 0.4.1

package/lib/src/genie.d.ts

@@ -0,0 +1,166 @@
+/**
+ * Genie tools for Mastra.
+ *
+ * Surfaces each configured Genie space as a small set of flat Mastra
+ * tools the calling agent drives directly - no inner orchestrator
+ * agent. The central agent decomposes user questions, picks which
+ * space to ask, streams the per-turn wire events (status, thinking,
+ * sql, rows) through `ctx.writer`, and composes the final reply.
+ * Rows are never fetched eagerly: the agent reads a statement's
+ * values only when it needs to reason about them, otherwise it embeds
+ * a `[data:<statement_id>]` marker in prose and lets the host UI
+ * resolve the data. Charts are minted asynchronously and referenced
+ * by `[chart:<chartId>]` markers so prose isn't blocked on chart
+ * generation; the host UI fetches the cached spec by id once ready.
+ * Space description and serialized-space lookups are available for
+ * grounding when the agent needs schema context.
+ *
+ * Each tool's `execute` pulls the per-request
+ * {@link WorkspaceClient} off `ctx.requestContext` (stamped by
+ * `MastraServer` under {@link MASTRA_USER_KEY}) and the per-call
+ * `writer` / `abortSignal` off `ctx`, so the tools are stateless
+ * across requests and the central agent owns the loop.
+ *
+ * The tools talk to Genie directly via `@dbx-tools/genie`
+ * (`genieEventChat`); statement-row fetching is delegated to
+ * {@link fetchStatementData} from `./statement.js`, which wraps
+ * the workspace `statementExecution.getStatement` API. AppKit's
+ * stock `genie` plugin is honored only for its `spaces` config
+ * so existing AppKit-style wiring keeps working without change.
+ *
+ * Suggested orchestration prompt for the central agent lives in
+ * {@link GENIE_INSTRUCTIONS}; compose it into the agent's own
+ * `instructions` when you want the canonical "how to drive the
+ * Genie tools" guidance.
+ *
+ * @module
+ */
+import { WorkspaceClient } from "@databricks/sdk-experimental";
+import { plugin } from "@dbx-tools/appkit";
+import type { MastraTools } from "./agents.js";
+import type { MastraPluginConfig } from "./config.js";
+/** Default alias used when a single unnamed Genie space is wired up. */
+export declare const DEFAULT_GENIE_ALIAS = "default";
+/** Per-space Genie agent configuration. */
+export interface GenieSpaceConfig {
+    /** Genie `space_id`. Required; resolves via `client.genie.getSpace`. */
+    spaceId: string;
+    /**
+     * Optional human-readable description appended to the per-space
+     * tool descriptions so the calling LLM has hints about *what
+     * data* this space covers (e.g. "orders, returns,
+     * fulfillment"). When omitted, only the space's own
+     * `description` (fetched on first use of `get_space_description`)
+     * is shown.
+     */
+    hint?: string;
+}
+/** Map of alias -> space config. Accepts either explicit objects or bare space ids. */
+export type GenieSpacesConfig = Record<string, GenieSpaceConfig | string>;
+/**
+ * Suggested orchestration prompt for the central agent that owns
+ * the Genie tools. Compose into your agent's `instructions` to
+ * get the canonical "decompose questions, ask Genie focused
+ * sub-questions, place data / chart markers in prose" behavior:
+ *
+ * ```ts
+ * createAgent({
+ *   instructions: `${myAgentInstructions}\n\n${GENIE_INSTRUCTIONS}`,
+ *   tools(plugins) {
+ *     return { ...plugins.genie?.toolkit() };
+ *   },
+ * });
+ * ```
+ *
+ * The prompt references the bare tool names (`ask_genie`,
+ * `get_space_description`, `get_space_serialized`,
+ * `get_statement`, `prepare_chart`) used for the single-space
+ * default alias. Multi-space deployments should write their own
+ * variant that names the suffixed per-space tools
+ * (e.g. `ask_genie_sales`).
+ */
+export declare const GENIE_INSTRUCTIONS: string;
+/**
+ * Normalize the {@link GenieSpacesConfig} record. Bare-string
+ * entries (`{ default: "01ef..." }`) get wrapped as
+ * `{ spaceId: "01ef..." }`; object entries pass through unchanged.
+ *
+ * @throws ConfigurationError when an alias is present but carries no space id
+ *   (`{ default: process.env.DATABRICKS_GENIE_SPACE_ID }` with the variable
+ *   unset). An alias that resolves to nothing is a wiring contradiction: the
+ *   agent would advertise no Genie tool for a space the deployment believes it
+ *   configured, so it fails at construction instead of going quiet.
+ */
+export declare function normalizeGenieSpaces(spaces: GenieSpacesConfig | Record<string, string | GenieSpaceConfig | undefined> | undefined): Record<string, GenieSpaceConfig>;
+/**
+ * Discover Genie space aliases from every supported source and
+ * merge them into a single record. Precedence (highest first):
+ *
+ *   1. {@link MastraPluginConfig.genieSpaces} on the `mastra(...)`
+ *      call. Explicit Mastra wiring always wins so users can
+ *      override AppKit's defaults per-agent.
+ *   2. AppKit `genie({ spaces: { ... } })` plugin instance. Lets
+ *      users keep using the existing AppKit config format
+ *      (`genie({ spaces: { sales: "...", ops: "..." } })`)
+ *      without restating the same record on the Mastra plugin.
+ *      Read off the live plugin instance via a structural cast
+ *      since `Plugin.config` is TS-protected (not runtime-private).
+ *   3. `DATABRICKS_GENIE_SPACE_ID` env var (registered under the
+ *      well-known `default` alias). Matches the AppKit `genie()`
+ *      plugin's fallback behavior so a bare `mastra()` + `genie()`
+ *      pair just works.
+ *
+ * Aliases collide cleanly: a higher-precedence source's value
+ * replaces a lower one's wholesale. A source that contributes zero
+ * aliases is skipped; a source that names an alias without a space
+ * id fails through {@link normalizeGenieSpaces}.
+ */
+export declare function resolveGenieSpaces(config: MastraPluginConfig, context: plugin.PluginContextLike | undefined): Record<string, GenieSpaceConfig>;
+/**
+ * Build the flat Mastra tools record for every configured Genie
+ * space. Two shared, space-agnostic tools (`get_statement`,
+ * `prepare_chart`) are registered once regardless of how many
+ * spaces are wired; the per-space tools (`ask_genie`,
+ * `get_space_description`, `get_space_serialized`) are suffixed
+ * with `_<alias>` for non-default aliases so multi-space
+ * deployments stay disambiguated.
+ *
+ * Returns a record keyed by tool id, ready to spread into the
+ * central `Agent`'s `tools` map (or surfaced via the
+ * `plugins.genie?.toolkit()` callback). Returns an empty record
+ * when `spaces` resolves to zero entries so the caller can spread
+ * safely.
+ */
+export declare function buildGenieTools(opts: {
+    spaces: GenieSpacesConfig | Record<string, GenieSpaceConfig>;
+    config: MastraPluginConfig;
+}): MastraTools;
+/**
+ * Plugin-toolkit adapter so the `plugins.genie?.toolkit()` lookup
+ * inside an agent's `tools(plugins)` callback returns the
+ * flat Genie tools record instead of throwing on missing plugin.
+ * Mirrors AppKit's `PluginToolkitProvider` shape.
+ */
+export declare function buildGenieToolkitProvider(opts: {
+    spaces: GenieSpacesConfig | Record<string, GenieSpaceConfig>;
+    config: MastraPluginConfig;
+}): {
+    toolkit(opts?: unknown): MastraTools;
+};
+/**
+ * Collect the curated starter questions across every resolved Genie
+ * space, deduped and capped. Each space's `sample_questions` are
+ * fetched once (cached for {@link SUGGESTION_CACHE_TTL_MS}) via
+ * {@link getGenieSpace} + {@link genieSampleQuestions}, then merged in
+ * alias-iteration order so a single-space app surfaces that space's
+ * questions and a multi-space app round-trips breadth-first up to the
+ * cap. A per-space fetch failure degrades to "no questions for that
+ * space" (logged, not thrown) so one unreachable space never blanks
+ * the whole list. Returns `[]` when no spaces are configured.
+ */
+export declare function collectSpaceSuggestions(opts: {
+    spaces: Record<string, GenieSpaceConfig>;
+    client: WorkspaceClient;
+    signal?: AbortSignal;
+    limit?: number;
+}): Promise<string[]>;