@dbx-tools/appkit-mastra 0.1.58 → 0.1.67
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 +41 -1
- package/dist/index.d.ts +1229 -17
- package/dist/index.js +3348 -19
- package/package.json +17 -28
- package/src/agents.ts +7 -1
- package/src/config.ts +64 -0
- package/src/genie.ts +9 -21
- package/{index.ts → src/index.ts} +7 -9
- package/src/mcp.ts +89 -0
- package/src/model.ts +9 -3
- package/src/plugin.ts +51 -8
- package/src/serving.ts +5 -9
- package/src/statement.ts +8 -36
- package/dist/src/agents.d.ts +0 -316
- package/dist/src/agents.js +0 -470
- package/dist/src/chart.d.ts +0 -153
- package/dist/src/chart.js +0 -570
- package/dist/src/config.d.ts +0 -295
- package/dist/src/config.js +0 -55
- package/dist/src/genie.d.ts +0 -161
- package/dist/src/genie.js +0 -973
- package/dist/src/history.d.ts +0 -95
- package/dist/src/history.js +0 -278
- package/dist/src/memory.d.ts +0 -109
- package/dist/src/memory.js +0 -253
- package/dist/src/model.d.ts +0 -52
- package/dist/src/model.js +0 -198
- package/dist/src/observability.d.ts +0 -64
- package/dist/src/observability.js +0 -83
- package/dist/src/plugin.d.ts +0 -177
- package/dist/src/plugin.js +0 -554
- package/dist/src/processor.d.ts +0 -8
- package/dist/src/processor.js +0 -40
- package/dist/src/processors/strip-stale-charts.d.ts +0 -32
- package/dist/src/processors/strip-stale-charts.js +0 -98
- package/dist/src/server.d.ts +0 -51
- package/dist/src/server.js +0 -152
- package/dist/src/serving.d.ts +0 -65
- package/dist/src/serving.js +0 -79
- package/dist/src/statement.d.ts +0 -66
- package/dist/src/statement.js +0 -111
- package/dist/src/writer.d.ts +0 -23
- package/dist/src/writer.js +0 -37
- package/dist/tsconfig.build.tsbuildinfo +0 -1
package/dist/src/memory.js
DELETED
|
@@ -1,253 +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
|
-
* Additionally, {@link MemoryBuilder.instanceStorage} returns a
|
|
18
|
-
* **Mastra-instance-level** `PostgresStore` (schema `mastra_instance`)
|
|
19
|
-
* used for workflow snapshots - the persistence layer
|
|
20
|
-
* `agent.resumeStream()` reads from when waking a suspended
|
|
21
|
-
* `requireApproval` tool call. Per-agent stores are not enough for
|
|
22
|
-
* this: workflow runs are scoped to the Mastra instance, not an
|
|
23
|
-
* individual agent's `Memory`.
|
|
24
|
-
*
|
|
25
|
-
* Plugin-level `config.storage` / `config.memory` act as the baseline
|
|
26
|
-
* (auto-defaulted to `true` in `plugin.ts` when the `lakebase` plugin
|
|
27
|
-
* is registered); per-agent settings cascade on top of that.
|
|
28
|
-
*/
|
|
29
|
-
import { getUsernameWithApiLookup } from "@databricks/appkit";
|
|
30
|
-
import { logUtils } from "@dbx-tools/shared";
|
|
31
|
-
import { fastembed } from "@mastra/fastembed";
|
|
32
|
-
import { Memory } from "@mastra/memory";
|
|
33
|
-
import { PgVector, PostgresStore } from "@mastra/pg";
|
|
34
|
-
import { randomUUID } from "node:crypto";
|
|
35
|
-
import { Pool } from "pg";
|
|
36
|
-
const log = logUtils.logger("mastra/memory");
|
|
37
|
-
/**
|
|
38
|
-
* Build a dedicated **service-principal** Lakebase pool for Mastra
|
|
39
|
-
* memory from the lakebase plugin's resolved SP pg config.
|
|
40
|
-
*
|
|
41
|
-
* The plugin's `exports().pool` is a `RoutingPool` that switches to
|
|
42
|
-
* the per-user (OBO) pool whenever a query runs inside an `asUser`
|
|
43
|
-
* scope - exactly the context the mastra plugin establishes around
|
|
44
|
-
* every chat turn. Memory (threads / messages + semantic recall) must
|
|
45
|
-
* instead always act as the app service principal: it owns the
|
|
46
|
-
* auto-created `mastra_*` schemas (a per-user role usually can't
|
|
47
|
-
* `CREATE SCHEMA`) and is shared across users, so it cannot inherit a
|
|
48
|
-
* request's OBO identity.
|
|
49
|
-
*
|
|
50
|
-
* `pgConfig` must be the plugin's `exports().getPgConfig()` evaluated
|
|
51
|
-
* **outside** any `asUser` scope (i.e. during setup), so it carries
|
|
52
|
-
* the SP connection target, OAuth token-refresh `password` callback,
|
|
53
|
-
* and any `lakebase({ pool })` tuning overrides - all of which this
|
|
54
|
-
* pool inherits. See the call site in `plugin.ts`.
|
|
55
|
-
*/
|
|
56
|
-
export async function createServicePrincipalPool(pgConfig) {
|
|
57
|
-
// `getPgConfig()` resolves the SP username synchronously from
|
|
58
|
-
// `PGUSER` / `DATABRICKS_CLIENT_ID`; fall back to the async API
|
|
59
|
-
// lookup (e.g. local dev authenticating via PAT) so the pool always
|
|
60
|
-
// has an identity to connect with.
|
|
61
|
-
const user = pgConfig.user ?? (await getUsernameWithApiLookup());
|
|
62
|
-
return new Pool({ ...pgConfig, user });
|
|
63
|
-
}
|
|
64
|
-
/**
|
|
65
|
-
* True when any plugin-level or per-agent setting could need the
|
|
66
|
-
* Lakebase pool. Used by `plugin.ts` to gate creation of the
|
|
67
|
-
* service-principal pool and the {@link MemoryBuilder} that consumes
|
|
68
|
-
* it; when false neither is built.
|
|
69
|
-
*/
|
|
70
|
-
export function needsLakebase(config) {
|
|
71
|
-
if (settingNeedsSharedPool(config.storage))
|
|
72
|
-
return true;
|
|
73
|
-
if (settingNeedsSharedPool(config.memory))
|
|
74
|
-
return true;
|
|
75
|
-
const defs = collectAgentDefinitions(config);
|
|
76
|
-
return defs.some((d) => settingNeedsSharedPool(d.storage) || settingNeedsSharedPool(d.memory));
|
|
77
|
-
}
|
|
78
|
-
/**
|
|
79
|
-
* Construct a per-agent {@link Memory} factory bound to the supplied
|
|
80
|
-
* service-principal pool (see {@link createServicePrincipalPool}).
|
|
81
|
-
* Caches the shared `PgVector` singleton (built on first need) so each
|
|
82
|
-
* agent build is O(1) after the first.
|
|
83
|
-
*/
|
|
84
|
-
export function createMemoryBuilder(config, servicePrincipalPool) {
|
|
85
|
-
return new MemoryBuilder(config, servicePrincipalPool);
|
|
86
|
-
}
|
|
87
|
-
/**
|
|
88
|
-
* Builds one `Memory` per agent against a shared service-principal
|
|
89
|
-
* Lakebase pool. Per-instance state keeps the shared `PgVector` alive
|
|
90
|
-
* across calls so registering N agents stays cheap.
|
|
91
|
-
*/
|
|
92
|
-
export class MemoryBuilder {
|
|
93
|
-
config;
|
|
94
|
-
servicePrincipalPool;
|
|
95
|
-
sharedVector;
|
|
96
|
-
constructor(config, servicePrincipalPool) {
|
|
97
|
-
this.config = config;
|
|
98
|
-
this.servicePrincipalPool = servicePrincipalPool;
|
|
99
|
-
}
|
|
100
|
-
/**
|
|
101
|
-
* Build a `Memory` for `agentId` after the plugin/agent cascade.
|
|
102
|
-
* Returns `undefined` when the agent has neither storage nor a
|
|
103
|
-
* vector store enabled - Mastra accepts a missing `memory` field
|
|
104
|
-
* and treats the agent as stateless.
|
|
105
|
-
*/
|
|
106
|
-
/**
|
|
107
|
-
* Build the Mastra-instance-level storage used for workflow
|
|
108
|
-
* snapshots. Returns `undefined` when plugin-level `storage` is
|
|
109
|
-
* disabled, in which case `agent.resumeStream()` (and therefore
|
|
110
|
-
* the `requireApproval` flow) will not be available.
|
|
111
|
-
*
|
|
112
|
-
* The store lives in a dedicated `mastra_instance` schema so it
|
|
113
|
-
* never collides with per-agent `mastra_<agentId>` namespaces.
|
|
114
|
-
* Workflow snapshots are not per-agent state; they belong to the
|
|
115
|
-
* `Mastra` instance that owns the workflow execution.
|
|
116
|
-
*/
|
|
117
|
-
instanceStorage() {
|
|
118
|
-
const setting = this.config.storage;
|
|
119
|
-
if (!setting)
|
|
120
|
-
return undefined;
|
|
121
|
-
if (typeof setting === "object") {
|
|
122
|
-
return new PostgresStore(withId(setting, "mastra-store__instance"));
|
|
123
|
-
}
|
|
124
|
-
return new PostgresStore({
|
|
125
|
-
id: "mastra-store__instance",
|
|
126
|
-
schemaName: "mastra_instance",
|
|
127
|
-
pool: this.servicePrincipalPool,
|
|
128
|
-
});
|
|
129
|
-
}
|
|
130
|
-
forAgent(agentId, def) {
|
|
131
|
-
const storageSetting = def.storage ?? this.config.storage;
|
|
132
|
-
const memorySetting = def.memory ?? this.config.memory;
|
|
133
|
-
const storage = this.buildStorage(agentId, storageSetting);
|
|
134
|
-
const vector = this.buildVector(memorySetting);
|
|
135
|
-
if (!storage && !vector) {
|
|
136
|
-
log.debug("agent:stateless", { agentId });
|
|
137
|
-
return undefined;
|
|
138
|
-
}
|
|
139
|
-
log.debug("agent:configured", {
|
|
140
|
-
agentId,
|
|
141
|
-
storage: storage !== undefined,
|
|
142
|
-
vector: vector !== undefined,
|
|
143
|
-
vectorMode: vector === undefined
|
|
144
|
-
? "off"
|
|
145
|
-
: typeof memorySetting === "object"
|
|
146
|
-
? "dedicated"
|
|
147
|
-
: "shared",
|
|
148
|
-
});
|
|
149
|
-
return new Memory({
|
|
150
|
-
...(storage ? { storage } : {}),
|
|
151
|
-
...(vector ? { vector, embedder: fastembed } : {}),
|
|
152
|
-
options: {
|
|
153
|
-
lastMessages: 10,
|
|
154
|
-
...(vector ? { semanticRecall: { topK: 3, messageRange: 2 } } : {}),
|
|
155
|
-
},
|
|
156
|
-
});
|
|
157
|
-
}
|
|
158
|
-
buildStorage(agentId, setting) {
|
|
159
|
-
if (!setting)
|
|
160
|
-
return undefined;
|
|
161
|
-
if (typeof setting === "boolean") {
|
|
162
|
-
return new PostgresStore({
|
|
163
|
-
id: `mastra-store__${agentId}`,
|
|
164
|
-
schemaName: `mastra_${agentId}`,
|
|
165
|
-
pool: this.servicePrincipalPool,
|
|
166
|
-
});
|
|
167
|
-
}
|
|
168
|
-
// Cast: `withId` guarantees `id` is set, but the distributive
|
|
169
|
-
// Omit + `id?: string` shape doesn't structurally narrow to the
|
|
170
|
-
// discriminated union members. Runtime shape is identical.
|
|
171
|
-
return new PostgresStore(withId(setting, `mastra-store__${agentId}`));
|
|
172
|
-
}
|
|
173
|
-
/**
|
|
174
|
-
* Resolve the agent's vector store. Cascade:
|
|
175
|
-
*
|
|
176
|
-
* - falsy: no vector.
|
|
177
|
-
* - `boolean` / `undefined-inheriting-true`: return the shared
|
|
178
|
-
* singleton (built lazily on first call). All agents that
|
|
179
|
-
* default-enable memory write into and recall from one index.
|
|
180
|
-
* - object: build a dedicated `PgVector` for this agent.
|
|
181
|
-
*/
|
|
182
|
-
buildVector(setting) {
|
|
183
|
-
if (!setting)
|
|
184
|
-
return undefined;
|
|
185
|
-
if (typeof setting === "boolean")
|
|
186
|
-
return this.getSharedVector();
|
|
187
|
-
return buildPgVector(setting);
|
|
188
|
-
}
|
|
189
|
-
getSharedVector() {
|
|
190
|
-
if (!this.sharedVector) {
|
|
191
|
-
this.sharedVector = buildSharedPgVector(this.servicePrincipalPool);
|
|
192
|
-
}
|
|
193
|
-
return this.sharedVector;
|
|
194
|
-
}
|
|
195
|
-
}
|
|
196
|
-
/**
|
|
197
|
-
* Build the shared `PgVector` that backs the default
|
|
198
|
-
* `def.memory === true` case across every agent.
|
|
199
|
-
*
|
|
200
|
-
* `PgVector`'s constructor accepts only connection-style configs
|
|
201
|
-
* (`HostConfig` / `ConnectionStringConfig` / `ClientConfig`); there is
|
|
202
|
-
* no `{ pool }` shorthand the way `PostgresStore` has one. Worse, the
|
|
203
|
-
* constructor synchronously kicks off a `cacheWarmupPromise` IIFE that
|
|
204
|
-
* calls `this.pool.connect()` before returning, so we can't cleanly
|
|
205
|
-
* hand it an inert config and patch the pool afterwards.
|
|
206
|
-
*
|
|
207
|
-
* The trick: pass illegal-but-validation-passing placeholders so the
|
|
208
|
-
* warmup's `net.connect()` rejects synchronously with `RangeError`
|
|
209
|
-
* (Node validates `0 <= port < 65536`). The IIFE's `catch {}` swallows
|
|
210
|
-
* it, no DNS lookup or TCP attempt happens, and we then swap
|
|
211
|
-
* `pgVector.pool` to the lakebase pool. Every subsequent `PgVector`
|
|
212
|
-
* method reads `this.pool` at call time, so all real I/O goes through
|
|
213
|
-
* the lakebase pool from then on. The placeholder pool is `.end()`'d
|
|
214
|
-
* so its socket book-keeping is released.
|
|
215
|
-
*/
|
|
216
|
-
function buildSharedPgVector(pool) {
|
|
217
|
-
const vector = new PgVector({
|
|
218
|
-
id: `pg${randomUUID()}`,
|
|
219
|
-
host: "-1",
|
|
220
|
-
port: -1,
|
|
221
|
-
database: "_",
|
|
222
|
-
user: "_",
|
|
223
|
-
password: "_",
|
|
224
|
-
});
|
|
225
|
-
const placeholder = vector.pool;
|
|
226
|
-
vector.pool = pool;
|
|
227
|
-
void placeholder.end().catch(() => undefined);
|
|
228
|
-
return vector;
|
|
229
|
-
}
|
|
230
|
-
/** Per-agent dedicated `PgVector` (rare; opt-in via object override). */
|
|
231
|
-
function buildPgVector(setting) {
|
|
232
|
-
return new PgVector(withId(setting, `pg-vector__${randomUUID()}`));
|
|
233
|
-
}
|
|
234
|
-
/** True when this setting requires the shared Lakebase pool. */
|
|
235
|
-
function settingNeedsSharedPool(setting) {
|
|
236
|
-
return setting === true;
|
|
237
|
-
}
|
|
238
|
-
/** Walk the three shapes of `config.agents` into a flat list. */
|
|
239
|
-
function collectAgentDefinitions(config) {
|
|
240
|
-
const agents = config.agents;
|
|
241
|
-
if (!agents)
|
|
242
|
-
return [];
|
|
243
|
-
if (Array.isArray(agents))
|
|
244
|
-
return agents;
|
|
245
|
-
if (typeof agents.instructions === "string") {
|
|
246
|
-
return [agents];
|
|
247
|
-
}
|
|
248
|
-
return Object.values(agents);
|
|
249
|
-
}
|
|
250
|
-
/** Fill in a default `id` when the caller didn't supply one. */
|
|
251
|
-
function withId(value, fallback) {
|
|
252
|
-
return value.id ? value : { ...value, id: fallback };
|
|
253
|
-
}
|
package/dist/src/model.d.ts
DELETED
|
@@ -1,52 +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
|
-
* This module only adds the Mastra-specific glue. The actual model
|
|
11
|
-
* selection - listing the workspace catalogue and resolving an
|
|
12
|
-
* explicit name / class / fallback chain to a real endpoint id - lives
|
|
13
|
-
* in `@dbx-tools/model` ({@link selectModel}) so non-Mastra consumers
|
|
14
|
-
* (e.g. a job that just needs a model name) can reuse it. Here we
|
|
15
|
-
* assemble the explicit ask from Mastra's request context (the
|
|
16
|
-
* per-request override under {@link MASTRA_MODEL_OVERRIDE_KEY}, the
|
|
17
|
-
* agent / plugin `modelId`, or `DATABRICKS_SERVING_ENDPOINT_NAME`),
|
|
18
|
-
* pass the plugin's fuzzy / class / fallback knobs through, and wrap
|
|
19
|
-
* the resolved id in the OpenAI-compatible provider config Mastra
|
|
20
|
-
* expects. Catalogue fetches fail loud: network / auth errors
|
|
21
|
-
* propagate so callers see the real SDK message.
|
|
22
|
-
*/
|
|
23
|
-
import { type ModelClass } from "@dbx-tools/model";
|
|
24
|
-
import type { MastraModelConfig } from "@mastra/core/llm";
|
|
25
|
-
import type { RequestContext } from "@mastra/core/request-context";
|
|
26
|
-
import { type MastraPluginConfig } from "./config.js";
|
|
27
|
-
export { classifyEndpoints, FALLBACK_MODEL_IDS, ModelClass, modelForClass, modelsForClass, } from "@dbx-tools/model";
|
|
28
|
-
/** Optional overrides accepted by {@link buildModel}. */
|
|
29
|
-
export interface BuildModelOverrides {
|
|
30
|
-
/**
|
|
31
|
-
* Static model id from the agent / plugin config (string sugar on
|
|
32
|
-
* `def.model` or `config.defaultModel`). Loses to the per-request
|
|
33
|
-
* override but wins over env / class / fallback.
|
|
34
|
-
*/
|
|
35
|
-
modelId?: string;
|
|
36
|
-
/**
|
|
37
|
-
* Chat capability class to resolve when no explicit model id is
|
|
38
|
-
* supplied. Used by internal agents (e.g. the chart planner asks for
|
|
39
|
-
* {@link ModelClass.ChatFast}) to express intent without pinning an
|
|
40
|
-
* endpoint name; the live catalogue is classified and the top
|
|
41
|
-
* available model in the class is chosen, falling back to the
|
|
42
|
-
* class's static list when the workspace has none.
|
|
43
|
-
*/
|
|
44
|
-
modelClass?: ModelClass;
|
|
45
|
-
}
|
|
46
|
-
/**
|
|
47
|
-
* Resolve a `MastraModelConfig` for the current agent step. Runs
|
|
48
|
-
* while `agent.stream` is inside the `asUser(req)` scope so tokens
|
|
49
|
-
* are user-scoped; outside an active user context the workspace
|
|
50
|
-
* client falls back to the service principal.
|
|
51
|
-
*/
|
|
52
|
-
export declare function buildModel(config: MastraPluginConfig, requestContext: RequestContext, overrides?: BuildModelOverrides): Promise<MastraModelConfig>;
|
package/dist/src/model.js
DELETED
|
@@ -1,198 +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
|
-
* This module only adds the Mastra-specific glue. The actual model
|
|
11
|
-
* selection - listing the workspace catalogue and resolving an
|
|
12
|
-
* explicit name / class / fallback chain to a real endpoint id - lives
|
|
13
|
-
* in `@dbx-tools/model` ({@link selectModel}) so non-Mastra consumers
|
|
14
|
-
* (e.g. a job that just needs a model name) can reuse it. Here we
|
|
15
|
-
* assemble the explicit ask from Mastra's request context (the
|
|
16
|
-
* per-request override under {@link MASTRA_MODEL_OVERRIDE_KEY}, the
|
|
17
|
-
* agent / plugin `modelId`, or `DATABRICKS_SERVING_ENDPOINT_NAME`),
|
|
18
|
-
* pass the plugin's fuzzy / class / fallback knobs through, and wrap
|
|
19
|
-
* the resolved id in the OpenAI-compatible provider config Mastra
|
|
20
|
-
* expects. Catalogue fetches fail loud: network / auth errors
|
|
21
|
-
* propagate so callers see the real SDK message.
|
|
22
|
-
*/
|
|
23
|
-
import { parseModelClass, selectModel } from "@dbx-tools/model";
|
|
24
|
-
import { commonUtils, logUtils, netUtils, stringUtils } from "@dbx-tools/shared";
|
|
25
|
-
import { MASTRA_USER_KEY } from "./config.js";
|
|
26
|
-
import { MASTRA_MODEL_OVERRIDE_KEY, resolveServingConfig } from "./serving.js";
|
|
27
|
-
export { classifyEndpoints, FALLBACK_MODEL_IDS, ModelClass, modelForClass, modelsForClass, } from "@dbx-tools/model";
|
|
28
|
-
/**
|
|
29
|
-
* Resolve a `MastraModelConfig` for the current agent step. Runs
|
|
30
|
-
* while `agent.stream` is inside the `asUser(req)` scope so tokens
|
|
31
|
-
* are user-scoped; outside an active user context the workspace
|
|
32
|
-
* client falls back to the service principal.
|
|
33
|
-
*/
|
|
34
|
-
export async function buildModel(config, requestContext, overrides = {}) {
|
|
35
|
-
void setupFetchInterceptor();
|
|
36
|
-
const user = requestContext.get(MASTRA_USER_KEY);
|
|
37
|
-
const clientConfig = user.executionContext.client.config;
|
|
38
|
-
const host = (await clientConfig.getHost()).toString();
|
|
39
|
-
const headers = new Headers();
|
|
40
|
-
await clientConfig.authenticate(headers);
|
|
41
|
-
// The OpenAI Node SDK appends paths like `/chat/completions` to whatever
|
|
42
|
-
// URL we hand it. Drop the trailing slash so the resulting URL stays
|
|
43
|
-
// well-formed (`/serving-endpoints/chat/completions`).
|
|
44
|
-
const url = new URL("/serving-endpoints", host).toString().replace(/\/$/, "");
|
|
45
|
-
const log = logUtils.logger(config);
|
|
46
|
-
const serving = resolveServingConfig(config);
|
|
47
|
-
const override = serving.allowOverride
|
|
48
|
-
? requestContext.get(MASTRA_MODEL_OVERRIDE_KEY)
|
|
49
|
-
: undefined;
|
|
50
|
-
// The override / agent default / env value can be either a concrete
|
|
51
|
-
// endpoint name or a model class slug ("chat-thinking" /
|
|
52
|
-
// "chat-balanced" / "chat-fast"). A class slug becomes a class intent
|
|
53
|
-
// (let the live catalogue pick the best model in that band); anything
|
|
54
|
-
// else is an explicit name fuzzy-matched against the catalogue. An
|
|
55
|
-
// internal `overrides.modelClass` (e.g. the chart planner) is the
|
|
56
|
-
// floor when nothing was requested.
|
|
57
|
-
const requested = override ?? overrides.modelId ?? process.env.DATABRICKS_SERVING_ENDPOINT_NAME;
|
|
58
|
-
const requestedClass = requested !== undefined ? parseModelClass(requested) : null;
|
|
59
|
-
const explicit = requestedClass === null ? requested : undefined;
|
|
60
|
-
const modelClass = requestedClass ?? overrides.modelClass;
|
|
61
|
-
const { modelId, source } = await selectModel(user.executionContext.client, host, {
|
|
62
|
-
...(explicit !== undefined ? { explicit } : {}),
|
|
63
|
-
fuzzy: serving.fuzzy,
|
|
64
|
-
threshold: serving.threshold,
|
|
65
|
-
...(modelClass !== undefined ? { modelClass } : {}),
|
|
66
|
-
fallbacks: serving.fallbacks,
|
|
67
|
-
ttlMs: serving.ttlMs,
|
|
68
|
-
});
|
|
69
|
-
log.debug("model selected", { modelId, source, requested });
|
|
70
|
-
return {
|
|
71
|
-
providerId: config.providerId ?? "databricks",
|
|
72
|
-
modelId,
|
|
73
|
-
url,
|
|
74
|
-
headers: Object.fromEntries(headers.entries()),
|
|
75
|
-
};
|
|
76
|
-
}
|
|
77
|
-
/** Path prefix that identifies a Databricks Model Serving REST call. */
|
|
78
|
-
const SERVING_ENDPOINTS_PATH_PREFIX = "/serving-endpoints/";
|
|
79
|
-
/**
|
|
80
|
-
* Install a single shared `globalThis.fetch` wrapper for every POST to
|
|
81
|
-
* `/serving-endpoints/...`. The wrapper does two things:
|
|
82
|
-
*
|
|
83
|
-
* 1. Rewrites the outgoing `messages` array to repair Mastra/AI SDK
|
|
84
|
-
* stream-replay quirks that Databricks-hosted Claude rejects (see
|
|
85
|
-
* {@link sanitizeServingMessages}).
|
|
86
|
-
* 2. At `LOG_LEVEL=debug`, dumps the (post-sanitize) JSON body so
|
|
87
|
-
* 4xx debugging doesn't have to fight AI SDK's `[Array]`
|
|
88
|
-
* formatter.
|
|
89
|
-
*
|
|
90
|
-
* Safe to call from any hot path: {@link commonUtils.memoize} ensures
|
|
91
|
-
* the wrapper is installed at most once per process, so subsequent
|
|
92
|
-
* calls collapse to a single cached promise even when
|
|
93
|
-
* {@link buildModel} fires on every agent step.
|
|
94
|
-
*/
|
|
95
|
-
const setupFetchInterceptor = commonUtils.memoize(() => {
|
|
96
|
-
const log = logUtils.logger("mastra/llm");
|
|
97
|
-
const original = globalThis.fetch.bind(globalThis);
|
|
98
|
-
globalThis.fetch = (async (input, init) => {
|
|
99
|
-
const url = netUtils.urlBuilder(input);
|
|
100
|
-
if (!url ||
|
|
101
|
-
!url.pathname.startsWith(SERVING_ENDPOINTS_PATH_PREFIX) ||
|
|
102
|
-
typeof init?.body !== "string") {
|
|
103
|
-
return original(input, init);
|
|
104
|
-
}
|
|
105
|
-
const rewritten = rewriteServingBody(init.body);
|
|
106
|
-
if (rewritten !== init.body) {
|
|
107
|
-
init = { ...init, body: rewritten };
|
|
108
|
-
}
|
|
109
|
-
try {
|
|
110
|
-
log.debug("POST", { url: url.toString(), body: JSON.parse(rewritten) });
|
|
111
|
-
}
|
|
112
|
-
catch {
|
|
113
|
-
log.debug("POST", { url: url.toString(), bodyType: "non-JSON" });
|
|
114
|
-
}
|
|
115
|
-
return original(input, init);
|
|
116
|
-
});
|
|
117
|
-
});
|
|
118
|
-
/**
|
|
119
|
-
* Parse, sanitize, and re-serialize a `/serving-endpoints/...` POST
|
|
120
|
-
* body. Returns the original string verbatim when the body is not
|
|
121
|
-
* JSON, has no `messages`, or no rewrite was needed; this lets the
|
|
122
|
-
* caller skip the allocation of a new `init` object in the common
|
|
123
|
-
* pass-through case.
|
|
124
|
-
*/
|
|
125
|
-
function rewriteServingBody(body) {
|
|
126
|
-
let parsed;
|
|
127
|
-
try {
|
|
128
|
-
parsed = JSON.parse(body);
|
|
129
|
-
}
|
|
130
|
-
catch {
|
|
131
|
-
return body;
|
|
132
|
-
}
|
|
133
|
-
if (!Array.isArray(parsed.messages))
|
|
134
|
-
return body;
|
|
135
|
-
const changed = sanitizeServingMessages(parsed.messages);
|
|
136
|
-
return changed ? JSON.stringify(parsed) : body;
|
|
137
|
-
}
|
|
138
|
-
/**
|
|
139
|
-
* Repair a Mastra/AI SDK message replay that Databricks-hosted Claude
|
|
140
|
-
* rejects with `"This model does not support assistant message
|
|
141
|
-
* prefill. The conversation must end with a user message."`.
|
|
142
|
-
*
|
|
143
|
-
* The bug pattern: when an assistant turn streams text *and* a
|
|
144
|
-
* `tool_call`, the AI SDK persists them as two separate assistant
|
|
145
|
-
* entries (text-only and tool-call-only). On the next agent step the
|
|
146
|
-
* tool-call entry is replayed *before* the tool result and the
|
|
147
|
-
* text entry is replayed *after* it, so the conversation ends with a
|
|
148
|
-
* trailing assistant text message. Anthropic interprets that as a
|
|
149
|
-
* prefill request and rejects it on Databricks (the upstream Bedrock
|
|
150
|
-
* route disallows prefill).
|
|
151
|
-
*
|
|
152
|
-
* Fix: when the last message is an assistant text with no `tool_calls`
|
|
153
|
-
* and the chain immediately before it is `assistant(tool_calls=...)`
|
|
154
|
-
* followed only by `tool(...)` results, fold the trailing text back
|
|
155
|
-
* into the `content` of that opening assistant and drop the duplicate.
|
|
156
|
-
* The result is the canonical OpenAI shape
|
|
157
|
-
* `[..., user, assistant(text + tool_calls), tool(...)]` which both
|
|
158
|
-
* Databricks Claude and every other endpoint accept.
|
|
159
|
-
*
|
|
160
|
-
* Mutates `messages` in place; returns `true` when something changed
|
|
161
|
-
* so the caller knows whether to re-serialize.
|
|
162
|
-
*/
|
|
163
|
-
function sanitizeServingMessages(messages) {
|
|
164
|
-
if (messages.length < 2)
|
|
165
|
-
return false;
|
|
166
|
-
const last = messages[messages.length - 1];
|
|
167
|
-
if (!last ||
|
|
168
|
-
last.role !== "assistant" ||
|
|
169
|
-
(last.tool_calls && last.tool_calls.length > 0)) {
|
|
170
|
-
return false;
|
|
171
|
-
}
|
|
172
|
-
// Walk back through any contiguous tool-result messages to find the
|
|
173
|
-
// assistant turn that opened this tool sequence.
|
|
174
|
-
let i = messages.length - 2;
|
|
175
|
-
while (i >= 0 && messages[i]?.role === "tool")
|
|
176
|
-
i--;
|
|
177
|
-
if (i < 0)
|
|
178
|
-
return false;
|
|
179
|
-
const opener = messages[i];
|
|
180
|
-
if (!opener ||
|
|
181
|
-
opener.role !== "assistant" ||
|
|
182
|
-
!opener.tool_calls ||
|
|
183
|
-
opener.tool_calls.length === 0) {
|
|
184
|
-
return false;
|
|
185
|
-
}
|
|
186
|
-
// `trimToNull` collapses the `typeof string && trimmed` dance and
|
|
187
|
-
// drops blank fragments before the `\n\n` join below, so the merge
|
|
188
|
-
// never introduces stray leading / trailing whitespace.
|
|
189
|
-
const merged = [
|
|
190
|
-
stringUtils.trimToNull(opener.content),
|
|
191
|
-
stringUtils.trimToNull(last.content),
|
|
192
|
-
]
|
|
193
|
-
.filter((s) => s !== null)
|
|
194
|
-
.join("\n\n");
|
|
195
|
-
opener.content = merged;
|
|
196
|
-
messages.pop();
|
|
197
|
-
return true;
|
|
198
|
-
}
|
|
@@ -1,64 +0,0 @@
|
|
|
1
|
-
/**
|
|
2
|
-
* Mastra observability wired through the same OTel pipeline AppKit's
|
|
3
|
-
* built-in plugins (e.g. `agents`) use, via `@mastra/otel-bridge`.
|
|
4
|
-
*
|
|
5
|
-
* How traces flow:
|
|
6
|
-
*
|
|
7
|
-
* 1. `@databricks/appkit` boots a global `NodeSDK` in
|
|
8
|
-
* `TelemetryManager.initialize()` (during `createApp`) when
|
|
9
|
-
* `OTEL_EXPORTER_OTLP_ENDPOINT` is set in the process env.
|
|
10
|
-
* 2. Every AppKit plugin span (e.g. the `agents` plugin's
|
|
11
|
-
* `executeStream`) is created via the global OTel tracer
|
|
12
|
-
* (`trace.getTracer(<plugin>)`), so it lands on that NodeSDK and
|
|
13
|
-
* is shipped through its OTLP exporter.
|
|
14
|
-
* 3. The Mastra `OtelBridge` ALSO creates real OTel spans on the same
|
|
15
|
-
* global tracer for every Mastra operation (agent runs, model
|
|
16
|
-
* calls, tool invocations, workflow steps). They inherit the
|
|
17
|
-
* ambient OTel context, so when Mastra is invoked from inside an
|
|
18
|
-
* AppKit HTTP span the trace stays connected.
|
|
19
|
-
*
|
|
20
|
-
* Net effect: Mastra spans get exactly the treatment AppKit's
|
|
21
|
-
* `agents` plugin gets. No custom OTLP pipeline lives in this
|
|
22
|
-
* package; the OTLP endpoint, headers, and resource attributes are
|
|
23
|
-
* driven by the standard OTel env vars
|
|
24
|
-
* (`OTEL_EXPORTER_OTLP_ENDPOINT`, `OTEL_EXPORTER_OTLP_HEADERS`,
|
|
25
|
-
* `OTEL_SERVICE_NAME`, `OTEL_RESOURCE_ATTRIBUTES`, ...) and consumed
|
|
26
|
-
* by AppKit's `TelemetryManager`. Set those once and both AppKit and
|
|
27
|
-
* Mastra spans end up at the same backend.
|
|
28
|
-
*
|
|
29
|
-
* When `OTEL_EXPORTER_OTLP_ENDPOINT` is unset the bridge's spans go
|
|
30
|
-
* to the global noop tracer, mirroring how the `agents` plugin
|
|
31
|
-
* silently no-ops in the same situation.
|
|
32
|
-
*/
|
|
33
|
-
import { Observability } from "@mastra/observability";
|
|
34
|
-
export interface BuildObservabilityOptions {
|
|
35
|
-
/**
|
|
36
|
-
* Service name attached to the Mastra `Observability` config. Used
|
|
37
|
-
* as the tracer scope name on bridged OTel spans (the `service.name`
|
|
38
|
-
* resource attribute is owned by AppKit's `TelemetryManager` instead
|
|
39
|
-
* - it reads `OTEL_SERVICE_NAME` / `DATABRICKS_APP_NAME` at
|
|
40
|
-
* `createApp` time).
|
|
41
|
-
*
|
|
42
|
-
* Defaults to project name then `"mastra"`.
|
|
43
|
-
*/
|
|
44
|
-
serviceName?: string;
|
|
45
|
-
/**
|
|
46
|
-
* `RequestContext` keys to extract as span metadata on every Mastra
|
|
47
|
-
* trace. Defaults to {@link TRACE_REQUEST_CONTEXT_KEYS} (user id,
|
|
48
|
-
* thread id, request id, environment, model override, ...).
|
|
49
|
-
*
|
|
50
|
-
* Supports dot notation for nested values per the Mastra docs.
|
|
51
|
-
*/
|
|
52
|
-
requestContextKeys?: readonly string[];
|
|
53
|
-
}
|
|
54
|
-
/**
|
|
55
|
-
* Build a Mastra `Observability` whose spans ride AppKit's global
|
|
56
|
-
* OTel pipeline via `@mastra/otel-bridge`.
|
|
57
|
-
*
|
|
58
|
-
* Returns `undefined` only if someone explicitly opts out in the
|
|
59
|
-
* future; today it always returns an `Observability` because the
|
|
60
|
-
* bridge degrades gracefully (no-op tracer) when no global OTel SDK
|
|
61
|
-
* is registered. Callers can spread `...(observability ? { observability } : {})`
|
|
62
|
-
* either way to stay forward-compatible.
|
|
63
|
-
*/
|
|
64
|
-
export declare function buildObservability(options?: BuildObservabilityOptions): Promise<Observability | undefined>;
|
|
@@ -1,83 +0,0 @@
|
|
|
1
|
-
/**
|
|
2
|
-
* Mastra observability wired through the same OTel pipeline AppKit's
|
|
3
|
-
* built-in plugins (e.g. `agents`) use, via `@mastra/otel-bridge`.
|
|
4
|
-
*
|
|
5
|
-
* How traces flow:
|
|
6
|
-
*
|
|
7
|
-
* 1. `@databricks/appkit` boots a global `NodeSDK` in
|
|
8
|
-
* `TelemetryManager.initialize()` (during `createApp`) when
|
|
9
|
-
* `OTEL_EXPORTER_OTLP_ENDPOINT` is set in the process env.
|
|
10
|
-
* 2. Every AppKit plugin span (e.g. the `agents` plugin's
|
|
11
|
-
* `executeStream`) is created via the global OTel tracer
|
|
12
|
-
* (`trace.getTracer(<plugin>)`), so it lands on that NodeSDK and
|
|
13
|
-
* is shipped through its OTLP exporter.
|
|
14
|
-
* 3. The Mastra `OtelBridge` ALSO creates real OTel spans on the same
|
|
15
|
-
* global tracer for every Mastra operation (agent runs, model
|
|
16
|
-
* calls, tool invocations, workflow steps). They inherit the
|
|
17
|
-
* ambient OTel context, so when Mastra is invoked from inside an
|
|
18
|
-
* AppKit HTTP span the trace stays connected.
|
|
19
|
-
*
|
|
20
|
-
* Net effect: Mastra spans get exactly the treatment AppKit's
|
|
21
|
-
* `agents` plugin gets. No custom OTLP pipeline lives in this
|
|
22
|
-
* package; the OTLP endpoint, headers, and resource attributes are
|
|
23
|
-
* driven by the standard OTel env vars
|
|
24
|
-
* (`OTEL_EXPORTER_OTLP_ENDPOINT`, `OTEL_EXPORTER_OTLP_HEADERS`,
|
|
25
|
-
* `OTEL_SERVICE_NAME`, `OTEL_RESOURCE_ATTRIBUTES`, ...) and consumed
|
|
26
|
-
* by AppKit's `TelemetryManager`. Set those once and both AppKit and
|
|
27
|
-
* Mastra spans end up at the same backend.
|
|
28
|
-
*
|
|
29
|
-
* When `OTEL_EXPORTER_OTLP_ENDPOINT` is unset the bridge's spans go
|
|
30
|
-
* to the global noop tracer, mirroring how the `agents` plugin
|
|
31
|
-
* silently no-ops in the same situation.
|
|
32
|
-
*/
|
|
33
|
-
import { logUtils, projectUtils } from "@dbx-tools/shared";
|
|
34
|
-
import { Observability } from "@mastra/observability";
|
|
35
|
-
import { OtelBridge } from "@mastra/otel-bridge";
|
|
36
|
-
import { TRACE_REQUEST_CONTEXT_KEYS } from "./config.js";
|
|
37
|
-
const log = logUtils.logger("mastra/observability");
|
|
38
|
-
const DEFAULT_SERVICE_NAME = "mastra";
|
|
39
|
-
/**
|
|
40
|
-
* Build a Mastra `Observability` whose spans ride AppKit's global
|
|
41
|
-
* OTel pipeline via `@mastra/otel-bridge`.
|
|
42
|
-
*
|
|
43
|
-
* Returns `undefined` only if someone explicitly opts out in the
|
|
44
|
-
* future; today it always returns an `Observability` because the
|
|
45
|
-
* bridge degrades gracefully (no-op tracer) when no global OTel SDK
|
|
46
|
-
* is registered. Callers can spread `...(observability ? { observability } : {})`
|
|
47
|
-
* either way to stay forward-compatible.
|
|
48
|
-
*/
|
|
49
|
-
export async function buildObservability(options) {
|
|
50
|
-
const serviceName = options?.serviceName ?? (await projectUtils.name()) ?? DEFAULT_SERVICE_NAME;
|
|
51
|
-
const requestContextKeys = [
|
|
52
|
-
...(options?.requestContextKeys ?? TRACE_REQUEST_CONTEXT_KEYS),
|
|
53
|
-
];
|
|
54
|
-
// The OTel HTTP exporter treats `OTEL_EXPORTER_OTLP_ENDPOINT` as a
|
|
55
|
-
// *base* URL and appends the signal path itself (e.g.
|
|
56
|
-
// `http://localhost:6006` -> `http://localhost:6006/v1/traces`). Log
|
|
57
|
-
// the resolved POST URL so misconfigurations (e.g. accidentally
|
|
58
|
-
// setting the base var to a `/v1/traces`-suffixed URL, which makes
|
|
59
|
-
// the SDK POST to `.../v1/traces/v1/traces` and Phoenix 404s) are
|
|
60
|
-
// obvious in startup output.
|
|
61
|
-
const otelBase = process.env.OTEL_EXPORTER_OTLP_ENDPOINT;
|
|
62
|
-
const otelTracesOverride = process.env.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT;
|
|
63
|
-
const resolvedTracesUrl = otelTracesOverride
|
|
64
|
-
? otelTracesOverride
|
|
65
|
-
: otelBase
|
|
66
|
-
? `${otelBase.replace(/\/+$/, "")}/v1/traces`
|
|
67
|
-
: undefined;
|
|
68
|
-
log.info("Mastra observability wired through OTel bridge", {
|
|
69
|
-
serviceName,
|
|
70
|
-
requestContextKeys,
|
|
71
|
-
otelBase: otelBase ?? "<unset>",
|
|
72
|
-
resolvedTracesUrl: resolvedTracesUrl ?? "<noop; OTLP endpoint unset>",
|
|
73
|
-
});
|
|
74
|
-
return new Observability({
|
|
75
|
-
configs: {
|
|
76
|
-
serviceName: {
|
|
77
|
-
serviceName,
|
|
78
|
-
bridge: new OtelBridge(),
|
|
79
|
-
requestContextKeys,
|
|
80
|
-
},
|
|
81
|
-
},
|
|
82
|
-
});
|
|
83
|
-
}
|