@dbx-tools/model 0.1.111 → 0.3.1
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 +134 -74
- package/index.ts +10 -0
- package/package.json +41 -20
- package/src/classes.ts +73 -0
- package/src/fallback.ts +73 -0
- package/src/resolve.ts +275 -0
- package/src/serving.ts +381 -0
- package/test/classes.test.ts +74 -0
- package/test/resolve.test.ts +169 -0
- package/tsconfig.json +41 -0
- package/dist/index.d.ts +0 -291
- package/dist/index.js +0 -576
package/src/resolve.ts
ADDED
|
@@ -0,0 +1,275 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* Workspace-aware model selection.
|
|
3
|
+
*
|
|
4
|
+
* Given a caller's intent - a search string, a capability {@link ModelClass}
|
|
5
|
+
* ceiling, both, or nothing - the toolkit returns matching endpoints ranked by
|
|
6
|
+
* match quality then class, or collapses to the single best id the workspace
|
|
7
|
+
* actually has, degrading from "best in range" down to the static fallback
|
|
8
|
+
* floor. Selection is chat-only: embedding endpoints surface only when
|
|
9
|
+
* `modelClass` is explicitly {@link ModelClass.Embedding}.
|
|
10
|
+
*
|
|
11
|
+
* Two shapes of selection, each in a pure form (over an endpoint list the
|
|
12
|
+
* caller already holds) and an I/O wrapper (that lists `/serving-endpoints`
|
|
13
|
+
* first): ranking, which returns a match- then class-ordered list, and
|
|
14
|
+
* single-selection, which collapses to one id plus how it was reached and
|
|
15
|
+
* layers the operator-pinned fallback / static-floor safety net on top. A chat
|
|
16
|
+
* `modelClass` acts as a ceiling: that band and the less-capable chat bands
|
|
17
|
+
* below it are eligible (see {@link classesAtOrBelow}), so a `chat-balanced`
|
|
18
|
+
* ask can fall to `chat-fast` but never escalate to `chat-thinking`.
|
|
19
|
+
*/
|
|
20
|
+
|
|
21
|
+
import {
|
|
22
|
+
classify,
|
|
23
|
+
model,
|
|
24
|
+
type ModelQuery,
|
|
25
|
+
type RankedModel,
|
|
26
|
+
type ServingEndpointSummary,
|
|
27
|
+
} from "@dbx-tools/shared-model";
|
|
28
|
+
import { object } from "@dbx-tools/shared-core";
|
|
29
|
+
|
|
30
|
+
import { CHAT_CLASS_ORDER, classesAtOrBelow, MODEL_CLASS_ORDER } from "./classes";
|
|
31
|
+
import { FALLBACK_MODEL_IDS, modelsForClass } from "./fallback";
|
|
32
|
+
import { listServingEndpoints, searchServingEndpoints, type WorkspaceClientLike } from "./serving";
|
|
33
|
+
|
|
34
|
+
type ModelClass = model.ModelClass;
|
|
35
|
+
|
|
36
|
+
/** Caller intent passed to {@link resolveModel}. */
|
|
37
|
+
export interface ResolveModelInput {
|
|
38
|
+
/**
|
|
39
|
+
* Explicit model id / loose name (per-request override, agent / plugin
|
|
40
|
+
* default, or env var). When set it wins over `modelClass` and `fallbacks`.
|
|
41
|
+
*/
|
|
42
|
+
explicit?: string;
|
|
43
|
+
/**
|
|
44
|
+
* Fuzzy-match an `explicit` name against the live catalogue so loose names
|
|
45
|
+
* like `"claude sonnet"` resolve. Default `true`. When `false` the explicit
|
|
46
|
+
* input is returned verbatim (Databricks surfaces the canonical 404 if it
|
|
47
|
+
* doesn't exist).
|
|
48
|
+
*/
|
|
49
|
+
fuzzy?: boolean;
|
|
50
|
+
/** Fuse.js threshold forwarded to the fuzzy `search` match ({@link searchServingEndpoints}). */
|
|
51
|
+
threshold?: number;
|
|
52
|
+
/**
|
|
53
|
+
* Chat capability class to resolve when no `explicit` id is given. The live
|
|
54
|
+
* catalogue is classified by its Foundation Model API scores and the top
|
|
55
|
+
* available model in the class (and the chat bands below it) wins, falling
|
|
56
|
+
* back to the class's small static list.
|
|
57
|
+
*/
|
|
58
|
+
modelClass?: ModelClass;
|
|
59
|
+
/**
|
|
60
|
+
* Operator-supplied fallback ids tried *first* in the no-explicit, no-class
|
|
61
|
+
* path (e.g. a regulated workspace pinned to an approved subset), ahead of
|
|
62
|
+
* the auto-classified catalogue.
|
|
63
|
+
*/
|
|
64
|
+
fallbacks?: readonly string[];
|
|
65
|
+
}
|
|
66
|
+
|
|
67
|
+
/** Outcome of {@link resolveModel}: the chosen id plus how it was reached. */
|
|
68
|
+
export interface ResolvedModelSelection {
|
|
69
|
+
modelId: string;
|
|
70
|
+
source: "explicit" | "fuzzy-match" | "class" | "fallback";
|
|
71
|
+
}
|
|
72
|
+
|
|
73
|
+
/** Intent + catalogue knobs passed to {@link selectModel}. */
|
|
74
|
+
export interface SelectModelInput extends ResolveModelInput {
|
|
75
|
+
/** TTL override for the cached `/serving-endpoints` listing, in ms. */
|
|
76
|
+
ttlMs?: number;
|
|
77
|
+
}
|
|
78
|
+
|
|
79
|
+
/** TTL override merged into a {@link ModelQuery} for {@link searchModels}. */
|
|
80
|
+
export interface SearchModelsInput extends ModelQuery {
|
|
81
|
+
/** TTL override for the cached `/serving-endpoints` listing, in ms. */
|
|
82
|
+
ttlMs?: number;
|
|
83
|
+
}
|
|
84
|
+
|
|
85
|
+
/**
|
|
86
|
+
* Round a Fuse score to the display precision so version siblings that match a
|
|
87
|
+
* token identically (e.g. `opus-4-7` vs `opus-4-8` for the query `"opus"`) tie
|
|
88
|
+
* on match and let the class / within-class rank decide - which is what
|
|
89
|
+
* surfaces the newer, higher-quality sibling.
|
|
90
|
+
*/
|
|
91
|
+
function matchBucket(score: number | undefined): number {
|
|
92
|
+
return Math.round((score ?? 0) * 1000);
|
|
93
|
+
}
|
|
94
|
+
|
|
95
|
+
/**
|
|
96
|
+
* Rank the live catalogue against a {@link ModelQuery}, best-first.
|
|
97
|
+
*
|
|
98
|
+
* Candidates are the classified endpoints in the eligible classes:
|
|
99
|
+
* {@link classesAtOrBelow} the requested `modelClass`, or - when none is given
|
|
100
|
+
* - the chat bands only ({@link CHAT_CLASS_ORDER}), so a general ask never
|
|
101
|
+
* surfaces an embedding endpoint. Each class bucket is already best-first from
|
|
102
|
+
* {@link classify.classifyEndpoints}. Ranking is **match then class**:
|
|
103
|
+
*
|
|
104
|
+
* 1. With a `search`, only endpoints matching it survive, ordered by match
|
|
105
|
+
* distance (bucketed via {@link matchBucket} so near-identical scores tie),
|
|
106
|
+
* then by class (more capable first), then by the stable within-class rank.
|
|
107
|
+
* 2. Without a `search`, the class-then-rank candidate order stands.
|
|
108
|
+
*
|
|
109
|
+
* A `limit` truncates the result. Returns `[]` when nothing is eligible or
|
|
110
|
+
* matches - callers layer their own fallback.
|
|
111
|
+
*/
|
|
112
|
+
export function rankModels(
|
|
113
|
+
endpoints: readonly ServingEndpointSummary[],
|
|
114
|
+
query: ModelQuery = {},
|
|
115
|
+
): RankedModel[] {
|
|
116
|
+
const classified = classify.classifyEndpoints(endpoints);
|
|
117
|
+
const eligible =
|
|
118
|
+
query.modelClass !== undefined ? classesAtOrBelow(query.modelClass) : CHAT_CLASS_ORDER;
|
|
119
|
+
|
|
120
|
+
// Flatten eligible classes in capability order, carrying each endpoint's
|
|
121
|
+
// class; bucket order is already best-first.
|
|
122
|
+
const candidates: RankedModel[] = [];
|
|
123
|
+
for (const modelClass of eligible) {
|
|
124
|
+
for (const endpoint of classified[modelClass]) candidates.push({ endpoint, modelClass });
|
|
125
|
+
}
|
|
126
|
+
|
|
127
|
+
const search = query.search?.trim();
|
|
128
|
+
let ranked: RankedModel[];
|
|
129
|
+
if (search) {
|
|
130
|
+
const scores = new Map<string, number>();
|
|
131
|
+
for (const match of searchServingEndpoints(
|
|
132
|
+
search,
|
|
133
|
+
candidates.map((c) => c.endpoint),
|
|
134
|
+
query.threshold !== undefined ? { threshold: query.threshold } : {},
|
|
135
|
+
)) {
|
|
136
|
+
scores.set(match.endpoint.name, match.score);
|
|
137
|
+
}
|
|
138
|
+
// `Array.prototype.sort` is stable, so endpoints equal on match and class
|
|
139
|
+
// keep their best-first within-class order.
|
|
140
|
+
ranked = candidates
|
|
141
|
+
.filter((c) => scores.has(c.endpoint.name))
|
|
142
|
+
.map((c) => ({ ...c, score: scores.get(c.endpoint.name) }))
|
|
143
|
+
.sort((a, b) => {
|
|
144
|
+
const byMatch = matchBucket(a.score) - matchBucket(b.score);
|
|
145
|
+
if (byMatch !== 0) return byMatch;
|
|
146
|
+
return MODEL_CLASS_ORDER.indexOf(a.modelClass) - MODEL_CLASS_ORDER.indexOf(b.modelClass);
|
|
147
|
+
});
|
|
148
|
+
} else {
|
|
149
|
+
ranked = candidates;
|
|
150
|
+
}
|
|
151
|
+
|
|
152
|
+
return query.limit !== undefined ? ranked.slice(0, Math.max(0, query.limit)) : ranked;
|
|
153
|
+
}
|
|
154
|
+
|
|
155
|
+
/**
|
|
156
|
+
* Rank a workspace's catalogue in one call: list its `/serving-endpoints`
|
|
157
|
+
* (cached) and run {@link rankModels} over the result. The list counterpart to
|
|
158
|
+
* {@link selectModel}, for a consumer that wants the full ranked set (a model
|
|
159
|
+
* picker, a CLI) rather than a single id. Catalogue fetches fail loud: network
|
|
160
|
+
* / auth errors propagate so the caller sees the real SDK message.
|
|
161
|
+
*
|
|
162
|
+
* @param host - Workspace host used as the cache key. Pass the value resolved
|
|
163
|
+
* from `client.config.getHost()`.
|
|
164
|
+
*/
|
|
165
|
+
export async function searchModels(
|
|
166
|
+
client: WorkspaceClientLike,
|
|
167
|
+
host: string,
|
|
168
|
+
input: SearchModelsInput = {},
|
|
169
|
+
): Promise<RankedModel[]> {
|
|
170
|
+
const endpoints = await listServingEndpoints(
|
|
171
|
+
client,
|
|
172
|
+
host,
|
|
173
|
+
input.ttlMs !== undefined ? { ttlMs: input.ttlMs } : {},
|
|
174
|
+
);
|
|
175
|
+
return rankModels(endpoints, input);
|
|
176
|
+
}
|
|
177
|
+
|
|
178
|
+
/**
|
|
179
|
+
* Resolve a model id for a workspace in one call: list its `/serving-endpoints`
|
|
180
|
+
* (cached) and run {@link resolveModel} over the result. This is the entry
|
|
181
|
+
* point for any consumer that holds a `WorkspaceClient` and just wants a usable
|
|
182
|
+
* model name - a Lakeflow job, a one-off script, or the Mastra plugin alike.
|
|
183
|
+
*
|
|
184
|
+
* Cheap exit: when an `explicit` name is given and `fuzzy` is off, the
|
|
185
|
+
* catalogue is never fetched - the name is returned verbatim and Databricks
|
|
186
|
+
* surfaces the canonical 404 if it doesn't exist. Catalogue fetches otherwise
|
|
187
|
+
* fail loud: network / auth errors propagate so the caller sees the real SDK
|
|
188
|
+
* message instead of a silent fallback.
|
|
189
|
+
*
|
|
190
|
+
* @param host - Workspace host used as the cache key. Pass the value resolved
|
|
191
|
+
* from `client.config.getHost()`.
|
|
192
|
+
*/
|
|
193
|
+
export async function selectModel(
|
|
194
|
+
client: WorkspaceClientLike,
|
|
195
|
+
host: string,
|
|
196
|
+
input: SelectModelInput = {},
|
|
197
|
+
): Promise<ResolvedModelSelection> {
|
|
198
|
+
if (input.explicit !== undefined && input.fuzzy === false) {
|
|
199
|
+
return { modelId: input.explicit, source: "explicit" };
|
|
200
|
+
}
|
|
201
|
+
const endpoints = await listServingEndpoints(client, host, {
|
|
202
|
+
...(input.ttlMs !== undefined ? { ttlMs: input.ttlMs } : {}),
|
|
203
|
+
});
|
|
204
|
+
return resolveModel(endpoints, input);
|
|
205
|
+
}
|
|
206
|
+
|
|
207
|
+
/**
|
|
208
|
+
* Resolve a single model id from the live catalogue and caller intent,
|
|
209
|
+
* delegating the live selection to {@link rankModels} with `limit: 1`.
|
|
210
|
+
*
|
|
211
|
+
* 1. **Explicit ask**: with `fuzzy` off, returned verbatim; otherwise
|
|
212
|
+
* fuzzy-ranked within the (optional) class ceiling and the best taken,
|
|
213
|
+
* falling back to the input verbatim when nothing matches.
|
|
214
|
+
* 2. **No explicit ask**: an operator-pinned `fallback` that exists in the live
|
|
215
|
+
* catalogue wins first; then the ranked live catalogue (class ceiling
|
|
216
|
+
* applied); then the static {@link FALLBACK_MODEL_IDS} floor when the
|
|
217
|
+
* catalogue yields nothing in range.
|
|
218
|
+
*/
|
|
219
|
+
export function resolveModel(
|
|
220
|
+
endpoints: readonly ServingEndpointSummary[],
|
|
221
|
+
input: ResolveModelInput = {},
|
|
222
|
+
): ResolvedModelSelection {
|
|
223
|
+
if (input.explicit !== undefined) {
|
|
224
|
+
if (input.fuzzy === false) {
|
|
225
|
+
return { modelId: input.explicit, source: "explicit" };
|
|
226
|
+
}
|
|
227
|
+
const [top] = rankModels(endpoints, buildQuery(input, input.explicit));
|
|
228
|
+
return { modelId: top?.endpoint.name ?? input.explicit, source: "fuzzy-match" };
|
|
229
|
+
}
|
|
230
|
+
|
|
231
|
+
// Operator-pinned fallbacks win when present and live (e.g. a regulated
|
|
232
|
+
// workspace restricted to an approved subset).
|
|
233
|
+
if (input.modelClass === undefined && input.fallbacks && input.fallbacks.length > 0) {
|
|
234
|
+
const present = new Set(endpoints.map((e) => e.name));
|
|
235
|
+
const pinned = input.fallbacks.find((id) => present.has(id));
|
|
236
|
+
if (pinned) return { modelId: pinned, source: "fallback" };
|
|
237
|
+
}
|
|
238
|
+
|
|
239
|
+
const source = input.modelClass !== undefined ? "class" : "fallback";
|
|
240
|
+
const [top] = rankModels(endpoints, buildQuery(input, undefined));
|
|
241
|
+
if (top) return { modelId: top.endpoint.name, source };
|
|
242
|
+
|
|
243
|
+
// Live catalogue yielded nothing in range: walk the static floor.
|
|
244
|
+
const floorSource =
|
|
245
|
+
input.modelClass !== undefined ? modelsForClass(input.modelClass) : (input.fallbacks ?? []);
|
|
246
|
+
const floor = object.sequence(floorSource).concat(FALLBACK_MODEL_IDS).distinct().toArray();
|
|
247
|
+
return { modelId: pickFirstAvailable(floor, endpoints), source };
|
|
248
|
+
}
|
|
249
|
+
|
|
250
|
+
/** Build a {@link ModelQuery} from {@link ResolveModelInput} for the `limit: 1` delegation. */
|
|
251
|
+
function buildQuery(input: ResolveModelInput, search: string | undefined): ModelQuery {
|
|
252
|
+
return {
|
|
253
|
+
...(search !== undefined ? { search } : {}),
|
|
254
|
+
...(input.modelClass !== undefined ? { modelClass: input.modelClass } : {}),
|
|
255
|
+
...(input.threshold !== undefined ? { threshold: input.threshold } : {}),
|
|
256
|
+
limit: 1,
|
|
257
|
+
};
|
|
258
|
+
}
|
|
259
|
+
|
|
260
|
+
/**
|
|
261
|
+
* Find the first id in `candidates` whose endpoint is present in `endpoints`.
|
|
262
|
+
* Returns the top candidate when the workspace has none of them so callers
|
|
263
|
+
* always get a string; an offline workspace then receives a clean 404 from
|
|
264
|
+
* Databricks instead of a malformed config.
|
|
265
|
+
*/
|
|
266
|
+
function pickFirstAvailable(
|
|
267
|
+
candidates: readonly string[],
|
|
268
|
+
endpoints: readonly ServingEndpointSummary[],
|
|
269
|
+
): string {
|
|
270
|
+
const present = new Set(endpoints.map((e) => e.name));
|
|
271
|
+
for (const candidate of candidates) {
|
|
272
|
+
if (present.has(candidate)) return candidate;
|
|
273
|
+
}
|
|
274
|
+
return candidates[0] ?? FALLBACK_MODEL_IDS[0]!;
|
|
275
|
+
}
|
package/src/serving.ts
ADDED
|
@@ -0,0 +1,381 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* Live Databricks Model Serving catalogue access.
|
|
3
|
+
*
|
|
4
|
+
* Lists the workspace's `/serving-endpoints` once per host and caches the
|
|
5
|
+
* result with a TTL via AppKit's `CacheManager`, with concurrent callers
|
|
6
|
+
* sharing one in-flight promise (the coalescing pattern of Python's
|
|
7
|
+
* `cachetools-async`). Surfaces each endpoint as a stable
|
|
8
|
+
* {@link ServingEndpointSummary} - including the Foundation Model API
|
|
9
|
+
* `quality` / `speed` / `cost` profile when present, the classified
|
|
10
|
+
* {@link ModelClass}, and (for embedding endpoints) the measured vector
|
|
11
|
+
* `dimension` - and snaps loose, human-typed names to real endpoint ids
|
|
12
|
+
* through `fuse.js` extended search so tokens like `"claude sonnet"` resolve to
|
|
13
|
+
* `databricks-claude-sonnet-4-6`.
|
|
14
|
+
*
|
|
15
|
+
* The class stamp and embedding dimension are computed once per cache load:
|
|
16
|
+
* every embedding endpoint is "pinged" in parallel and the resulting vector
|
|
17
|
+
* length recorded, so the cost is paid on a cache miss, not per read. The ping
|
|
18
|
+
* is best-effort - a failure logs at debug and leaves `dimension` unset rather
|
|
19
|
+
* than failing the whole listing.
|
|
20
|
+
*/
|
|
21
|
+
|
|
22
|
+
import { error, log, string } from "@dbx-tools/shared-core";
|
|
23
|
+
import {
|
|
24
|
+
classify,
|
|
25
|
+
display,
|
|
26
|
+
model,
|
|
27
|
+
type ModelProfile,
|
|
28
|
+
type ServingEndpointSummary,
|
|
29
|
+
} from "@dbx-tools/shared-model";
|
|
30
|
+
import { appkit } from "@dbx-tools/appkit";
|
|
31
|
+
import { CacheManager } from "@databricks/appkit";
|
|
32
|
+
import Fuse from "fuse.js";
|
|
33
|
+
|
|
34
|
+
import { MODEL_CLASS_ORDER } from "./classes";
|
|
35
|
+
|
|
36
|
+
const { ModelClass } = model;
|
|
37
|
+
|
|
38
|
+
const logger = log.logger("model/serving");
|
|
39
|
+
|
|
40
|
+
/**
|
|
41
|
+
* Structural type for the Databricks workspace client, re-exported so the rest
|
|
42
|
+
* of this package can keep importing it from here. See
|
|
43
|
+
* `appkit.WorkspaceClientLike` (node-appkit) for the canonical definition.
|
|
44
|
+
*/
|
|
45
|
+
export type WorkspaceClientLike = appkit.WorkspaceClientLike;
|
|
46
|
+
|
|
47
|
+
/** Default TTL for the in-memory endpoint cache. Matches the Databricks SDK's session lifetime budget. */
|
|
48
|
+
export const DEFAULT_MODEL_CACHE_TTL_MS = 5 * 60 * 1000;
|
|
49
|
+
|
|
50
|
+
/** Default Fuse.js score threshold below which a fuzzy match is accepted. */
|
|
51
|
+
export const DEFAULT_FUZZY_THRESHOLD = 0.4;
|
|
52
|
+
|
|
53
|
+
/** Cache key parts under which endpoint listings are stored. */
|
|
54
|
+
const CACHE_KEY_NAMESPACE = "serving-endpoints";
|
|
55
|
+
|
|
56
|
+
/**
|
|
57
|
+
* Stable `userKey` arg for AppKit's `CacheManager.getOrExecute`. Endpoint
|
|
58
|
+
* visibility is effectively workspace-scoped (we cache by host in the key
|
|
59
|
+
* parts), so a single shared key lets every user of the same workspace share
|
|
60
|
+
* one cached fetch and coalesce on the in-flight promise. Permissions can
|
|
61
|
+
* differ in theory, but the Foundation Model API catalogue is the same view for
|
|
62
|
+
* every caller.
|
|
63
|
+
*/
|
|
64
|
+
const SHARED_USER_KEY = "model-shared";
|
|
65
|
+
|
|
66
|
+
/** Options for {@link listServingEndpoints}. */
|
|
67
|
+
export interface ListServingEndpointsOptions {
|
|
68
|
+
/**
|
|
69
|
+
* Override the default cache TTL for this call, in milliseconds. Forwarded to
|
|
70
|
+
* `CacheManager` as seconds.
|
|
71
|
+
*/
|
|
72
|
+
ttlMs?: number;
|
|
73
|
+
}
|
|
74
|
+
|
|
75
|
+
/**
|
|
76
|
+
* List Model Serving endpoints for the workspace owning `client`, routed
|
|
77
|
+
* through AppKit's `CacheManager`. The manager gives us everything
|
|
78
|
+
* `cachetools.TTLCache` provides plus what `cachetools-async` adds on top:
|
|
79
|
+
* per-entry TTL, in-flight request coalescing (concurrent callers share one
|
|
80
|
+
* fetch via the manager's internal `inFlightRequests` map), bounded size,
|
|
81
|
+
* telemetry spans (`cache.getOrExecute`), and optional Lakebase persistence so
|
|
82
|
+
* the catalogue survives restarts when the lakebase plugin is wired up.
|
|
83
|
+
*
|
|
84
|
+
* Returns plain {@link ServingEndpointSummary} objects (a stable subset of the
|
|
85
|
+
* SDK type) so cache hits never expose stale SDK internals. Errors from
|
|
86
|
+
* `CacheManager` or the SDK fetch propagate to the caller - we don't swallow
|
|
87
|
+
* them so users see the real auth / network issue.
|
|
88
|
+
*
|
|
89
|
+
* @param host - Workspace host used as the cache key. Pass the value resolved
|
|
90
|
+
* from `client.config.getHost()` so multi-host apps share one entry per
|
|
91
|
+
* workspace.
|
|
92
|
+
* @param options.ttlMs - Override the default TTL just for this call. Forwarded
|
|
93
|
+
* to `CacheManager` as seconds.
|
|
94
|
+
*/
|
|
95
|
+
export async function listServingEndpoints(
|
|
96
|
+
client: WorkspaceClientLike,
|
|
97
|
+
host: string,
|
|
98
|
+
options: ListServingEndpointsOptions = {},
|
|
99
|
+
): Promise<ServingEndpointSummary[]> {
|
|
100
|
+
const ttlSec = Math.max(1, Math.round((options.ttlMs ?? DEFAULT_MODEL_CACHE_TTL_MS) / 1000));
|
|
101
|
+
return CacheManager.getInstanceSync().getOrExecute(
|
|
102
|
+
[CACHE_KEY_NAMESPACE, host],
|
|
103
|
+
() => fetchEndpoints(client),
|
|
104
|
+
SHARED_USER_KEY,
|
|
105
|
+
{ ttl: ttlSec },
|
|
106
|
+
);
|
|
107
|
+
}
|
|
108
|
+
|
|
109
|
+
/**
|
|
110
|
+
* List the workspace's serving endpoints as minimal
|
|
111
|
+
* {@link ServingEndpointSummary} objects straight from the SDK: no caching, and
|
|
112
|
+
* none of the cache-load enrichment ({@link listServingEndpoints} adds the
|
|
113
|
+
* {@link ModelClass} stamp and the embedding-dimension probe). Use this for a
|
|
114
|
+
* one-shot, dependency-light listing - e.g. a CLI that only needs names/tasks
|
|
115
|
+
* for fuzzy resolution and doesn't want AppKit's `CacheManager` or the
|
|
116
|
+
* per-embedding ping cost. Prefer {@link listServingEndpoints} for the cached,
|
|
117
|
+
* enriched view.
|
|
118
|
+
*/
|
|
119
|
+
export async function listServingEndpointsUncached(
|
|
120
|
+
client: WorkspaceClientLike,
|
|
121
|
+
): Promise<ServingEndpointSummary[]> {
|
|
122
|
+
const out: ServingEndpointSummary[] = [];
|
|
123
|
+
for await (const ep of client.servingEndpoints.list()) {
|
|
124
|
+
if (!ep.name) continue;
|
|
125
|
+
const profile = extractProfile(ep);
|
|
126
|
+
out.push({
|
|
127
|
+
name: ep.name,
|
|
128
|
+
// Prefer a Databricks-provided human name (a display-name tag or an
|
|
129
|
+
// external-model name); else derive a title-cased label from the id.
|
|
130
|
+
displayName: display.toModelDisplayName(ep.name, providedDisplayName(ep)),
|
|
131
|
+
...(ep.task !== undefined ? { task: ep.task } : {}),
|
|
132
|
+
...(ep.state?.ready !== undefined ? { state: String(ep.state.ready) } : {}),
|
|
133
|
+
...(ep.description !== undefined ? { description: ep.description } : {}),
|
|
134
|
+
...(profile ? { profile } : {}),
|
|
135
|
+
});
|
|
136
|
+
}
|
|
137
|
+
return out;
|
|
138
|
+
}
|
|
139
|
+
|
|
140
|
+
/**
|
|
141
|
+
* Pull a Databricks-provided human name off a serving endpoint, if any:
|
|
142
|
+
* a `display_name` / `displayName` / `name` tag first, then an
|
|
143
|
+
* external-model `name` from the served-entity config. Returns `null`
|
|
144
|
+
* when the endpoint carries no explicit name (the common case for
|
|
145
|
+
* Foundation Model API endpoints), so the caller falls back to deriving
|
|
146
|
+
* one from the id.
|
|
147
|
+
*/
|
|
148
|
+
function providedDisplayName(ep: {
|
|
149
|
+
tags?: Array<{ key: string; value?: string }>;
|
|
150
|
+
config?: { served_entities?: Array<{ external_model?: { name?: string } }> };
|
|
151
|
+
}): string | null {
|
|
152
|
+
const tag = ep.tags?.find(
|
|
153
|
+
(t) => t.key === "display_name" || t.key === "displayName" || t.key === "name",
|
|
154
|
+
);
|
|
155
|
+
const fromTag = string.trimToNull(tag?.value);
|
|
156
|
+
if (fromTag) return fromTag;
|
|
157
|
+
for (const entity of ep.config?.served_entities ?? []) {
|
|
158
|
+
const external = string.trimToNull(entity.external_model?.name);
|
|
159
|
+
if (external) return external;
|
|
160
|
+
}
|
|
161
|
+
return null;
|
|
162
|
+
}
|
|
163
|
+
|
|
164
|
+
async function fetchEndpoints(client: WorkspaceClientLike): Promise<ServingEndpointSummary[]> {
|
|
165
|
+
const startedAt = Date.now();
|
|
166
|
+
const out = await listServingEndpointsUncached(client);
|
|
167
|
+
stampModelClasses(out);
|
|
168
|
+
await measureEmbeddingDimensions(client, out);
|
|
169
|
+
logger.debug("listed", { count: out.length, elapsedMs: Date.now() - startedAt });
|
|
170
|
+
return out;
|
|
171
|
+
}
|
|
172
|
+
|
|
173
|
+
/**
|
|
174
|
+
* Stamp each summary's {@link ServingEndpointSummary.class} from the relative
|
|
175
|
+
* classification of the whole set. Mutates `summaries` in place. Endpoints the
|
|
176
|
+
* classifier doesn't recognize (custom, unscored, non-LLM) are left without a
|
|
177
|
+
* class.
|
|
178
|
+
*/
|
|
179
|
+
function stampModelClasses(summaries: ServingEndpointSummary[]): void {
|
|
180
|
+
const buckets = classify.classifyEndpoints(summaries);
|
|
181
|
+
const classOf = new Map<string, model.ModelClass>();
|
|
182
|
+
for (const cls of MODEL_CLASS_ORDER) {
|
|
183
|
+
for (const ep of buckets[cls]) classOf.set(ep.name, cls);
|
|
184
|
+
}
|
|
185
|
+
for (const summary of summaries) {
|
|
186
|
+
const cls = classOf.get(summary.name);
|
|
187
|
+
if (cls !== undefined) summary.class = cls;
|
|
188
|
+
}
|
|
189
|
+
}
|
|
190
|
+
|
|
191
|
+
/**
|
|
192
|
+
* Measure the embedding vector dimension of every {@link ModelClass.Embedding}
|
|
193
|
+
* endpoint by pinging it once, all in parallel. Mutates the matching summaries
|
|
194
|
+
* in place with the resulting `dimension`. Runs only on a cache miss (it's
|
|
195
|
+
* called from {@link fetchEndpoints}), so the probe cost is amortized across the
|
|
196
|
+
* cached TTL window. Per-endpoint failures are swallowed (logged at warn) so
|
|
197
|
+
* one unreachable embedding model never fails the listing.
|
|
198
|
+
*/
|
|
199
|
+
async function measureEmbeddingDimensions(
|
|
200
|
+
client: WorkspaceClientLike,
|
|
201
|
+
summaries: ServingEndpointSummary[],
|
|
202
|
+
): Promise<void> {
|
|
203
|
+
await Promise.all(
|
|
204
|
+
summaries
|
|
205
|
+
.filter((s) => s.class === ModelClass.Embedding)
|
|
206
|
+
.map(async (summary) => {
|
|
207
|
+
const dimension = await pingEmbeddingDimension(client, summary.name);
|
|
208
|
+
if (dimension !== undefined) summary.dimension = dimension;
|
|
209
|
+
}),
|
|
210
|
+
);
|
|
211
|
+
}
|
|
212
|
+
|
|
213
|
+
/**
|
|
214
|
+
* Best-effort embedding dimension probe: query `name` with a tiny `"ping"`
|
|
215
|
+
* input and return the length of the returned vector. Returns `undefined` (and
|
|
216
|
+
* logs at warn) when the endpoint can't be queried or returns no vector - the
|
|
217
|
+
* dimension is informational, never required.
|
|
218
|
+
*/
|
|
219
|
+
async function pingEmbeddingDimension(
|
|
220
|
+
client: WorkspaceClientLike,
|
|
221
|
+
name: string,
|
|
222
|
+
): Promise<number | undefined> {
|
|
223
|
+
try {
|
|
224
|
+
const response = await client.servingEndpoints.query({ name, input: "ping" });
|
|
225
|
+
const dimension = response.data?.[0]?.embedding?.length;
|
|
226
|
+
if (typeof dimension === "number" && dimension > 0) return dimension;
|
|
227
|
+
logger.warn("embedding ping returned no vector", { name });
|
|
228
|
+
return undefined;
|
|
229
|
+
} catch (err) {
|
|
230
|
+
logger.warn("embedding ping failed", { name, error: error.errorMessage(err) });
|
|
231
|
+
return undefined;
|
|
232
|
+
}
|
|
233
|
+
}
|
|
234
|
+
|
|
235
|
+
/**
|
|
236
|
+
* Pull the Foundation Model API `quality` / `speed` / `cost` scores off a
|
|
237
|
+
* serving-endpoint listing entry. Databricks returns these under
|
|
238
|
+
* `config.served_entities[].foundation_model.ai_gateway_model_profile`
|
|
239
|
+
* (snake_case at the wire level, preserved verbatim by the SDK), but the field
|
|
240
|
+
* is newer than the typed `FoundationModel` interface, so we read it through a
|
|
241
|
+
* structural cast. Returns `undefined` when no served entity carries a profile
|
|
242
|
+
* (custom models, embeddings, and brand-new endpoints that Databricks has not
|
|
243
|
+
* scored yet).
|
|
244
|
+
*/
|
|
245
|
+
function extractProfile(ep: unknown): ModelProfile | undefined {
|
|
246
|
+
const entities = (
|
|
247
|
+
ep as {
|
|
248
|
+
config?: {
|
|
249
|
+
served_entities?: Array<{
|
|
250
|
+
foundation_model?: {
|
|
251
|
+
ai_gateway_model_profile?: {
|
|
252
|
+
quality?: number;
|
|
253
|
+
speed?: number;
|
|
254
|
+
cost?: number;
|
|
255
|
+
};
|
|
256
|
+
};
|
|
257
|
+
}>;
|
|
258
|
+
};
|
|
259
|
+
}
|
|
260
|
+
).config?.served_entities;
|
|
261
|
+
if (!entities) return undefined;
|
|
262
|
+
for (const entity of entities) {
|
|
263
|
+
const raw = entity.foundation_model?.ai_gateway_model_profile;
|
|
264
|
+
if (!raw) continue;
|
|
265
|
+
const profile: ModelProfile = {};
|
|
266
|
+
if (Number.isFinite(raw.quality)) profile.quality = raw.quality;
|
|
267
|
+
if (Number.isFinite(raw.speed)) profile.speed = raw.speed;
|
|
268
|
+
if (Number.isFinite(raw.cost)) profile.cost = raw.cost;
|
|
269
|
+
if (Object.keys(profile).length > 0) return profile;
|
|
270
|
+
}
|
|
271
|
+
return undefined;
|
|
272
|
+
}
|
|
273
|
+
|
|
274
|
+
/**
|
|
275
|
+
* Force-evict cached endpoint listings via AppKit's `CacheManager`. With a
|
|
276
|
+
* `host` deletes that one workspace's entry; without one clears every cache
|
|
277
|
+
* entry on the manager (since `CacheManager` doesn't expose a namespace-scoped
|
|
278
|
+
* clear, this is the brute-force path - fine for tests, avoid in steady-state
|
|
279
|
+
* code).
|
|
280
|
+
*/
|
|
281
|
+
export async function clearServingEndpointsCache(host?: string): Promise<void> {
|
|
282
|
+
const cache = CacheManager.getInstanceSync();
|
|
283
|
+
if (host) {
|
|
284
|
+
const key = cache.generateKey([CACHE_KEY_NAMESPACE, host], SHARED_USER_KEY);
|
|
285
|
+
await cache.delete(key);
|
|
286
|
+
} else {
|
|
287
|
+
await cache.clear();
|
|
288
|
+
}
|
|
289
|
+
}
|
|
290
|
+
|
|
291
|
+
/**
|
|
292
|
+
* Result of fuzzy-resolving a user-supplied model name against the live
|
|
293
|
+
* endpoint list. `score` is Fuse.js's distance (`0` is exact, `1` is no match);
|
|
294
|
+
* `matched` is `false` when the score exceeds the configured threshold so
|
|
295
|
+
* callers can fall back to the original input (Databricks will then return a
|
|
296
|
+
* clean 404).
|
|
297
|
+
*/
|
|
298
|
+
export interface ResolvedModel {
|
|
299
|
+
modelId: string;
|
|
300
|
+
matched: boolean;
|
|
301
|
+
score?: number;
|
|
302
|
+
}
|
|
303
|
+
|
|
304
|
+
/** Options accepted by {@link resolveModelId} / {@link searchServingEndpoints}. */
|
|
305
|
+
export interface ResolveModelOptions {
|
|
306
|
+
/** Fuse.js threshold (0 = exact, 1 = anything). Default `0.4`. */
|
|
307
|
+
threshold?: number;
|
|
308
|
+
}
|
|
309
|
+
|
|
310
|
+
/** A serving endpoint paired with its fuzzy-match distance for a query. */
|
|
311
|
+
export interface ScoredEndpoint {
|
|
312
|
+
endpoint: ServingEndpointSummary;
|
|
313
|
+
/** Fuse.js distance: `0` is exact, `1` is no match. */
|
|
314
|
+
score: number;
|
|
315
|
+
}
|
|
316
|
+
|
|
317
|
+
/**
|
|
318
|
+
* Fuzzy-rank endpoints by how closely their `name` matches `input`, best
|
|
319
|
+
* (lowest score) first, keeping only those within `threshold`:
|
|
320
|
+
*
|
|
321
|
+
* 1. An exact name match short-circuits to a single `score: 0` result.
|
|
322
|
+
* 2. Otherwise the input is tokenized (dashes / underscores / spaces become
|
|
323
|
+
* separators) and fed through Fuse.js extended search, which AND-s each
|
|
324
|
+
* token with fuzzy matching enabled - the "tokenized fuzzy match" a caller
|
|
325
|
+
* reaches for when they type `"claude sonnet"` instead of the full endpoint
|
|
326
|
+
* name.
|
|
327
|
+
*
|
|
328
|
+
* Returns `[]` for an empty endpoint list or when `input` tokenizes to nothing,
|
|
329
|
+
* so callers fall back to the raw input and let Databricks surface a clean 404.
|
|
330
|
+
* This multi-result core is shared by {@link resolveModelId} (single best) and
|
|
331
|
+
* the ranked `rankModels` selector.
|
|
332
|
+
*/
|
|
333
|
+
export function searchServingEndpoints(
|
|
334
|
+
input: string,
|
|
335
|
+
endpoints: readonly ServingEndpointSummary[],
|
|
336
|
+
options: ResolveModelOptions = {},
|
|
337
|
+
): ScoredEndpoint[] {
|
|
338
|
+
if (endpoints.length === 0) return [];
|
|
339
|
+
for (const ep of endpoints) {
|
|
340
|
+
if (ep.name === input) return [{ endpoint: ep, score: 0 }];
|
|
341
|
+
}
|
|
342
|
+
const threshold = options.threshold ?? DEFAULT_FUZZY_THRESHOLD;
|
|
343
|
+
// Fuse 7.3 has no built-in tokenize hook; in extended search, space-separated
|
|
344
|
+
// tokens are AND-ed with fuzzy matching enabled. We lean on the shared
|
|
345
|
+
// tokenizer so the splitting rules stay consistent with the rest of the toolkit.
|
|
346
|
+
const query = Array.from(
|
|
347
|
+
string.tokenizeWithOptions({ lowerCase: true, camelCase: false }, input),
|
|
348
|
+
).join(" ");
|
|
349
|
+
if (!query) return [];
|
|
350
|
+
const fuse = new Fuse(endpoints, {
|
|
351
|
+
keys: ["name"],
|
|
352
|
+
threshold,
|
|
353
|
+
ignoreLocation: true,
|
|
354
|
+
includeScore: true,
|
|
355
|
+
useExtendedSearch: true,
|
|
356
|
+
isCaseSensitive: false,
|
|
357
|
+
});
|
|
358
|
+
return fuse
|
|
359
|
+
.search(query)
|
|
360
|
+
.filter((r) => (r.score ?? 0) <= threshold)
|
|
361
|
+
.map((r) => ({ endpoint: r.item, score: r.score ?? 0 }));
|
|
362
|
+
}
|
|
363
|
+
|
|
364
|
+
/**
|
|
365
|
+
* Snap a user-supplied model name to the single closest configured serving
|
|
366
|
+
* endpoint via {@link searchServingEndpoints}. Returns the input unchanged with
|
|
367
|
+
* `matched: false` when nothing scores within the threshold (or the catalogue
|
|
368
|
+
* is empty), so a deliberate model id is never silently rewritten to a
|
|
369
|
+
* similar-looking neighbour and the upstream call surfaces the canonical 404.
|
|
370
|
+
*/
|
|
371
|
+
export function resolveModelId(
|
|
372
|
+
input: string,
|
|
373
|
+
endpoints: readonly ServingEndpointSummary[],
|
|
374
|
+
options: ResolveModelOptions = {},
|
|
375
|
+
): ResolvedModel {
|
|
376
|
+
const [best] = searchServingEndpoints(input, endpoints, options);
|
|
377
|
+
if (best) {
|
|
378
|
+
return { modelId: best.endpoint.name, matched: true, score: best.score };
|
|
379
|
+
}
|
|
380
|
+
return { modelId: input, matched: false };
|
|
381
|
+
}
|