@dbx-tools/appkit-model-gateway 0.9.36

package/src/registry.ts

@@ -0,0 +1,225 @@
+/**
+ * Request-scoped Databricks model discovery for the AppKit gateway.
+ *
+ * @module
+ */
+
+import { getExecutionContext, type WorkspaceClient } from "@databricks/appkit";
+import { metadata, modelCatalog, policy, resolve as modelResolve } from "@dbx-tools/model";
+import { log } from "@dbx-tools/shared-core";
+import {
+  ModelClass,
+  type ResolveModelInput,
+  type ServingEndpointSummary,
+} from "@dbx-tools/shared-model/contracts";
+import type {
+  ModelCapabilities,
+  ModelCapabilityOverride,
+  ModelTarget,
+} from "@dbx-tools/shared-model-gateway";
+
+const DEFAULT_CACHE_TTL_MS = 60_000;
+const logger = log.logger("appkit/model-gateway/registry");
+
+/** Resolution controls passed through to the model policy owner. */
+export type ModelRegistryResolveOptions = Pick<ResolveModelInput, "modelClass" | "requiresTools">;
+
+/** Dynamic model catalogue used by request routing and model-list responses. */
+export interface ModelRegistry {
+  list(): Promise<ModelTarget[]>;
+  search(query: string): Promise<ModelTarget[]>;
+  resolve(model: string, options?: ModelRegistryResolveOptions): Promise<ModelTarget | undefined>;
+  refresh(): Promise<void>;
+}
+
+/** Configuration for {@link DatabricksModelRegistry}. */
+export interface ModelRegistryOptions {
+  readonly ttlMs?: number;
+  readonly overrides?: readonly ModelCapabilityOverride[];
+}
+
+interface RegistryContext {
+  readonly client: WorkspaceClient;
+  readonly host: string;
+  readonly identity: string;
+}
+
+/** Model registry backed by the active AppKit execution context. */
+export class DatabricksModelRegistry implements ModelRegistry {
+  private readonly ttlMs: number;
+  private readonly overrides: readonly ModelCapabilityOverride[];
+
+  constructor(options: ModelRegistryOptions = {}) {
+    this.ttlMs = positiveTtl(options.ttlMs);
+    this.overrides = options.overrides ?? [];
+  }
+
+  async list(): Promise<ModelTarget[]> {
+    const context = await registryContext();
+    logger.debug("loading catalogue", { host: context.host, ttlMs: this.ttlMs });
+    const endpoints = await modelCatalog.listServingEndpoints(context.client, context.host, {
+      cacheIdentity: context.identity,
+      ttlMs: this.ttlMs,
+    });
+    const targets = endpoints.map((endpoint) => this.target(endpoint));
+    logger.debug("loaded catalogue", {
+      endpointCount: endpoints.length,
+      host: context.host,
+      targetCount: targets.length,
+    });
+    return targets;
+  }
+
+  async resolve(
+    model: string,
+    options: ModelRegistryResolveOptions = {},
+  ): Promise<ModelTarget | undefined> {
+    const requested = unqualifiedModel(model);
+    const targets = await this.list();
+    const exact = targets.find((target) =>
+      target.aliases.some(
+        (alias) => alias.localeCompare(requested, undefined, { sensitivity: "accent" }) === 0,
+      ),
+    );
+    if (exact && matchesResolveOptions(exact, options)) {
+      logger.debug("resolved exact model", { requested, resolved: exact.id });
+      return exact;
+    }
+
+    const endpoints = targets
+      .map((target) => target.endpoint)
+      .filter((endpoint): endpoint is ServingEndpointSummary => endpoint !== undefined);
+    const resolved = modelResolve.resolveModel(endpoints, {
+      explicit: requested,
+      ...options,
+    });
+    const target = targets.find((candidate) => candidate.id === resolved.modelId);
+    logger.debug("resolved fuzzy model", {
+      requested,
+      resolved: target?.id,
+      source: resolved.source,
+    });
+    return target;
+  }
+
+  async search(query: string): Promise<ModelTarget[]> {
+    const targets = await this.list();
+    const byId = new Map(targets.map((target) => [target.id, target]));
+    const endpoints = targets
+      .map((target) => target.endpoint)
+      .filter((endpoint): endpoint is ServingEndpointSummary => endpoint !== undefined);
+    const ranked = [
+      ...modelResolve.lookupModels(endpoints, { search: query }),
+      ...modelResolve.lookupModels(endpoints, {
+        search: query,
+        modelClass: ModelClass.Embedding,
+      }),
+    ].sort((left, right) => (left.score ?? 0) - (right.score ?? 0));
+    return ranked
+      .map(({ endpoint }) => byId.get(endpoint.name))
+      .filter((target): target is ModelTarget => target !== undefined)
+      .filter(
+        (target, index, all) => all.findIndex((candidate) => candidate.id === target.id) === index,
+      );
+  }
+
+  async refresh(): Promise<void> {
+    const context = await registryContext();
+    logger.debug("clearing catalogue", { host: context.host });
+    await modelCatalog.clearServingEndpointsCache(context.host, context.identity);
+    await modelCatalog.listServingEndpoints(context.client, context.host, {
+      cacheIdentity: context.identity,
+      ttlMs: this.ttlMs,
+    });
+    logger.debug("refreshed catalogue", { host: context.host });
+  }
+
+  private target(endpoint: ServingEndpointSummary): ModelTarget {
+    const published = metadata.modelCapabilitiesFor(endpoint);
+    const family = endpoint.family?.toLowerCase();
+    const openAiFamily = family === "gpt" || family === "openai" || endpoint.name.includes("codex");
+    const modelServiceName = normalizedModelServiceName(endpoint.modelServiceName);
+    const embeddings =
+      endpoint.class === ModelClass.Embedding || endpoint.task === "llm/v1/embeddings";
+    const capabilities: ModelCapabilities = {
+      responses: !embeddings && (openAiFamily || published.responses),
+      openResponses: !embeddings && !openAiFamily,
+      chat: !embeddings,
+      anthropic: family === "claude",
+      embeddings,
+      aiGatewayCodex:
+        !embeddings && Boolean(modelServiceName) && family !== "claude" && family !== "gemini",
+      tools: endpoint.supportsTools === true,
+      reasoning: policy.modelReasoningEfforts(endpoint.name).length > 0,
+      streaming: true,
+      parallelTools: endpoint.supportsTools === true,
+      customTools: published.applyPatch || endpoint.name.includes("codex"),
+      structuredOutput: true,
+      webSearch: published.webSearch,
+    };
+    const aliases = new Set(
+      [
+        endpoint.name,
+        endpoint.displayName,
+        endpoint.modelServiceName,
+        modelServiceName,
+        ...Object.values(endpoint.serviceNames ?? {}),
+      ].filter((value): value is string => Boolean(value?.trim())),
+    );
+    const override = this.overrides.find(
+      (candidate) =>
+        candidate.model === endpoint.name ||
+        [...aliases].some((alias) => alias === candidate.model),
+    );
+    return {
+      id: endpoint.name,
+      aliases: [...aliases],
+      displayName: endpoint.displayName ?? endpoint.name,
+      ...(endpoint.family ? { family: endpoint.family } : {}),
+      ...(modelServiceName ? { modelServiceName } : {}),
+      endpoint,
+      capabilities: { ...capabilities, ...override?.capabilities },
+      reasoningEfforts: endpoint.reasoningEfforts ?? [],
+    };
+  }
+}
+
+async function registryContext(): Promise<RegistryContext> {
+  const context = getExecutionContext();
+  const host = (await context.client.config.getHost()).toString();
+  return {
+    client: context.client,
+    host,
+    identity: "userId" in context ? context.userId : context.serviceUserId,
+  };
+}
+
+function unqualifiedModel(model: string): string {
+  return model.trim().replace(/^(?:dbx|databricks)\//i, "");
+}
+
+function positiveTtl(value = DEFAULT_CACHE_TTL_MS): number {
+  if (!Number.isFinite(value) || value <= 0) {
+    throw new Error("Model registry ttlMs must be a positive finite number");
+  }
+  return value;
+}
+
+function normalizedModelServiceName(value: string | undefined): string | undefined {
+  const normalized = value?.trim().replace(/^databricks\//, "");
+  if (!normalized) return undefined;
+  if (normalized.startsWith("system.ai.databricks-")) {
+    return `system.ai.${normalized.slice("system.ai.databricks-".length)}`;
+  }
+  if (normalized.startsWith("databricks-")) {
+    return `system.ai.${normalized.slice("databricks-".length)}`;
+  }
+  return normalized;
+}
+
+function matchesResolveOptions(target: ModelTarget, options: ModelRegistryResolveOptions): boolean {
+  if (options.requiresTools && !target.capabilities.tools) return false;
+  if (options.modelClass === ModelClass.Embedding) return target.capabilities.embeddings;
+  if (options.modelClass !== undefined) return target.capabilities.chat;
+  return true;
+}

package/src/router.ts

@@ -0,0 +1,174 @@
+/**
+ * Pure capability-driven route selection for model-gateway requests.
+ *
+ * @module
+ */
+
+import type {
+  GatewayRoute,
+  ModelTarget,
+  RequestedFeatures,
+  ResolveRouteInput,
+  UpstreamProtocol,
+} from "@dbx-tools/shared-model-gateway";
+
+/** Error returned when preserving a requested protocol feature is impossible. */
+export class UnsupportedGatewayFeatureError extends Error {
+  readonly status = 400;
+
+  constructor(readonly features: readonly string[]) {
+    super(`The selected model cannot preserve requested features: ${features.join(", ")}`);
+    this.name = "UnsupportedGatewayFeatureError";
+  }
+}
+
+/** Select the lowest-overhead upstream that preserves the client contract. */
+export function resolveRoute(input: ResolveRouteInput): GatewayRoute {
+  const { clientProtocol, requestedModel, features, target } = input;
+  if (clientProtocol === "openai-embeddings") {
+    if (!target.capabilities.embeddings) {
+      throw new UnsupportedGatewayFeatureError(["embeddings"]);
+    }
+    return directRoute(input, "databricks-embeddings", target.id);
+  }
+  if (clientProtocol === "openai-responses") {
+    rejectStatefulResponses(features);
+    if (
+      isCodexOriginator(input.originator) &&
+      targetsModelService(requestedModel, target) &&
+      target.capabilities.aiGatewayCodex &&
+      supportsDirectResponses(target, features)
+    ) {
+      return directRoute(input, "databricks-ai-gateway-codex", target.modelServiceName!);
+    }
+    if (target.capabilities.responses && supportsDirectResponses(target, features)) {
+      return directRoute(input, "databricks-responses", target.id);
+    }
+    if (
+      target.capabilities.aiGatewayCodex &&
+      target.modelServiceName &&
+      supportsDirectResponses(target, features)
+    ) {
+      return directRoute(input, "databricks-ai-gateway-codex", target.modelServiceName);
+    }
+    if (target.capabilities.openResponses && supportsOpenResponses(target, features)) {
+      return directRoute(input, "databricks-open-responses", target.id);
+    }
+    return translatedRoute(input);
+  }
+
+  if (clientProtocol === "openai-chat") {
+    if (target.capabilities.chat && supportsTools(target, features)) {
+      return directRoute(input, "databricks-chat", target.id);
+    }
+    return translatedRoute(input);
+  }
+
+  if (target.capabilities.anthropic && supportsTools(target, features)) {
+    return directRoute(input, "databricks-anthropic", target.id);
+  }
+  return translatedRoute(input);
+}
+
+/** Detect route-affecting features from an OpenAI or Anthropic request body. */
+export function requestedFeatures(body: Readonly<Record<string, unknown>>): RequestedFeatures {
+  const tools = Array.isArray(body.tools) ? body.tools : [];
+  const toolTypes = tools
+    .filter((tool): tool is Record<string, unknown> => isRecord(tool))
+    .map((tool) => tool.type);
+  const text = isRecord(body.text) ? body.text : {};
+  const format = isRecord(text.format) ? text.format.type : text.format;
+  return {
+    background: body.background === true,
+    customTools: toolTypes.some((type) => type === "custom"),
+    parallelTools: body.parallel_tool_calls === true,
+    previousResponse:
+      typeof body.previous_response_id === "string" && body.previous_response_id.length > 0,
+    reasoning: isRecord(body.reasoning) || body.reasoning_effort !== undefined,
+    storage: body.store === true,
+    structuredOutput: format !== undefined && format !== "text",
+    tools: tools.length > 0,
+    unsupportedOpenResponsesTools: toolTypes.some(
+      (type) => type !== "function" && type !== "web_search" && type !== "web_search_preview",
+    ),
+    webSearch: toolTypes.some((type) => type === "web_search" || type === "web_search_preview"),
+  };
+}
+
+/** Return whether an Originator header selects Codex compatibility behavior. */
+export function isCodexOriginator(originator: string | undefined): boolean {
+  return originator?.toLowerCase().includes("codex") === true;
+}
+
+function directRoute(
+  input: ResolveRouteInput,
+  upstreamProtocol: Exclude<UpstreamProtocol, "ai-sdk">,
+  upstreamModel: string,
+): GatewayRoute {
+  return {
+    clientProtocol: input.clientProtocol,
+    upstreamProtocol,
+    target: input.target,
+    upstreamModel,
+    translateRequest: false,
+    translateResponse: false,
+  };
+}
+
+function translatedRoute(input: ResolveRouteInput): GatewayRoute {
+  return {
+    clientProtocol: input.clientProtocol,
+    upstreamProtocol: "ai-sdk",
+    target: input.target,
+    upstreamModel: input.target.id,
+    translateRequest: true,
+    translateResponse: true,
+  };
+}
+
+function supportsDirectResponses(target: ModelTarget, features: RequestedFeatures): boolean {
+  return (
+    supportsTools(target, features) &&
+    (!features.customTools || target.capabilities.customTools) &&
+    (!features.structuredOutput || target.capabilities.structuredOutput)
+  );
+}
+
+function supportsOpenResponses(target: ModelTarget, features: RequestedFeatures): boolean {
+  return (
+    supportsTools(target, features) &&
+    !features.customTools &&
+    !features.previousResponse &&
+    !features.storage &&
+    !features.background &&
+    !features.unsupportedOpenResponsesTools &&
+    (!features.webSearch || target.capabilities.webSearch) &&
+    (!features.structuredOutput || target.capabilities.structuredOutput)
+  );
+}
+
+function supportsTools(target: ModelTarget, features: RequestedFeatures): boolean {
+  return (
+    (!features.tools || target.capabilities.tools) &&
+    (!features.parallelTools || target.capabilities.parallelTools)
+  );
+}
+
+function rejectStatefulResponses(features: RequestedFeatures): void {
+  const unsupported = [
+    ...(features.previousResponse ? ["previous_response_id"] : []),
+    ...(features.storage ? ["store"] : []),
+    ...(features.background ? ["background"] : []),
+  ];
+  if (unsupported.length > 0) throw new UnsupportedGatewayFeatureError(unsupported);
+}
+
+function targetsModelService(requestedModel: string, target: ModelTarget): boolean {
+  if (!target.modelServiceName) return false;
+  const unqualified = requestedModel.trim().replace(/^(?:dbx|databricks)\//i, "");
+  return unqualified === target.modelServiceName;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return value !== null && typeof value === "object" && !Array.isArray(value);
+}

package/src/routes.ts

@@ -0,0 +1,181 @@
+/**
+ * Express adapters for the shared model-gateway service.
+ *
+ * @module
+ */
+
+import { Readable } from "node:stream";
+import { pipeline } from "node:stream/promises";
+import { log } from "@dbx-tools/shared-core";
+import type { ClientProtocol } from "@dbx-tools/shared-model-gateway";
+import type express from "express";
+
+import { ModelGateway, ModelGatewayError } from "./gateway.ts";
+import { UnsupportedGatewayFeatureError } from "./router.ts";
+
+const logger = log.logger("appkit/model-gateway/http");
+
+/** Send the lightweight compatibility health response. */
+export function sendHealth(response: express.Response): void {
+  response.json({ ready: true });
+}
+
+/** Send the active principal's dynamic model catalogue. */
+export async function sendModels(
+  gateway: ModelGateway,
+  request: express.Request,
+  response: express.Response,
+): Promise<void> {
+  const startedAt = Date.now();
+  logger.debug("models request", requestContext(request));
+  try {
+    const payload = await gateway.models(
+      request.header("originator"),
+      optionalSearch(request.query.search),
+    );
+    response.json(payload);
+    logger.debug("models response", {
+      ...requestContext(request),
+      elapsedMs: Date.now() - startedAt,
+      status: response.statusCode,
+    });
+  } catch (error) {
+    logger.error("models request failed", {
+      ...requestContext(request),
+      elapsedMs: Date.now() - startedAt,
+      error,
+    });
+    sendError("openai-responses", response, error);
+  }
+}
+
+/** Execute and stream one model request through Express with cancellation. */
+export async function sendInference(
+  gateway: ModelGateway,
+  protocol: ClientProtocol,
+  request: express.Request,
+  response: express.Response,
+): Promise<void> {
+  const startedAt = Date.now();
+  const controller = new AbortController();
+  const abort = () => {
+    if (response.writableEnded) return;
+    logger.debug("client disconnected", {
+      ...requestContext(request),
+      elapsedMs: Date.now() - startedAt,
+      protocol,
+    });
+    controller.abort(new Error("Model gateway client disconnected"));
+  };
+  request.once("aborted", abort);
+  response.once("close", abort);
+  try {
+    const body = record(request.body);
+    logger.debug("inference request", {
+      ...requestContext(request),
+      model: typeof body.model === "string" ? body.model : undefined,
+      protocol,
+      stream: body.stream === true,
+    });
+    const upstream = await gateway.inference(
+      protocol,
+      body,
+      requestHeaders(request),
+      controller.signal,
+    );
+    response.status(upstream.status);
+    upstream.headers.forEach((value, name) => response.setHeader(name, value));
+    if (!upstream.body) {
+      response.end();
+      logger.debug("inference response", {
+        ...requestContext(request),
+        elapsedMs: Date.now() - startedAt,
+        protocol,
+        status: upstream.status,
+        upstreamProtocol: upstream.headers.get("x-dbx-tools-upstream"),
+      });
+      return;
+    }
+    await pipeline(Readable.fromWeb(upstream.body as never), response);
+    logger.debug("inference stream completed", {
+      ...requestContext(request),
+      elapsedMs: Date.now() - startedAt,
+      protocol,
+      status: upstream.status,
+      upstreamProtocol: upstream.headers.get("x-dbx-tools-upstream"),
+    });
+  } catch (error) {
+    if (controller.signal.aborted) {
+      response.destroy();
+      return;
+    }
+    logger.error("inference request failed", {
+      ...requestContext(request),
+      elapsedMs: Date.now() - startedAt,
+      error,
+      protocol,
+    });
+    if (response.headersSent) {
+      response.destroy(error instanceof Error ? error : undefined);
+      return;
+    }
+    sendError(protocol, response, error);
+  } finally {
+    request.off("aborted", abort);
+    response.off("close", abort);
+  }
+}
+
+function sendError(protocol: ClientProtocol, response: express.Response, error: unknown): void {
+  const status =
+    error instanceof ModelGatewayError || error instanceof UnsupportedGatewayFeatureError
+      ? error.status
+      : 500;
+  const message = error instanceof Error ? error.message : "Model gateway request failed";
+  if (protocol === "anthropic-messages") {
+    response.status(status).json({
+      type: "error",
+      error: { type: status < 500 ? "invalid_request_error" : "api_error", message },
+    });
+    return;
+  }
+  response.status(status).json({
+    error: {
+      message,
+      type: status < 500 ? "invalid_request_error" : "api_error",
+      code: status,
+    },
+  });
+}
+
+function requestHeaders(request: express.Request): Headers {
+  const headers = new Headers();
+  for (const [name, raw] of Object.entries(request.headers)) {
+    if (typeof raw === "string") headers.set(name, raw);
+    else if (Array.isArray(raw)) headers.set(name, raw.join(", "));
+  }
+  return headers;
+}
+
+function record(value: unknown): Record<string, unknown> {
+  if (value !== null && typeof value === "object" && !Array.isArray(value)) {
+    return value as Record<string, unknown>;
+  }
+  throw new ModelGatewayError(400, "Request body must be a JSON object");
+}
+
+function optionalSearch(value: unknown): string | undefined {
+  if (value === undefined) return undefined;
+  if (typeof value !== "string") {
+    throw new ModelGatewayError(400, "search must be a string");
+  }
+  return value;
+}
+
+function requestContext(request: express.Request): Record<string, unknown> {
+  return {
+    method: request.method,
+    path: request.path,
+    requestId: request.header("x-request-id") ?? request.header("request-id"),
+  };
+}