@dbx-tools/model 0.9.5 → 0.9.6

package/package.json

@@ -25,11 +25,10 @@
   },
   "dependencies": {
     "@databricks/appkit": "0.81.0",
-    "@dbx-tools/appkit": "^0.9.5",
-    "@dbx-tools/model-rs": "^0.9.5",
-    "@dbx-tools/shared-core": "^0.9.5",
-    "@dbx-tools/shared-model": "^0.9.5",
-    "fuse.js": "^7.4.2"
+    "@dbx-tools/appkit": "^0.9.6",
+    "@dbx-tools/model-rs": "^0.9.6",
+    "@dbx-tools/shared-core": "^0.9.6",
+    "@dbx-tools/shared-model": "^0.9.6"
   },
   "main": "./lib/index.js",
   "license": "Apache-2.0",
@@ -45,7 +44,7 @@
       "./package.json": "./package.json"
     }
   },
-  "version": "0.9.5",
+  "version": "0.9.6",
   "types": "./lib/index.d.ts",
   "type": "module",
   "exports": {

package/src/_native.ts

@@ -0,0 +1,174 @@
+import {
+  ModelClass as NativeModelClass,
+  ReasoningEffort as NativeReasoningEffort,
+  normalizeServingEndpointsJson,
+  rankModels as rankModelsWithRust,
+  type ModelQuery as NativeModelQuery,
+  type ServingEndpointSummary as NativeServingEndpointSummary,
+} from "@dbx-tools/model-rs";
+import {
+  model,
+  type ModelQuery,
+  type RankedModel,
+  type ReasoningEffort,
+  type ServingEndpointSummary,
+} from "@dbx-tools/shared-model";
+
+type ModelClass = model.ModelClass;
+
+const MODEL_CLASS_TO_NATIVE: Readonly<Record<ModelClass, NativeModelClass>> = {
+  [model.ModelClass.ChatThinking]: NativeModelClass.ChatThinking,
+  [model.ModelClass.ChatBalanced]: NativeModelClass.ChatBalanced,
+  [model.ModelClass.ChatFast]: NativeModelClass.ChatFast,
+  [model.ModelClass.Embedding]: NativeModelClass.Embedding,
+};
+
+const REASONING_EFFORT_TO_NATIVE: Readonly<Record<string, NativeReasoningEffort>> = {
+  none: NativeReasoningEffort.None,
+  minimal: NativeReasoningEffort.Minimal,
+  low: NativeReasoningEffort.Low,
+  medium: NativeReasoningEffort.Medium,
+  high: NativeReasoningEffort.High,
+  xhigh: NativeReasoningEffort.Xhigh,
+  max: NativeReasoningEffort.Max,
+};
+
+const REASONING_EFFORT_FROM_NATIVE: Readonly<Record<NativeReasoningEffort, ReasoningEffort>> = {
+  [NativeReasoningEffort.None]: "none",
+  [NativeReasoningEffort.Minimal]: "minimal",
+  [NativeReasoningEffort.Low]: "low",
+  [NativeReasoningEffort.Medium]: "medium",
+  [NativeReasoningEffort.High]: "high",
+  [NativeReasoningEffort.Xhigh]: "xhigh",
+  [NativeReasoningEffort.Max]: "max",
+};
+
+interface NativeRankingOptions {
+  readonly includeDeprecated?: boolean;
+  readonly modelClass?: ModelClass;
+  readonly task?: string;
+}
+
+/** Rank public endpoint records through the Rust-owned model policy. */
+export function rankEndpoints(
+  endpoints: readonly ServingEndpointSummary[],
+  query: ModelQuery = {},
+  options: NativeRankingOptions = {},
+): RankedModel[] {
+  const endpointsByName = new Map(endpoints.map((endpoint) => [endpoint.name, endpoint]));
+  return rankModelsWithRust(
+    endpoints.map((endpoint) => toNativeEndpoint(endpoint, options)),
+    toNativeQuery(query, options.includeDeprecated ?? false),
+  ).map((ranked) => {
+    const endpoint = endpointsByName.get(ranked.endpoint.name);
+    if (!endpoint) {
+      throw new Error(`Rust model ranking returned unknown endpoint "${ranked.endpoint.name}"`);
+    }
+    return {
+      endpoint,
+      modelClass: fromNativeClass(ranked.modelClass),
+      ...(ranked.score !== undefined ? { score: ranked.score } : {}),
+    };
+  });
+}
+
+/** Classify chat and embedding endpoints through the Rust-owned policy. */
+export function classifyEndpointClasses(
+  endpoints: readonly ServingEndpointSummary[],
+): ReadonlyMap<string, ModelClass> {
+  const chat = rankEndpoints(endpoints, {}, { includeDeprecated: true });
+  const embeddings = rankEndpoints(
+    endpoints,
+    { modelClass: model.ModelClass.Embedding },
+    { includeDeprecated: true },
+  );
+  return new Map(
+    [...chat, ...embeddings].map((ranked) => [ranked.endpoint.name, ranked.modelClass]),
+  );
+}
+
+/** Normalize serialized SDK endpoint records through the Rust-owned catalogue policy. */
+export function normalizeEndpoints(endpoints: readonly unknown[]): ServingEndpointSummary[] {
+  return normalizeServingEndpointsJson(JSON.stringify({ endpoints })).map(fromNativeEndpoint);
+}
+
+function toNativeQuery(query: ModelQuery, includeDeprecated: boolean): NativeModelQuery {
+  return {
+    search: query.search,
+    modelClass:
+      query.modelClass === undefined ? undefined : MODEL_CLASS_TO_NATIVE[query.modelClass],
+    requiresTools: query.requiresTools,
+    includeDeprecated: includeDeprecated || undefined,
+    limit: query.limit,
+    threshold: query.threshold,
+  };
+}
+
+function toNativeEndpoint(
+  endpoint: ServingEndpointSummary,
+  options: NativeRankingOptions,
+): NativeServingEndpointSummary {
+  const modelClass = options.modelClass ?? endpoint.class;
+  return {
+    name: endpoint.name,
+    displayName: endpoint.displayName,
+    family: endpoint.family,
+    task: options.task ?? endpoint.task,
+    state: endpoint.state,
+    description: endpoint.description,
+    supportsTools: endpoint.supportsTools,
+    profile: endpoint.profile,
+    modelClass: modelClass === undefined ? undefined : MODEL_CLASS_TO_NATIVE[modelClass],
+    serviceNames: new Map(Object.entries(endpoint.serviceNames ?? {})),
+    modelServiceName: endpoint.modelServiceName,
+    reasoningEfforts: (endpoint.reasoningEfforts ?? []).map((effort) => {
+      const native = REASONING_EFFORT_TO_NATIVE[effort];
+      if (native === undefined) throw new Error(`Unknown reasoning effort "${effort}"`);
+      return native;
+    }),
+    status: { deprecated: endpoint.status?.deprecated ?? false },
+    dimension: endpoint.dimension,
+  };
+}
+
+function fromNativeEndpoint(endpoint: NativeServingEndpointSummary): ServingEndpointSummary {
+  return {
+    name: endpoint.name,
+    ...(endpoint.displayName !== undefined ? { displayName: endpoint.displayName } : {}),
+    ...(endpoint.family !== undefined ? { family: endpoint.family } : {}),
+    ...(endpoint.task !== undefined ? { task: endpoint.task } : {}),
+    ...(endpoint.state !== undefined ? { state: endpoint.state } : {}),
+    ...(endpoint.description !== undefined ? { description: endpoint.description } : {}),
+    ...(endpoint.supportsTools !== undefined ? { supportsTools: endpoint.supportsTools } : {}),
+    ...(endpoint.profile !== undefined ? { profile: endpoint.profile } : {}),
+    ...(endpoint.modelClass !== undefined ? { class: fromNativeClass(endpoint.modelClass) } : {}),
+    ...(endpoint.serviceNames.size > 0
+      ? { serviceNames: Object.fromEntries(endpoint.serviceNames) }
+      : {}),
+    ...(endpoint.modelServiceName !== undefined
+      ? { modelServiceName: endpoint.modelServiceName }
+      : {}),
+    ...(endpoint.reasoningEfforts.length > 0
+      ? {
+          reasoningEfforts: endpoint.reasoningEfforts.map(
+            (effort) => REASONING_EFFORT_FROM_NATIVE[effort],
+          ),
+        }
+      : {}),
+    status: { deprecated: endpoint.status.deprecated },
+    ...(endpoint.dimension !== undefined ? { dimension: endpoint.dimension } : {}),
+  };
+}
+
+function fromNativeClass(modelClass: NativeModelClass): ModelClass {
+  switch (modelClass) {
+    case NativeModelClass.ChatThinking:
+      return model.ModelClass.ChatThinking;
+    case NativeModelClass.ChatBalanced:
+      return model.ModelClass.ChatBalanced;
+    case NativeModelClass.ChatFast:
+      return model.ModelClass.ChatFast;
+    case NativeModelClass.Embedding:
+      return model.ModelClass.Embedding;
+  }
+}

package/src/fallback.ts

@@ -5,27 +5,27 @@
  * 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 chat {@link ModelClass} by the shared
- * {@link classify.classifyByFamily} heuristic (no classes are hard-coded here)
+ * each bucketed into a chat {@link ModelClass} by the Rust family policy
  * 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 nothing
  * to fall back to and must not assume a stale, baked-in model list; it consumes
- * the live `/models` response instead. The pure classifier
- * ({@link classify.classifyEndpoints}) is what the client shares.
+ * the classified live `/models` response instead.
  *
  * @module
  */
 
-import { classify, model, type FamilyClass } from "@dbx-tools/shared-model";
+import { model, type ServingEndpointSummary } from "@dbx-tools/shared-model";
+
+import { rankEndpoints } from "./_native.ts";
 
 type ModelClass = model.ModelClass;
 const { ModelClass } = model;
 
 /**
  * Small, last-resort set of well-known Foundation Model API endpoint names,
- * ordered best-first within each class by {@link classify.classifyByFamily}.
+ * ordered best-first within each class by the Rust model ranker.
  * 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. Classes are not hard-coded here
@@ -43,19 +43,23 @@ const FALLBACK_MODEL_NAMES: readonly string[] = [
   "databricks-meta-llama-3-1-8b-instruct",
 ];
 
+const FALLBACK_ENDPOINTS: readonly ServingEndpointSummary[] = FALLBACK_MODEL_NAMES.map((name) => ({
+  name,
+  task: "llm/v1/chat",
+}));
+
 /**
  * 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
- * {@link classify.classifyEndpoints}. Returns `[]` for
+ * the Rust-classified workspace catalogue. Returns `[]` for
  * {@link ModelClass.Embedding} (the floor is chat-only).
  */
 export function modelsForClass(cls: ModelClass): readonly string[] {
-  return FALLBACK_MODEL_NAMES.map((name) => ({ name, c: classify.classifyByFamily(name) }))
-    .filter((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);
+  return rankEndpoints(FALLBACK_ENDPOINTS, { modelClass: cls })
+    .filter((ranked) => ranked.modelClass === cls)
+    .map((ranked) => ranked.endpoint.name);
 }
 
 /** Top static fallback model id for a chat class. */

package/src/{serving.ts → model-catalog.ts}

@@ -1,15 +1,15 @@
 /**
  * Live Databricks Model Serving catalogue access.
  *
- * Lists the workspace's `/serving-endpoints` once per host and caches the
- * result with a TTL via AppKit's `CacheManager`, with concurrent callers
- * sharing one in-flight promise (the coalescing pattern of Python's
- * `cachetools-async`). Surfaces each endpoint as a stable
+ * Lists the workspace's `/serving-endpoints` and, when the caller supplies a
+ * trusted credential identity, caches the result per host and identity with a
+ * TTL via AppKit's `CacheManager`. Concurrent callers in the same scope share
+ * one in-flight promise. Surfaces each endpoint as a stable
  * {@link ServingEndpointSummary} - including the Foundation Model API
  * `quality` / `speed` / `cost` profile when present, the classified
  * {@link ModelClass}, and (for embedding endpoints) the measured vector
  * `dimension` - and snaps loose, human-typed names to real endpoint ids
- * through `fuse.js` extended search so tokens like `"claude sonnet"` resolve to
+ * through the Rust model ranker so tokens like `"claude sonnet"` resolve to
  * `databricks-claude-sonnet-4-6`.
  *
  * The class stamp and embedding dimension are computed once per cache load:
@@ -21,20 +21,13 @@
  * @module
  */
 
+import { createHash } from "node:crypto";
 import { CacheManager } from "@databricks/appkit";
 import { appkit } from "@dbx-tools/appkit";
-import { error, log, string } from "@dbx-tools/shared-core";
-import {
-  classify,
-  display,
-  model,
-  type ModelProfile,
-  type ServingEndpointSummary,
-} from "@dbx-tools/shared-model";
-import Fuse from "fuse.js";
+import { errorUtils, log } from "@dbx-tools/shared-core";
+import { model, type ServingEndpointSummary } from "@dbx-tools/shared-model";
 
-import { MODEL_CLASS_ORDER } from "./classes.ts";
-import { modelFamily, modelReasoningEfforts } from "./policy.ts";
+import { classifyEndpointClasses, normalizeEndpoints, rankEndpoints } from "./_native.ts";
 
 const { ModelClass } = model;
 
@@ -50,22 +43,12 @@ export type WorkspaceClientLike = Pick<appkit.WorkspaceClientLike, "servingEndpo
 /** Default TTL for the in-memory endpoint cache. Matches the Databricks SDK's session lifetime budget. */
 export const DEFAULT_MODEL_CACHE_TTL_MS = 5 * 60 * 1000;
 
-/** Default Fuse.js score threshold below which a fuzzy match is accepted. */
+/** Default Rust distance threshold below which a fuzzy match is accepted. */
 export const DEFAULT_FUZZY_THRESHOLD = 0.4;
 
 /** Cache key parts under which endpoint listings are stored. */
 const CACHE_KEY_NAMESPACE = "serving-endpoints";
 
-/**
- * Stable `userKey` arg for AppKit's `CacheManager.getOrExecute`. Endpoint
- * visibility is effectively workspace-scoped (we cache by host in the key
- * parts), so a single shared key lets every user of the same workspace share
- * one cached fetch and coalesce on the in-flight promise. Permissions can
- * differ in theory, but the Foundation Model API catalogue is the same view for
- * every caller.
- */
-const SHARED_USER_KEY = "model-shared";
-
 /** Options for {@link listServingEndpoints}. */
 export interface ListServingEndpointsOptions {
   /**
@@ -73,6 +56,20 @@ export interface ListServingEndpointsOptions {
    * `CacheManager` as seconds.
    */
   ttlMs?: number;
+  /**
+   * Trusted identity of the credential used by `client`.
+   *
+   * The value is hashed before it enters the cache key and is never logged.
+   * Omit it when the caller cannot prove the credential identity; that call
+   * bypasses shared caching so one principal can never read another's
+   * catalogue.
+   */
+  cacheIdentity?: string;
+}
+
+/** Hash an opaque credential identity before handing it to CacheManager. */
+function cacheUserKey(identity: string): string {
+  return `model-${createHash("sha256").update(identity).digest("hex")}`;
 }
 
 /**
@@ -100,11 +97,12 @@ export async function listServingEndpoints(
   host: string,
   options: ListServingEndpointsOptions = {},
 ): Promise<ServingEndpointSummary[]> {
+  if (!options.cacheIdentity) return fetchEndpoints(client);
   const ttlSec = Math.max(1, Math.round((options.ttlMs ?? DEFAULT_MODEL_CACHE_TTL_MS) / 1000));
   return CacheManager.getInstanceSync().getOrExecute(
     [CACHE_KEY_NAMESPACE, host],
     () => fetchEndpoints(client),
-    SHARED_USER_KEY,
+    cacheUserKey(options.cacheIdentity),
     { ttl: ttlSec },
   );
 }
@@ -122,51 +120,11 @@ export async function listServingEndpoints(
 export async function listServingEndpointsUncached(
   client: WorkspaceClientLike,
 ): Promise<ServingEndpointSummary[]> {
-  const out: ServingEndpointSummary[] = [];
+  const endpoints: unknown[] = [];
   for await (const ep of client.servingEndpoints.list()) {
-    if (!ep.name) continue;
-    const profile = extractProfile(ep);
-    const family = modelFamily(ep.name);
-    const reasoningEfforts = modelReasoningEfforts(ep.name);
-    out.push({
-      name: ep.name,
-      // Prefer a Databricks-provided human name (a display-name tag or an
-      // external-model name); else derive a title-cased label from the id.
-      displayName: display.toModelDisplayName(ep.name, providedDisplayName(ep)),
-      ...(family !== undefined ? { family } : {}),
-      ...(ep.task !== undefined ? { task: ep.task } : {}),
-      ...(ep.state?.ready !== undefined ? { state: String(ep.state.ready) } : {}),
-      ...(ep.description !== undefined ? { description: ep.description } : {}),
-      supportsTools: classify.supportsToolsByFamily(ep.name),
-      ...(profile ? { profile } : {}),
-      ...(reasoningEfforts.length > 0 ? { reasoningEfforts } : {}),
-    });
-  }
-  return out;
-}
-
-/**
- * Pull a Databricks-provided human name off a serving endpoint, if any:
- * a `display_name` / `displayName` / `name` tag first, then an
- * external-model `name` from the served-entity config. Returns `null`
- * when the endpoint carries no explicit name (the common case for
- * Foundation Model API endpoints), so the caller falls back to deriving
- * one from the id.
- */
-function providedDisplayName(ep: {
-  tags?: Array<{ key: string; value?: string }>;
-  config?: { served_entities?: Array<{ external_model?: { name?: string } }> };
-}): string | null {
-  const tag = ep.tags?.find(
-    (t) => t.key === "display_name" || t.key === "displayName" || t.key === "name",
-  );
-  const fromTag = string.trimToNull(tag?.value);
-  if (fromTag) return fromTag;
-  for (const entity of ep.config?.served_entities ?? []) {
-    const external = string.trimToNull(entity.external_model?.name);
-    if (external) return external;
+    endpoints.push(ep);
   }
-  return null;
+  return normalizeEndpoints(endpoints);
 }
 
 async function fetchEndpoints(client: WorkspaceClientLike): Promise<ServingEndpointSummary[]> {
@@ -185,11 +143,7 @@ async function fetchEndpoints(client: WorkspaceClientLike): Promise<ServingEndpo
  * class.
  */
 function stampModelClasses(summaries: ServingEndpointSummary[]): void {
-  const buckets = classify.classifyEndpoints(summaries);
-  const classOf = new Map<string, model.ModelClass>();
-  for (const cls of MODEL_CLASS_ORDER) {
-    for (const ep of buckets[cls]) classOf.set(ep.name, cls);
-  }
+  const classOf = classifyEndpointClasses(summaries);
   for (const summary of summaries) {
     const cls = classOf.get(summary.name);
     if (cls !== undefined) summary.class = cls;
@@ -235,61 +189,24 @@ async function pingEmbeddingDimension(
     logger.warn("embedding ping returned no vector", { name });
     return undefined;
   } catch (err) {
-    logger.warn("embedding ping failed", { name, error: error.errorMessage(err) });
+    logger.warn("embedding ping failed", { name, error: errorUtils.errorMessage(err) });
     return undefined;
   }
 }
 
 /**
- * Pull the Foundation Model API `quality` / `speed` / `cost` scores off a
- * serving-endpoint listing entry. Databricks returns these under
- * `config.served_entities[].foundation_model.ai_gateway_model_profile`
- * (snake_case at the wire level, preserved verbatim by the SDK), but the field
- * is newer than the typed `FoundationModel` interface, so we read it through a
- * structural cast. Returns `undefined` when no served entity carries a profile
- * (custom models, embeddings, and brand-new endpoints that Databricks has not
- * scored yet).
+ * Force-evict cached endpoint listings via AppKit's `CacheManager`. A host and
+ * identity delete only that principal's workspace entry. Because CacheManager
+ * cannot enumerate a namespace, omitting either value clears the complete
+ * manager and should be reserved for tests or administrative refreshes.
  */
-function extractProfile(ep: unknown): ModelProfile | undefined {
-  const entities = (
-    ep as {
-      config?: {
-        served_entities?: Array<{
-          foundation_model?: {
-            ai_gateway_model_profile?: {
-              quality?: number;
-              speed?: number;
-              cost?: number;
-            };
-          };
-        }>;
-      };
-    }
-  ).config?.served_entities;
-  if (!entities) return undefined;
-  for (const entity of entities) {
-    const raw = entity.foundation_model?.ai_gateway_model_profile;
-    if (!raw) continue;
-    const profile: ModelProfile = {};
-    if (Number.isFinite(raw.quality)) profile.quality = raw.quality;
-    if (Number.isFinite(raw.speed)) profile.speed = raw.speed;
-    if (Number.isFinite(raw.cost)) profile.cost = raw.cost;
-    if (Object.keys(profile).length > 0) return profile;
-  }
-  return undefined;
-}
-
-/**
- * Force-evict cached endpoint listings via AppKit's `CacheManager`. With a
- * `host` deletes that one workspace's entry; without one clears every cache
- * entry on the manager (since `CacheManager` doesn't expose a namespace-scoped
- * clear, this is the brute-force path - fine for tests, avoid in steady-state
- * code).
- */
-export async function clearServingEndpointsCache(host?: string): Promise<void> {
+export async function clearServingEndpointsCache(
+  host?: string,
+  cacheIdentity?: string,
+): Promise<void> {
   const cache = CacheManager.getInstanceSync();
-  if (host) {
-    const key = cache.generateKey([CACHE_KEY_NAMESPACE, host], SHARED_USER_KEY);
+  if (host && cacheIdentity) {
+    const key = cache.generateKey([CACHE_KEY_NAMESPACE, host], cacheUserKey(cacheIdentity));
     await cache.delete(key);
   } else {
     await cache.clear();
@@ -298,7 +215,7 @@ export async function clearServingEndpointsCache(host?: string): Promise<void> {
 
 /**
  * Result of fuzzy-resolving a user-supplied model name against the live
- * endpoint list. `score` is Fuse.js's distance (`0` is exact, `1` is no match);
+ * endpoint list. `score` is the Rust ranker's distance (`0` is exact, `1` is no match);
  * `matched` is `false` when the score exceeds the configured threshold so
  * callers can fall back to the original input (Databricks will then return a
  * clean 404).
@@ -311,7 +228,7 @@ export interface ResolvedModel {
 
 /** Options accepted by {@link resolveModelId} / {@link searchServingEndpoints}. */
 export interface ResolveModelOptions {
-  /** Fuse.js threshold (0 = exact, 1 = anything). Default `0.4`. */
+  /** Fuzzy distance threshold (0 = exact, 1 = anything). Default `0.4`. */
   threshold?: number;
   /** Require a model verified for a complete function-tool round-trip. */
   requiresTools?: boolean;
@@ -320,7 +237,7 @@ export interface ResolveModelOptions {
 /** A serving endpoint paired with its fuzzy-match distance for a query. */
 export interface ScoredEndpoint {
   endpoint: ServingEndpointSummary;
-  /** Fuse.js distance: `0` is exact, `1` is no match. */
+  /** Fuzzy distance: `0` is exact, `1` is no match. */
   score: number;
 }
 
@@ -329,11 +246,8 @@ export interface ScoredEndpoint {
  * (lowest score) first, keeping only those within `threshold`:
  *
  * 1. An exact name match short-circuits to a single `score: 0` result.
- * 2. Otherwise the input is tokenized (dashes / underscores / spaces become
- *    separators) and fed through Fuse.js extended search, which AND-s each
- *    token with fuzzy matching enabled - the "tokenized fuzzy match" a caller
- *    reaches for when they type `"claude sonnet"` instead of the full endpoint
- *    name.
+ * 2. Otherwise Rust tokenizes the input and scores every endpoint through the
+ *    same fuzzy-distance policy used by {@link lookupModels}.
  *
  * Returns `[]` for an empty endpoint list or when `input` tokenizes to nothing,
  * so callers fall back to the raw input and let Databricks surface a clean 404.
@@ -345,30 +259,20 @@ export function searchServingEndpoints(
   endpoints: readonly ServingEndpointSummary[],
   options: ResolveModelOptions = {},
 ): ScoredEndpoint[] {
-  if (endpoints.length === 0) return [];
-  for (const ep of endpoints) {
-    if (ep.name === input) return [{ endpoint: ep, score: 0 }];
-  }
-  const threshold = options.threshold ?? DEFAULT_FUZZY_THRESHOLD;
-  // Fuse 7.3 has no built-in tokenize hook; in extended search, space-separated
-  // tokens are AND-ed with fuzzy matching enabled. We lean on the shared
-  // tokenizer so the splitting rules stay consistent with the rest of the toolkit.
-  const query = Array.from(
-    string.tokenizeWithOptions({ lowerCase: true, camelCase: false }, input),
-  ).join(" ");
-  if (!query) return [];
-  const fuse = new Fuse(endpoints, {
-    keys: ["name"],
-    threshold,
-    ignoreLocation: true,
-    includeScore: true,
-    useExtendedSearch: true,
-    isCaseSensitive: false,
-  });
-  return fuse
-    .search(query)
-    .filter((r) => (r.score ?? 0) <= threshold)
-    .map((r) => ({ endpoint: r.item, score: r.score ?? 0 }));
+  if (endpoints.length === 0 || !input.trim()) return [];
+  return rankEndpoints(
+    endpoints,
+    {
+      search: input,
+      ...(options.requiresTools !== undefined ? { requiresTools: options.requiresTools } : {}),
+      threshold: options.threshold ?? DEFAULT_FUZZY_THRESHOLD,
+    },
+    {
+      includeDeprecated: true,
+      modelClass: model.ModelClass.ChatBalanced,
+      task: "llm/v1/chat",
+    },
+  ).map((ranked) => ({ endpoint: ranked.endpoint, score: ranked.score ?? 0 }));
 }
 
 /**

package/src/policy.ts

@@ -7,8 +7,14 @@
 import {
   modelFamily as modelFamilyWithRust,
   reasoningEffortNamesByFamily,
+  supportsToolsByFamily as supportsToolsByFamilyWithRust,
 } from "@dbx-tools/model-rs";
-import type { ReasoningEffort } from "@dbx-tools/shared-model";
+import {
+  model,
+  ReasoningEffortSchema,
+  type ReasoningEffort,
+  type ServingEndpointSummary,
+} from "@dbx-tools/shared-model";
 
 /** Return the normalized family parsed by the Rust model-name policy. */
 export function modelFamily(name: string): string | undefined {
@@ -17,5 +23,18 @@ export function modelFamily(name: string): string | undefined {
 
 /** Return the reasoning efforts accepted by a model family. */
 export function modelReasoningEfforts(name: string): ReasoningEffort[] {
-  return reasoningEffortNamesByFamily(name);
+  return ReasoningEffortSchema.array().parse(reasoningEffortNamesByFamily(name));
+}
+
+/** Return whether Rust policy verifies a complete tool-calling round trip. */
+export function modelSupportsTools(name: string): boolean {
+  return supportsToolsByFamilyWithRust(name);
+}
+
+/** Return whether a discovered endpoint can be used for a tool-calling chat. */
+export function endpointSupportsTools(endpoint: ServingEndpointSummary): boolean {
+  const embedding =
+    endpoint.task === "llm/v1/embeddings" || endpoint.class === model.ModelClass.Embedding;
+  const chat = !embedding && (endpoint.task === "llm/v1/chat" || endpoint.class !== undefined);
+  return chat && (endpoint.supportsTools ?? modelSupportsTools(endpoint.name));
 }