@dbx-tools/appkit-model-gateway 0.9.36

package/package.json

@@ -0,0 +1,150 @@
+{
+  "name": "@dbx-tools/appkit-model-gateway",
+  "description": "AppKit model gateway for Databricks OpenAI, Anthropic, Codex, and AI Gateway APIs",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/reggie-db/dbx-tools.git",
+    "directory": "packages/js/node/appkit-model-gateway"
+  },
+  "scripts": {
+    "build": "projen build",
+    "compile": "projen compile",
+    "default": "projen default",
+    "package": "projen package",
+    "post-compile": "projen post-compile",
+    "pre-compile": "projen pre-compile",
+    "prepack": "projen prepack",
+    "test": "projen test",
+    "watch": "projen watch",
+    "projen": "projen"
+  },
+  "devDependencies": {
+    "@types/express": "^5.0.5",
+    "@types/json-schema": "^7",
+    "@types/node": "^24.6.0",
+    "openai": "^6.16.0",
+    "typescript": "^5.9.3",
+    "vitest": "3.2.4"
+  },
+  "dependencies": {
+    "@ai-sdk/anthropic": "^4.0.0",
+    "@ai-sdk/open-responses": "^2.0.58",
+    "@ai-sdk/openai-compatible": "^3.0.62",
+    "@databricks/appkit": "0.81.0",
+    "@dbx-tools/databricks": "^0.9.36",
+    "@dbx-tools/model": "^0.9.36",
+    "@dbx-tools/shared-core": "^0.9.36",
+    "@dbx-tools/shared-model": "^0.9.36",
+    "@dbx-tools/shared-model-gateway": "^0.9.36",
+    "ai": "^7.0.0",
+    "zod": "4.3.6"
+  },
+  "main": "./lib/index.js",
+  "license": "Apache-2.0",
+  "publishConfig": {
+    "access": "public",
+    "main": "./lib/index.js",
+    "types": "./lib/index.d.ts",
+    "exports": {
+      ".": {
+        "types": "./lib/index.d.ts",
+        "default": "./lib/index.js"
+      },
+      "./gateway": {
+        "types": "./lib/src/gateway.d.ts",
+        "default": "./lib/src/gateway.js"
+      },
+      "./models": {
+        "types": "./lib/src/models.d.ts",
+        "default": "./lib/src/models.js"
+      },
+      "./plugin": {
+        "types": "./lib/src/plugin.d.ts",
+        "default": "./lib/src/plugin.js"
+      },
+      "./reasoning-adapt": {
+        "types": "./lib/src/reasoning-adapt.d.ts",
+        "default": "./lib/src/reasoning-adapt.js"
+      },
+      "./registry": {
+        "types": "./lib/src/registry.d.ts",
+        "default": "./lib/src/registry.js"
+      },
+      "./router": {
+        "types": "./lib/src/router.d.ts",
+        "default": "./lib/src/router.js"
+      },
+      "./routes": {
+        "types": "./lib/src/routes.d.ts",
+        "default": "./lib/src/routes.js"
+      },
+      "./translation": {
+        "types": "./lib/src/translation.d.ts",
+        "default": "./lib/src/translation.js"
+      },
+      "./transport": {
+        "types": "./lib/src/transport.d.ts",
+        "default": "./lib/src/transport.js"
+      },
+      "./package.json": "./package.json"
+    }
+  },
+  "version": "0.9.36",
+  "types": "./lib/index.d.ts",
+  "dbxToolsConfig": {
+    "tags": [
+      "node"
+    ]
+  },
+  "type": "module",
+  "exports": {
+    ".": {
+      "types": "./lib/index.d.ts",
+      "default": "./lib/index.js"
+    },
+    "./gateway": {
+      "types": "./lib/src/gateway.d.ts",
+      "default": "./lib/src/gateway.js"
+    },
+    "./models": {
+      "types": "./lib/src/models.d.ts",
+      "default": "./lib/src/models.js"
+    },
+    "./plugin": {
+      "types": "./lib/src/plugin.d.ts",
+      "default": "./lib/src/plugin.js"
+    },
+    "./reasoning-adapt": {
+      "types": "./lib/src/reasoning-adapt.d.ts",
+      "default": "./lib/src/reasoning-adapt.js"
+    },
+    "./registry": {
+      "types": "./lib/src/registry.d.ts",
+      "default": "./lib/src/registry.js"
+    },
+    "./router": {
+      "types": "./lib/src/router.d.ts",
+      "default": "./lib/src/router.js"
+    },
+    "./routes": {
+      "types": "./lib/src/routes.d.ts",
+      "default": "./lib/src/routes.js"
+    },
+    "./translation": {
+      "types": "./lib/src/translation.d.ts",
+      "default": "./lib/src/translation.js"
+    },
+    "./transport": {
+      "types": "./lib/src/transport.d.ts",
+      "default": "./lib/src/transport.js"
+    },
+    "./package.json": "./package.json"
+  },
+  "files": [
+    "index.ts",
+    "src",
+    "dist/plugins",
+    "lib"
+  ],
+  "//": "~~ Generated by projen. To modify, edit .projenrc.js and run \"bunx projen\"."
+}

package/src/gateway.ts

@@ -0,0 +1,167 @@
+/**
+ * Protocol-independent gateway service used by plugin and standalone routes.
+ *
+ * @module
+ */
+
+import { log } from "@dbx-tools/shared-core";
+import { ModelClass } from "@dbx-tools/shared-model";
+import type {
+  ClientProtocol,
+  GatewayRoute,
+  ModelCapabilityOverride,
+  ModelListResponse,
+} from "@dbx-tools/shared-model-gateway";
+
+import { listModelsPayload } from "./models.ts";
+import {
+  DatabricksModelRegistry,
+  type ModelRegistry,
+  type ModelRegistryOptions,
+} from "./registry.ts";
+import {
+  adaptInferenceReasoning,
+  learnAndAdaptReasoningRetry,
+} from "./reasoning-adapt.ts";
+import { isCodexOriginator, requestedFeatures, resolveRoute } from "./router.ts";
+import { translateGatewayRequest } from "./translation.ts";
+import { fetchDatabricks, gatewayResponseHeaders } from "./transport.ts";
+
+const logger = log.logger("appkit/model-gateway");
+
+/** Model gateway runtime configuration. */
+export interface ModelGatewayOptions {
+  readonly cacheTtlMs?: number;
+  readonly overrides?: readonly ModelCapabilityOverride[];
+}
+
+/** Error carrying an HTTP status and protocol-safe message. */
+export class ModelGatewayError extends Error {
+  constructor(
+    readonly status: number,
+    message: string,
+  ) {
+    super(message);
+    this.name = "ModelGatewayError";
+  }
+}
+
+/** AppKit-scoped model gateway with dynamic model discovery and streaming inference. */
+export class ModelGateway {
+  private readonly registry: ModelRegistry;
+
+  constructor(options: ModelGatewayOptions = {}, registry?: ModelRegistry) {
+    const registryOptions: ModelRegistryOptions = {
+      ...(options.cacheTtlMs !== undefined ? { ttlMs: options.cacheTtlMs } : {}),
+      ...(options.overrides ? { overrides: options.overrides } : {}),
+    };
+    this.registry = registry ?? new DatabricksModelRegistry(registryOptions);
+  }
+
+  /** Return the live OpenAI model list, including Codex metadata when requested. */
+  async models(originator?: string, search?: string): Promise<ModelListResponse> {
+    const codex = isCodexOriginator(originator);
+    logger.debug("listing models", { codex, search });
+    const targets = search?.trim()
+      ? await this.registry.search(search.trim())
+      : await this.registry.list();
+    logger.debug("listed models", { codex, count: targets.length });
+    return listModelsPayload(targets, codex);
+  }
+
+  /** Resolve and execute one inference request. */
+  async inference(
+    protocol: ClientProtocol,
+    body: Readonly<Record<string, unknown>>,
+    headers: Headers,
+    signal: AbortSignal,
+  ): Promise<Response> {
+    const requestedModel = requiredModel(body.model);
+    logger.debug("resolving request", {
+      clientProtocol: protocol,
+      model: requestedModel,
+      originator: headers.get("originator") ?? undefined,
+      stream: body.stream === true,
+    });
+    const features = requestedFeatures(body);
+    const target = await this.registry.resolve(requestedModel, {
+      ...(protocol === "openai-embeddings" ? { modelClass: ModelClass.Embedding } : {}),
+      ...(features.tools ? { requiresTools: true } : {}),
+    });
+    if (!target) throw new ModelGatewayError(404, `Model not found: ${requestedModel}`);
+    const route = resolveRoute({
+      clientProtocol: protocol,
+      requestedModel,
+      originator: headers.get("originator") ?? undefined,
+      features,
+      target,
+    });
+    logger.debug("routed", {
+      clientProtocol: protocol,
+      model: target.id,
+      upstreamProtocol: route.upstreamProtocol,
+    });
+
+    const adapted = adaptInferenceReasoning(target.id, body);
+    let response = await this.forward(route, adapted.body, headers, signal);
+    const retryBody = await learnAndAdaptReasoningRetry({
+      model: target.id,
+      body,
+      response,
+      previousWireEffort: adapted.wireEffort,
+    });
+    if (retryBody) {
+      await response.body?.cancel().catch(() => undefined);
+      response = await this.forward(route, retryBody, headers, signal);
+    }
+
+    logResponse(route.upstreamProtocol, target.id, response);
+    if (route.upstreamProtocol === "ai-sdk") {
+      response.headers.set("x-dbx-tools-upstream", route.upstreamProtocol);
+      return response;
+    }
+    const forwarded = gatewayResponseHeaders(response.headers);
+    forwarded.set("x-dbx-tools-upstream", route.upstreamProtocol);
+    return new Response(response.body, {
+      status: response.status,
+      statusText: response.statusText,
+      headers: forwarded,
+    });
+  }
+
+  /** Refresh the active principal's model catalogue. */
+  refresh(): Promise<void> {
+    logger.debug("refreshing model catalogue");
+    return this.registry.refresh();
+  }
+
+  private async forward(
+    route: GatewayRoute,
+    body: Readonly<Record<string, unknown>>,
+    headers: Headers,
+    signal: AbortSignal,
+  ): Promise<Response> {
+    if (route.upstreamProtocol === "ai-sdk") {
+      return translateGatewayRequest(route, body, signal);
+    }
+    return fetchDatabricks({ route, body, headers, signal });
+  }
+}
+
+function requiredModel(value: unknown): string {
+  if (typeof value !== "string" || !value.trim()) {
+    throw new ModelGatewayError(400, "Request body must include a non-empty model");
+  }
+  return value.trim();
+}
+
+function logResponse(upstreamProtocol: string, model: string, response: Response): void {
+  const details = {
+    contentType: response.headers.get("content-type"),
+    model,
+    status: response.status,
+    upstreamProtocol,
+  };
+  if (response.ok) logger.debug("upstream response", details);
+  else logger.warn("upstream request failed", details);
+}

package/src/manifest.json

@@ -0,0 +1,13 @@
+{
+  "$schema": "https://databricks.github.io/appkit/schemas/plugin-manifest.schema.json",
+  "name": "modelGateway",
+  "displayName": "Model Gateway",
+  "description": "OpenAI, Anthropic, Codex, and AI Gateway compatibility backed by Databricks",
+  "author": "dbx-tools",
+  "license": "Apache-2.0",
+  "stability": "beta",
+  "resources": {
+    "required": [],
+    "optional": []
+  }
+}

package/src/models.ts

@@ -0,0 +1,152 @@
+/**
+ * OpenAI and Codex model-list payloads derived from the live registry.
+ *
+ * @module
+ */
+
+import { isChatClass } from "@dbx-tools/model/classes";
+import type {
+  CodexModel,
+  CodexModelListResponse,
+  ModelListResponse,
+  ModelTarget,
+  OpenAIModel,
+  OpenAIModelListResponse,
+} from "@dbx-tools/shared-model-gateway";
+
+const CODEX_BASE_INSTRUCTIONS =
+  "You are a coding agent. Follow the user's instructions and use the available tools to work in the current repository.";
+
+const REASONING_DESCRIPTIONS: Readonly<Record<string, string>> = {
+  none: "Disable explicit reasoning",
+  minimal: "Use the smallest available reasoning budget",
+  low: "Use a low reasoning budget",
+  medium: "Use a medium reasoning budget",
+  high: "Use a high reasoning budget",
+  xhigh: "Use an extra-high reasoning budget",
+  max: "Use the largest available reasoning budget",
+};
+
+/** Build the combined OpenAI and optional Codex model catalogue. */
+export function listModelsPayload(
+  targets: readonly ModelTarget[],
+  includeCodex: true,
+): CodexModelListResponse;
+export function listModelsPayload(
+  targets: readonly ModelTarget[],
+  includeCodex: false,
+): OpenAIModelListResponse;
+export function listModelsPayload(
+  targets: readonly ModelTarget[],
+  includeCodex: boolean,
+): ModelListResponse;
+export function listModelsPayload(
+  targets: readonly ModelTarget[],
+  includeCodex: boolean,
+): ModelListResponse {
+  const available = targets
+    .filter((target) => target.endpoint?.status?.deprecated !== true)
+    .filter((target) => !includeCodex || hasChatClass(target))
+    .slice()
+    .sort(compareCatalogueTargets);
+  if (includeCodex) {
+    return {
+      models: available.map((target, index) => codexModel(target, index + 1)),
+    };
+  }
+  return {
+    object: "list",
+    data: available.map(openAiModel),
+  };
+}
+
+function openAiModel(target: ModelTarget): OpenAIModel {
+  return {
+    id: target.id,
+    object: "model",
+    created: 0,
+    owned_by: "databricks",
+    name: target.displayName,
+    task: target.endpoint?.task,
+    status: target.endpoint?.status ?? { deprecated: false },
+    capabilities: {
+      tools: target.capabilities.tools,
+      reasoning: target.reasoningEfforts,
+      responses: target.capabilities.responses,
+      open_responses: target.capabilities.openResponses,
+      anthropic: target.capabilities.anthropic,
+      embeddings: target.capabilities.embeddings,
+      ai_gateway_codex: target.capabilities.aiGatewayCodex,
+      streaming: target.capabilities.streaming,
+    },
+  };
+}
+
+function codexModel(target: ModelTarget, priority: number): CodexModel {
+  const model =
+    target.capabilities.aiGatewayCodex && target.modelServiceName
+      ? target.modelServiceName
+      : target.id;
+  const efforts = target.reasoningEfforts;
+  return {
+    slug: `databricks/${model}`,
+    display_name: target.displayName,
+    description: target.endpoint?.description ?? "Databricks model service",
+    base_instructions: CODEX_BASE_INSTRUCTIONS,
+    status: target.endpoint?.status ?? { deprecated: false },
+    supported_reasoning_levels: efforts.map((effort) => ({
+      effort,
+      description: REASONING_DESCRIPTIONS[effort] ?? `Use the ${effort} reasoning budget`,
+    })),
+    ...(efforts.length > 0
+      ? { default_reasoning_level: efforts.includes("medium") ? "medium" : efforts[0] }
+      : {}),
+    shell_type: "unified_exec",
+    visibility: "list",
+    supported_in_api: true,
+    priority,
+    availability_nux: null,
+    upgrade: null,
+    support_verbosity: false,
+    default_verbosity: null,
+    apply_patch_tool_type: target.capabilities.customTools ? "freeform" : null,
+    truncation_policy: { mode: "tokens", limit: 128_000 },
+    context_window: null,
+    experimental_supported_tools: [],
+    input_modalities: ["text"],
+    web_search_tool_type: "text",
+    supports_search_tool: false,
+    supports_image_detail_original: false,
+  };
+}
+
+/**
+ * Order listed models by family A-Z, then display name. Targets with no family
+ * follow every named family, still sorted by display name.
+ */
+function compareCatalogueTargets(left: ModelTarget, right: ModelTarget): number {
+  const leftFamily = catalogueFamily(left);
+  const rightFamily = catalogueFamily(right);
+  if (Boolean(leftFamily) !== Boolean(rightFamily)) return leftFamily ? -1 : 1;
+  if (leftFamily && rightFamily) {
+    const familyOrder = leftFamily.localeCompare(rightFamily, undefined, { sensitivity: "base" });
+    if (familyOrder !== 0) return familyOrder;
+  }
+  return catalogueName(left).localeCompare(catalogueName(right), undefined, {
+    sensitivity: "base",
+  });
+}
+
+function catalogueFamily(target: ModelTarget): string | undefined {
+  const family = target.family?.trim() || target.endpoint?.family?.trim();
+  return family || undefined;
+}
+
+function catalogueName(target: ModelTarget): string {
+  return target.displayName.trim() || target.id;
+}
+
+function hasChatClass(target: ModelTarget): boolean {
+  const modelClass = target.endpoint?.class;
+  return modelClass !== undefined && isChatClass(modelClass);
+}

package/src/plugin.ts

@@ -0,0 +1,95 @@
+/**
+ * AppKit model-gateway plugin.
+ *
+ * @module
+ */
+
+import {
+  defineManifest,
+  Plugin,
+  toPlugin,
+  type BasePluginConfig,
+  type IAppRouter,
+} from "@databricks/appkit";
+import type { ClientProtocol } from "@dbx-tools/shared-model-gateway";
+import type express from "express";
+
+import { ModelGateway, type ModelGatewayOptions } from "./gateway.ts";
+import manifest from "./manifest.json" with { type: "json" };
+import { sendHealth, sendInference, sendModels } from "./routes.ts";
+
+/** AppKit configuration for the model-gateway plugin. */
+export interface ModelGatewayPluginConfig extends BasePluginConfig, ModelGatewayOptions {}
+
+/** AppKit plugin exposing dynamic model discovery and compatibility routes. */
+export class ModelGatewayPlugin extends Plugin<ModelGatewayPluginConfig> {
+  static manifest = defineManifest<"modelGateway">(manifest);
+
+  private readonly gateway: ModelGateway;
+
+  constructor(config: ModelGatewayPluginConfig = {}) {
+    super(config);
+    this.gateway = new ModelGateway({
+      ...(config.cacheTtlMs !== undefined ? { cacheTtlMs: config.cacheTtlMs } : {}),
+      ...(config.overrides ? { overrides: config.overrides } : {}),
+    });
+  }
+
+  override injectRoutes(router: IAppRouter): void {
+    this.route(router, {
+      name: "health",
+      method: "get",
+      path: "/healthz",
+      handler: async (_request, response) => sendHealth(response),
+    });
+    this.route(router, {
+      name: "models",
+      method: "get",
+      path: "/v1/models",
+      handler: (request, response) => this.models(request, response),
+    });
+    for (const [name, path, protocol] of [
+      ["chatCompletions", "/v1/chat/completions", "openai-chat"],
+      ["responses", "/v1/responses", "openai-responses"],
+      ["messages", "/v1/messages", "anthropic-messages"],
+      ["embeddings", "/v1/embeddings", "openai-embeddings"],
+    ] as const) {
+      this.route(router, {
+        name,
+        method: "post",
+        path,
+        handler: (request, response) => this.inference(protocol, request, response),
+      });
+    }
+  }
+
+  /** Send the active principal's model catalogue. */
+  async models(request: express.Request, response: express.Response): Promise<void> {
+    await sendModels(this.gateway, request, response);
+  }
+
+  /** Execute one compatibility request in the active AppKit identity scope. */
+  async inference(
+    protocol: ClientProtocol,
+    request: express.Request,
+    response: express.Response,
+  ): Promise<void> {
+    await sendInference(this.gateway, protocol, request, response);
+  }
+
+  /** Refresh the active principal's model catalogue. */
+  async refresh(): Promise<void> {
+    await this.gateway.refresh();
+  }
+
+  override exports() {
+    return {
+      models: this.models,
+      inference: this.inference,
+      refresh: this.refresh,
+    };
+  }
+}
+
+/** AppKit plugin factory for the Databricks model gateway. */
+export const modelGateway = toPlugin(ModelGatewayPlugin);