@dbx-tools/model 0.1.58 → 0.1.68

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.
@@ -1,122 +0,0 @@
1
- /**
2
- * Live Databricks Model Serving catalogue access.
3
- *
4
- * Lists the workspace's `/serving-endpoints` once per host and caches
5
- * the result with a TTL via AppKit's `CacheManager`, with concurrent
6
- * callers sharing one in-flight promise (the coalescing pattern of
7
- * Python's `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"`
13
- * resolve to `databricks-claude-sonnet-4-6`.
14
- *
15
- * The class stamp and embedding dimension are computed once per cache
16
- * load: every embedding endpoint is "pinged" in parallel and the
17
- * resulting vector length recorded, so the cost is paid on a cache miss,
18
- * not per read. The ping is best-effort - a failure logs at debug and
19
- * leaves `dimension` unset rather than failing the whole listing.
20
- */
21
- import { type ServingEndpointSummary } from "@dbx-tools/model-shared";
22
- import type { appkitUtils } from "@dbx-tools/shared";
23
- /**
24
- * Structural type for the Databricks workspace client, re-exported from
25
- * `@dbx-tools/shared` so the rest of this package can keep importing it
26
- * from here. See `appkitUtils.WorkspaceClientLike` for the canonical
27
- * definition.
28
- */
29
- export type WorkspaceClientLike = appkitUtils.WorkspaceClientLike;
30
- /** Default TTL for the in-memory endpoint cache. Matches the Databricks SDK's session lifetime budget. */
31
- export declare const DEFAULT_MODEL_CACHE_TTL_MS: number;
32
- /** Default Fuse.js score threshold below which a fuzzy match is accepted. */
33
- export declare const DEFAULT_FUZZY_THRESHOLD = 0.4;
34
- /** Options for {@link listServingEndpoints}. */
35
- export interface ListServingEndpointsOptions {
36
- /**
37
- * Override the default cache TTL for this call, in milliseconds.
38
- * Forwarded to `CacheManager` as seconds.
39
- */
40
- ttlMs?: number;
41
- }
42
- /**
43
- * List Model Serving endpoints for the workspace owning `client`,
44
- * routed through AppKit's `CacheManager`. The manager gives us
45
- * everything `cachetools.TTLCache` provides plus what
46
- * `cachetools-async` adds on top: per-entry TTL, in-flight request
47
- * coalescing (concurrent callers share one fetch via the manager's
48
- * internal `inFlightRequests` map), bounded size, telemetry spans
49
- * (`cache.getOrExecute`), and optional Lakebase persistence so the
50
- * catalogue survives restarts when the lakebase plugin is wired up.
51
- *
52
- * Returns plain {@link ServingEndpointSummary} objects (a stable
53
- * subset of the SDK type) so cache hits never expose stale SDK
54
- * internals. Errors from `CacheManager` or the SDK fetch propagate
55
- * to the caller - we don't swallow them so users see the real
56
- * auth / network issue.
57
- *
58
- * @param host - Workspace host used as the cache key. Pass the value
59
- * resolved from `client.config.getHost()` so multi-host apps share
60
- * one entry per workspace.
61
- * @param options.ttlMs - Override the default TTL just for this call.
62
- * Forwarded to `CacheManager` as seconds.
63
- */
64
- export declare function listServingEndpoints(client: WorkspaceClientLike, host: string, options?: ListServingEndpointsOptions): Promise<ServingEndpointSummary[]>;
65
- /**
66
- * Force-evict cached endpoint listings via AppKit's `CacheManager`.
67
- * With a `host` deletes that one workspace's entry; without one
68
- * clears every cache entry on the manager (since `CacheManager`
69
- * doesn't expose a namespace-scoped clear, this is the brute-force
70
- * path - fine for tests, avoid in steady-state code).
71
- */
72
- export declare function clearServingEndpointsCache(host?: string): Promise<void>;
73
- /**
74
- * Result of fuzzy-resolving a user-supplied model name against the
75
- * live endpoint list. `score` is Fuse.js's distance (`0` is exact,
76
- * `1` is no match); `matched` is `false` when the score exceeds the
77
- * configured threshold so callers can fall back to the original
78
- * input (Databricks will then return a clean 404).
79
- */
80
- export interface ResolvedModel {
81
- modelId: string;
82
- matched: boolean;
83
- score?: number;
84
- }
85
- /** Options accepted by {@link resolveModelId} / {@link searchServingEndpoints}. */
86
- export interface ResolveModelOptions {
87
- /** Fuse.js threshold (0 = exact, 1 = anything). Default `0.4`. */
88
- threshold?: number;
89
- }
90
- /** A serving endpoint paired with its fuzzy-match distance for a query. */
91
- export interface ScoredEndpoint {
92
- endpoint: ServingEndpointSummary;
93
- /** Fuse.js distance: `0` is exact, `1` is no match. */
94
- score: number;
95
- }
96
- /**
97
- * Fuzzy-rank endpoints by how closely their `name` matches `input`,
98
- * best (lowest score) first, keeping only those within `threshold`:
99
- *
100
- * 1. An exact name match short-circuits to a single `score: 0` result.
101
- * 2. Otherwise the input is tokenized (dashes / underscores / spaces
102
- * become separators) and fed through Fuse.js extended search, which
103
- * AND-s each token with fuzzy matching enabled - the "tokenized
104
- * fuzzy match" a caller reaches for when they type `"claude sonnet"`
105
- * instead of the full endpoint name.
106
- *
107
- * Returns `[]` for an empty endpoint list or when `input` tokenizes to
108
- * nothing, so callers fall back to the raw input and let Databricks
109
- * surface a clean 404. This multi-result core is shared by
110
- * {@link resolveModelId} (single best) and the ranked `rankModels`
111
- * selector.
112
- */
113
- export declare function searchServingEndpoints(input: string, endpoints: readonly ServingEndpointSummary[], options?: ResolveModelOptions): ScoredEndpoint[];
114
- /**
115
- * Snap a user-supplied model name to the single closest configured
116
- * serving endpoint via {@link searchServingEndpoints}. Returns the
117
- * input unchanged with `matched: false` when nothing scores within the
118
- * threshold (or the catalogue is empty), so a deliberate model id is
119
- * never silently rewritten to a similar-looking neighbour and the
120
- * upstream call surfaces the canonical 404.
121
- */
122
- export declare function resolveModelId(input: string, endpoints: readonly ServingEndpointSummary[], options?: ResolveModelOptions): ResolvedModel;
@@ -1,251 +0,0 @@
1
- /**
2
- * Live Databricks Model Serving catalogue access.
3
- *
4
- * Lists the workspace's `/serving-endpoints` once per host and caches
5
- * the result with a TTL via AppKit's `CacheManager`, with concurrent
6
- * callers sharing one in-flight promise (the coalescing pattern of
7
- * Python's `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"`
13
- * resolve to `databricks-claude-sonnet-4-6`.
14
- *
15
- * The class stamp and embedding dimension are computed once per cache
16
- * load: every embedding endpoint is "pinged" in parallel and the
17
- * resulting vector length recorded, so the cost is paid on a cache miss,
18
- * not per read. The ping is best-effort - a failure logs at debug and
19
- * leaves `dimension` unset rather than failing the whole listing.
20
- */
21
- import { CacheManager } from "@databricks/appkit";
22
- import { classifyEndpoints, ModelClass, } from "@dbx-tools/model-shared";
23
- import { commonUtils, logUtils, stringUtils } from "@dbx-tools/shared";
24
- import Fuse from "fuse.js";
25
- import { MODEL_CLASS_ORDER } from "./classes.js";
26
- const log = logUtils.logger("model/serving");
27
- /** Default TTL for the in-memory endpoint cache. Matches the Databricks SDK's session lifetime budget. */
28
- export const DEFAULT_MODEL_CACHE_TTL_MS = 5 * 60 * 1000;
29
- /** Default Fuse.js score threshold below which a fuzzy match is accepted. */
30
- export const DEFAULT_FUZZY_THRESHOLD = 0.4;
31
- /** Cache key parts under which endpoint listings are stored. */
32
- const CACHE_KEY_NAMESPACE = "serving-endpoints";
33
- /**
34
- * Stable `userKey` arg for AppKit's `CacheManager.getOrExecute`.
35
- * Endpoint visibility is effectively workspace-scoped (we cache by
36
- * host in the key parts), so a single shared key lets every user of
37
- * the same workspace share one cached fetch and coalesce on the
38
- * in-flight promise. Permissions can differ in theory, but the
39
- * Foundation Model API catalogue is the same view for every caller.
40
- */
41
- const SHARED_USER_KEY = "model-shared";
42
- /**
43
- * List Model Serving endpoints for the workspace owning `client`,
44
- * routed through AppKit's `CacheManager`. The manager gives us
45
- * everything `cachetools.TTLCache` provides plus what
46
- * `cachetools-async` adds on top: per-entry TTL, in-flight request
47
- * coalescing (concurrent callers share one fetch via the manager's
48
- * internal `inFlightRequests` map), bounded size, telemetry spans
49
- * (`cache.getOrExecute`), and optional Lakebase persistence so the
50
- * catalogue survives restarts when the lakebase plugin is wired up.
51
- *
52
- * Returns plain {@link ServingEndpointSummary} objects (a stable
53
- * subset of the SDK type) so cache hits never expose stale SDK
54
- * internals. Errors from `CacheManager` or the SDK fetch propagate
55
- * to the caller - we don't swallow them so users see the real
56
- * auth / network issue.
57
- *
58
- * @param host - Workspace host used as the cache key. Pass the value
59
- * resolved from `client.config.getHost()` so multi-host apps share
60
- * one entry per workspace.
61
- * @param options.ttlMs - Override the default TTL just for this call.
62
- * Forwarded to `CacheManager` as seconds.
63
- */
64
- export async function listServingEndpoints(client, host, options = {}) {
65
- const ttlSec = Math.max(1, Math.round((options.ttlMs ?? DEFAULT_MODEL_CACHE_TTL_MS) / 1000));
66
- return CacheManager.getInstanceSync().getOrExecute([CACHE_KEY_NAMESPACE, host], () => fetchEndpoints(client), SHARED_USER_KEY, { ttl: ttlSec });
67
- }
68
- async function fetchEndpoints(client) {
69
- const startedAt = Date.now();
70
- const out = [];
71
- for await (const ep of client.servingEndpoints.list()) {
72
- if (!ep.name)
73
- continue;
74
- const profile = extractProfile(ep);
75
- out.push({
76
- name: ep.name,
77
- ...(ep.task !== undefined ? { task: ep.task } : {}),
78
- ...(ep.state?.ready !== undefined ? { state: String(ep.state.ready) } : {}),
79
- ...(ep.description !== undefined ? { description: ep.description } : {}),
80
- ...(profile ? { profile } : {}),
81
- });
82
- }
83
- stampModelClasses(out);
84
- await measureEmbeddingDimensions(client, out);
85
- log.debug("listed", { count: out.length, elapsedMs: Date.now() - startedAt });
86
- return out;
87
- }
88
- /**
89
- * Stamp each summary's {@link ServingEndpointSummary.class} from the
90
- * relative classification of the whole set. Mutates `summaries` in
91
- * place. Endpoints the classifier doesn't recognize (custom, unscored,
92
- * non-LLM) are left without a class.
93
- */
94
- function stampModelClasses(summaries) {
95
- const buckets = classifyEndpoints(summaries);
96
- const classOf = new Map();
97
- for (const cls of MODEL_CLASS_ORDER) {
98
- for (const ep of buckets[cls])
99
- classOf.set(ep.name, cls);
100
- }
101
- for (const summary of summaries) {
102
- const cls = classOf.get(summary.name);
103
- if (cls !== undefined)
104
- summary.class = cls;
105
- }
106
- }
107
- /**
108
- * Measure the embedding vector dimension of every
109
- * {@link ModelClass.Embedding} endpoint by pinging it once, all in
110
- * parallel. Mutates the matching summaries in place with the resulting
111
- * `dimension`. Runs only on a cache miss (it's called from
112
- * {@link fetchEndpoints}), so the probe cost is amortized across the
113
- * cached TTL window. Per-endpoint failures are swallowed (logged at
114
- * warn) so one unreachable embedding model never fails the listing.
115
- */
116
- async function measureEmbeddingDimensions(client, summaries) {
117
- await Promise.all(summaries
118
- .filter((s) => s.class === ModelClass.Embedding)
119
- .map(async (summary) => {
120
- const dimension = await pingEmbeddingDimension(client, summary.name);
121
- if (dimension !== undefined)
122
- summary.dimension = dimension;
123
- }));
124
- }
125
- /**
126
- * Best-effort embedding dimension probe: query `name` with a tiny
127
- * `"ping"` input and return the length of the returned vector. Returns
128
- * `undefined` (and logs at warn) when the endpoint can't be queried or
129
- * returns no vector - the dimension is informational, never required.
130
- */
131
- async function pingEmbeddingDimension(client, name) {
132
- try {
133
- const response = await client.servingEndpoints.query({ name, input: "ping" });
134
- const dimension = response.data?.[0]?.embedding?.length;
135
- if (typeof dimension === "number" && dimension > 0)
136
- return dimension;
137
- log.warn("embedding ping returned no vector", { name });
138
- return undefined;
139
- }
140
- catch (err) {
141
- log.warn("embedding ping failed", { name, error: commonUtils.errorMessage(err) });
142
- return undefined;
143
- }
144
- }
145
- /**
146
- * Pull the Foundation Model API `quality` / `speed` / `cost` scores
147
- * off a serving-endpoint listing entry. Databricks returns these
148
- * under `config.served_entities[].foundation_model.ai_gateway_model_profile`
149
- * (snake_case at the wire level, preserved verbatim by the SDK), but
150
- * the field is newer than the typed `FoundationModel` interface, so we
151
- * read it through a structural cast. Returns `undefined` when no
152
- * served entity carries a profile (custom models, embeddings, and
153
- * brand-new endpoints that Databricks has not scored yet).
154
- */
155
- function extractProfile(ep) {
156
- const entities = ep.config?.served_entities;
157
- if (!entities)
158
- return undefined;
159
- for (const entity of entities) {
160
- const raw = entity.foundation_model?.ai_gateway_model_profile;
161
- if (!raw)
162
- continue;
163
- const profile = {};
164
- if (Number.isFinite(raw.quality))
165
- profile.quality = raw.quality;
166
- if (Number.isFinite(raw.speed))
167
- profile.speed = raw.speed;
168
- if (Number.isFinite(raw.cost))
169
- profile.cost = raw.cost;
170
- if (Object.keys(profile).length > 0)
171
- return profile;
172
- }
173
- return undefined;
174
- }
175
- /**
176
- * Force-evict cached endpoint listings via AppKit's `CacheManager`.
177
- * With a `host` deletes that one workspace's entry; without one
178
- * clears every cache entry on the manager (since `CacheManager`
179
- * doesn't expose a namespace-scoped clear, this is the brute-force
180
- * path - fine for tests, avoid in steady-state code).
181
- */
182
- export async function clearServingEndpointsCache(host) {
183
- const cache = CacheManager.getInstanceSync();
184
- if (host) {
185
- const key = cache.generateKey([CACHE_KEY_NAMESPACE, host], SHARED_USER_KEY);
186
- await cache.delete(key);
187
- }
188
- else {
189
- await cache.clear();
190
- }
191
- }
192
- /**
193
- * Fuzzy-rank endpoints by how closely their `name` matches `input`,
194
- * best (lowest score) first, keeping only those within `threshold`:
195
- *
196
- * 1. An exact name match short-circuits to a single `score: 0` result.
197
- * 2. Otherwise the input is tokenized (dashes / underscores / spaces
198
- * become separators) and fed through Fuse.js extended search, which
199
- * AND-s each token with fuzzy matching enabled - the "tokenized
200
- * fuzzy match" a caller reaches for when they type `"claude sonnet"`
201
- * instead of the full endpoint name.
202
- *
203
- * Returns `[]` for an empty endpoint list or when `input` tokenizes to
204
- * nothing, so callers fall back to the raw input and let Databricks
205
- * surface a clean 404. This multi-result core is shared by
206
- * {@link resolveModelId} (single best) and the ranked `rankModels`
207
- * selector.
208
- */
209
- export function searchServingEndpoints(input, endpoints, options = {}) {
210
- if (endpoints.length === 0)
211
- return [];
212
- for (const ep of endpoints) {
213
- if (ep.name === input)
214
- return [{ endpoint: ep, score: 0 }];
215
- }
216
- const threshold = options.threshold ?? DEFAULT_FUZZY_THRESHOLD;
217
- // Fuse 7.3 has no built-in tokenize hook; in extended search,
218
- // space-separated tokens are AND-ed with fuzzy matching enabled. We
219
- // lean on the shared tokenizer so the splitting rules stay
220
- // consistent with the rest of the toolkit.
221
- const query = Array.from(stringUtils.tokenizeWithOptions({ lowerCase: true, camelCase: false }, input)).join(" ");
222
- if (!query)
223
- return [];
224
- const fuse = new Fuse(endpoints, {
225
- keys: ["name"],
226
- threshold,
227
- ignoreLocation: true,
228
- includeScore: true,
229
- useExtendedSearch: true,
230
- isCaseSensitive: false,
231
- });
232
- return fuse
233
- .search(query)
234
- .filter((r) => (r.score ?? 0) <= threshold)
235
- .map((r) => ({ endpoint: r.item, score: r.score ?? 0 }));
236
- }
237
- /**
238
- * Snap a user-supplied model name to the single closest configured
239
- * serving endpoint via {@link searchServingEndpoints}. Returns the
240
- * input unchanged with `matched: false` when nothing scores within the
241
- * threshold (or the catalogue is empty), so a deliberate model id is
242
- * never silently rewritten to a similar-looking neighbour and the
243
- * upstream call surfaces the canonical 404.
244
- */
245
- export function resolveModelId(input, endpoints, options = {}) {
246
- const [best] = searchServingEndpoints(input, endpoints, options);
247
- if (best) {
248
- return { modelId: best.endpoint.name, matched: true, score: best.score };
249
- }
250
- return { modelId: input, matched: false };
251
- }