@dbx-tools/appkit-mastra 0.1.5 → 0.1.10
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +394 -0
- package/index.ts +46 -37
- package/package.json +58 -47
- package/src/agents.ts +233 -62
- package/src/chart.ts +555 -285
- package/src/config.ts +282 -18
- package/src/filesystems.ts +1090 -0
- package/src/genie.ts +1004 -601
- package/src/history.ts +178 -79
- package/src/mcp.ts +105 -0
- package/src/memory.ts +129 -74
- package/src/mlflow.ts +149 -0
- package/src/model.ts +80 -408
- package/src/observability.ts +144 -0
- package/src/pagination.ts +34 -0
- package/src/plugin.ts +573 -59
- package/src/processors.ts +168 -0
- package/src/rest.ts +67 -0
- package/src/server.ts +243 -45
- package/src/serving-sanitize.ts +167 -0
- package/src/serving.ts +27 -224
- package/src/statement.ts +89 -0
- package/src/storage-schema.ts +41 -0
- package/src/summarize.ts +176 -0
- package/src/threads.ts +338 -0
- package/src/workspaces.ts +346 -0
- package/src/writer.ts +44 -0
- package/tsconfig.json +41 -0
- package/dist/index.d.ts +0 -19
- package/dist/index.js +0 -19
- package/dist/src/agents.d.ts +0 -306
- package/dist/src/agents.js +0 -393
- package/dist/src/chart.d.ts +0 -104
- package/dist/src/chart.js +0 -375
- package/dist/src/config.d.ts +0 -170
- package/dist/src/config.js +0 -12
- package/dist/src/genie.d.ts +0 -116
- package/dist/src/genie.js +0 -594
- package/dist/src/history.d.ts +0 -67
- package/dist/src/history.js +0 -158
- package/dist/src/memory.d.ts +0 -79
- package/dist/src/memory.js +0 -197
- package/dist/src/model.d.ts +0 -159
- package/dist/src/model.js +0 -423
- package/dist/src/plugin.d.ts +0 -130
- package/dist/src/plugin.js +0 -255
- package/dist/src/render-chart-route.d.ts +0 -33
- package/dist/src/render-chart-route.js +0 -120
- package/dist/src/server.d.ts +0 -46
- package/dist/src/server.js +0 -113
- package/dist/src/serving.d.ts +0 -156
- package/dist/src/serving.js +0 -214
- package/src/render-chart-route.ts +0 -141
package/dist/src/history.js
DELETED
|
@@ -1,158 +0,0 @@
|
|
|
1
|
-
/**
|
|
2
|
-
* Thread history loader exposed as a Mastra custom API route.
|
|
3
|
-
*
|
|
4
|
-
* Backed entirely by native Mastra: looks up the active agent by id,
|
|
5
|
-
* asks its `Memory` instance to `recall` a page of `MastraDBMessage`s,
|
|
6
|
-
* and converts the result to AI SDK V5 `UIMessage`s with the official
|
|
7
|
-
* {@link toAISdkV5Messages} helper from `@mastra/ai-sdk/ui`. No direct
|
|
8
|
-
* database reads.
|
|
9
|
-
*
|
|
10
|
-
* The route is registered through {@link historyRoute} as a Mastra
|
|
11
|
-
* `registerApiRoute` so it sits in the same dispatcher pipeline as
|
|
12
|
-
* `chatRoute`. That means the `MastraServer` auth middleware (in
|
|
13
|
-
* `./server.ts`) has already populated `RequestContext` with
|
|
14
|
-
* `MASTRA_THREAD_ID_KEY` and `MASTRA_RESOURCE_ID_KEY` by the time
|
|
15
|
-
* the handler runs - no cookie or user lookups happen here, and the
|
|
16
|
-
* session-cookie logic stays the single source of truth in `server.ts`.
|
|
17
|
-
*/
|
|
18
|
-
import { toAISdkV5Messages } from "@mastra/ai-sdk/ui";
|
|
19
|
-
import { MASTRA_RESOURCE_ID_KEY, MASTRA_THREAD_ID_KEY, } from "@mastra/core/request-context";
|
|
20
|
-
import { registerApiRoute } from "@mastra/core/server";
|
|
21
|
-
/** Default history page size; matches the Mastra storage default. */
|
|
22
|
-
const DEFAULT_PER_PAGE = 20;
|
|
23
|
-
/** Hard cap so a misbehaving client can't fetch the whole thread at once. */
|
|
24
|
-
const MAX_PER_PAGE = 200;
|
|
25
|
-
/**
|
|
26
|
-
* Fetch a page of UI-formatted messages for a thread.
|
|
27
|
-
*
|
|
28
|
-
* Uses the agent's resolved `Memory` (`getMemory()`) so per-agent
|
|
29
|
-
* storage namespaces (`mastra_<agentId>` schemas) and any future
|
|
30
|
-
* memory-side filters apply automatically. When the agent has no
|
|
31
|
-
* memory configured the response is a successful empty page so
|
|
32
|
-
* callers don't have to special-case stateless agents.
|
|
33
|
-
*
|
|
34
|
-
* Pagination is descending-by-default: page 0 is the most recent
|
|
35
|
-
* page, page 1 the page before that, etc. The returned `uiMessages`
|
|
36
|
-
* are always re-sorted into chronological order (oldest -> newest)
|
|
37
|
-
* so the client can prepend them above the existing transcript
|
|
38
|
-
* without sorting locally.
|
|
39
|
-
*/
|
|
40
|
-
export async function loadHistory(opts) {
|
|
41
|
-
const perPage = clampPerPage(opts.perPage);
|
|
42
|
-
const page = Math.max(0, Math.trunc(opts.page ?? 0));
|
|
43
|
-
const memory = await opts.agent.getMemory();
|
|
44
|
-
if (!memory) {
|
|
45
|
-
return { uiMessages: [], page, perPage, total: 0, hasMore: false };
|
|
46
|
-
}
|
|
47
|
-
const result = await memory.recall({
|
|
48
|
-
threadId: opts.threadId,
|
|
49
|
-
...(opts.resourceId ? { resourceId: opts.resourceId } : {}),
|
|
50
|
-
page,
|
|
51
|
-
perPage,
|
|
52
|
-
orderBy: {
|
|
53
|
-
field: "createdAt",
|
|
54
|
-
direction: opts.ascending ? "ASC" : "DESC",
|
|
55
|
-
},
|
|
56
|
-
});
|
|
57
|
-
const chronological = sortChronological(result.messages);
|
|
58
|
-
const uiMessages = toAISdkV5Messages(chronological);
|
|
59
|
-
return {
|
|
60
|
-
uiMessages,
|
|
61
|
-
page,
|
|
62
|
-
perPage,
|
|
63
|
-
total: result.total,
|
|
64
|
-
hasMore: result.hasMore,
|
|
65
|
-
};
|
|
66
|
-
}
|
|
67
|
-
/**
|
|
68
|
-
* Register a `GET <path>` Mastra custom API route that returns a page
|
|
69
|
-
* of AI SDK V5 `UIMessage`s for the caller's current thread.
|
|
70
|
-
*
|
|
71
|
-
* Modeled after `chatRoute` from `@mastra/ai-sdk`: pass `agent` for a
|
|
72
|
-
* fixed-agent mount, or include `:agentId` in the path for dynamic
|
|
73
|
-
* routing. Pairs cleanly with the AppKit Mastra plugin's chat route
|
|
74
|
-
* layout (`/route/chat` + `/route/chat/:agentId`).
|
|
75
|
-
*
|
|
76
|
-
* The handler reads `threadId` and `resourceId` from `RequestContext`
|
|
77
|
-
* (populated upstream by `MastraServer.registerAuthMiddleware`), so
|
|
78
|
-
* no cookie or user lookups happen here.
|
|
79
|
-
*/
|
|
80
|
-
export function historyRoute(options) {
|
|
81
|
-
const { path } = options;
|
|
82
|
-
const fixedAgent = "agent" in options ? options.agent : undefined;
|
|
83
|
-
if (!fixedAgent && !path.includes(":agentId")) {
|
|
84
|
-
throw new Error("historyRoute path must include `:agentId` or `agent` must be passed explicitly");
|
|
85
|
-
}
|
|
86
|
-
return registerApiRoute(path, {
|
|
87
|
-
method: "GET",
|
|
88
|
-
handler: async (c) => {
|
|
89
|
-
const mastra = c.get("mastra");
|
|
90
|
-
const requestContext = c.get("requestContext");
|
|
91
|
-
const agentId = fixedAgent ?? c.req.param("agentId");
|
|
92
|
-
if (!agentId) {
|
|
93
|
-
return c.json({ error: "agentId is required" }, 400);
|
|
94
|
-
}
|
|
95
|
-
const agent = mastra.getAgentById(agentId);
|
|
96
|
-
if (!agent) {
|
|
97
|
-
return c.json({ error: `Unknown agent "${agentId}"` }, 404);
|
|
98
|
-
}
|
|
99
|
-
const threadId = requestContext.get(MASTRA_THREAD_ID_KEY);
|
|
100
|
-
if (!threadId) {
|
|
101
|
-
return c.json({ error: "thread id missing from request context" }, 400);
|
|
102
|
-
}
|
|
103
|
-
const resourceId = requestContext.get(MASTRA_RESOURCE_ID_KEY);
|
|
104
|
-
const payload = await loadHistory({
|
|
105
|
-
agent,
|
|
106
|
-
threadId,
|
|
107
|
-
...(resourceId ? { resourceId } : {}),
|
|
108
|
-
page: parseIntParam(c.req.query("page")),
|
|
109
|
-
perPage: parseIntParam(c.req.query("perPage")),
|
|
110
|
-
});
|
|
111
|
-
return c.json(payload);
|
|
112
|
-
},
|
|
113
|
-
});
|
|
114
|
-
}
|
|
115
|
-
/** Coerce / clamp `perPage`; falls back to the page-size default. */
|
|
116
|
-
function clampPerPage(value) {
|
|
117
|
-
if (value === undefined || Number.isNaN(value))
|
|
118
|
-
return DEFAULT_PER_PAGE;
|
|
119
|
-
const n = Math.trunc(value);
|
|
120
|
-
if (n <= 0)
|
|
121
|
-
return DEFAULT_PER_PAGE;
|
|
122
|
-
return Math.min(n, MAX_PER_PAGE);
|
|
123
|
-
}
|
|
124
|
-
/**
|
|
125
|
-
* Sort messages oldest-first by `createdAt`, falling back to whatever
|
|
126
|
-
* order the storage returned them in. The native `recall` call honors
|
|
127
|
-
* `orderBy` but doesn't guarantee a stable secondary sort, so we
|
|
128
|
-
* normalize here before handing the page to the AI SDK converter.
|
|
129
|
-
*/
|
|
130
|
-
function sortChronological(messages) {
|
|
131
|
-
return [...messages].sort((a, b) => {
|
|
132
|
-
const ta = toEpoch(a.createdAt);
|
|
133
|
-
const tb = toEpoch(b.createdAt);
|
|
134
|
-
return ta - tb;
|
|
135
|
-
});
|
|
136
|
-
}
|
|
137
|
-
function toEpoch(value) {
|
|
138
|
-
if (value instanceof Date)
|
|
139
|
-
return value.getTime();
|
|
140
|
-
if (typeof value === "string" || typeof value === "number") {
|
|
141
|
-
const parsed = new Date(value).getTime();
|
|
142
|
-
return Number.isNaN(parsed) ? 0 : parsed;
|
|
143
|
-
}
|
|
144
|
-
return 0;
|
|
145
|
-
}
|
|
146
|
-
/**
|
|
147
|
-
* Coerce a Hono query value into a non-negative integer. Returns
|
|
148
|
-
* `undefined` for empty / non-numeric / negative inputs so
|
|
149
|
-
* {@link loadHistory} can apply its built-in defaults.
|
|
150
|
-
*/
|
|
151
|
-
function parseIntParam(value) {
|
|
152
|
-
if (!value)
|
|
153
|
-
return undefined;
|
|
154
|
-
const n = Number(value);
|
|
155
|
-
if (!Number.isFinite(n) || n < 0)
|
|
156
|
-
return undefined;
|
|
157
|
-
return Math.trunc(n);
|
|
158
|
-
}
|
package/dist/src/memory.d.ts
DELETED
|
@@ -1,79 +0,0 @@
|
|
|
1
|
-
/**
|
|
2
|
-
* Lakebase-backed Mastra memory wiring.
|
|
3
|
-
*
|
|
4
|
-
* Provides a {@link MemoryBuilder} that mints one `Memory` per agent
|
|
5
|
-
* with two independent knobs:
|
|
6
|
-
*
|
|
7
|
-
* - **Storage** (threads / messages via `PostgresStore`): defaults to
|
|
8
|
-
* **per-agent** namespacing via `schemaName: "mastra_<agentId>"` so
|
|
9
|
-
* conversation history stays isolated between agents in the same
|
|
10
|
-
* database. `PostgresStore` auto-creates the schema with
|
|
11
|
-
* `CREATE SCHEMA IF NOT EXISTS` on init.
|
|
12
|
-
* - **Memory** (semantic recall via `PgVector`): defaults to a single
|
|
13
|
-
* **shared** instance across every agent. Cross-agent recall on one
|
|
14
|
-
* index is almost always what users want; opt into per-agent recall
|
|
15
|
-
* by passing a {@link MastraMemoryConfigOverride} on the agent.
|
|
16
|
-
*
|
|
17
|
-
* Plugin-level `config.storage` / `config.memory` act as the baseline
|
|
18
|
-
* (auto-defaulted to `true` in `plugin.ts` when the `lakebase` plugin
|
|
19
|
-
* is registered); per-agent settings cascade on top of that.
|
|
20
|
-
*/
|
|
21
|
-
import { lakebase } from "@databricks/appkit";
|
|
22
|
-
import { pluginUtils } from "@dbx-tools/appkit-shared";
|
|
23
|
-
import { Memory } from "@mastra/memory";
|
|
24
|
-
import type { MastraAgentDefinition } from "./agents.js";
|
|
25
|
-
import type { MastraPluginConfig } from "./config.js";
|
|
26
|
-
/** Pool handle returned by the AppKit `lakebase` plugin `exports().pool`. */
|
|
27
|
-
export type LakebasePool = ReturnType<InstanceType<ReturnType<typeof lakebase>["plugin"]>["exports"]>["pool"];
|
|
28
|
-
/**
|
|
29
|
-
* True when any plugin-level or per-agent setting could need the
|
|
30
|
-
* Lakebase pool. Used by `plugin.ts` to gate pool acquisition; the
|
|
31
|
-
* builder also acquires lazily so missed cases still fail with a
|
|
32
|
-
* clear lakebase-not-registered error.
|
|
33
|
-
*/
|
|
34
|
-
export declare function needsLakebase(config: MastraPluginConfig): boolean;
|
|
35
|
-
/**
|
|
36
|
-
* Look up the `lakebase` plugin and return its managed `pg.Pool`.
|
|
37
|
-
* Throws when the sibling plugin is not registered; enabling
|
|
38
|
-
* `storage` / `memory` without lakebase is a wiring bug, not a runtime
|
|
39
|
-
* condition we can recover from.
|
|
40
|
-
*/
|
|
41
|
-
export declare function resolveLakebasePool(context: pluginUtils.PluginContextLike | undefined, caller: MastraPluginConfig): LakebasePool;
|
|
42
|
-
/**
|
|
43
|
-
* Construct a per-agent {@link Memory} factory. Caches the shared
|
|
44
|
-
* `PgVector` singleton (built on first need) and the lazily-resolved
|
|
45
|
-
* Lakebase pool so each agent build is O(1) after the first.
|
|
46
|
-
*/
|
|
47
|
-
export declare function createMemoryBuilder(config: MastraPluginConfig, context: pluginUtils.PluginContextLike | undefined): MemoryBuilder;
|
|
48
|
-
/**
|
|
49
|
-
* Builds one `Memory` per agent. Per-instance state keeps the shared
|
|
50
|
-
* `PgVector` and the resolved Lakebase pool alive across calls so
|
|
51
|
-
* registering N agents stays cheap.
|
|
52
|
-
*/
|
|
53
|
-
export declare class MemoryBuilder {
|
|
54
|
-
private readonly config;
|
|
55
|
-
private readonly context;
|
|
56
|
-
private sharedVector;
|
|
57
|
-
private pool;
|
|
58
|
-
constructor(config: MastraPluginConfig, context: pluginUtils.PluginContextLike | undefined);
|
|
59
|
-
/**
|
|
60
|
-
* Build a `Memory` for `agentId` after the plugin/agent cascade.
|
|
61
|
-
* Returns `undefined` when the agent has neither storage nor a
|
|
62
|
-
* vector store enabled - Mastra accepts a missing `memory` field
|
|
63
|
-
* and treats the agent as stateless.
|
|
64
|
-
*/
|
|
65
|
-
forAgent(agentId: string, def: MastraAgentDefinition): Memory | undefined;
|
|
66
|
-
private buildStorage;
|
|
67
|
-
/**
|
|
68
|
-
* Resolve the agent's vector store. Cascade:
|
|
69
|
-
*
|
|
70
|
-
* - falsy: no vector.
|
|
71
|
-
* - `boolean` / `undefined-inheriting-true`: return the shared
|
|
72
|
-
* singleton (built lazily on first call). All agents that
|
|
73
|
-
* default-enable memory write into and recall from one index.
|
|
74
|
-
* - object: build a dedicated `PgVector` for this agent.
|
|
75
|
-
*/
|
|
76
|
-
private buildVector;
|
|
77
|
-
private getSharedVector;
|
|
78
|
-
private requirePool;
|
|
79
|
-
}
|
package/dist/src/memory.js
DELETED
|
@@ -1,197 +0,0 @@
|
|
|
1
|
-
/**
|
|
2
|
-
* Lakebase-backed Mastra memory wiring.
|
|
3
|
-
*
|
|
4
|
-
* Provides a {@link MemoryBuilder} that mints one `Memory` per agent
|
|
5
|
-
* with two independent knobs:
|
|
6
|
-
*
|
|
7
|
-
* - **Storage** (threads / messages via `PostgresStore`): defaults to
|
|
8
|
-
* **per-agent** namespacing via `schemaName: "mastra_<agentId>"` so
|
|
9
|
-
* conversation history stays isolated between agents in the same
|
|
10
|
-
* database. `PostgresStore` auto-creates the schema with
|
|
11
|
-
* `CREATE SCHEMA IF NOT EXISTS` on init.
|
|
12
|
-
* - **Memory** (semantic recall via `PgVector`): defaults to a single
|
|
13
|
-
* **shared** instance across every agent. Cross-agent recall on one
|
|
14
|
-
* index is almost always what users want; opt into per-agent recall
|
|
15
|
-
* by passing a {@link MastraMemoryConfigOverride} on the agent.
|
|
16
|
-
*
|
|
17
|
-
* Plugin-level `config.storage` / `config.memory` act as the baseline
|
|
18
|
-
* (auto-defaulted to `true` in `plugin.ts` when the `lakebase` plugin
|
|
19
|
-
* is registered); per-agent settings cascade on top of that.
|
|
20
|
-
*/
|
|
21
|
-
import { lakebase } from "@databricks/appkit";
|
|
22
|
-
import { pluginUtils } from "@dbx-tools/appkit-shared";
|
|
23
|
-
import { fastembed } from "@mastra/fastembed";
|
|
24
|
-
import { Memory } from "@mastra/memory";
|
|
25
|
-
import { PgVector, PostgresStore } from "@mastra/pg";
|
|
26
|
-
import { randomUUID } from "node:crypto";
|
|
27
|
-
/**
|
|
28
|
-
* True when any plugin-level or per-agent setting could need the
|
|
29
|
-
* Lakebase pool. Used by `plugin.ts` to gate pool acquisition; the
|
|
30
|
-
* builder also acquires lazily so missed cases still fail with a
|
|
31
|
-
* clear lakebase-not-registered error.
|
|
32
|
-
*/
|
|
33
|
-
export function needsLakebase(config) {
|
|
34
|
-
if (settingNeedsSharedPool(config.storage))
|
|
35
|
-
return true;
|
|
36
|
-
if (settingNeedsSharedPool(config.memory))
|
|
37
|
-
return true;
|
|
38
|
-
const defs = collectAgentDefinitions(config);
|
|
39
|
-
return defs.some((d) => settingNeedsSharedPool(d.storage) || settingNeedsSharedPool(d.memory));
|
|
40
|
-
}
|
|
41
|
-
/**
|
|
42
|
-
* Look up the `lakebase` plugin and return its managed `pg.Pool`.
|
|
43
|
-
* Throws when the sibling plugin is not registered; enabling
|
|
44
|
-
* `storage` / `memory` without lakebase is a wiring bug, not a runtime
|
|
45
|
-
* condition we can recover from.
|
|
46
|
-
*/
|
|
47
|
-
export function resolveLakebasePool(context, caller) {
|
|
48
|
-
return pluginUtils.require(context, lakebase, caller).exports().pool;
|
|
49
|
-
}
|
|
50
|
-
/**
|
|
51
|
-
* Construct a per-agent {@link Memory} factory. Caches the shared
|
|
52
|
-
* `PgVector` singleton (built on first need) and the lazily-resolved
|
|
53
|
-
* Lakebase pool so each agent build is O(1) after the first.
|
|
54
|
-
*/
|
|
55
|
-
export function createMemoryBuilder(config, context) {
|
|
56
|
-
return new MemoryBuilder(config, context);
|
|
57
|
-
}
|
|
58
|
-
/**
|
|
59
|
-
* Builds one `Memory` per agent. Per-instance state keeps the shared
|
|
60
|
-
* `PgVector` and the resolved Lakebase pool alive across calls so
|
|
61
|
-
* registering N agents stays cheap.
|
|
62
|
-
*/
|
|
63
|
-
export class MemoryBuilder {
|
|
64
|
-
config;
|
|
65
|
-
context;
|
|
66
|
-
sharedVector;
|
|
67
|
-
pool;
|
|
68
|
-
constructor(config, context) {
|
|
69
|
-
this.config = config;
|
|
70
|
-
this.context = context;
|
|
71
|
-
}
|
|
72
|
-
/**
|
|
73
|
-
* Build a `Memory` for `agentId` after the plugin/agent cascade.
|
|
74
|
-
* Returns `undefined` when the agent has neither storage nor a
|
|
75
|
-
* vector store enabled - Mastra accepts a missing `memory` field
|
|
76
|
-
* and treats the agent as stateless.
|
|
77
|
-
*/
|
|
78
|
-
forAgent(agentId, def) {
|
|
79
|
-
const storageSetting = def.storage ?? this.config.storage;
|
|
80
|
-
const memorySetting = def.memory ?? this.config.memory;
|
|
81
|
-
const storage = this.buildStorage(agentId, storageSetting);
|
|
82
|
-
const vector = this.buildVector(memorySetting);
|
|
83
|
-
if (!storage && !vector)
|
|
84
|
-
return undefined;
|
|
85
|
-
return new Memory({
|
|
86
|
-
...(storage ? { storage } : {}),
|
|
87
|
-
...(vector ? { vector, embedder: fastembed } : {}),
|
|
88
|
-
options: {
|
|
89
|
-
lastMessages: 10,
|
|
90
|
-
...(vector
|
|
91
|
-
? { semanticRecall: { topK: 3, messageRange: 2 } }
|
|
92
|
-
: {}),
|
|
93
|
-
},
|
|
94
|
-
});
|
|
95
|
-
}
|
|
96
|
-
buildStorage(agentId, setting) {
|
|
97
|
-
if (!setting)
|
|
98
|
-
return undefined;
|
|
99
|
-
if (typeof setting === "boolean") {
|
|
100
|
-
return new PostgresStore({
|
|
101
|
-
id: `mastra-store__${agentId}`,
|
|
102
|
-
schemaName: `mastra_${agentId}`,
|
|
103
|
-
pool: this.requirePool(),
|
|
104
|
-
});
|
|
105
|
-
}
|
|
106
|
-
// Cast: `withId` guarantees `id` is set, but the distributive
|
|
107
|
-
// Omit + `id?: string` shape doesn't structurally narrow to the
|
|
108
|
-
// discriminated union members. Runtime shape is identical.
|
|
109
|
-
return new PostgresStore(withId(setting, `mastra-store__${agentId}`));
|
|
110
|
-
}
|
|
111
|
-
/**
|
|
112
|
-
* Resolve the agent's vector store. Cascade:
|
|
113
|
-
*
|
|
114
|
-
* - falsy: no vector.
|
|
115
|
-
* - `boolean` / `undefined-inheriting-true`: return the shared
|
|
116
|
-
* singleton (built lazily on first call). All agents that
|
|
117
|
-
* default-enable memory write into and recall from one index.
|
|
118
|
-
* - object: build a dedicated `PgVector` for this agent.
|
|
119
|
-
*/
|
|
120
|
-
buildVector(setting) {
|
|
121
|
-
if (!setting)
|
|
122
|
-
return undefined;
|
|
123
|
-
if (typeof setting === "boolean")
|
|
124
|
-
return this.getSharedVector();
|
|
125
|
-
return buildPgVector(setting);
|
|
126
|
-
}
|
|
127
|
-
getSharedVector() {
|
|
128
|
-
if (!this.sharedVector) {
|
|
129
|
-
this.sharedVector = buildSharedPgVector(this.requirePool());
|
|
130
|
-
}
|
|
131
|
-
return this.sharedVector;
|
|
132
|
-
}
|
|
133
|
-
requirePool() {
|
|
134
|
-
if (!this.pool) {
|
|
135
|
-
this.pool = resolveLakebasePool(this.context, this.config);
|
|
136
|
-
}
|
|
137
|
-
return this.pool;
|
|
138
|
-
}
|
|
139
|
-
}
|
|
140
|
-
/**
|
|
141
|
-
* Build the shared `PgVector` that backs the default
|
|
142
|
-
* `def.memory === true` case across every agent.
|
|
143
|
-
*
|
|
144
|
-
* `PgVector`'s constructor accepts only connection-style configs
|
|
145
|
-
* (`HostConfig` / `ConnectionStringConfig` / `ClientConfig`); there is
|
|
146
|
-
* no `{ pool }` shorthand the way `PostgresStore` has one. Worse, the
|
|
147
|
-
* constructor synchronously kicks off a `cacheWarmupPromise` IIFE that
|
|
148
|
-
* calls `this.pool.connect()` before returning, so we can't cleanly
|
|
149
|
-
* hand it an inert config and patch the pool afterwards.
|
|
150
|
-
*
|
|
151
|
-
* The trick: pass illegal-but-validation-passing placeholders so the
|
|
152
|
-
* warmup's `net.connect()` rejects synchronously with `RangeError`
|
|
153
|
-
* (Node validates `0 <= port < 65536`). The IIFE's `catch {}` swallows
|
|
154
|
-
* it, no DNS lookup or TCP attempt happens, and we then swap
|
|
155
|
-
* `pgVector.pool` to the lakebase pool. Every subsequent `PgVector`
|
|
156
|
-
* method reads `this.pool` at call time, so all real I/O goes through
|
|
157
|
-
* the lakebase pool from then on. The placeholder pool is `.end()`'d
|
|
158
|
-
* so its socket book-keeping is released.
|
|
159
|
-
*/
|
|
160
|
-
function buildSharedPgVector(pool) {
|
|
161
|
-
const vector = new PgVector({
|
|
162
|
-
id: `pg${randomUUID()}`,
|
|
163
|
-
host: "-1",
|
|
164
|
-
port: -1,
|
|
165
|
-
database: "_",
|
|
166
|
-
user: "_",
|
|
167
|
-
password: "_",
|
|
168
|
-
});
|
|
169
|
-
const placeholder = vector.pool;
|
|
170
|
-
vector.pool = pool;
|
|
171
|
-
void placeholder.end().catch(() => undefined);
|
|
172
|
-
return vector;
|
|
173
|
-
}
|
|
174
|
-
/** Per-agent dedicated `PgVector` (rare; opt-in via object override). */
|
|
175
|
-
function buildPgVector(setting) {
|
|
176
|
-
return new PgVector(withId(setting, `pg-vector__${randomUUID()}`));
|
|
177
|
-
}
|
|
178
|
-
/** True when this setting requires the shared Lakebase pool. */
|
|
179
|
-
function settingNeedsSharedPool(setting) {
|
|
180
|
-
return setting === true;
|
|
181
|
-
}
|
|
182
|
-
/** Walk the three shapes of `config.agents` into a flat list. */
|
|
183
|
-
function collectAgentDefinitions(config) {
|
|
184
|
-
const agents = config.agents;
|
|
185
|
-
if (!agents)
|
|
186
|
-
return [];
|
|
187
|
-
if (Array.isArray(agents))
|
|
188
|
-
return agents;
|
|
189
|
-
if (typeof agents.instructions === "string") {
|
|
190
|
-
return [agents];
|
|
191
|
-
}
|
|
192
|
-
return Object.values(agents);
|
|
193
|
-
}
|
|
194
|
-
/** Fill in a default `id` when the caller didn't supply one. */
|
|
195
|
-
function withId(value, fallback) {
|
|
196
|
-
return value.id ? value : { ...value, id: fallback };
|
|
197
|
-
}
|
package/dist/src/model.d.ts
DELETED
|
@@ -1,159 +0,0 @@
|
|
|
1
|
-
/**
|
|
2
|
-
* Databricks Model Serving resolver for Mastra agents.
|
|
3
|
-
*
|
|
4
|
-
* Each agent step calls {@link buildModel} with the active
|
|
5
|
-
* `RequestContext`. The user stamped by `MastraServer` carries an
|
|
6
|
-
* AppKit `WorkspaceClient`; we ask it for the workspace host and a
|
|
7
|
-
* fresh bearer header, then point Mastra's OpenAI-compatible provider
|
|
8
|
-
* at `/serving-endpoints` on that host.
|
|
9
|
-
*
|
|
10
|
-
* Model id resolution walks three sources before falling back to the
|
|
11
|
-
* hard-coded default, **in this priority order**:
|
|
12
|
-
*
|
|
13
|
-
* 1. Per-request override stashed by the auth middleware under
|
|
14
|
-
* {@link MASTRA_MODEL_OVERRIDE_KEY} (header / query / body).
|
|
15
|
-
* 2. The static `modelId` passed in by the agent / plugin (string
|
|
16
|
-
* sugar on `def.model` or `config.defaultModel`).
|
|
17
|
-
* 3. `DATABRICKS_SERVING_ENDPOINT_NAME` env var.
|
|
18
|
-
* 4. {@link FALLBACK_MODEL_ID}.
|
|
19
|
-
*
|
|
20
|
-
* Whatever wins is then fuzzy-matched against the live
|
|
21
|
-
* `/serving-endpoints` list ({@link listServingEndpoints}) so loose
|
|
22
|
-
* names like `"claude sonnet"` resolve to the real endpoint name.
|
|
23
|
-
* Fuzzy matching is best-effort: when the workspace client throws
|
|
24
|
-
* (network blip, expired token at cache-fill time) we fall back to
|
|
25
|
-
* the input verbatim and let Databricks return the canonical error.
|
|
26
|
-
*/
|
|
27
|
-
import type { MastraModelConfig } from "@mastra/core/llm";
|
|
28
|
-
import type { RequestContext } from "@mastra/core/request-context";
|
|
29
|
-
import { type MastraPluginConfig } from "./config.js";
|
|
30
|
-
/**
|
|
31
|
-
* Capability tiers for Databricks Foundation Model API endpoints.
|
|
32
|
-
*
|
|
33
|
-
* - {@link ModelTier.Thinking}: deepest reasoning / "thinking" models
|
|
34
|
-
* (Claude Opus, GPT-5.5 Pro, Gemini Pro, Llama 4 Maverick, etc).
|
|
35
|
-
* Highest cost and latency; reserve for hard multi-step reasoning.
|
|
36
|
-
* - {@link ModelTier.Balanced}: cost/latency sweet spot for general
|
|
37
|
-
* agent work (Claude Sonnet, GPT-5.x, Gemini Flash, Llama 3.3 70B).
|
|
38
|
-
* The right default for most agents.
|
|
39
|
-
* - {@link ModelTier.Fast}: cheap and quick; classification, routing,
|
|
40
|
-
* tool-arg extraction, simple summarisation (Claude Haiku, GPT-5
|
|
41
|
-
* mini/nano, Gemini Flash Lite, GPT-OSS 20B, Llama 3.1 8B).
|
|
42
|
-
*
|
|
43
|
-
* String enum so the value is the slug we use in cache keys, logs,
|
|
44
|
-
* and as the value users see in serialized configs.
|
|
45
|
-
*/
|
|
46
|
-
export declare enum ModelTier {
|
|
47
|
-
Thinking = "thinking",
|
|
48
|
-
Balanced = "balanced",
|
|
49
|
-
Fast = "fast"
|
|
50
|
-
}
|
|
51
|
-
/**
|
|
52
|
-
* Catalogue of Databricks-hosted Foundation Model API endpoints,
|
|
53
|
-
* grouped by capability {@link ModelTier} and then by provider. Each
|
|
54
|
-
* inner array is priority-ordered (most powerful first within the
|
|
55
|
-
* same provider+tier).
|
|
56
|
-
*
|
|
57
|
-
* Provider buckets:
|
|
58
|
-
*
|
|
59
|
-
* - `claude`: Anthropic Claude family (closed; flagship reasoning).
|
|
60
|
-
* - `gpt`: OpenAI GPT-5 family (closed; "ChatGPT" on Databricks FMAPI).
|
|
61
|
-
* - `gemini`: Google Gemini family (closed; multimodal + web-search).
|
|
62
|
-
* - `openSource`: open-weights models (widest regional / SKU availability).
|
|
63
|
-
*
|
|
64
|
-
* The list is curated by hand; refresh from the Databricks "supported
|
|
65
|
-
* foundation models" doc when new endpoints land.
|
|
66
|
-
*/
|
|
67
|
-
export declare const MODEL_CATALOG: {
|
|
68
|
-
readonly thinking: {
|
|
69
|
-
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"];
|
|
70
|
-
readonly gpt: readonly ["databricks-gpt-5-5-pro"];
|
|
71
|
-
readonly gemini: readonly ["databricks-gemini-3-1-pro", "databricks-gemini-3-pro", "databricks-gemini-2-5-pro"];
|
|
72
|
-
readonly openSource: readonly ["databricks-llama-4-maverick", "databricks-gpt-oss-120b", "databricks-meta-llama-3-1-405b-instruct"];
|
|
73
|
-
};
|
|
74
|
-
readonly balanced: {
|
|
75
|
-
readonly claude: readonly ["databricks-claude-sonnet-4-6", "databricks-claude-sonnet-4-5", "databricks-claude-sonnet-4"];
|
|
76
|
-
readonly gpt: readonly ["databricks-gpt-5-5", "databricks-gpt-5-4", "databricks-gpt-5-2", "databricks-gpt-5-1", "databricks-gpt-5"];
|
|
77
|
-
readonly gemini: readonly ["databricks-gemini-3-5-flash", "databricks-gemini-3-flash", "databricks-gemini-2-5-flash"];
|
|
78
|
-
readonly openSource: readonly ["databricks-meta-llama-3-3-70b-instruct", "databricks-qwen3-next-80b-a3b-instruct", "databricks-qwen35-122b-a10b"];
|
|
79
|
-
};
|
|
80
|
-
readonly fast: {
|
|
81
|
-
readonly claude: readonly ["databricks-claude-haiku-4-5"];
|
|
82
|
-
readonly gpt: readonly ["databricks-gpt-5-4-mini", "databricks-gpt-5-4-nano", "databricks-gpt-5-mini", "databricks-gpt-5-nano"];
|
|
83
|
-
readonly gemini: readonly ["databricks-gemini-3-1-flash-lite"];
|
|
84
|
-
readonly openSource: readonly ["databricks-gpt-oss-20b", "databricks-gemma-3-12b", "databricks-meta-llama-3-1-8b-instruct"];
|
|
85
|
-
};
|
|
86
|
-
};
|
|
87
|
-
/**
|
|
88
|
-
* Priority-ordered model ids for a single capability {@link ModelTier},
|
|
89
|
-
* interleaved across providers so a workspace missing the top Claude
|
|
90
|
-
* still lands on a flagship GPT / Gemini on the next probe.
|
|
91
|
-
*
|
|
92
|
-
* Provider order within the interleave: Claude, GPT, Gemini, then the
|
|
93
|
-
* open-weights tail appended verbatim as the universal floor (widest
|
|
94
|
-
* regional availability).
|
|
95
|
-
*
|
|
96
|
-
* @example
|
|
97
|
-
* ```ts
|
|
98
|
-
* mastra({
|
|
99
|
-
* defaultModelFallbacks: modelsForTier(ModelTier.Fast),
|
|
100
|
-
* });
|
|
101
|
-
* ```
|
|
102
|
-
*/
|
|
103
|
-
export declare function modelsForTier(tier: ModelTier): readonly string[];
|
|
104
|
-
/**
|
|
105
|
-
* Top model id at the given {@link ModelTier}. Sync; the agent-step
|
|
106
|
-
* resolver fuzzy-matches it against the workspace catalogue at call
|
|
107
|
-
* time, so this works even when the literal top pick isn't deployed.
|
|
108
|
-
*
|
|
109
|
-
* Use when wiring a tier-appropriate model into an agent definition:
|
|
110
|
-
*
|
|
111
|
-
* @example
|
|
112
|
-
* ```ts
|
|
113
|
-
* const classifier = createAgent({
|
|
114
|
-
* instructions: "Classify this email",
|
|
115
|
-
* model: modelForTier(ModelTier.Fast), // cheap, quick
|
|
116
|
-
* });
|
|
117
|
-
*
|
|
118
|
-
* const planner = createAgent({
|
|
119
|
-
* instructions: "Plan a multi-step migration",
|
|
120
|
-
* model: modelForTier(ModelTier.Thinking), // deep reasoning
|
|
121
|
-
* });
|
|
122
|
-
* ```
|
|
123
|
-
*/
|
|
124
|
-
export declare function modelForTier(tier: ModelTier): string;
|
|
125
|
-
/**
|
|
126
|
-
* Last-resort model ids used when neither `config.defaultModel`,
|
|
127
|
-
* per-agent `model`, nor `DATABRICKS_SERVING_ENDPOINT_NAME` is set.
|
|
128
|
-
*
|
|
129
|
-
* Walked in order at resolve time: the first id whose endpoint is
|
|
130
|
-
* actually present in the workspace's `/serving-endpoints` listing
|
|
131
|
-
* wins. Workspaces vary - not every region / SKU has every model,
|
|
132
|
-
* and the list of Foundation Model APIs evolves quickly - so the
|
|
133
|
-
* resolver degrades all the way from "best thinking model" down to
|
|
134
|
-
* "smallest commodity Llama" before giving up.
|
|
135
|
-
*
|
|
136
|
-
* Built by chaining the per-tier interleaves (Thinking -> Balanced
|
|
137
|
-
* -> Fast); within each tier the providers are round-robin-zipped
|
|
138
|
-
* (Claude, GPT, Gemini, then open-weights tail). Override the entire
|
|
139
|
-
* list via `MastraPluginConfig.defaultModelFallbacks` (e.g. to pin a
|
|
140
|
-
* regulated workspace to a specific approved subset, or to bias the
|
|
141
|
-
* priority toward a particular tier).
|
|
142
|
-
*/
|
|
143
|
-
export declare const FALLBACK_MODEL_IDS: readonly string[];
|
|
144
|
-
/** Optional overrides accepted by {@link buildModel}. */
|
|
145
|
-
export interface BuildModelOverrides {
|
|
146
|
-
/**
|
|
147
|
-
* Static model id from the agent / plugin config (string sugar on
|
|
148
|
-
* `def.model` or `config.defaultModel`). Loses to the per-request
|
|
149
|
-
* override but wins over env / fallback.
|
|
150
|
-
*/
|
|
151
|
-
modelId?: string;
|
|
152
|
-
}
|
|
153
|
-
/**
|
|
154
|
-
* Resolve a `MastraModelConfig` for the current agent step. Runs
|
|
155
|
-
* while `agent.stream` is inside the `asUser(req)` scope so tokens
|
|
156
|
-
* are user-scoped; outside an active user context the workspace
|
|
157
|
-
* client falls back to the service principal.
|
|
158
|
-
*/
|
|
159
|
-
export declare function buildModel(config: MastraPluginConfig, requestContext: RequestContext, overrides?: BuildModelOverrides): Promise<MastraModelConfig>;
|