@dbx-tools/appkit-model-gateway 0.9.189 → 0.9.193

package/src/models.ts

@@ -12,10 +12,11 @@ import {
   ModelFamily,
   modelFamily,
   modelFamilyRank,
+  modelMetadataFor,
   versionTuple,
 } from "@dbx-tools/model";
 import { ModelClass, ReasoningEffort } from "@dbx-tools/shared-model";
-import { SMART_ROUTER_MODEL_ID } from "@dbx-tools/shared-model-gateway";
+import { AI_DECIDE_MODEL_ID, SMART_ROUTER_MODEL_ID } from "@dbx-tools/shared-model-gateway";
 import type {
   CodexModel,
   CodexModelListResponse,
@@ -29,6 +30,7 @@ import type {
 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 CODEX_SMART_ROUTER_PRIORITY = 9_000;
+const CODEX_AI_DECIDE_PRIORITY = 9_001;
 const CODEX_MODEL_PRIORITY_BASE = 10_000;
 const CODEX_FAMILY_ORDER = [
   ModelFamily.Gpt,
@@ -52,6 +54,8 @@ export function listModelsPayload(
   preserveOrder?: boolean,
   preferSystemModels?: boolean,
   includeSmartRouter?: boolean,
+  includeAiDecide?: boolean,
+  includeMetadata?: boolean,
 ): CodexModelListResponse;
 export function listModelsPayload(
   targets: readonly ModelTarget[],
@@ -59,6 +63,8 @@ export function listModelsPayload(
   preserveOrder?: boolean,
   preferSystemModels?: boolean,
   includeSmartRouter?: boolean,
+  includeAiDecide?: boolean,
+  includeMetadata?: boolean,
 ): OpenAIModelListResponse;
 export function listModelsPayload(
   targets: readonly ModelTarget[],
@@ -66,6 +72,8 @@ export function listModelsPayload(
   preserveOrder?: boolean,
   preferSystemModels?: boolean,
   includeSmartRouter?: boolean,
+  includeAiDecide?: boolean,
+  includeMetadata?: boolean,
 ): ModelListResponse;
 export function listModelsPayload(
   targets: readonly ModelTarget[],
@@ -73,6 +81,8 @@ export function listModelsPayload(
   preserveOrder = false,
   preferSystemModels = false,
   includeSmartRouter = false,
+  includeAiDecide = false,
+  includeMetadata = false,
 ): ModelListResponse {
   const available = targets
     .filter((target) => target.endpoint?.status?.deprecated !== true)
@@ -85,6 +95,9 @@ export function listModelsPayload(
     );
   }
   const smartRouterCandidates = available.filter((target) => target.capabilities.chat);
+  if (includeAiDecide && smartRouterCandidates.length > 0) {
+    available.unshift(aiDecideModelTarget(smartRouterCandidates));
+  }
   if (includeSmartRouter && smartRouterCandidates.length > 0) {
     available.unshift(smartRouterModelTarget(smartRouterCandidates));
   }
@@ -96,29 +109,35 @@ export function listModelsPayload(
           target,
           target.id === SMART_ROUTER_MODEL_ID
             ? CODEX_SMART_ROUTER_PRIORITY
-            : CODEX_MODEL_PRIORITY_BASE + modelIndex++,
+            : target.id === AI_DECIDE_MODEL_ID
+              ? CODEX_AI_DECIDE_PRIORITY
+              : CODEX_MODEL_PRIORITY_BASE + modelIndex++,
           preferSystemModels,
+          includeMetadata,
         ),
       ),
     };
   }
   return {
     object: "list",
-    data: available.map(openAiModel),
+    data: available.map((target) => openAiModel(target, includeMetadata)),
   };
 }
 
-function openAiModel(target: ModelTarget): OpenAIModel {
+function openAiModel(target: ModelTarget, includeMetadata: boolean): OpenAIModel {
   return {
     id: target.id,
     object: "model",
     created: 0,
     owned_by: "databricks",
     name: target.displayName,
+    model_source: target.endpoint?.modelSource ?? "custom",
     task: target.endpoint?.task,
     status: target.endpoint?.status ?? { deprecated: false },
+    ...(includeMetadata ? { metadata: modelMetadataFor(target.endpoint ?? target.id) } : {}),
     capabilities: {
       smart_router: target.id === SMART_ROUTER_MODEL_ID,
+      ai_decide: target.id === AI_DECIDE_MODEL_ID,
       tools: target.capabilities.tools,
       reasoning: target.reasoningEfforts,
       responses: target.capabilities.responses,
@@ -134,6 +153,46 @@ function openAiModel(target: ModelTarget): OpenAIModel {
   };
 }
 
+function aiDecideModelTarget(candidates: readonly ModelTarget[]): ModelTarget {
+  const has = (capability: keyof ModelCapabilities) =>
+    candidates.some((target) => target.capabilities[capability]);
+  const reasoningEfforts = Object.values(ReasoningEffort).filter((effort) =>
+    candidates.some((target) => target.reasoningEfforts.includes(effort)),
+  );
+  return {
+    id: AI_DECIDE_MODEL_ID,
+    aliases: [AI_DECIDE_MODEL_ID],
+    displayName: "AI Decide",
+    endpoint: {
+      name: AI_DECIDE_MODEL_ID,
+      displayName: "AI Decide",
+      description: "Uses Databricks ai_decide to choose a sufficient foundation model.",
+      task: "llm/v1/chat",
+      class: ModelClass.ChatBalanced,
+      modelSource: "foundation",
+      status: { deprecated: false },
+    },
+    capabilities: {
+      responses: true,
+      openResponses: true,
+      chat: true,
+      anthropic: true,
+      embeddings: false,
+      aiGatewayCodex: true,
+      tools: has("tools"),
+      reasoning: reasoningEfforts.length > 0,
+      streaming: true,
+      parallelTools: has("parallelTools"),
+      customTools: has("customTools"),
+      structuredOutput: has("structuredOutput"),
+      webSearch: has("webSearch"),
+      imageInput: has("imageInput"),
+      applyPatch: has("applyPatch"),
+    },
+    reasoningEfforts,
+  };
+}
+
 function smartRouterModelTarget(candidates: readonly ModelTarget[]): ModelTarget {
   const has = (capability: keyof ModelCapabilities) =>
     candidates.some((target) => target.capabilities[capability]);
@@ -151,6 +210,7 @@ function smartRouterModelTarget(candidates: readonly ModelTarget[]): ModelTarget
         "Uses Databricks task_v3 Smart Router to select a compatible system model for each turn.",
       task: "llm/v1/chat",
       class: ModelClass.ChatBalanced,
+      modelSource: "foundation",
       status: { deprecated: false },
     },
     capabilities: {
@@ -187,6 +247,7 @@ function codexModel(
   target: ModelTarget,
   priority: number,
   preferSystemModels: boolean,
+  includeMetadata: boolean,
 ): CodexModel {
   const efforts = target.reasoningEfforts;
   return {
@@ -195,6 +256,7 @@ function codexModel(
     description: target.endpoint?.description ?? "Databricks model service",
     base_instructions: CODEX_BASE_INSTRUCTIONS,
     status: target.endpoint?.status ?? { deprecated: false },
+    ...(includeMetadata ? { metadata: modelMetadataFor(target.endpoint ?? target.id) } : {}),
     supported_reasoning_levels: efforts.map((effort) => ({
       effort,
       description: REASONING_DESCRIPTIONS[effort] ?? `Use the ${effort} reasoning budget`,

package/src/reasoning-adapt.ts

@@ -14,7 +14,7 @@ import {
   adaptRequestReasoning,
   parseReasoningLevels,
 } from "@dbx-tools/model";
-import { log } from "@dbx-tools/shared-core";
+import { isRecord, log } from "@dbx-tools/shared-core";
 
 const logger = log.logger("appkit/model-gateway/reasoning");
 const INVALID_ENCRYPTED_REASONING_PATTERN =
@@ -123,7 +123,3 @@ function stripRequestReasoning(
   delete stripped.reasoning_effort;
   return stripped;
 }
-
-function isRecord(value: unknown): value is Readonly<Record<string, unknown>> {
-  return typeof value === "object" && value !== null && !Array.isArray(value);
-}

package/src/routes.ts

@@ -10,6 +10,7 @@ import { log, toError } from "@dbx-tools/shared-core";
 import type { ClientProtocol } from "@dbx-tools/shared-model-gateway";
 import type express from "express";
 
+import { AiDecideRouterError } from "./ai-decide.ts";
 import { ModelGateway, ModelGatewayError } from "./gateway.ts";
 import { UnsupportedGatewayFeatureError } from "./router.ts";
 import { SmartRouterError } from "./smart-router.ts";
@@ -33,6 +34,8 @@ export async function sendModels(
     const payload = await gateway.models(
       request.header("originator"),
       optionalSearch(request.query.search),
+      optionalModelSource(request.query.model_source ?? request.query.modelSource),
+      optionalBoolean(request.query.metadata, "metadata"),
     );
     response.json(payload);
     logger.debug("models response", {
@@ -130,6 +133,7 @@ export async function sendInference(
 function sendError(protocol: ClientProtocol, response: express.Response, error: unknown): void {
   const status =
     error instanceof ModelGatewayError ||
+    error instanceof AiDecideRouterError ||
     error instanceof SmartRouterError ||
     error instanceof UnsupportedGatewayFeatureError
       ? error.status
@@ -175,6 +179,18 @@ function optionalSearch(value: unknown): string | undefined {
   return value;
 }
 
+function optionalModelSource(value: unknown): "foundation" | "custom" | "all" {
+  if (value === undefined) return "foundation";
+  if (value === "foundation" || value === "custom" || value === "all") return value;
+  throw new ModelGatewayError(400, "model_source must be foundation, custom, or all");
+}
+
+function optionalBoolean(value: unknown, name: string): boolean {
+  if (value === undefined || value === "false") return false;
+  if (value === "true") return true;
+  throw new ModelGatewayError(400, `${name} must be true or false`);
+}
+
 function requestContext(request: express.Request): Record<string, unknown> {
   return {
     method: request.method,

package/src/smart-router.ts

@@ -22,10 +22,9 @@ import {
   type SmartRouterObservation,
   type SmartRouterRequest,
 } from "@dbx-tools/shared-model-gateway";
-import { Tiktoken } from "js-tiktoken/lite";
-import o200kBase from "js-tiktoken/ranks/o200k_base";
 import { z } from "zod";
 
+import { normalizeModelContext, type NormalizedModelContext } from "./model-context.ts";
 import { resolveRoute } from "./router.ts";
 
 const logger = log.logger("appkit/model-gateway/smart-router");
@@ -35,7 +34,6 @@ const PROBE_CACHE_NAMESPACE = "model-gateway-smart-router-probe";
 const PROBE_CACHE_VERSION = 1;
 const PROBE_TASK = "Check whether Smart Router can select an available model.";
 const NEGATIVE_PROBE_TTL_MS = 30_000;
-const CONTEXT_ENCODING = new Tiktoken(o200kBase);
 
 /** Runtime settings for Smart Router. */
 export interface SmartRouterSettings {
@@ -109,11 +107,7 @@ export interface SmartRouterResult {
 }
 
 /** Bounded task context extracted from an arbitrary supported input shape. */
-export interface SmartRouterContext {
-  readonly continuation: boolean;
-  readonly logicalTurn: string;
-  readonly task: string;
-}
+export type SmartRouterContext = NormalizedModelContext;
 
 /** Error carrying an HTTP status for Smart Router validation or availability failures. */
 export class SmartRouterError extends Error {
@@ -210,7 +204,7 @@ export class SmartRouterService {
     candidates: readonly SmartRouterCandidate<T>[],
     options: SmartRouterSelectOptions = {},
   ): Promise<SmartRouterSelection<T>> {
-    const task = this.context.build(input);
+    const task = await this.context.build(input);
     if (!task) throw new SmartRouterError(400, "Smart Router input contains no routable text");
     const client = options.client ?? getExecutionContext().client;
     try {
@@ -450,18 +444,13 @@ export class SmartRouterContextBuilder {
   constructor(private readonly tokenBudget = 4_096) {}
 
   /** Return only the bounded task prompt required by routes:select. */
-  build(input: unknown): string {
-    return this.analyze(input).task;
+  async build(input: unknown): Promise<string> {
+    return (await this.analyze(input)).task;
   }
 
   /** Return bounded task text plus logical-turn metadata used for pinning. */
-  analyze(input: unknown): SmartRouterContext {
-    const normalized = parseJsonString(input);
-    const events = conversationEvents(normalized);
-    if (events.length > 0) return contextFromEvents(events, this.tokenBudget);
-    const values = collectTaskStrings(normalized);
-    const task = truncateText(values.join("\n"), this.tokenBudget, "start");
-    return { task, logicalTurn: task, continuation: false };
+  analyze(input: unknown): Promise<SmartRouterContext> {
+    return normalizeModelContext(input, { tokenLimit: this.tokenBudget });
   }
 }
 
@@ -520,7 +509,7 @@ export class SmartModelRouter {
       );
     }
 
-    const context = this.service.context.analyze(input.body);
+    const context = await this.service.context.analyze(input.body);
     if (!context.task) {
       throw new SmartRouterError(400, "Smart Router request contains no task text");
     }
@@ -670,174 +659,6 @@ function requestedReasoningEffort(
   return parsed.data;
 }
 
-interface ConversationEvent {
-  readonly role: "assistant" | "system" | "tool" | "user";
-  readonly text: string;
-}
-
-function contextFromEvents(
-  events: readonly ConversationEvent[],
-  tokenBudget: number,
-): SmartRouterContext {
-  const latestUserIndex = lastEventIndex(events, "user");
-  const taskEvent = latestUserIndex >= 0 ? events[latestUserIndex] : events.at(-1);
-  const task = truncateText(taskEvent?.text ?? "", tokenBudget, "start");
-  let remaining = Math.max(0, tokenBudget - countTokens(task));
-  const context: string[] = [];
-  for (
-    let index = (taskEvent ? events.indexOf(taskEvent) : events.length) - 1;
-    index >= 0;
-    index -= 1
-  ) {
-    const event = events[index];
-    if (!event || !event.text || event.role === "tool") continue;
-    const value = `${event.role}: ${event.text}`;
-    const retained = truncateText(value, remaining, "end");
-    if (!retained) break;
-    context.unshift(retained);
-    remaining -= countTokens(retained);
-  }
-  const combined = [...context, task].filter(Boolean).join("\n");
-  const logicalEvents = latestUserIndex >= 0 ? events.slice(0, latestUserIndex + 1) : events;
-  return {
-    task: combined,
-    logicalTurn: logicalEvents.map((event) => `${event.role}:${event.text}`).join("\n"),
-    continuation:
-      latestUserIndex >= 0 &&
-      events.slice(latestUserIndex + 1).some((event) => event.role !== "system"),
-  };
-}
-
-function conversationEvents(input: unknown): ConversationEvent[] {
-  if (typeof input === "string") return input.trim() ? [{ role: "user", text: input.trim() }] : [];
-  if (Array.isArray(input)) return input.flatMap(messageEvent);
-  if (!isRecord(input)) return [];
-  for (const key of ["messages", "input"] as const) {
-    const value = input[key];
-    if (typeof value === "string" && value.trim()) return [{ role: "user", text: value.trim() }];
-    if (Array.isArray(value)) {
-      const events = value.flatMap(messageEvent);
-      if (events.length > 0) return events;
-    }
-  }
-  return [];
-}
-
-function messageEvent(value: unknown): ConversationEvent[] {
-  if (typeof value === "string") return value.trim() ? [{ role: "user", text: value.trim() }] : [];
-  if (!isRecord(value)) return [];
-  const role = messageRole(value);
-  if (!role) return [];
-  const text = contentText(value.content ?? value.text ?? value.input_text ?? value.output_text);
-  return text ? [{ role, text }] : [];
-}
-
-function messageRole(
-  value: Readonly<Record<string, unknown>>,
-): ConversationEvent["role"] | undefined {
-  if (value.role === "user" || value.role === "assistant" || value.role === "system") {
-    return value.role;
-  }
-  if (
-    value.role === "tool" ||
-    value.type === "function_call_output" ||
-    value.type === "tool_result"
-  ) {
-    return "tool";
-  }
-  if (value.type === "function_call" || value.type === "reasoning") return "assistant";
-  if (value.type === "message" && value.role === undefined) return "user";
-  return undefined;
-}
-
-function contentText(value: unknown): string {
-  if (typeof value === "string") return value.trim();
-  if (Array.isArray(value)) return value.map(contentText).filter(Boolean).join("\n");
-  if (!isRecord(value)) return "";
-  for (const key of ["text", "input_text", "output_text", "content"] as const) {
-    const text = contentText(value[key]);
-    if (text) return text;
-  }
-  return "";
-}
-
-const TASK_KEYS = [
-  "prompt",
-  "task",
-  "query",
-  "question",
-  "request",
-  "description",
-  "message",
-  "text",
-  "input",
-] as const;
-const IGNORED_KEYS = new Set([
-  "api_key",
-  "authorization",
-  "arguments",
-  "credentials",
-  "headers",
-  "metadata",
-  "password",
-  "secret",
-  "token",
-  "tools",
-]);
-
-function collectTaskStrings(input: unknown): string[] {
-  if (typeof input === "string") return input.trim() ? [input.trim()] : [];
-  const preferred: string[] = [];
-  if (isRecord(input)) {
-    for (const key of TASK_KEYS) preferred.push(...stringLeaves(input[key], key, 0));
-  }
-  return preferred.length > 0 ? preferred.slice(0, 32) : stringLeaves(input, "", 0).slice(0, 32);
-}
-
-function stringLeaves(value: unknown, key: string, depth: number): string[] {
-  if (depth > 6 || IGNORED_KEYS.has(key.toLowerCase())) return [];
-  if (typeof value === "string") return value.trim() ? [value.trim()] : [];
-  if (Array.isArray(value)) return value.flatMap((item) => stringLeaves(item, key, depth + 1));
-  if (!isRecord(value)) return [];
-  return Object.entries(value).flatMap(([childKey, child]) =>
-    stringLeaves(child, childKey, depth + 1),
-  );
-}
-
-function parseJsonString(input: unknown): unknown {
-  if (typeof input !== "string") return input;
-  const trimmed = input.trim();
-  if (!trimmed || (!trimmed.startsWith("{") && !trimmed.startsWith("["))) return input;
-  try {
-    return JSON.parse(trimmed) as unknown;
-  } catch {
-    return input;
-  }
-}
-
-function countTokens(text: string): number {
-  return CONTEXT_ENCODING.encode(text, [], []).length;
-}
-
-function truncateText(text: string, tokenBudget: number, edge: "start" | "end"): string {
-  if (tokenBudget <= 0 || !text) return "";
-  const tokens = CONTEXT_ENCODING.encode(text, [], []);
-  if (tokens.length <= tokenBudget) return text;
-  return CONTEXT_ENCODING.decode(
-    edge === "start" ? tokens.slice(0, tokenBudget) : tokens.slice(-tokenBudget),
-  );
-}
-
-function lastEventIndex(
-  events: readonly ConversationEvent[],
-  role: ConversationEvent["role"],
-): number {
-  for (let index = events.length - 1; index >= 0; index -= 1) {
-    if (events[index]?.role === role) return index;
-  }
-  return -1;
-}
-
 function uniqueCandidates<T>(
   candidates: readonly SmartRouterCandidate<T>[],
 ): SmartRouterCandidate<T>[] {