@dbx-tools/appkit-model-gateway 0.9.58 → 0.9.61

package/src/registry.ts

@@ -5,7 +5,13 @@
  */
 
 import { getExecutionContext, type WorkspaceClient } from "@databricks/appkit";
-import { metadata, modelCatalog, policy, resolve as modelResolve } from "@dbx-tools/model";
+import {
+  metadata,
+  modelCatalog,
+  policy,
+  resolve as modelResolve,
+  type ListServingEndpointsOptions,
+} from "@dbx-tools/model";
 import { log } from "@dbx-tools/shared-core";
 import {
   ModelClass,
@@ -36,31 +42,36 @@ export interface ModelRegistry {
 export interface ModelRegistryOptions {
   readonly ttlMs?: number;
   readonly overrides?: readonly ModelCapabilityOverride[];
+  /** Discover models with this client instead of the AppKit execution context. */
+  readonly client?: WorkspaceClient;
 }
 
 interface RegistryContext {
   readonly client: WorkspaceClient;
   readonly host: string;
-  readonly identity: 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[];
+  private readonly client: WorkspaceClient | undefined;
 
   constructor(options: ModelRegistryOptions = {}) {
     this.ttlMs = positiveTtl(options.ttlMs);
     this.overrides = options.overrides ?? [];
+    this.client = options.client;
   }
 
   async list(): Promise<ModelTarget[]> {
-    const context = await registryContext();
+    const context = await this.context();
     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 endpoints = await modelCatalog.listServingEndpoints(
+      context.client,
+      context.host,
+      catalogueOptions(context, this.ttlMs),
+    );
     const targets = endpoints.map((endpoint) => this.target(endpoint));
     logger.debug("loaded catalogue", {
       endpointCount: endpoints.length,
@@ -124,13 +135,16 @@ export class DatabricksModelRegistry implements ModelRegistry {
   }
 
   async refresh(): Promise<void> {
-    const context = await registryContext();
+    const context = await this.context();
     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,
-    });
+    if (context.identity) {
+      await modelCatalog.clearServingEndpointsCache(context.host, context.identity);
+    }
+    await modelCatalog.listServingEndpoints(
+      context.client,
+      context.host,
+      catalogueOptions(context, this.ttlMs),
+    );
     logger.debug("refreshed catalogue", { host: context.host });
   }
 
@@ -182,6 +196,21 @@ export class DatabricksModelRegistry implements ModelRegistry {
       reasoningEfforts: endpoint.reasoningEfforts ?? [],
     };
   }
+
+  private async context(): Promise<RegistryContext> {
+    if (this.client) {
+      const host = (await this.client.config.getHost()).toString();
+      return { client: this.client, host };
+    }
+    return registryContext();
+  }
+}
+
+function catalogueOptions(context: RegistryContext, ttlMs: number): ListServingEndpointsOptions {
+  return {
+    ttlMs,
+    ...(context.identity ? { cacheIdentity: context.identity } : {}),
+  };
 }
 
 async function registryContext(): Promise<RegistryContext> {

package/src/router.ts

@@ -4,6 +4,7 @@
  * @module
  */
 
+import { object } from "@dbx-tools/shared-core";
 import type {
   GatewayRoute,
   ModelTarget,
@@ -35,21 +36,16 @@ export function resolveRoute(input: ResolveRouteInput): GatewayRoute {
     rejectStatefulResponses(features);
     if (
       isCodexOriginator(input.originator) &&
-      targetsModelService(requestedModel, target) &&
-      target.capabilities.aiGatewayCodex &&
-      supportsDirectResponses(target, features)
+      canUseAiGatewayCodex(target, features) &&
+      !requestedServingEndpoint(requestedModel, target)
     ) {
       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 (canUseAiGatewayCodex(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);
@@ -74,17 +70,17 @@ export function resolveRoute(input: ResolveRouteInput): GatewayRoute {
 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))
+    .filter((tool): tool is Record<string, unknown> => object.isRecord(tool))
     .map((tool) => tool.type);
-  const text = isRecord(body.text) ? body.text : {};
-  const format = isRecord(text.format) ? text.format.type : text.format;
+  const text = object.isRecord(body.text) ? body.text : {};
+  const format = object.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,
+    reasoning: object.isRecord(body.reasoning) || body.reasoning_effort !== undefined,
     storage: body.store === true,
     structuredOutput: format !== undefined && format !== "text",
     tools: tools.length > 0,
@@ -163,12 +159,14 @@ function rejectStatefulResponses(features: RequestedFeatures): void {
   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 canUseAiGatewayCodex(target: ModelTarget, features: RequestedFeatures): boolean {
+  return (
+    Boolean(target.capabilities.aiGatewayCodex && target.modelServiceName) &&
+    supportsDirectResponses(target, features)
+  );
 }
 
-function isRecord(value: unknown): value is Record<string, unknown> {
-  return value !== null && typeof value === "object" && !Array.isArray(value);
+/** Serving-endpoint ids stay on Databricks Responses; aliases and model-service slugs use Codex. */
+function requestedServingEndpoint(requestedModel: string, target: ModelTarget): boolean {
+  return requestedModel.trim().replace(/^(?:dbx|databricks)\//i, "") === target.id;
 }

package/src/transport.ts

@@ -7,6 +7,7 @@
 import { getExecutionContext } from "@databricks/appkit";
 import { workspaceClient } from "@dbx-tools/databricks";
 import { log } from "@dbx-tools/shared-core";
+import { openaiChat } from "@dbx-tools/shared-model";
 import { upstreamUrl, type GatewayRoute } from "@dbx-tools/shared-model-gateway";
 
 const logger = log.logger("appkit/model-gateway/transport");
@@ -60,7 +61,7 @@ export async function fetchDatabricks(request: DatabricksTransportRequest): Prom
   await client.config.authenticate(headers);
   const url = upstreamUrl(host, request.route);
   const startedAt = Date.now();
-  const body = JSON.stringify({ ...request.body, model: request.route.upstreamModel });
+  const body = JSON.stringify(directDatabricksRequestBody(request.route, request.body));
   logger.debug("sending upstream request", {
     model: request.route.upstreamModel,
     path: new URL(url).pathname,
@@ -119,6 +120,20 @@ export async function fetchDatabricks(request: DatabricksTransportRequest): Prom
   }
 }
 
+/** Remove client compatibility fields rejected by direct Databricks protocols. */
+export function directDatabricksRequestBody(
+  route: Pick<GatewayRoute, "upstreamProtocol" | "upstreamModel">,
+  body: Readonly<Record<string, unknown>>,
+): Record<string, unknown> {
+  const sanitized: Record<string, unknown> = { ...body, model: route.upstreamModel };
+  if (route.upstreamProtocol === "databricks-chat") {
+    openaiChat.stripUnsupportedChatFields(sanitized);
+  } else {
+    delete sanitized.parallel_tool_calls;
+  }
+  return sanitized;
+}
+
 /** Copy only response headers that belong to the public model protocol. */
 export function gatewayResponseHeaders(upstream: Headers): Headers {
   const headers = new Headers();