@dbx-tools/model 0.9.21 → 0.9.29

package/src/model-client.ts

@@ -0,0 +1,311 @@
+/**
+ * Authentication-aware Databricks Model Serving client.
+ *
+ * This is the portable entry point used by generated Python bindings. It keeps
+ * authentication, endpoint discovery, caching, normalization, ranking, route
+ * selection, and published metadata in Node while exposing one small client.
+ *
+ * @module
+ */
+import {
+  DatabricksClient,
+  type AuthClient,
+  type DatabricksAuthDependencies,
+  type DatabricksAuthOptions,
+  type DatabricksProfileSummary,
+} from "@dbx-tools/auth";
+import * as log from "@dbx-tools/shared-core/log";
+import * as object from "@dbx-tools/shared-core/object";
+import {
+  ModelClass,
+  type ModelQuery,
+  type RankedModel,
+  type ServingEndpointSummary,
+} from "@dbx-tools/shared-model/contracts";
+
+import { normalizeEndpoints } from "./_ranking.ts";
+import {
+  lookupModels,
+  type ResolveModelInput,
+  type ResolvedModelSelection,
+  resolveModel,
+} from "./_selection.ts";
+import { chatCompletionsUrl, invocationsUrl, responsesUpstreamUrl } from "./invoke.ts";
+import { type ModelMetadata, modelMetadataFor } from "./metadata.ts";
+import { modelServingApi } from "./policy.ts";
+
+const logger = log.logger("model/client");
+
+/** Default lifetime for one live endpoint catalogue. */
+export const DEFAULT_MODEL_CLIENT_CACHE_TTL_MS = 5 * 60 * 1000;
+
+/** Construction options for {@link createModelClient}. */
+export interface ModelClientOptions {
+  /** Databricks authentication and profile selection options. */
+  auth?: DatabricksAuthOptions;
+  /** In-memory endpoint catalogue TTL in milliseconds. */
+  cacheTtlMs?: number;
+}
+
+/** Secret-free runtime identity and cache configuration. */
+export interface ModelClientStatus {
+  readonly profile: string;
+  readonly host: string;
+  readonly principal: string;
+  readonly workspaceId?: string;
+  readonly cacheTtlMs: number;
+}
+
+/** Serving protocol selected for one resolved model. */
+export type ModelProtocol = "chat" | "responses" | "embeddings";
+
+/** Model selection plus route controls. */
+export interface ModelRouteInput extends ResolveModelInput {
+  /** Force a live endpoint catalogue refresh before selection. */
+  refresh?: boolean;
+  /** Requested upstream protocol. Defaults from endpoint task and model policy. */
+  protocol?: ModelProtocol;
+  /** Whether authentication may run an interactive login. */
+  login?: boolean;
+}
+
+/** Fully resolved Databricks route for a model request. */
+export interface ModelRoute {
+  readonly modelId: string;
+  readonly endpoint?: ServingEndpointSummary;
+  readonly source: ResolvedModelSelection["source"];
+  readonly protocol: ModelProtocol;
+  readonly host: string;
+  readonly apiBase: string;
+  readonly url: string;
+  readonly headers: Readonly<Record<string, string>>;
+  readonly metadata: ModelMetadata;
+}
+
+/** Stable model facade shared by Node and generated language bindings. */
+export interface ModelClient {
+  listModels(refresh?: boolean): Promise<ServingEndpointSummary[]>;
+  searchModels(query?: ModelQuery, refresh?: boolean): Promise<RankedModel[]>;
+  resolveModel(input?: ResolveModelInput, refresh?: boolean): Promise<ResolvedModelSelection>;
+  route(input?: ModelRouteInput): Promise<ModelRoute>;
+  authenticate(login?: boolean): Promise<Record<string, string>>;
+  metadata(model: string | ServingEndpointSummary): ModelMetadata;
+  status(): ModelClientStatus;
+  profile(name?: string): DatabricksProfileSummary;
+  listProfiles(refresh?: boolean): DatabricksProfileSummary[];
+}
+
+interface DatabricksModelClient {
+  readonly auth: AuthClient;
+  host(): string;
+  principal(): string;
+  workspaceId(): string | undefined;
+  request(path: string): Promise<unknown>;
+}
+
+interface CatalogueEntry {
+  readonly promise: Promise<readonly ServingEndpointSummary[]>;
+  expiresAt: number;
+}
+
+const catalogueCache = new Map<string, CatalogueEntry>();
+
+class DefaultModelClient implements ModelClient {
+  constructor(
+    private readonly client: DatabricksModelClient,
+    private readonly cacheTtlMs: number,
+  ) {}
+
+  async listModels(refresh = false): Promise<ServingEndpointSummary[]> {
+    return cloneEndpoints(await this.catalogue(refresh));
+  }
+
+  async searchModels(query: ModelQuery = {}, refresh = false): Promise<RankedModel[]> {
+    return lookupModels(await this.catalogue(refresh), query).map(({ endpoint, ...ranked }) => ({
+      ...ranked,
+      endpoint: cloneEndpoint(endpoint),
+    }));
+  }
+
+  async resolveModel(
+    input: ResolveModelInput = {},
+    refresh = false,
+  ): Promise<ResolvedModelSelection> {
+    const catalogue = await this.catalogue(refresh);
+    const resolved = resolveModel(catalogue, input);
+    if (isUnmatchedExplicit(catalogue, input, resolved)) {
+      return resolveModel(await this.catalogue(true), input);
+    }
+    return resolved;
+  }
+
+  async route(input: ModelRouteInput = {}): Promise<ModelRoute> {
+    const { refresh = false, protocol: requestedProtocol, login, ...selection } = input;
+    const effectiveSelection =
+      requestedProtocol === "embeddings" && selection.modelClass === undefined
+        ? { ...selection, modelClass: ModelClass.Embedding }
+        : selection;
+    const catalogue = await this.catalogue(refresh);
+    let resolved = resolveModel(catalogue, effectiveSelection);
+    if (isUnmatchedExplicit(catalogue, effectiveSelection, resolved)) {
+      resolved = resolveModel(await this.catalogue(true), effectiveSelection);
+    }
+    const endpoint = (await this.catalogue(false)).find(
+      (candidate) => candidate.name === resolved.modelId,
+    );
+    const protocol = requestedProtocol ?? inferProtocol(endpoint, resolved.modelId);
+    const host = this.client.host();
+    const url = routeUrl(host, resolved.modelId, protocol);
+    const headers = await this.client.auth.requestHeadersForUrl(url, login);
+    return {
+      modelId: resolved.modelId,
+      ...(endpoint ? { endpoint: cloneEndpoint(endpoint) } : {}),
+      source: resolved.source,
+      protocol,
+      host,
+      apiBase: new URL("serving-endpoints", `${host.replace(/\/$/, "")}/`).toString(),
+      url,
+      headers,
+      metadata: modelMetadataFor(endpoint ?? resolved.modelId),
+    };
+  }
+
+  authenticate(login?: boolean): Promise<Record<string, string>> {
+    return this.client.auth.authenticate(login);
+  }
+
+  metadata(model: string | ServingEndpointSummary): ModelMetadata {
+    return modelMetadataFor(model);
+  }
+
+  status(): ModelClientStatus {
+    return {
+      profile: this.client.auth.status().profile,
+      host: this.client.host(),
+      principal: this.client.principal(),
+      ...(this.client.workspaceId() ? { workspaceId: this.client.workspaceId() } : {}),
+      cacheTtlMs: this.cacheTtlMs,
+    };
+  }
+
+  profile(name?: string): DatabricksProfileSummary {
+    return this.client.auth.profile(name);
+  }
+
+  listProfiles(refresh = false): DatabricksProfileSummary[] {
+    return this.client.auth.listProfiles(refresh);
+  }
+
+  private catalogue(refresh: boolean): Promise<readonly ServingEndpointSummary[]> {
+    const key = catalogueKey(this.client);
+    if (refresh) catalogueCache.delete(key);
+    const current = catalogueCache.get(key);
+    if (current && Date.now() < current.expiresAt) return current.promise;
+    const promise = this.fetchCatalogue();
+    const entry: CatalogueEntry = { promise, expiresAt: Number.POSITIVE_INFINITY };
+    catalogueCache.set(key, entry);
+    void promise.then(
+      () => {
+        entry.expiresAt = Date.now() + this.cacheTtlMs;
+      },
+      () => {
+        if (catalogueCache.get(key) === entry) catalogueCache.delete(key);
+      },
+    );
+    return promise;
+  }
+
+  private async fetchCatalogue(): Promise<readonly ServingEndpointSummary[]> {
+    const startedAt = Date.now();
+    const response = await this.client.request("/api/2.0/serving-endpoints");
+    if (!object.isRecord(response) || !Array.isArray(response.endpoints)) {
+      throw new Error("Databricks serving-endpoints response is missing an endpoints array");
+    }
+    const endpoints = normalizeEndpoints(response.endpoints);
+    logger.debug("listed", {
+      count: endpoints.length,
+      host: this.client.host(),
+      elapsedMs: Date.now() - startedAt,
+    });
+    return endpoints;
+  }
+}
+
+/** Create an authentication-aware Databricks model client. */
+export async function createModelClient(options: ModelClientOptions = {}): Promise<ModelClient> {
+  const cacheTtlMs = validateCacheTtl(options.cacheTtlMs);
+  return new DefaultModelClient(await DatabricksClient.create(options.auth), cacheTtlMs);
+}
+
+/** @internal Construct a model client around a test or host-owned Databricks client. */
+export function createModelClientWithDatabricksClient(
+  client: DatabricksModelClient,
+  cacheTtlMs = DEFAULT_MODEL_CLIENT_CACHE_TTL_MS,
+): ModelClient {
+  return new DefaultModelClient(client, validateCacheTtl(cacheTtlMs));
+}
+
+/** @internal Create a model client with injectable auth host capabilities. */
+export async function createModelClientWithDependencies(
+  options: ModelClientOptions,
+  dependencies: DatabricksAuthDependencies,
+): Promise<ModelClient> {
+  return new DefaultModelClient(
+    await DatabricksClient.create(options.auth, dependencies),
+    validateCacheTtl(options.cacheTtlMs),
+  );
+}
+
+function validateCacheTtl(value = DEFAULT_MODEL_CLIENT_CACHE_TTL_MS): number {
+  if (!Number.isFinite(value) || value <= 0) {
+    throw new Error("Model client cacheTtlMs must be a positive finite number");
+  }
+  return value;
+}
+
+function catalogueKey(client: DatabricksModelClient): string {
+  return [client.host(), client.workspaceId() ?? "", client.principal()].join("\u0000");
+}
+
+function inferProtocol(
+  endpoint: ServingEndpointSummary | undefined,
+  modelId: string,
+): ModelProtocol {
+  if (endpoint?.task === "llm/v1/embeddings") return "embeddings";
+  return modelServingApi(modelId);
+}
+
+function routeUrl(host: string, modelId: string, protocol: ModelProtocol): string {
+  if (protocol === "responses") return responsesUpstreamUrl(host, modelId);
+  if (protocol === "chat" && modelServingApi(modelId) === "responses") {
+    return chatCompletionsUrl(host);
+  }
+  return invocationsUrl(host, modelId);
+}
+
+function isUnmatchedExplicit(
+  catalogue: readonly ServingEndpointSummary[],
+  input: ResolveModelInput,
+  resolved: ResolvedModelSelection,
+): boolean {
+  return (
+    input.explicit !== undefined &&
+    input.fuzzy !== false &&
+    resolved.modelId === input.explicit &&
+    !catalogue.some((endpoint) => endpoint.name === input.explicit)
+  );
+}
+
+function cloneEndpoints(endpoints: readonly ServingEndpointSummary[]): ServingEndpointSummary[] {
+  return endpoints.map(cloneEndpoint);
+}
+
+function cloneEndpoint(endpoint: ServingEndpointSummary): ServingEndpointSummary {
+  return {
+    ...endpoint,
+    ...(endpoint.profile ? { profile: { ...endpoint.profile } } : {}),
+    ...(endpoint.serviceNames ? { serviceNames: { ...endpoint.serviceNames } } : {}),
+    ...(endpoint.reasoningEfforts ? { reasoningEfforts: [...endpoint.reasoningEfforts] } : {}),
+    ...(endpoint.status ? { status: { ...endpoint.status } } : {}),
+  };
+}

package/src/policy.ts

@@ -1,40 +1,197 @@
-/**
- * Browser-safe model metadata projected from the Rust policy bindings.
- *
- * @module
- */
+/** TypeScript-owned model-family and serving capability policy. */
 
 import {
-  modelFamily as modelFamilyWithRust,
-  reasoningEffortNamesByFamily,
-  supportsToolsByFamily as supportsToolsByFamilyWithRust,
-} from "@dbx-tools/model-rs";
-import {
-  model,
-  ReasoningEffortSchema,
+  ModelClass,
   type ReasoningEffort,
   type ServingEndpointSummary,
-} from "@dbx-tools/shared-model";
+} from "@dbx-tools/shared-model/contracts";
+
+import { supportsToolsByFamily } from "./classify.ts";
+
+const FAMILIES = [
+  "bge",
+  "claude",
+  "deepseek",
+  "gemini",
+  "gemma",
+  "glm",
+  "gpt",
+  "grok",
+  "gte",
+  "inkling",
+  "kimi",
+  "llama",
+  "qwen",
+] as const;
 
-/** Return the normalized family parsed by the Rust model-name policy. */
+const STANDARD: ReasoningEffort[] = ["low", "medium", "high"];
+const GPT_5_6: ReasoningEffort[] = ["none", "low", "medium", "high", "xhigh", "max"];
+const GPT_5_5_PRO: ReasoningEffort[] = ["medium", "high", "xhigh"];
+const CLAUDE: ReasoningEffort[] = ["none", "minimal", "low", "medium", "high", "xhigh", "max"];
+const GEMINI: ReasoningEffort[] = ["minimal", "low", "medium", "high"];
+
+interface ParsedModelName {
+  family: (typeof FAMILIES)[number];
+  version: number[];
+  model: string[];
+}
+
+/** Return the normalized family parsed from a model identity. */
 export function modelFamily(name: string): string | undefined {
-  return modelFamilyWithRust(name);
+  return parseModelName(name)?.family;
+}
+
+/** Normalize a recognized model identity into family, version, and variant tokens. */
+export function modelSearchQuery(name: string): string | undefined {
+  const parsed = parseModelName(name);
+  if (!parsed) return undefined;
+  return [parsed.family, ...parsed.version.map(String), ...parsed.model].join(" ");
 }
 
 /** Return the reasoning efforts accepted by a model family. */
 export function modelReasoningEfforts(name: string): ReasoningEffort[] {
-  return ReasoningEffortSchema.array().parse(reasoningEffortNamesByFamily(name));
+  const normalized = name.toLowerCase();
+  const parsed = parseModelName(name);
+  let efforts: ReasoningEffort[] = [];
+  if (parsed?.family === "gpt" && parsed.version.length > 0) {
+    const [major = 0, minor = 0] = parsed.version;
+    if (major >= 5) {
+      efforts =
+        major === 5 && minor === 5 && parsed.model.includes("pro")
+          ? GPT_5_5_PRO
+          : major === 5 && minor === 6
+            ? GPT_5_6
+            : STANDARD;
+    }
+  }
+  if (parsed?.family === "gpt" && parsed.model.includes("oss")) efforts = STANDARD;
+  if (parsed?.family === "claude") {
+    const [major = 0, minor = 0] = parsed.version;
+    if (major > 3 || (major === 3 && minor >= 7)) efforts = CLAUDE;
+  }
+  if (parsed?.family === "gemini") efforts = GEMINI;
+  if (normalized.includes("codex") || /(?:^|[-_./])o(?:1|3|4)(?:[-_./]|$)/i.test(name)) {
+    efforts = STANDARD;
+  }
+  return efforts;
 }
 
-/** Return whether Rust policy verifies a complete tool-calling round trip. */
+/** Return whether policy verifies a complete tool-calling round trip. */
 export function modelSupportsTools(name: string): boolean {
-  return supportsToolsByFamilyWithRust(name);
+  return supportsToolsByFamily(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;
+    endpoint.task === "llm/v1/embeddings" || endpoint.class === ModelClass.Embedding;
   const chat = !embedding && (endpoint.task === "llm/v1/chat" || endpoint.class !== undefined);
   return chat && (endpoint.supportsTools ?? modelSupportsTools(endpoint.name));
 }
+
+/** Return whether a model requires Databricks' native Responses endpoint. */
+export function isResponsesOnly(name: string): boolean {
+  if (name.toLowerCase().includes("codex")) return true;
+  const parsed = parseModelName(name);
+  if (!parsed || parsed.family !== "gpt" || parsed.model.includes("oss")) return false;
+  const [major = 0, minor = 0] = parsed.version;
+  return major > 5 || (major === 5 && minor >= 4);
+}
+
+/** Return the Databricks inference protocol required by a model. */
+export function modelServingApi(name: string): "chat" | "responses" {
+  return isResponsesOnly(name) ? "responses" : "chat";
+}
+
+/** Return the Chat Completions effort required when function tools are present. */
+export function chatToolReasoningEffort(name: string): ReasoningEffort | undefined {
+  const parsed = parseModelName(name);
+  if (!parsed || parsed.family !== "gpt") return undefined;
+  const [major = 0, minor = 0] = parsed.version;
+  return major === 5 && minor === 6 ? "none" : undefined;
+}
+
+/** Derive provider-specific model names from a Databricks model identity. */
+export function modelServiceNames(name: string): Record<string, string> {
+  const parsed = parseModelName(name);
+  if (!parsed) return {};
+  const version = parsed.version.join(".");
+  const suffix = parsed.model.length ? `-${parsed.model.join("-")}` : "";
+  switch (parsed.family) {
+    case "gpt":
+      return parsed.model.includes("oss") ? {} : { openai: `gpt-${version}${suffix}` };
+    case "claude": {
+      const [variant, ...rest] = parsed.model;
+      return variant
+        ? { anthropic: `claude-${variant}-${version}${rest.length ? `-${rest.join("-")}` : ""}` }
+        : {};
+    }
+    case "gemini":
+    case "gemma":
+      return { google: `${parsed.family}-${version}${suffix}` };
+    case "glm":
+      return { zhipu: `glm-${version}${suffix}` };
+    case "grok":
+      return { xai: `grok-${version}${suffix}` };
+    case "llama":
+      return { meta: `llama-${version}${suffix}` };
+    case "qwen":
+      return { alibaba: `qwen${version}${suffix}` };
+    case "deepseek":
+      return { deepseek: `deepseek-v${parsed.version[0] ?? ""}${suffix}` };
+    case "kimi":
+      return { moonshot: `kimi-k${parsed.version[0] ?? ""}${suffix}` };
+    default:
+      return {};
+  }
+}
+
+function parseModelName(name: string): ParsedModelName | undefined {
+  const tokens = name.toLowerCase().match(/[a-z0-9]+/g) ?? [];
+  let familyIndex = -1;
+  let family: ParsedModelName["family"] | undefined;
+  for (let index = 0; index < tokens.length; index += 1) {
+    const token = tokens[index]!;
+    const compactQwen = /^qwen(\d+)$/.exec(token);
+    if (compactQwen) {
+      familyIndex = index;
+      family = "qwen";
+      break;
+    }
+    if (FAMILIES.includes(token as ParsedModelName["family"])) {
+      familyIndex = index;
+      family = token as ParsedModelName["family"];
+      break;
+    }
+  }
+  if (!family || familyIndex < 0) return undefined;
+  const remainder = tokens.slice(familyIndex + 1);
+  const numeric = (value: string) => /^\d+$/.test(value);
+  let versionStart = remainder.findIndex(numeric);
+  if (family === "qwen" && /^qwen\d+$/.test(tokens[familyIndex]!)) {
+    versionStart = 0;
+    remainder.unshift(tokens[familyIndex]!.slice(4));
+  }
+  if (family === "deepseek") {
+    const prefixed = remainder.findIndex((value) => /^v\d+$/.test(value));
+    if (prefixed >= 0) remainder[prefixed] = remainder[prefixed]!.slice(1);
+    versionStart = remainder.findIndex(numeric);
+  }
+  if (family === "kimi") {
+    const prefixed = remainder.findIndex((value) => /^k\d+$/.test(value));
+    if (prefixed >= 0) remainder[prefixed] = remainder[prefixed]!.slice(1);
+    versionStart = remainder.findIndex(numeric);
+  }
+  if (versionStart < 0) return { family, version: [], model: remainder };
+  const version: number[] = [];
+  let cursor = versionStart;
+  while (cursor < remainder.length && version.length < 2 && numeric(remainder[cursor]!)) {
+    version.push(Number(remainder[cursor]));
+    cursor += 1;
+  }
+  return {
+    family,
+    version,
+    model: [...remainder.slice(0, versionStart), ...remainder.slice(cursor)],
+  };
+}

package/src/python.ts

@@ -0,0 +1,2 @@
+/** Portable entry point used by generated Python bindings. */
+export { createModelClient } from "./model-client.ts";