@dbx-tools/appkit-model-gateway 0.9.60 → 0.9.64

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,
@@ -28,7 +34,7 @@ export type ModelRegistryResolveOptions = Pick<ResolveModelInput, "modelClass" |
 export interface ModelRegistry {
   list(): Promise<ModelTarget[]>;
   search(query: string): Promise<ModelTarget[]>;
-  resolve(model: string, options?: ModelRegistryResolveOptions): Promise<ModelTarget | undefined>;
+  resolve(model?: string, options?: ModelRegistryResolveOptions): Promise<ModelTarget | undefined>;
   refresh(): Promise<void>;
 }
 
@@ -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,
@@ -71,16 +82,19 @@ export class DatabricksModelRegistry implements ModelRegistry {
   }
 
   async resolve(
-    model: string,
+    model: string | undefined,
     options: ModelRegistryResolveOptions = {},
   ): Promise<ModelTarget | undefined> {
-    const requested = unqualifiedModel(model);
+    const requested = model === undefined ? undefined : unqualifiedModel(model);
     const targets = await this.list();
-    const exact = targets.find((target) =>
-      target.aliases.some(
-        (alias) => alias.localeCompare(requested, undefined, { sensitivity: "accent" }) === 0,
-      ),
-    );
+    const exact =
+      requested === undefined
+        ? undefined
+        : 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;
@@ -90,7 +104,7 @@ export class DatabricksModelRegistry implements ModelRegistry {
       .map((target) => target.endpoint)
       .filter((endpoint): endpoint is ServingEndpointSummary => endpoint !== undefined);
     const resolved = modelResolve.resolveModel(endpoints, {
-      explicit: requested,
+      ...(requested !== undefined ? { explicit: requested } : {}),
       ...options,
     });
     const target = targets.find((candidate) => candidate.id === resolved.modelId);
@@ -124,13 +138,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 +199,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

@@ -36,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);
@@ -164,8 +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)
+  );
+}
+
+/** 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();