@dbx-tools/model 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/src/serving.ts CHANGED
@@ -109,10 +109,19 @@ export async function listServingEndpoints(
109
109
  );
110
110
  }
111
111
 
112
- async function fetchEndpoints(
112
+ /**
113
+ * List the workspace's serving endpoints as minimal
114
+ * {@link ServingEndpointSummary} objects straight from the SDK: no
115
+ * caching, and none of the cache-load enrichment ({@link listServingEndpoints}
116
+ * adds the {@link ModelClass} stamp and the embedding-dimension probe).
117
+ * Use this for a one-shot, dependency-light listing - e.g. a CLI that
118
+ * only needs names/tasks for fuzzy resolution and doesn't want AppKit's
119
+ * `CacheManager` or the per-embedding ping cost. Prefer
120
+ * {@link listServingEndpoints} for the cached, enriched view.
121
+ */
122
+ export async function listServingEndpointsUncached(
113
123
  client: WorkspaceClientLike,
114
124
  ): Promise<ServingEndpointSummary[]> {
115
- const startedAt = Date.now();
116
125
  const out: ServingEndpointSummary[] = [];
117
126
  for await (const ep of client.servingEndpoints.list()) {
118
127
  if (!ep.name) continue;
@@ -125,6 +134,14 @@ async function fetchEndpoints(
125
134
  ...(profile ? { profile } : {}),
126
135
  });
127
136
  }
137
+ return out;
138
+ }
139
+
140
+ async function fetchEndpoints(
141
+ client: WorkspaceClientLike,
142
+ ): Promise<ServingEndpointSummary[]> {
143
+ const startedAt = Date.now();
144
+ const out = await listServingEndpointsUncached(client);
128
145
  stampModelClasses(out);
129
146
  await measureEmbeddingDimensions(client, out);
130
147
  log.debug("listed", { count: out.length, elapsedMs: Date.now() - startedAt });
@@ -1,51 +0,0 @@
1
- /**
2
- * Model-class ordering for the model service.
3
- *
4
- * `@dbx-tools/model-shared` declares the {@link ModelClass} values and
5
- * their schema; the *behavior* over those values - the chat capability
6
- * order, the "this class and below" ceiling, and coercing loose request
7
- * input to a class - lives here in the service so the shared
8
- * wire-format surface stays purely declarative.
9
- */
10
- import { ModelClass } from "@dbx-tools/model-shared";
11
- /**
12
- * Chat capability ladder in descending order - most capable
13
- * ({@link ModelClass.ChatThinking}) first, least
14
- * ({@link ModelClass.ChatFast}) last. {@link ModelClass.Embedding} is a
15
- * separate modality and is deliberately absent: the ceiling never spans
16
- * it. This is the order used to filter "this class and below" and to
17
- * break ranking ties toward the more capable model.
18
- */
19
- export declare const CHAT_CLASS_ORDER: readonly ModelClass[];
20
- /**
21
- * Every class in display order: the chat ladder followed by
22
- * {@link ModelClass.Embedding}. Used when flattening a full
23
- * classification (e.g. stamping the class onto each cached endpoint).
24
- */
25
- export declare const MODEL_CLASS_ORDER: readonly ModelClass[];
26
- /** Whether `cls` is one of the chat capability bands (vs. embedding). */
27
- export declare function isChatClass(cls: ModelClass): boolean;
28
- /**
29
- * Coerce an arbitrary value (query string, header, body field) to a
30
- * {@link ModelClass}, returning `null` when it isn't a known class.
31
- * Lets a route accept a class request without throwing on junk input.
32
- *
33
- * Matches the full slug first (`"chat-fast"`, `"embedding"`), then -
34
- * for a bare chat band - retries with a `chat-` prefix, so the shorthand
35
- * `"thinking"` / `"balanced"` / `"fast"` resolves to the corresponding
36
- * `chat-*` class.
37
- */
38
- export declare function parseModelClass(value: unknown): ModelClass | null;
39
- /**
40
- * Classes at or below `cls` in chat capability: the class itself plus
41
- * every less-capable chat band, in {@link CHAT_CLASS_ORDER}. Treats a
42
- * chat `cls` as a ceiling - requesting {@link ModelClass.ChatBalanced}
43
- * yields `[ChatBalanced, ChatFast]`, never
44
- * {@link ModelClass.ChatThinking} - so a class ask can degrade downward
45
- * to a smaller chat model but never escalate to a larger one.
46
- *
47
- * {@link ModelClass.Embedding} is its own modality, not a rung on the
48
- * chat ladder, so it yields just `[Embedding]`. An unrecognized class
49
- * yields the full chat ladder.
50
- */
51
- export declare function classesAtOrBelow(cls: ModelClass): ModelClass[];
@@ -1,71 +0,0 @@
1
- /**
2
- * Model-class ordering for the model service.
3
- *
4
- * `@dbx-tools/model-shared` declares the {@link ModelClass} values and
5
- * their schema; the *behavior* over those values - the chat capability
6
- * order, the "this class and below" ceiling, and coercing loose request
7
- * input to a class - lives here in the service so the shared
8
- * wire-format surface stays purely declarative.
9
- */
10
- import { ModelClass, ModelClassSchema } from "@dbx-tools/model-shared";
11
- /**
12
- * Chat capability ladder in descending order - most capable
13
- * ({@link ModelClass.ChatThinking}) first, least
14
- * ({@link ModelClass.ChatFast}) last. {@link ModelClass.Embedding} is a
15
- * separate modality and is deliberately absent: the ceiling never spans
16
- * it. This is the order used to filter "this class and below" and to
17
- * break ranking ties toward the more capable model.
18
- */
19
- export const CHAT_CLASS_ORDER = [
20
- ModelClass.ChatThinking,
21
- ModelClass.ChatBalanced,
22
- ModelClass.ChatFast,
23
- ];
24
- /**
25
- * Every class in display order: the chat ladder followed by
26
- * {@link ModelClass.Embedding}. Used when flattening a full
27
- * classification (e.g. stamping the class onto each cached endpoint).
28
- */
29
- export const MODEL_CLASS_ORDER = [
30
- ...CHAT_CLASS_ORDER,
31
- ModelClass.Embedding,
32
- ];
33
- /** Whether `cls` is one of the chat capability bands (vs. embedding). */
34
- export function isChatClass(cls) {
35
- return CHAT_CLASS_ORDER.includes(cls);
36
- }
37
- /**
38
- * Coerce an arbitrary value (query string, header, body field) to a
39
- * {@link ModelClass}, returning `null` when it isn't a known class.
40
- * Lets a route accept a class request without throwing on junk input.
41
- *
42
- * Matches the full slug first (`"chat-fast"`, `"embedding"`), then -
43
- * for a bare chat band - retries with a `chat-` prefix, so the shorthand
44
- * `"thinking"` / `"balanced"` / `"fast"` resolves to the corresponding
45
- * `chat-*` class.
46
- */
47
- export function parseModelClass(value) {
48
- const exact = ModelClassSchema.safeParse(value);
49
- if (exact.success)
50
- return exact.data;
51
- const prefixed = ModelClassSchema.safeParse(`chat-${value}`);
52
- return prefixed.success ? prefixed.data : null;
53
- }
54
- /**
55
- * Classes at or below `cls` in chat capability: the class itself plus
56
- * every less-capable chat band, in {@link CHAT_CLASS_ORDER}. Treats a
57
- * chat `cls` as a ceiling - requesting {@link ModelClass.ChatBalanced}
58
- * yields `[ChatBalanced, ChatFast]`, never
59
- * {@link ModelClass.ChatThinking} - so a class ask can degrade downward
60
- * to a smaller chat model but never escalate to a larger one.
61
- *
62
- * {@link ModelClass.Embedding} is its own modality, not a rung on the
63
- * chat ladder, so it yields just `[Embedding]`. An unrecognized class
64
- * yields the full chat ladder.
65
- */
66
- export function classesAtOrBelow(cls) {
67
- if (cls === ModelClass.Embedding)
68
- return [ModelClass.Embedding];
69
- const index = CHAT_CLASS_ORDER.indexOf(cls);
70
- return index < 0 ? [...CHAT_CLASS_ORDER] : [...CHAT_CLASS_ORDER.slice(index)];
71
- }
@@ -1,36 +0,0 @@
1
- /**
2
- * Server-side offline fallback opinion for model selection.
3
- *
4
- * When the live `/serving-endpoints` catalogue can't be read at all -
5
- * no user token, the service principal can't list, or the workspace is
6
- * unreachable - the resolver still has to name *some* endpoint. This
7
- * module holds that floor: a small, hard-coded set of well-known
8
- * Foundation Model API endpoint names, each bucketed into a chat
9
- * {@link ModelClass} by the shared {@link classifyByFamily} heuristic
10
- * (no classes are hard-coded here) and ordered best-first.
11
- *
12
- * This is deliberately *server-only*. A browser client never talks to
13
- * Databricks directly - it always goes through this server - so it has
14
- * nothing to fall back to and must not assume a stale, baked-in model
15
- * list; it consumes the live `/models` response instead. The pure
16
- * classifier ({@link classifyEndpoints}) is what the client shares.
17
- */
18
- import { ModelClass } from "@dbx-tools/model-shared";
19
- /**
20
- * Static fallback model ids for a chat class, drawn from the small
21
- * built-in {@link FALLBACK_MODEL_NAMES} list and ordered best-first by
22
- * family rank. Sync and workspace-independent: this is the *fallback
23
- * opinion* used to seed default lists or when the live catalogue is
24
- * unreachable - live resolution prefers `classifyEndpoints`. Returns
25
- * `[]` for {@link ModelClass.Embedding} (the floor is chat-only).
26
- */
27
- export declare function modelsForClass(cls: ModelClass): readonly string[];
28
- /** Top static fallback model id for a chat class. */
29
- export declare function modelForClass(cls: ModelClass): string;
30
- /**
31
- * Priority-ordered fallback chain (ChatThinking -> ChatBalanced ->
32
- * ChatFast) over the small built-in list. The floor walked at resolve
33
- * time when no agent / plugin / env / request model is set *and* the
34
- * live catalogue yields nothing.
35
- */
36
- export declare const FALLBACK_MODEL_IDS: readonly string[];
@@ -1,66 +0,0 @@
1
- /**
2
- * Server-side offline fallback opinion for model selection.
3
- *
4
- * When the live `/serving-endpoints` catalogue can't be read at all -
5
- * no user token, the service principal can't list, or the workspace is
6
- * unreachable - the resolver still has to name *some* endpoint. This
7
- * module holds that floor: a small, hard-coded set of well-known
8
- * Foundation Model API endpoint names, each bucketed into a chat
9
- * {@link ModelClass} by the shared {@link classifyByFamily} heuristic
10
- * (no classes are hard-coded here) and ordered best-first.
11
- *
12
- * This is deliberately *server-only*. A browser client never talks to
13
- * Databricks directly - it always goes through this server - so it has
14
- * nothing to fall back to and must not assume a stale, baked-in model
15
- * list; it consumes the live `/models` response instead. The pure
16
- * classifier ({@link classifyEndpoints}) is what the client shares.
17
- */
18
- import { classifyByFamily, ModelClass, } from "@dbx-tools/model-shared";
19
- /**
20
- * Small, last-resort set of well-known Foundation Model API endpoint
21
- * names, ordered best-first within each class by {@link classifyByFamily}.
22
- * Used only as the floor when the live `/serving-endpoints` catalogue
23
- * can't be read at resolve time; the live, score-driven classification
24
- * supersedes it whenever the workspace listing is available. Classes are
25
- * not hard-coded here - each name is classified by family heuristic.
26
- */
27
- const FALLBACK_MODEL_NAMES = [
28
- "databricks-claude-opus-4-8",
29
- "databricks-gpt-5-5-pro",
30
- "databricks-gemini-3-1-pro",
31
- "databricks-claude-sonnet-4-6",
32
- "databricks-gpt-5-5",
33
- "databricks-meta-llama-3-3-70b-instruct",
34
- "databricks-claude-haiku-4-5",
35
- "databricks-gpt-5-nano",
36
- "databricks-meta-llama-3-1-8b-instruct",
37
- ];
38
- /**
39
- * Static fallback model ids for a chat class, drawn from the small
40
- * built-in {@link FALLBACK_MODEL_NAMES} list and ordered best-first by
41
- * family rank. Sync and workspace-independent: this is the *fallback
42
- * opinion* used to seed default lists or when the live catalogue is
43
- * unreachable - live resolution prefers `classifyEndpoints`. Returns
44
- * `[]` for {@link ModelClass.Embedding} (the floor is chat-only).
45
- */
46
- export function modelsForClass(cls) {
47
- return FALLBACK_MODEL_NAMES.map((name) => ({ name, c: classifyByFamily(name) }))
48
- .filter((x) => x.c !== null && x.c.class === cls)
49
- .sort((a, b) => b.c.rank - a.c.rank)
50
- .map((x) => x.name);
51
- }
52
- /** Top static fallback model id for a chat class. */
53
- export function modelForClass(cls) {
54
- return modelsForClass(cls)[0];
55
- }
56
- /**
57
- * Priority-ordered fallback chain (ChatThinking -> ChatBalanced ->
58
- * ChatFast) over the small built-in list. The floor walked at resolve
59
- * time when no agent / plugin / env / request model is set *and* the
60
- * live catalogue yields nothing.
61
- */
62
- export const FALLBACK_MODEL_IDS = [
63
- ...modelsForClass(ModelClass.ChatThinking),
64
- ...modelsForClass(ModelClass.ChatBalanced),
65
- ...modelsForClass(ModelClass.ChatFast),
66
- ];
@@ -1,139 +0,0 @@
1
- /**
2
- * Workspace-aware model selection.
3
- *
4
- * Given a caller's intent - a search string, a capability
5
- * {@link ModelClass} ceiling, both, or nothing - the toolkit returns
6
- * matching endpoints ranked by match quality then class, or collapses
7
- * to the single best id the workspace actually has, degrading from
8
- * "best in range" down to the static fallback floor. Selection is
9
- * chat-only: embedding endpoints surface only when `modelClass` is
10
- * explicitly {@link ModelClass.Embedding}.
11
- *
12
- * Ranking entry points (return a ranked list):
13
- *
14
- * - {@link rankModels} is the pure ranker over an endpoint list you
15
- * already hold. A chat `modelClass` acts as a ceiling: only that
16
- * band and the less-capable chat bands below it are eligible (see
17
- * {@link classesAtOrBelow}), so a `chat-balanced` ask can fall to
18
- * `chat-fast` but never escalate to `chat-thinking`.
19
- * - {@link searchModels} is the I/O wrapper: hand it a
20
- * `WorkspaceClient` and it lists `/serving-endpoints` (cached) then
21
- * ranks in one step.
22
- *
23
- * Single-selection entry points (return one id + how it was reached):
24
- *
25
- * - {@link resolveModel} is the pure selector; it delegates to
26
- * {@link rankModels} with `limit: 1` and adds the operator-pinned
27
- * fallback / static-floor safety net.
28
- * - {@link selectModel} is the I/O wrapper over {@link resolveModel}.
29
- */
30
- import { ModelClass, type ModelQuery, type RankedModel, type ServingEndpointSummary } from "@dbx-tools/model-shared";
31
- import { type WorkspaceClientLike } from "./serving.js";
32
- /** Caller intent passed to {@link resolveModel}. */
33
- export interface ResolveModelInput {
34
- /**
35
- * Explicit model id / loose name (per-request override, agent /
36
- * plugin default, or env var). When set it wins over `modelClass`
37
- * and `fallbacks`.
38
- */
39
- explicit?: string;
40
- /**
41
- * Fuzzy-match an `explicit` name against the live catalogue so loose
42
- * names like `"claude sonnet"` resolve. Default `true`. When `false`
43
- * the explicit input is returned verbatim (Databricks surfaces the
44
- * canonical 404 if it doesn't exist).
45
- */
46
- fuzzy?: boolean;
47
- /** Fuse.js threshold forwarded to {@link resolveModelId}. */
48
- threshold?: number;
49
- /**
50
- * Chat capability class to resolve when no `explicit` id is given.
51
- * The live catalogue is classified by its Foundation Model API scores
52
- * and the top available model in the class (and the chat bands below
53
- * it) wins, falling back to the class's small static list.
54
- */
55
- modelClass?: ModelClass;
56
- /**
57
- * Operator-supplied fallback ids tried *first* in the no-explicit,
58
- * no-class path (e.g. a regulated workspace pinned to an approved
59
- * subset), ahead of the auto-classified catalogue.
60
- */
61
- fallbacks?: readonly string[];
62
- }
63
- /** Outcome of {@link resolveModel}: the chosen id plus how it was reached. */
64
- export interface ResolvedModelSelection {
65
- modelId: string;
66
- source: "explicit" | "fuzzy-match" | "class" | "fallback";
67
- }
68
- /** Intent + catalogue knobs passed to {@link selectModel}. */
69
- export interface SelectModelInput extends ResolveModelInput {
70
- /** TTL override for the cached `/serving-endpoints` listing, in ms. */
71
- ttlMs?: number;
72
- }
73
- /** TTL override merged into a {@link ModelQuery} for {@link searchModels}. */
74
- export interface SearchModelsInput extends ModelQuery {
75
- /** TTL override for the cached `/serving-endpoints` listing, in ms. */
76
- ttlMs?: number;
77
- }
78
- /**
79
- * Rank the live catalogue against a {@link ModelQuery}, best-first.
80
- *
81
- * Candidates are the classified endpoints in the eligible classes:
82
- * {@link classesAtOrBelow} the requested `modelClass`, or - when none is
83
- * given - the chat bands only ({@link CHAT_CLASS_ORDER}), so a general
84
- * ask never surfaces an embedding endpoint. Each class bucket is
85
- * already best-first from {@link classifyEndpoints}. Ranking is **match
86
- * then class**:
87
- *
88
- * 1. With a `search`, only endpoints matching it survive, ordered by
89
- * match distance (bucketed via {@link matchBucket} so near-identical
90
- * scores tie), then by class (more capable first), then by the stable
91
- * within-class rank.
92
- * 2. Without a `search`, the class-then-rank candidate order stands.
93
- *
94
- * A `limit` truncates the result. Returns `[]` when nothing is
95
- * eligible or matches - callers layer their own fallback.
96
- */
97
- export declare function rankModels(endpoints: readonly ServingEndpointSummary[], query?: ModelQuery): RankedModel[];
98
- /**
99
- * Rank a workspace's catalogue in one call: list its
100
- * `/serving-endpoints` (cached) and run {@link rankModels} over the
101
- * result. The list counterpart to {@link selectModel}, for a consumer
102
- * that wants the full ranked set (a model picker, a CLI) rather than a
103
- * single id. Catalogue fetches fail loud: network / auth errors
104
- * propagate so the caller sees the real SDK message.
105
- *
106
- * @param host - Workspace host used as the cache key. Pass the value
107
- * resolved from `client.config.getHost()`.
108
- */
109
- export declare function searchModels(client: WorkspaceClientLike, host: string, input?: SearchModelsInput): Promise<RankedModel[]>;
110
- /**
111
- * Resolve a model id for a workspace in one call: list its
112
- * `/serving-endpoints` (cached) and run {@link resolveModel} over the
113
- * result. This is the entry point for any consumer that holds a
114
- * `WorkspaceClient` and just wants a usable model name - a Lakeflow
115
- * job, a one-off script, or the Mastra plugin alike.
116
- *
117
- * Cheap exit: when an `explicit` name is given and `fuzzy` is off, the
118
- * catalogue is never fetched - the name is returned verbatim and
119
- * Databricks surfaces the canonical 404 if it doesn't exist. Catalogue
120
- * fetches otherwise fail loud: network / auth errors propagate so the
121
- * caller sees the real SDK message instead of a silent fallback.
122
- *
123
- * @param host - Workspace host used as the cache key. Pass the value
124
- * resolved from `client.config.getHost()`.
125
- */
126
- export declare function selectModel(client: WorkspaceClientLike, host: string, input?: SelectModelInput): Promise<ResolvedModelSelection>;
127
- /**
128
- * Resolve a single model id from the live catalogue and caller intent,
129
- * delegating the live selection to {@link rankModels} with `limit: 1`.
130
- *
131
- * 1. **Explicit ask**: with `fuzzy` off, returned verbatim; otherwise
132
- * fuzzy-ranked within the (optional) class ceiling and the best taken,
133
- * falling back to the input verbatim when nothing matches.
134
- * 2. **No explicit ask**: an operator-pinned `fallback` that exists in
135
- * the live catalogue wins first; then the ranked live catalogue
136
- * (class ceiling applied); then the static {@link FALLBACK_MODEL_IDS}
137
- * floor when the catalogue yields nothing in range.
138
- */
139
- export declare function resolveModel(endpoints: readonly ServingEndpointSummary[], input?: ResolveModelInput): ResolvedModelSelection;
@@ -1,203 +0,0 @@
1
- /**
2
- * Workspace-aware model selection.
3
- *
4
- * Given a caller's intent - a search string, a capability
5
- * {@link ModelClass} ceiling, both, or nothing - the toolkit returns
6
- * matching endpoints ranked by match quality then class, or collapses
7
- * to the single best id the workspace actually has, degrading from
8
- * "best in range" down to the static fallback floor. Selection is
9
- * chat-only: embedding endpoints surface only when `modelClass` is
10
- * explicitly {@link ModelClass.Embedding}.
11
- *
12
- * Ranking entry points (return a ranked list):
13
- *
14
- * - {@link rankModels} is the pure ranker over an endpoint list you
15
- * already hold. A chat `modelClass` acts as a ceiling: only that
16
- * band and the less-capable chat bands below it are eligible (see
17
- * {@link classesAtOrBelow}), so a `chat-balanced` ask can fall to
18
- * `chat-fast` but never escalate to `chat-thinking`.
19
- * - {@link searchModels} is the I/O wrapper: hand it a
20
- * `WorkspaceClient` and it lists `/serving-endpoints` (cached) then
21
- * ranks in one step.
22
- *
23
- * Single-selection entry points (return one id + how it was reached):
24
- *
25
- * - {@link resolveModel} is the pure selector; it delegates to
26
- * {@link rankModels} with `limit: 1` and adds the operator-pinned
27
- * fallback / static-floor safety net.
28
- * - {@link selectModel} is the I/O wrapper over {@link resolveModel}.
29
- */
30
- import { classifyEndpoints, ModelClass, } from "@dbx-tools/model-shared";
31
- import { CHAT_CLASS_ORDER, classesAtOrBelow, MODEL_CLASS_ORDER } from "./classes.js";
32
- import { FALLBACK_MODEL_IDS, modelsForClass } from "./fallback.js";
33
- import { listServingEndpoints, searchServingEndpoints, } from "./serving.js";
34
- /**
35
- * Round a Fuse score to the display precision so version siblings that
36
- * match a token identically (e.g. `opus-4-7` vs `opus-4-8` for the
37
- * query `"opus"`) tie on match and let the class / within-class rank
38
- * decide - which is what surfaces the newer, higher-quality sibling.
39
- */
40
- function matchBucket(score) {
41
- return Math.round((score ?? 0) * 1000);
42
- }
43
- /**
44
- * Rank the live catalogue against a {@link ModelQuery}, best-first.
45
- *
46
- * Candidates are the classified endpoints in the eligible classes:
47
- * {@link classesAtOrBelow} the requested `modelClass`, or - when none is
48
- * given - the chat bands only ({@link CHAT_CLASS_ORDER}), so a general
49
- * ask never surfaces an embedding endpoint. Each class bucket is
50
- * already best-first from {@link classifyEndpoints}. Ranking is **match
51
- * then class**:
52
- *
53
- * 1. With a `search`, only endpoints matching it survive, ordered by
54
- * match distance (bucketed via {@link matchBucket} so near-identical
55
- * scores tie), then by class (more capable first), then by the stable
56
- * within-class rank.
57
- * 2. Without a `search`, the class-then-rank candidate order stands.
58
- *
59
- * A `limit` truncates the result. Returns `[]` when nothing is
60
- * eligible or matches - callers layer their own fallback.
61
- */
62
- export function rankModels(endpoints, query = {}) {
63
- const classified = classifyEndpoints(endpoints);
64
- const eligible = query.modelClass !== undefined
65
- ? classesAtOrBelow(query.modelClass)
66
- : CHAT_CLASS_ORDER;
67
- // Flatten eligible classes in capability order, carrying each
68
- // endpoint's class; bucket order is already best-first.
69
- const candidates = [];
70
- for (const modelClass of eligible) {
71
- for (const endpoint of classified[modelClass])
72
- candidates.push({ endpoint, modelClass });
73
- }
74
- const search = query.search?.trim();
75
- let ranked;
76
- if (search) {
77
- const scores = new Map();
78
- for (const match of searchServingEndpoints(search, candidates.map((c) => c.endpoint), query.threshold !== undefined ? { threshold: query.threshold } : {})) {
79
- scores.set(match.endpoint.name, match.score);
80
- }
81
- // `Array.prototype.sort` is stable, so endpoints equal on match and
82
- // class keep their best-first within-class order.
83
- ranked = candidates
84
- .filter((c) => scores.has(c.endpoint.name))
85
- .map((c) => ({ ...c, score: scores.get(c.endpoint.name) }))
86
- .sort((a, b) => {
87
- const byMatch = matchBucket(a.score) - matchBucket(b.score);
88
- if (byMatch !== 0)
89
- return byMatch;
90
- return (MODEL_CLASS_ORDER.indexOf(a.modelClass) -
91
- MODEL_CLASS_ORDER.indexOf(b.modelClass));
92
- });
93
- }
94
- else {
95
- ranked = candidates;
96
- }
97
- return query.limit !== undefined ? ranked.slice(0, Math.max(0, query.limit)) : ranked;
98
- }
99
- /**
100
- * Rank a workspace's catalogue in one call: list its
101
- * `/serving-endpoints` (cached) and run {@link rankModels} over the
102
- * result. The list counterpart to {@link selectModel}, for a consumer
103
- * that wants the full ranked set (a model picker, a CLI) rather than a
104
- * single id. Catalogue fetches fail loud: network / auth errors
105
- * propagate so the caller sees the real SDK message.
106
- *
107
- * @param host - Workspace host used as the cache key. Pass the value
108
- * resolved from `client.config.getHost()`.
109
- */
110
- export async function searchModels(client, host, input = {}) {
111
- const endpoints = await listServingEndpoints(client, host, input.ttlMs !== undefined ? { ttlMs: input.ttlMs } : {});
112
- return rankModels(endpoints, input);
113
- }
114
- /**
115
- * Resolve a model id for a workspace in one call: list its
116
- * `/serving-endpoints` (cached) and run {@link resolveModel} over the
117
- * result. This is the entry point for any consumer that holds a
118
- * `WorkspaceClient` and just wants a usable model name - a Lakeflow
119
- * job, a one-off script, or the Mastra plugin alike.
120
- *
121
- * Cheap exit: when an `explicit` name is given and `fuzzy` is off, the
122
- * catalogue is never fetched - the name is returned verbatim and
123
- * Databricks surfaces the canonical 404 if it doesn't exist. Catalogue
124
- * fetches otherwise fail loud: network / auth errors propagate so the
125
- * caller sees the real SDK message instead of a silent fallback.
126
- *
127
- * @param host - Workspace host used as the cache key. Pass the value
128
- * resolved from `client.config.getHost()`.
129
- */
130
- export async function selectModel(client, host, input = {}) {
131
- if (input.explicit !== undefined && input.fuzzy === false) {
132
- return { modelId: input.explicit, source: "explicit" };
133
- }
134
- const endpoints = await listServingEndpoints(client, host, {
135
- ...(input.ttlMs !== undefined ? { ttlMs: input.ttlMs } : {}),
136
- });
137
- return resolveModel(endpoints, input);
138
- }
139
- /**
140
- * Resolve a single model id from the live catalogue and caller intent,
141
- * delegating the live selection to {@link rankModels} with `limit: 1`.
142
- *
143
- * 1. **Explicit ask**: with `fuzzy` off, returned verbatim; otherwise
144
- * fuzzy-ranked within the (optional) class ceiling and the best taken,
145
- * falling back to the input verbatim when nothing matches.
146
- * 2. **No explicit ask**: an operator-pinned `fallback` that exists in
147
- * the live catalogue wins first; then the ranked live catalogue
148
- * (class ceiling applied); then the static {@link FALLBACK_MODEL_IDS}
149
- * floor when the catalogue yields nothing in range.
150
- */
151
- export function resolveModel(endpoints, input = {}) {
152
- if (input.explicit !== undefined) {
153
- if (input.fuzzy === false) {
154
- return { modelId: input.explicit, source: "explicit" };
155
- }
156
- const [top] = rankModels(endpoints, buildQuery(input, input.explicit));
157
- return { modelId: top?.endpoint.name ?? input.explicit, source: "fuzzy-match" };
158
- }
159
- // Operator-pinned fallbacks win when present and live (e.g. a
160
- // regulated workspace restricted to an approved subset).
161
- if (input.modelClass === undefined && input.fallbacks && input.fallbacks.length > 0) {
162
- const present = new Set(endpoints.map((e) => e.name));
163
- const pinned = input.fallbacks.find((id) => present.has(id));
164
- if (pinned)
165
- return { modelId: pinned, source: "fallback" };
166
- }
167
- const source = input.modelClass !== undefined ? "class" : "fallback";
168
- const [top] = rankModels(endpoints, buildQuery(input, undefined));
169
- if (top)
170
- return { modelId: top.endpoint.name, source };
171
- // Live catalogue yielded nothing in range: walk the static floor.
172
- const floor = input.modelClass !== undefined
173
- ? dedupe([...modelsForClass(input.modelClass), ...FALLBACK_MODEL_IDS])
174
- : dedupe([...(input.fallbacks ?? []), ...FALLBACK_MODEL_IDS]);
175
- return { modelId: pickFirstAvailable(floor, endpoints), source };
176
- }
177
- /** Build a {@link ModelQuery} from {@link ResolveModelInput} for the `limit: 1` delegation. */
178
- function buildQuery(input, search) {
179
- return {
180
- ...(search !== undefined ? { search } : {}),
181
- ...(input.modelClass !== undefined ? { modelClass: input.modelClass } : {}),
182
- ...(input.threshold !== undefined ? { threshold: input.threshold } : {}),
183
- limit: 1,
184
- };
185
- }
186
- /** Drop duplicate ids while preserving first-seen order. */
187
- function dedupe(ids) {
188
- return [...new Set(ids)];
189
- }
190
- /**
191
- * Find the first id in `candidates` whose endpoint is present in
192
- * `endpoints`. Returns the top candidate when the workspace has none
193
- * of them so callers always get a string; an offline workspace then
194
- * receives a clean 404 from Databricks instead of a malformed config.
195
- */
196
- function pickFirstAvailable(candidates, endpoints) {
197
- const present = new Set(endpoints.map((e) => e.name));
198
- for (const candidate of candidates) {
199
- if (present.has(candidate))
200
- return candidate;
201
- }
202
- return candidates[0] ?? FALLBACK_MODEL_IDS[0];
203
- }