@dbx-tools/appkit-mastra 0.1.0

package/dist/src/memory.js

@@ -0,0 +1,197 @@
+/**
+ * Lakebase-backed Mastra memory wiring.
+ *
+ * Provides a {@link MemoryBuilder} that mints one `Memory` per agent
+ * with two independent knobs:
+ *
+ * - **Storage** (threads / messages via `PostgresStore`): defaults to
+ *   **per-agent** namespacing via `schemaName: "mastra_<agentId>"` so
+ *   conversation history stays isolated between agents in the same
+ *   database. `PostgresStore` auto-creates the schema with
+ *   `CREATE SCHEMA IF NOT EXISTS` on init.
+ * - **Memory** (semantic recall via `PgVector`): defaults to a single
+ *   **shared** instance across every agent. Cross-agent recall on one
+ *   index is almost always what users want; opt into per-agent recall
+ *   by passing a {@link MastraMemoryConfigOverride} on the agent.
+ *
+ * Plugin-level `config.storage` / `config.memory` act as the baseline
+ * (auto-defaulted to `true` in `plugin.ts` when the `lakebase` plugin
+ * is registered); per-agent settings cascade on top of that.
+ */
+import { lakebase } from "@databricks/appkit";
+import { pluginUtils } from "@dbx-tools/appkit-shared";
+import { fastembed } from "@mastra/fastembed";
+import { Memory } from "@mastra/memory";
+import { PgVector, PostgresStore } from "@mastra/pg";
+import { randomUUID } from "node:crypto";
+/**
+ * True when any plugin-level or per-agent setting could need the
+ * Lakebase pool. Used by `plugin.ts` to gate pool acquisition; the
+ * builder also acquires lazily so missed cases still fail with a
+ * clear lakebase-not-registered error.
+ */
+export function needsLakebase(config) {
+    if (settingNeedsSharedPool(config.storage))
+        return true;
+    if (settingNeedsSharedPool(config.memory))
+        return true;
+    const defs = collectAgentDefinitions(config);
+    return defs.some((d) => settingNeedsSharedPool(d.storage) || settingNeedsSharedPool(d.memory));
+}
+/**
+ * Look up the `lakebase` plugin and return its managed `pg.Pool`.
+ * Throws when the sibling plugin is not registered; enabling
+ * `storage` / `memory` without lakebase is a wiring bug, not a runtime
+ * condition we can recover from.
+ */
+export function resolveLakebasePool(context, caller) {
+    return pluginUtils.require(context, lakebase, caller).exports().pool;
+}
+/**
+ * Construct a per-agent {@link Memory} factory. Caches the shared
+ * `PgVector` singleton (built on first need) and the lazily-resolved
+ * Lakebase pool so each agent build is O(1) after the first.
+ */
+export function createMemoryBuilder(config, context) {
+    return new MemoryBuilder(config, context);
+}
+/**
+ * Builds one `Memory` per agent. Per-instance state keeps the shared
+ * `PgVector` and the resolved Lakebase pool alive across calls so
+ * registering N agents stays cheap.
+ */
+export class MemoryBuilder {
+    config;
+    context;
+    sharedVector;
+    pool;
+    constructor(config, context) {
+        this.config = config;
+        this.context = context;
+    }
+    /**
+     * Build a `Memory` for `agentId` after the plugin/agent cascade.
+     * Returns `undefined` when the agent has neither storage nor a
+     * vector store enabled - Mastra accepts a missing `memory` field
+     * and treats the agent as stateless.
+     */
+    forAgent(agentId, def) {
+        const storageSetting = def.storage ?? this.config.storage;
+        const memorySetting = def.memory ?? this.config.memory;
+        const storage = this.buildStorage(agentId, storageSetting);
+        const vector = this.buildVector(memorySetting);
+        if (!storage && !vector)
+            return undefined;
+        return new Memory({
+            ...(storage ? { storage } : {}),
+            ...(vector ? { vector, embedder: fastembed } : {}),
+            options: {
+                lastMessages: 10,
+                ...(vector
+                    ? { semanticRecall: { topK: 3, messageRange: 2 } }
+                    : {}),
+            },
+        });
+    }
+    buildStorage(agentId, setting) {
+        if (!setting)
+            return undefined;
+        if (typeof setting === "boolean") {
+            return new PostgresStore({
+                id: `mastra-store__${agentId}`,
+                schemaName: `mastra_${agentId}`,
+                pool: this.requirePool(),
+            });
+        }
+        // Cast: `withId` guarantees `id` is set, but the distributive
+        // Omit + `id?: string` shape doesn't structurally narrow to the
+        // discriminated union members. Runtime shape is identical.
+        return new PostgresStore(withId(setting, `mastra-store__${agentId}`));
+    }
+    /**
+     * Resolve the agent's vector store. Cascade:
+     *
+     * - falsy: no vector.
+     * - `boolean` / `undefined-inheriting-true`: return the shared
+     *   singleton (built lazily on first call). All agents that
+     *   default-enable memory write into and recall from one index.
+     * - object: build a dedicated `PgVector` for this agent.
+     */
+    buildVector(setting) {
+        if (!setting)
+            return undefined;
+        if (typeof setting === "boolean")
+            return this.getSharedVector();
+        return buildPgVector(setting);
+    }
+    getSharedVector() {
+        if (!this.sharedVector) {
+            this.sharedVector = buildSharedPgVector(this.requirePool());
+        }
+        return this.sharedVector;
+    }
+    requirePool() {
+        if (!this.pool) {
+            this.pool = resolveLakebasePool(this.context, this.config);
+        }
+        return this.pool;
+    }
+}
+/**
+ * Build the shared `PgVector` that backs the default
+ * `def.memory === true` case across every agent.
+ *
+ * `PgVector`'s constructor accepts only connection-style configs
+ * (`HostConfig` / `ConnectionStringConfig` / `ClientConfig`); there is
+ * no `{ pool }` shorthand the way `PostgresStore` has one. Worse, the
+ * constructor synchronously kicks off a `cacheWarmupPromise` IIFE that
+ * calls `this.pool.connect()` before returning, so we can't cleanly
+ * hand it an inert config and patch the pool afterwards.
+ *
+ * The trick: pass illegal-but-validation-passing placeholders so the
+ * warmup's `net.connect()` rejects synchronously with `RangeError`
+ * (Node validates `0 <= port < 65536`). The IIFE's `catch {}` swallows
+ * it, no DNS lookup or TCP attempt happens, and we then swap
+ * `pgVector.pool` to the lakebase pool. Every subsequent `PgVector`
+ * method reads `this.pool` at call time, so all real I/O goes through
+ * the lakebase pool from then on. The placeholder pool is `.end()`'d
+ * so its socket book-keeping is released.
+ */
+function buildSharedPgVector(pool) {
+    const vector = new PgVector({
+        id: `pg${randomUUID()}`,
+        host: "-1",
+        port: -1,
+        database: "_",
+        user: "_",
+        password: "_",
+    });
+    const placeholder = vector.pool;
+    vector.pool = pool;
+    void placeholder.end().catch(() => undefined);
+    return vector;
+}
+/** Per-agent dedicated `PgVector` (rare; opt-in via object override). */
+function buildPgVector(setting) {
+    return new PgVector(withId(setting, `pg-vector__${randomUUID()}`));
+}
+/** True when this setting requires the shared Lakebase pool. */
+function settingNeedsSharedPool(setting) {
+    return setting === true;
+}
+/** Walk the three shapes of `config.agents` into a flat list. */
+function collectAgentDefinitions(config) {
+    const agents = config.agents;
+    if (!agents)
+        return [];
+    if (Array.isArray(agents))
+        return agents;
+    if (typeof agents.instructions === "string") {
+        return [agents];
+    }
+    return Object.values(agents);
+}
+/** Fill in a default `id` when the caller didn't supply one. */
+function withId(value, fallback) {
+    return value.id ? value : { ...value, id: fallback };
+}

package/dist/src/model.d.ts

@@ -0,0 +1,159 @@
+/**
+ * Databricks Model Serving resolver for Mastra agents.
+ *
+ * Each agent step calls {@link buildModel} with the active
+ * `RequestContext`. The user stamped by `MastraServer` carries an
+ * AppKit `WorkspaceClient`; we ask it for the workspace host and a
+ * fresh bearer header, then point Mastra's OpenAI-compatible provider
+ * at `/serving-endpoints` on that host.
+ *
+ * Model id resolution walks three sources before falling back to the
+ * hard-coded default, **in this priority order**:
+ *
+ *   1. Per-request override stashed by the auth middleware under
+ *      {@link MASTRA_MODEL_OVERRIDE_KEY} (header / query / body).
+ *   2. The static `modelId` passed in by the agent / plugin (string
+ *      sugar on `def.model` or `config.defaultModel`).
+ *   3. `DATABRICKS_SERVING_ENDPOINT_NAME` env var.
+ *   4. {@link FALLBACK_MODEL_ID}.
+ *
+ * Whatever wins is then fuzzy-matched against the live
+ * `/serving-endpoints` list ({@link listServingEndpoints}) so loose
+ * names like `"claude sonnet"` resolve to the real endpoint name.
+ * Fuzzy matching is best-effort: when the workspace client throws
+ * (network blip, expired token at cache-fill time) we fall back to
+ * the input verbatim and let Databricks return the canonical error.
+ */
+import type { MastraModelConfig } from "@mastra/core/llm";
+import type { RequestContext } from "@mastra/core/request-context";
+import { type MastraPluginConfig } from "./config.js";
+/**
+ * Capability tiers for Databricks Foundation Model API endpoints.
+ *
+ * - {@link ModelTier.Thinking}: deepest reasoning / "thinking" models
+ *   (Claude Opus, GPT-5.5 Pro, Gemini Pro, Llama 4 Maverick, etc).
+ *   Highest cost and latency; reserve for hard multi-step reasoning.
+ * - {@link ModelTier.Balanced}: cost/latency sweet spot for general
+ *   agent work (Claude Sonnet, GPT-5.x, Gemini Flash, Llama 3.3 70B).
+ *   The right default for most agents.
+ * - {@link ModelTier.Fast}: cheap and quick; classification, routing,
+ *   tool-arg extraction, simple summarisation (Claude Haiku, GPT-5
+ *   mini/nano, Gemini Flash Lite, GPT-OSS 20B, Llama 3.1 8B).
+ *
+ * String enum so the value is the slug we use in cache keys, logs,
+ * and as the value users see in serialized configs.
+ */
+export declare enum ModelTier {
+    Thinking = "thinking",
+    Balanced = "balanced",
+    Fast = "fast"
+}
+/**
+ * Catalogue of Databricks-hosted Foundation Model API endpoints,
+ * grouped by capability {@link ModelTier} and then by provider. Each
+ * inner array is priority-ordered (most powerful first within the
+ * same provider+tier).
+ *
+ * Provider buckets:
+ *
+ * - `claude`: Anthropic Claude family (closed; flagship reasoning).
+ * - `gpt`: OpenAI GPT-5 family (closed; "ChatGPT" on Databricks FMAPI).
+ * - `gemini`: Google Gemini family (closed; multimodal + web-search).
+ * - `openSource`: open-weights models (widest regional / SKU availability).
+ *
+ * The list is curated by hand; refresh from the Databricks "supported
+ * foundation models" doc when new endpoints land.
+ */
+export declare const MODEL_CATALOG: {
+    readonly thinking: {
+        readonly claude: readonly ["databricks-claude-opus-4-8", "databricks-claude-opus-4-7", "databricks-claude-opus-4-6", "databricks-claude-opus-4-5", "databricks-claude-opus-4-1"];
+        readonly gpt: readonly ["databricks-gpt-5-5-pro"];
+        readonly gemini: readonly ["databricks-gemini-3-1-pro", "databricks-gemini-3-pro", "databricks-gemini-2-5-pro"];
+        readonly openSource: readonly ["databricks-llama-4-maverick", "databricks-gpt-oss-120b", "databricks-meta-llama-3-1-405b-instruct"];
+    };
+    readonly balanced: {
+        readonly claude: readonly ["databricks-claude-sonnet-4-6", "databricks-claude-sonnet-4-5", "databricks-claude-sonnet-4"];
+        readonly gpt: readonly ["databricks-gpt-5-5", "databricks-gpt-5-4", "databricks-gpt-5-2", "databricks-gpt-5-1", "databricks-gpt-5"];
+        readonly gemini: readonly ["databricks-gemini-3-5-flash", "databricks-gemini-3-flash", "databricks-gemini-2-5-flash"];
+        readonly openSource: readonly ["databricks-meta-llama-3-3-70b-instruct", "databricks-qwen3-next-80b-a3b-instruct", "databricks-qwen35-122b-a10b"];
+    };
+    readonly fast: {
+        readonly claude: readonly ["databricks-claude-haiku-4-5"];
+        readonly gpt: readonly ["databricks-gpt-5-4-mini", "databricks-gpt-5-4-nano", "databricks-gpt-5-mini", "databricks-gpt-5-nano"];
+        readonly gemini: readonly ["databricks-gemini-3-1-flash-lite"];
+        readonly openSource: readonly ["databricks-gpt-oss-20b", "databricks-gemma-3-12b", "databricks-meta-llama-3-1-8b-instruct"];
+    };
+};
+/**
+ * Priority-ordered model ids for a single capability {@link ModelTier},
+ * interleaved across providers so a workspace missing the top Claude
+ * still lands on a flagship GPT / Gemini on the next probe.
+ *
+ * Provider order within the interleave: Claude, GPT, Gemini, then the
+ * open-weights tail appended verbatim as the universal floor (widest
+ * regional availability).
+ *
+ * @example
+ * ```ts
+ * mastra({
+ *   defaultModelFallbacks: modelsForTier(ModelTier.Fast),
+ * });
+ * ```
+ */
+export declare function modelsForTier(tier: ModelTier): readonly string[];
+/**
+ * Top model id at the given {@link ModelTier}. Sync; the agent-step
+ * resolver fuzzy-matches it against the workspace catalogue at call
+ * time, so this works even when the literal top pick isn't deployed.
+ *
+ * Use when wiring a tier-appropriate model into an agent definition:
+ *
+ * @example
+ * ```ts
+ * const classifier = createAgent({
+ *   instructions: "Classify this email",
+ *   model: modelForTier(ModelTier.Fast),  // cheap, quick
+ * });
+ *
+ * const planner = createAgent({
+ *   instructions: "Plan a multi-step migration",
+ *   model: modelForTier(ModelTier.Thinking),  // deep reasoning
+ * });
+ * ```
+ */
+export declare function modelForTier(tier: ModelTier): string;
+/**
+ * Last-resort model ids used when neither `config.defaultModel`,
+ * per-agent `model`, nor `DATABRICKS_SERVING_ENDPOINT_NAME` is set.
+ *
+ * Walked in order at resolve time: the first id whose endpoint is
+ * actually present in the workspace's `/serving-endpoints` listing
+ * wins. Workspaces vary - not every region / SKU has every model,
+ * and the list of Foundation Model APIs evolves quickly - so the
+ * resolver degrades all the way from "best thinking model" down to
+ * "smallest commodity Llama" before giving up.
+ *
+ * Built by chaining the per-tier interleaves (Thinking -> Balanced
+ * -> Fast); within each tier the providers are round-robin-zipped
+ * (Claude, GPT, Gemini, then open-weights tail). Override the entire
+ * list via `MastraPluginConfig.defaultModelFallbacks` (e.g. to pin a
+ * regulated workspace to a specific approved subset, or to bias the
+ * priority toward a particular tier).
+ */
+export declare const FALLBACK_MODEL_IDS: readonly string[];
+/** Optional overrides accepted by {@link buildModel}. */
+export interface BuildModelOverrides {
+    /**
+     * Static model id from the agent / plugin config (string sugar on
+     * `def.model` or `config.defaultModel`). Loses to the per-request
+     * override but wins over env / fallback.
+     */
+    modelId?: string;
+}
+/**
+ * Resolve a `MastraModelConfig` for the current agent step. Runs
+ * while `agent.stream` is inside the `asUser(req)` scope so tokens
+ * are user-scoped; outside an active user context the workspace
+ * client falls back to the service principal.
+ */
+export declare function buildModel(config: MastraPluginConfig, requestContext: RequestContext, overrides?: BuildModelOverrides): Promise<MastraModelConfig>;