@dbx-tools/model 0.1.42 → 0.1.48

package/index.ts

@@ -4,7 +4,7 @@
  * Bundles the package's Node-side Model Serving access (cached
  * `/serving-endpoints` listing plus fuzzy name resolution), the
  * workspace-aware {@link resolveModel} selector, and the server-only
- * offline fallback opinion (`FALLBACK_MODEL_IDS` / `modelsForTier`)
+ * offline fallback opinion (`FALLBACK_MODEL_IDS` / `modelsForClass`)
  * with a re-export of the pure `@dbx-tools/model-shared` surface
  * (capability tiers, the score profile, the endpoint descriptor, and
  * the tier classifier) so a single `from "@dbx-tools/model"` import
@@ -17,6 +17,7 @@
  */
 
 export * from "@dbx-tools/model-shared";
+export * from "./src/classes.js";
 export * from "./src/fallback.js";
 export * from "./src/resolve.js";
 export * from "./src/serving.js";

package/package.json

@@ -9,10 +9,10 @@
     }
   },
   "name": "@dbx-tools/model",
-  "version": "0.1.42",
+  "version": "0.1.48",
   "dependencies": {
-    "@dbx-tools/model-shared": "0.1.42",
-    "@dbx-tools/shared": "0.1.42",
+    "@dbx-tools/model-shared": "0.1.48",
+    "@dbx-tools/shared": "0.1.48",
     "fuse.js": "^7.0.0"
   },
   "module": "index.ts",

package/src/classes.ts

@@ -0,0 +1,75 @@
+/**
+ * Model-class ordering for the model service.
+ *
+ * `@dbx-tools/model-shared` declares the {@link ModelClass} values and
+ * their schema; the *behavior* over those values - the chat capability
+ * order, the "this class and below" ceiling, and coercing loose request
+ * input to a class - lives here in the service so the shared
+ * wire-format surface stays purely declarative.
+ */
+
+import { ModelClass, ModelClassSchema } from "@dbx-tools/model-shared";
+
+/**
+ * Chat capability ladder in descending order - most capable
+ * ({@link ModelClass.ChatThinking}) first, least
+ * ({@link ModelClass.ChatFast}) last. {@link ModelClass.Embedding} is a
+ * separate modality and is deliberately absent: the ceiling never spans
+ * it. This is the order used to filter "this class and below" and to
+ * break ranking ties toward the more capable model.
+ */
+export const CHAT_CLASS_ORDER: readonly ModelClass[] = [
+  ModelClass.ChatThinking,
+  ModelClass.ChatBalanced,
+  ModelClass.ChatFast,
+];
+
+/**
+ * Every class in display order: the chat ladder followed by
+ * {@link ModelClass.Embedding}. Used when flattening a full
+ * classification (e.g. stamping the class onto each cached endpoint).
+ */
+export const MODEL_CLASS_ORDER: readonly ModelClass[] = [
+  ...CHAT_CLASS_ORDER,
+  ModelClass.Embedding,
+];
+
+/** Whether `cls` is one of the chat capability bands (vs. embedding). */
+export function isChatClass(cls: ModelClass): boolean {
+  return CHAT_CLASS_ORDER.includes(cls);
+}
+
+/**
+ * Coerce an arbitrary value (query string, header, body field) to a
+ * {@link ModelClass}, returning `null` when it isn't a known class.
+ * Lets a route accept a class request without throwing on junk input.
+ *
+ * Matches the full slug first (`"chat-fast"`, `"embedding"`), then -
+ * for a bare chat band - retries with a `chat-` prefix, so the shorthand
+ * `"thinking"` / `"balanced"` / `"fast"` resolves to the corresponding
+ * `chat-*` class.
+ */
+export function parseModelClass(value: unknown): ModelClass | null {
+  const exact = ModelClassSchema.safeParse(value);
+  if (exact.success) return exact.data;
+  const prefixed = ModelClassSchema.safeParse(`chat-${value}`);
+  return prefixed.success ? prefixed.data : null;
+}
+
+/**
+ * Classes at or below `cls` in chat capability: the class itself plus
+ * every less-capable chat band, in {@link CHAT_CLASS_ORDER}. Treats a
+ * chat `cls` as a ceiling - requesting {@link ModelClass.ChatBalanced}
+ * yields `[ChatBalanced, ChatFast]`, never
+ * {@link ModelClass.ChatThinking} - so a class ask can degrade downward
+ * to a smaller chat model but never escalate to a larger one.
+ *
+ * {@link ModelClass.Embedding} is its own modality, not a rung on the
+ * chat ladder, so it yields just `[Embedding]`. An unrecognized class
+ * yields the full chat ladder.
+ */
+export function classesAtOrBelow(cls: ModelClass): ModelClass[] {
+  if (cls === ModelClass.Embedding) return [ModelClass.Embedding];
+  const index = CHAT_CLASS_ORDER.indexOf(cls);
+  return index < 0 ? [...CHAT_CLASS_ORDER] : [...CHAT_CLASS_ORDER.slice(index)];
+}

package/src/fallback.ts

@@ -5,9 +5,9 @@
  * no user token, the service principal can't list, or the workspace is
  * unreachable - the resolver still has to name *some* endpoint. This
  * module holds that floor: a small, hard-coded set of well-known
- * Foundation Model API endpoint names, each bucketed into a
- * {@link ModelTier} by the shared {@link classifyByFamily} heuristic
- * (no tiers are hard-coded here) and ordered best-first.
+ * Foundation Model API endpoint names, each bucketed into a chat
+ * {@link ModelClass} by the shared {@link classifyByFamily} heuristic
+ * (no classes are hard-coded here) and ordered best-first.
  *
  * This is deliberately *server-only*. A browser client never talks to
  * Databricks directly - it always goes through this server - so it has
@@ -16,14 +16,18 @@
  * classifier ({@link classifyEndpoints}) is what the client shares.
  */
 
-import { classifyByFamily, type FamilyClass, ModelTier } from "@dbx-tools/model-shared";
+import {
+  classifyByFamily,
+  type FamilyClass,
+  ModelClass,
+} from "@dbx-tools/model-shared";
 
 /**
  * Small, last-resort set of well-known Foundation Model API endpoint
- * names, ordered best-first within each tier by {@link classifyByFamily}.
+ * names, ordered best-first within each class by {@link classifyByFamily}.
  * Used only as the floor when the live `/serving-endpoints` catalogue
  * can't be read at resolve time; the live, score-driven classification
- * supersedes it whenever the workspace listing is available. Tiers are
+ * supersedes it whenever the workspace listing is available. Classes are
  * not hard-coded here - each name is classified by family heuristic.
  */
 const FALLBACK_MODEL_NAMES: readonly string[] = [
@@ -39,34 +43,35 @@ const FALLBACK_MODEL_NAMES: readonly string[] = [
 ];
 
 /**
- * Static fallback model ids for a tier, drawn from the small built-in
- * {@link FALLBACK_MODEL_NAMES} list and ordered best-first by family
- * rank. Sync and workspace-independent: this is the *fallback opinion*
- * used to seed default lists or when the live catalogue is unreachable
- * - live resolution prefers `classifyEndpoints`.
+ * Static fallback model ids for a chat class, drawn from the small
+ * built-in {@link FALLBACK_MODEL_NAMES} list and ordered best-first by
+ * family rank. Sync and workspace-independent: this is the *fallback
+ * opinion* used to seed default lists or when the live catalogue is
+ * unreachable - live resolution prefers `classifyEndpoints`. Returns
+ * `[]` for {@link ModelClass.Embedding} (the floor is chat-only).
  */
-export function modelsForTier(tier: ModelTier): readonly string[] {
+export function modelsForClass(cls: ModelClass): readonly string[] {
   return FALLBACK_MODEL_NAMES.map((name) => ({ name, c: classifyByFamily(name) }))
     .filter(
-      (x): x is { name: string; c: FamilyClass } => x.c !== null && x.c.tier === tier,
+      (x): x is { name: string; c: FamilyClass } => x.c !== null && x.c.class === cls,
     )
     .sort((a, b) => b.c.rank - a.c.rank)
     .map((x) => x.name);
 }
 
-/** Top static fallback model id for a tier. */
-export function modelForTier(tier: ModelTier): string {
-  return modelsForTier(tier)[0]!;
+/** Top static fallback model id for a chat class. */
+export function modelForClass(cls: ModelClass): string {
+  return modelsForClass(cls)[0]!;
 }
 
 /**
- * Priority-ordered fallback chain (Thinking -> Balanced -> Fast) over
- * the small built-in list. The floor walked at resolve time when no
- * agent / plugin / env / request model is set *and* the live
- * catalogue yields nothing.
+ * Priority-ordered fallback chain (ChatThinking -> ChatBalanced ->
+ * ChatFast) over the small built-in list. The floor walked at resolve
+ * time when no agent / plugin / env / request model is set *and* the
+ * live catalogue yields nothing.
  */
 export const FALLBACK_MODEL_IDS: readonly string[] = [
-  ...modelsForTier(ModelTier.Thinking),
-  ...modelsForTier(ModelTier.Balanced),
-  ...modelsForTier(ModelTier.Fast),
+  ...modelsForClass(ModelClass.ChatThinking),
+  ...modelsForClass(ModelClass.ChatBalanced),
+  ...modelsForClass(ModelClass.ChatFast),
 ];

package/src/resolve.ts

@@ -1,29 +1,46 @@
 /**
  * Workspace-aware model selection.
  *
- * Given a caller's intent (an explicit name, a capability
- * {@link ModelTier}, or nothing), the resolver returns the endpoint id
- * to use - always one the workspace actually has, degrading from "best
- * in tier" down to the static fallback floor. Two entry points:
+ * Given a caller's intent - a search string, a capability
+ * {@link ModelClass} ceiling, both, or nothing - the toolkit returns
+ * matching endpoints ranked by match quality then class, or collapses
+ * to the single best id the workspace actually has, degrading from
+ * "best in range" down to the static fallback floor. Selection is
+ * chat-only: embedding endpoints surface only when `modelClass` is
+ * explicitly {@link ModelClass.Embedding}.
  *
- *   - {@link selectModel} is the high-level, I/O-doing call: hand it a
- *     `WorkspaceClient` and intent and it lists `/serving-endpoints`
- *     (cached) and resolves in one step. This is what a non-Mastra
- *     consumer (e.g. a job that just needs a model name) reaches for.
- *   - {@link resolveModel} is the pure form over an endpoint list you
- *     already hold; `selectModel` is a thin wrapper over it.
+ * Ranking entry points (return a ranked list):
+ *
+ *   - {@link rankModels} is the pure ranker over an endpoint list you
+ *     already hold. A chat `modelClass` acts as a ceiling: only that
+ *     band and the less-capable chat bands below it are eligible (see
+ *     {@link classesAtOrBelow}), so a `chat-balanced` ask can fall to
+ *     `chat-fast` but never escalate to `chat-thinking`.
+ *   - {@link searchModels} is the I/O wrapper: hand it a
+ *     `WorkspaceClient` and it lists `/serving-endpoints` (cached) then
+ *     ranks in one step.
+ *
+ * Single-selection entry points (return one id + how it was reached):
+ *
+ *   - {@link resolveModel} is the pure selector; it delegates to
+ *     {@link rankModels} with `limit: 1` and adds the operator-pinned
+ *     fallback / static-floor safety net.
+ *   - {@link selectModel} is the I/O wrapper over {@link resolveModel}.
  */
 
 import {
   classifyEndpoints,
-  ModelTier,
+  ModelClass,
+  type ModelQuery,
+  type RankedModel,
   type ServingEndpointSummary,
 } from "@dbx-tools/model-shared";
 
-import { FALLBACK_MODEL_IDS, modelsForTier } from "./fallback.js";
+import { CHAT_CLASS_ORDER, classesAtOrBelow, MODEL_CLASS_ORDER } from "./classes.js";
+import { FALLBACK_MODEL_IDS, modelsForClass } from "./fallback.js";
 import {
   listServingEndpoints,
-  resolveModelId,
+  searchServingEndpoints,
   type WorkspaceClientLike,
 } from "./serving.js";
 
@@ -31,8 +48,8 @@ import {
 export interface ResolveModelInput {
   /**
    * Explicit model id / loose name (per-request override, agent /
-   * plugin default, or env var). When set it wins over `tier` and
-   * `fallbacks`.
+   * plugin default, or env var). When set it wins over `modelClass`
+   * and `fallbacks`.
    */
   explicit?: string;
   /**
@@ -45,15 +62,15 @@ export interface ResolveModelInput {
   /** Fuse.js threshold forwarded to {@link resolveModelId}. */
   threshold?: number;
   /**
-   * Capability tier to resolve when no `explicit` id is given. The
-   * live catalogue is classified by its Foundation Model API scores
-   * and the top available model in the tier wins, falling back to the
-   * tier's small static list.
+   * Chat capability class to resolve when no `explicit` id is given.
+   * The live catalogue is classified by its Foundation Model API scores
+   * and the top available model in the class (and the chat bands below
+   * it) wins, falling back to the class's small static list.
    */
-  tier?: ModelTier;
+  modelClass?: ModelClass;
   /**
    * Operator-supplied fallback ids tried *first* in the no-explicit,
-   * no-tier path (e.g. a regulated workspace pinned to an approved
+   * no-class path (e.g. a regulated workspace pinned to an approved
    * subset), ahead of the auto-classified catalogue.
    */
   fallbacks?: readonly string[];
@@ -62,7 +79,7 @@ export interface ResolveModelInput {
 /** Outcome of {@link resolveModel}: the chosen id plus how it was reached. */
 export interface ResolvedModelSelection {
   modelId: string;
-  source: "explicit" | "fuzzy-match" | "tier" | "fallback";
+  source: "explicit" | "fuzzy-match" | "class" | "fallback";
 }
 
 /** Intent + catalogue knobs passed to {@link selectModel}. */
@@ -71,6 +88,114 @@ export interface SelectModelInput extends ResolveModelInput {
   ttlMs?: number;
 }
 
+/** TTL override merged into a {@link ModelQuery} for {@link searchModels}. */
+export interface SearchModelsInput extends ModelQuery {
+  /** TTL override for the cached `/serving-endpoints` listing, in ms. */
+  ttlMs?: number;
+}
+
+/**
+ * Round a Fuse score to the display precision so version siblings that
+ * match a token identically (e.g. `opus-4-7` vs `opus-4-8` for the
+ * query `"opus"`) tie on match and let the class / within-class rank
+ * decide - which is what surfaces the newer, higher-quality sibling.
+ */
+function matchBucket(score: number | undefined): number {
+  return Math.round((score ?? 0) * 1000);
+}
+
+/**
+ * Rank the live catalogue against a {@link ModelQuery}, best-first.
+ *
+ * Candidates are the classified endpoints in the eligible classes:
+ * {@link classesAtOrBelow} the requested `modelClass`, or - when none is
+ * given - the chat bands only ({@link CHAT_CLASS_ORDER}), so a general
+ * ask never surfaces an embedding endpoint. Each class bucket is
+ * already best-first from {@link classifyEndpoints}. Ranking is **match
+ * then class**:
+ *
+ * 1. With a `search`, only endpoints matching it survive, ordered by
+ *    match distance (bucketed via {@link matchBucket} so near-identical
+ *    scores tie), then by class (more capable first), then by the stable
+ *    within-class rank.
+ * 2. Without a `search`, the class-then-rank candidate order stands.
+ *
+ * A `limit` truncates the result. Returns `[]` when nothing is
+ * eligible or matches - callers layer their own fallback.
+ */
+export function rankModels(
+  endpoints: readonly ServingEndpointSummary[],
+  query: ModelQuery = {},
+): RankedModel[] {
+  const classified = classifyEndpoints(endpoints);
+  const eligible =
+    query.modelClass !== undefined
+      ? classesAtOrBelow(query.modelClass)
+      : CHAT_CLASS_ORDER;
+
+  // Flatten eligible classes in capability order, carrying each
+  // endpoint's class; bucket order is already best-first.
+  const candidates: RankedModel[] = [];
+  for (const modelClass of eligible) {
+    for (const endpoint of classified[modelClass])
+      candidates.push({ endpoint, modelClass });
+  }
+
+  const search = query.search?.trim();
+  let ranked: RankedModel[];
+  if (search) {
+    const scores = new Map<string, number>();
+    for (const match of searchServingEndpoints(
+      search,
+      candidates.map((c) => c.endpoint),
+      query.threshold !== undefined ? { threshold: query.threshold } : {},
+    )) {
+      scores.set(match.endpoint.name, match.score);
+    }
+    // `Array.prototype.sort` is stable, so endpoints equal on match and
+    // class keep their best-first within-class order.
+    ranked = candidates
+      .filter((c) => scores.has(c.endpoint.name))
+      .map((c) => ({ ...c, score: scores.get(c.endpoint.name) }))
+      .sort((a, b) => {
+        const byMatch = matchBucket(a.score) - matchBucket(b.score);
+        if (byMatch !== 0) return byMatch;
+        return (
+          MODEL_CLASS_ORDER.indexOf(a.modelClass) -
+          MODEL_CLASS_ORDER.indexOf(b.modelClass)
+        );
+      });
+  } else {
+    ranked = candidates;
+  }
+
+  return query.limit !== undefined ? ranked.slice(0, Math.max(0, query.limit)) : ranked;
+}
+
+/**
+ * Rank a workspace's catalogue in one call: list its
+ * `/serving-endpoints` (cached) and run {@link rankModels} over the
+ * result. The list counterpart to {@link selectModel}, for a consumer
+ * that wants the full ranked set (a model picker, a CLI) rather than a
+ * single id. Catalogue fetches fail loud: network / auth errors
+ * propagate so the caller sees the real SDK message.
+ *
+ * @param host - Workspace host used as the cache key. Pass the value
+ *   resolved from `client.config.getHost()`.
+ */
+export async function searchModels(
+  client: WorkspaceClientLike,
+  host: string,
+  input: SearchModelsInput = {},
+): Promise<RankedModel[]> {
+  const endpoints = await listServingEndpoints(
+    client,
+    host,
+    input.ttlMs !== undefined ? { ttlMs: input.ttlMs } : {},
+  );
+  return rankModels(endpoints, input);
+}
+
 /**
  * Resolve a model id for a workspace in one call: list its
  * `/serving-endpoints` (cached) and run {@link resolveModel} over the
@@ -102,16 +227,16 @@ export async function selectModel(
 }
 
 /**
- * Resolve a model id from the live catalogue and caller intent.
+ * Resolve a single model id from the live catalogue and caller intent,
+ * delegating the live selection to {@link rankModels} with `limit: 1`.
  *
  * 1. **Explicit ask**: with `fuzzy` off, returned verbatim; otherwise
- *    snapped to the closest live endpoint via {@link resolveModelId}.
- * 2. **Tier intent**: classify the live catalogue and return the top
- *    available model in that tier, then the tier's static list, then
- *    the {@link FALLBACK_MODEL_IDS} floor.
- * 3. **Neither**: walk `fallbacks` first, then the live
- *    score-classified catalogue (Thinking -> Balanced -> Fast), then
- *    the static floor, and return the first id actually present.
+ *    fuzzy-ranked within the (optional) class ceiling and the best taken,
+ *    falling back to the input verbatim when nothing matches.
+ * 2. **No explicit ask**: an operator-pinned `fallback` that exists in
+ *    the live catalogue wins first; then the ranked live catalogue
+ *    (class ceiling applied); then the static {@link FALLBACK_MODEL_IDS}
+ *    floor when the catalogue yields nothing in range.
  */
 export function resolveModel(
   endpoints: readonly ServingEndpointSummary[],
@@ -121,54 +246,38 @@ export function resolveModel(
     if (input.fuzzy === false) {
       return { modelId: input.explicit, source: "explicit" };
     }
-    const { modelId } = resolveModelId(input.explicit, endpoints, {
-      threshold: input.threshold,
-    });
-    return { modelId, source: "fuzzy-match" };
+    const [top] = rankModels(endpoints, buildQuery(input, input.explicit));
+    return { modelId: top?.endpoint.name ?? input.explicit, source: "fuzzy-match" };
   }
 
-  const classified = classifyEndpoints(endpoints);
-  const chain =
-    input.tier !== undefined
-      ? tierChain(classified, input.tier)
-      : defaultChain(classified, input.fallbacks);
-  return {
-    modelId: pickFirstAvailable(chain, endpoints),
-    source: input.tier !== undefined ? "tier" : "fallback",
-  };
-}
+  // Operator-pinned fallbacks win when present and live (e.g. a
+  // regulated workspace restricted to an approved subset).
+  if (input.modelClass === undefined && input.fallbacks && input.fallbacks.length > 0) {
+    const present = new Set(endpoints.map((e) => e.name));
+    const pinned = input.fallbacks.find((id) => present.has(id));
+    if (pinned) return { modelId: pinned, source: "fallback" };
+  }
 
-/**
- * Candidate chain for an explicit {@link ModelTier} ask: live models
- * classified into that tier (best-first), then the tier's small static
- * list, then the full fallback floor as a last resort.
- */
-function tierChain(
-  classified: Record<ModelTier, ServingEndpointSummary[]>,
-  tier: ModelTier,
-): string[] {
-  return dedupe([
-    ...classified[tier].map((e) => e.name),
-    ...modelsForTier(tier),
-    ...FALLBACK_MODEL_IDS,
-  ]);
+  const source = input.modelClass !== undefined ? "class" : "fallback";
+  const [top] = rankModels(endpoints, buildQuery(input, undefined));
+  if (top) return { modelId: top.endpoint.name, source };
+
+  // Live catalogue yielded nothing in range: walk the static floor.
+  const floor =
+    input.modelClass !== undefined
+      ? dedupe([...modelsForClass(input.modelClass), ...FALLBACK_MODEL_IDS])
+      : dedupe([...(input.fallbacks ?? []), ...FALLBACK_MODEL_IDS]);
+  return { modelId: pickFirstAvailable(floor, endpoints), source };
 }
 
-/**
- * Candidate chain for the no-explicit default: any operator-supplied
- * `fallbacks` first, then the live score-classified catalogue in
- * descending tier order, then the static fallback floor.
- */
-function defaultChain(
-  classified: Record<ModelTier, ServingEndpointSummary[]>,
-  fallbacks: readonly string[] | undefined,
-): string[] {
-  const live = [
-    ...classified[ModelTier.Thinking],
-    ...classified[ModelTier.Balanced],
-    ...classified[ModelTier.Fast],
-  ].map((e) => e.name);
-  return dedupe([...(fallbacks ?? []), ...live, ...FALLBACK_MODEL_IDS]);
+/** Build a {@link ModelQuery} from {@link ResolveModelInput} for the `limit: 1` delegation. */
+function buildQuery(input: ResolveModelInput, search: string | undefined): ModelQuery {
+  return {
+    ...(search !== undefined ? { search } : {}),
+    ...(input.modelClass !== undefined ? { modelClass: input.modelClass } : {}),
+    ...(input.threshold !== undefined ? { threshold: input.threshold } : {}),
+    limit: 1,
+  };
 }
 
 /** Drop duplicate ids while preserving first-seen order. */