@dbx-tools/appkit-model-gateway 0.9.190 → 0.9.194

package/src/gateway.ts

@@ -9,10 +9,12 @@ import { log } from "@dbx-tools/shared-core";
 import {
   ModelClass,
   ModelClassSchema,
+  type ModelSource,
   type ServingEndpointSummary,
   type ModelClass as ModelClassType,
 } from "@dbx-tools/shared-model";
 import type {
+  AiDecideObservation,
   ClientProtocol,
   GatewayRoute,
   ModelCapabilityOverride,
@@ -21,6 +23,11 @@ import type {
   SmartRouterMode,
 } from "@dbx-tools/shared-model-gateway";
 import {
+  AI_DECIDE_CONFIDENCE_HEADER,
+  AI_DECIDE_MODEL_ID,
+  AI_DECIDE_SELECTED_MODEL_HEADER,
+  AI_DECIDE_STRATEGY_HEADER,
+  AiDecideStrategySchema,
   MODEL_GATEWAY_DEFAULTS,
   SMART_ROUTER_HARNESS_HEADER,
   SMART_ROUTER_MODEL_ID,
@@ -30,9 +37,11 @@ import {
   SMART_ROUTER_SERVED_TIER_HEADER,
   SMART_ROUTER_SOURCE_HEADER,
   SMART_ROUTER_TASK_TIER_HEADER,
+  isAiDecideModelId,
   isSmartRouterModelId,
 } from "@dbx-tools/shared-model-gateway";
 
+import { AiDecideRouterService } from "./ai-decide.ts";
 import { listModelsPayload } from "./models.ts";
 import { adaptInferenceReasoning, learnAndAdaptReasoningRetry } from "./reasoning-adapt.ts";
 import {
@@ -64,6 +73,7 @@ export interface ModelGatewayOptions {
   readonly smartRouterTimeoutMs?: number;
   readonly smartRouterService?: SmartRouterService;
   readonly smartRouterStore?: SmartRouterStore;
+  readonly aiDecideService?: AiDecideRouterService;
   readonly overrides?: readonly ModelCapabilityOverride[];
   readonly model?: string;
   readonly modelClass?: ModelClassType;
@@ -85,6 +95,7 @@ export class ModelGatewayError extends Error {
 export class ModelGateway {
   private readonly registry: ModelRegistry;
   private readonly smartRouter?: SmartModelRouter;
+  private readonly aiDecide: AiDecideRouterService;
   private readonly smartRouterMode: SmartRouterMode;
   private readonly model?: string;
   private readonly modelClass?: ModelClassType;
@@ -99,6 +110,7 @@ export class ModelGateway {
         : {}),
     };
     this.registry = registry ?? new DatabricksModelRegistry(registryOptions);
+    this.aiDecide = options.aiDecideService ?? new AiDecideRouterService();
     this.smartRouterMode = options.smartRouter ?? MODEL_GATEWAY_DEFAULTS.smartRouter;
     if (this.smartRouterMode !== "false") {
       this.smartRouter = new SmartModelRouter(
@@ -126,22 +138,43 @@ export class ModelGateway {
   }
 
   /** Return the live OpenAI model list, including Codex metadata when requested. */
-  async models(originator?: string, search?: string): Promise<ModelListResponse> {
+  async models(
+    originator?: string,
+    search?: string,
+    modelSource: ModelSource | "all" = "foundation",
+    includeMetadata = false,
+  ): Promise<ModelListResponse> {
     const codex = isCodexOriginator(originator);
     logger.debug("listing models", { codex, search });
     const query = search?.trim();
     const smartRouterSearch =
       query !== undefined && SMART_ROUTER_MODEL_ID.includes(query.toLowerCase());
+    const aiDecideSearch = query !== undefined && AI_DECIDE_MODEL_ID.includes(query.toLowerCase());
+    const discovered =
+      query && !smartRouterSearch && !aiDecideSearch
+        ? await this.registry.search(query)
+        : await this.registry.list();
     const targets =
-      query && !smartRouterSearch ? await this.registry.search(query) : await this.registry.list();
+      modelSource === "all"
+        ? discovered
+        : discovered.filter((target) => target.endpoint?.modelSource === modelSource);
     const smartRouterTargets =
       this.smartRouterMode === "true" && query && !smartRouterSearch
-        ? await this.registry.list()
+        ? (await this.registry.list()).filter(
+            (target) => modelSource === "all" || target.endpoint?.modelSource === modelSource,
+          )
         : targets;
     const shouldProbeSmartRouter =
       Boolean(this.smartRouter) && (this.smartRouterMode === "true" || !query || smartRouterSearch);
     const smartRouterEnabled =
       shouldProbeSmartRouter && (await this.smartRouter!.isEnabled(smartRouterTargets));
+    const aiDecideEnabled =
+      (!query || aiDecideSearch) &&
+      (await this.aiDecide.isEnabled(
+        targets
+          .map((target) => target.endpoint)
+          .filter((endpoint): endpoint is ServingEndpointSummary => endpoint !== undefined),
+      ));
     if (this.smartRouterMode === "true" && !smartRouterEnabled) {
       throw new ModelGatewayError(
         503,
@@ -155,8 +188,12 @@ export class ModelGateway {
       query !== undefined && query.length > 0,
       this.preferSystemModels,
       smartRouterEnabled && (!query || smartRouterSearch),
+      aiDecideEnabled && (!query || aiDecideSearch),
+      includeMetadata,
     );
-    return smartRouterSearch ? onlySmartRouter(payload) : payload;
+    return smartRouterSearch || aiDecideSearch
+      ? onlyVirtualModel(payload, smartRouterSearch ? SMART_ROUTER_MODEL_ID : AI_DECIDE_MODEL_ID)
+      : payload;
   }
 
   /** Resolve and execute one inference request. */
@@ -178,6 +215,7 @@ export class ModelGateway {
     );
     logger.debug("resolving request", {
       smartRouter: selection.smartRouter === true,
+      aiDecide: selection.aiDecide === true,
       clientProtocol: protocol,
       model: selection.model,
       modelClass: selection.modelClass,
@@ -185,8 +223,11 @@ export class ModelGateway {
       stream: sanitizedBody.stream === true,
     });
     const features = requestedFeatures(sanitizedBody);
-    if (selection.smartRouter && protocol === "openai-embeddings") {
-      throw new ModelGatewayError(400, "Model smart-router is available only for chat requests");
+    if ((selection.smartRouter || selection.aiDecide) && protocol === "openai-embeddings") {
+      throw new ModelGatewayError(
+        400,
+        "Virtual model routers are available only for chat requests",
+      );
     }
     const smartRouter =
       selection.smartRouter && this.smartRouter
@@ -203,8 +244,29 @@ export class ModelGateway {
     if (selection.smartRouter && !smartRouter) {
       throw new ModelGatewayError(404, "Databricks Smart Router is disabled");
     }
+    const registryTargets = smartRouter ? undefined : await this.registry.list();
+    const aiDecide = selection.aiDecide
+      ? await this.aiDecide.select(
+          sanitizedBody,
+          (registryTargets ?? [])
+            .map((candidate) => candidate.endpoint)
+            .filter((endpoint): endpoint is ServingEndpointSummary => endpoint !== undefined),
+          undefined,
+          {
+            strategy: aiDecideStrategy(headers),
+            filters: {
+              ...(features.tools ? { requiresTools: true } : {}),
+              ...(features.structuredOutput ? { requiresStructuredOutput: true } : {}),
+            },
+            signal,
+          },
+        )
+      : undefined;
     const target =
       smartRouter?.target ??
+      (aiDecide
+        ? registryTargets?.find((candidate) => candidate.id === aiDecide.endpoint.name)
+        : undefined) ??
       (await this.registry.resolve(selection.model, {
         ...(protocol === "openai-embeddings" ? { modelClass: ModelClass.Embedding } : {}),
         ...(selection.modelClass ? { modelClass: selection.modelClass } : {}),
@@ -220,7 +282,7 @@ export class ModelGateway {
     }
     const route = resolveRoute({
       clientProtocol: protocol,
-      ...(!smartRouter && selection.model ? { requestedModel: selection.model } : {}),
+      ...(!smartRouter && !aiDecide && selection.model ? { requestedModel: selection.model } : {}),
       originator: headers.get("originator") ?? undefined,
       preferSystemModels: this.preferSystemModels,
       features,
@@ -255,11 +317,13 @@ export class ModelGateway {
       if (smartRouter) {
         applySmartRouterObservationHeaders(response.headers, smartRouter.observation);
       }
+      if (aiDecide) applyAiDecideObservationHeaders(response.headers, aiDecide.observation);
       return response;
     }
     const forwarded = gatewayResponseHeaders(response.headers);
     forwarded.set("x-dbx-tools-upstream", route.upstreamProtocol);
     if (smartRouter) applySmartRouterObservationHeaders(forwarded, smartRouter.observation);
+    if (aiDecide) applyAiDecideObservationHeaders(forwarded, aiDecide.observation);
     return new Response(response.body, {
       status: response.status,
       statusText: response.statusText,
@@ -330,6 +394,7 @@ export function gatewayModelSelection(
   protocol?: ClientProtocol,
 ): {
   readonly smartRouter?: boolean;
+  readonly aiDecide?: boolean;
   readonly model?: string;
   readonly modelClass?: ModelClassType;
 } {
@@ -342,10 +407,16 @@ export function gatewayModelSelection(
   if (model && parsedClass?.data) {
     throw new ModelGatewayError(400, "Model and model class are mutually exclusive");
   }
-  if (model) return isSmartRouterModelId(model) ? { smartRouter: true } : { model };
+  if (model) {
+    if (isSmartRouterModelId(model)) return { smartRouter: true };
+    if (isAiDecideModelId(model)) return { aiDecide: true };
+    return { model };
+  }
   if (parsedClass?.data) return { modelClass: parsedClass.data };
   if (defaults.model) {
-    return isSmartRouterModelId(defaults.model) ? { smartRouter: true } : { model: defaults.model };
+    if (isSmartRouterModelId(defaults.model)) return { smartRouter: true };
+    if (isAiDecideModelId(defaults.model)) return { aiDecide: true };
+    return { model: defaults.model };
   }
   if (defaults.modelClass) return { modelClass: defaults.modelClass };
   if (protocol === "openai-embeddings") return { modelClass: ModelClass.Embedding };
@@ -368,6 +439,19 @@ function applySmartRouterObservationHeaders(
   if (profilerStatus) headers.set(SMART_ROUTER_PROFILER_STATUS_HEADER, profilerStatus);
 }
 
+function applyAiDecideObservationHeaders(headers: Headers, observation: AiDecideObservation): void {
+  headers.set(AI_DECIDE_SELECTED_MODEL_HEADER, observation.selectedModel);
+  headers.set(AI_DECIDE_STRATEGY_HEADER, observation.strategy);
+  if (observation.confidence !== undefined) {
+    headers.set(AI_DECIDE_CONFIDENCE_HEADER, String(observation.confidence));
+  }
+}
+
+function aiDecideStrategy(headers: Headers) {
+  const value = headers.get(AI_DECIDE_STRATEGY_HEADER)?.trim().toLowerCase();
+  return value ? AiDecideStrategySchema.parse(value) : "balanced";
+}
+
 function stringParam(
   params: Readonly<Record<string, unknown>> | undefined,
   key: string,
@@ -376,17 +460,15 @@ function stringParam(
   return typeof value === "string" && value ? value : undefined;
 }
 
-function onlySmartRouter(payload: ModelListResponse): ModelListResponse {
+function onlyVirtualModel(payload: ModelListResponse, modelId: string): ModelListResponse {
   if ("models" in payload) {
     return {
-      models: payload.models.filter((model) =>
-        model.slug.toLowerCase().endsWith(`/${SMART_ROUTER_MODEL_ID}`),
-      ),
+      models: payload.models.filter((model) => model.slug.toLowerCase().endsWith(`/${modelId}`)),
     };
   }
   return {
     object: "list",
-    data: payload.data.filter((model) => model.id === SMART_ROUTER_MODEL_ID),
+    data: payload.data.filter((model) => model.id === modelId),
   };
 }
 

package/src/model-context.ts

@@ -0,0 +1,306 @@
+/**
+ * Forgiving model-context normalization for routing and suggestion generation.
+ *
+ * @module
+ */
+
+import { isRecord } from "@dbx-tools/shared-core";
+import {
+  convertToModelMessages,
+  modelMessageSchema,
+  pruneMessages,
+  safeValidateUIMessages,
+  type ModelMessage,
+} from "ai";
+import { Tiktoken } from "js-tiktoken/lite";
+import o200kBase from "js-tiktoken/ranks/o200k_base";
+
+const CONTEXT_ENCODING = new Tiktoken(o200kBase);
+const DEFAULT_TOKEN_LIMIT = 4_096;
+
+/** One useful conversational event retained from arbitrary caller input. */
+export interface ModelContextEvent {
+  readonly role: "assistant" | "system" | "tool" | "user";
+  readonly text: string;
+}
+
+/** Controls for converting arbitrary values into bounded model context. */
+export interface NormalizeModelContextOptions {
+  /** Maximum retained tokens. Defaults to 4096. */
+  readonly tokenLimit?: number;
+  /** Remove reasoning, tool calls/results, metadata, and other non-conversational noise. */
+  readonly removeNonUsefulMessages?: boolean;
+}
+
+/** Bounded context plus logical-turn metadata used by model routers. */
+export interface NormalizedModelContext {
+  readonly continuation: boolean;
+  readonly events: readonly ModelContextEvent[];
+  readonly logicalTurn: string;
+  readonly task: string;
+  readonly tokenCount: number;
+}
+
+/**
+ * Normalize strings, JSON, AI SDK messages, Responses input, and task-shaped objects.
+ *
+ * AI SDK UI/model message parsing is preferred. The fallback handles external wire
+ * shapes and arbitrary JSON without retaining credentials, metadata, or tool arguments.
+ */
+export async function normalizeModelContext(
+  input: unknown,
+  options: NormalizeModelContextOptions = {},
+): Promise<NormalizedModelContext> {
+  const tokenLimit = Math.max(0, options.tokenLimit ?? DEFAULT_TOKEN_LIMIT);
+  const removeNonUsefulMessages = options.removeNonUsefulMessages ?? true;
+  const normalized = parseJsonString(input);
+  const conversation = conversationValue(normalized);
+  const sdkMessages = await aiSdkModelMessages(conversation);
+  const events = sdkMessages
+    ? modelMessageEvents(sdkMessages, removeNonUsefulMessages)
+    : conversationEvents(conversation, removeNonUsefulMessages);
+  const rawEvents = conversationEvents(conversation, false);
+  const stateEvents = rawEvents.length > 0 ? rawEvents : events;
+  if (events.length > 0) return contextFromEvents(events, tokenLimit, stateEvents);
+
+  const values = collectTaskStrings(normalized);
+  const task = truncateText(values.join("\n"), tokenLimit, "start");
+  return {
+    task,
+    logicalTurn: task,
+    continuation: false,
+    events: task ? [{ role: "user", text: task }] : [],
+    tokenCount: countTokens(task),
+  };
+}
+
+async function aiSdkModelMessages(input: unknown): Promise<ModelMessage[] | undefined> {
+  if (!Array.isArray(input) || input.length === 0) return undefined;
+  const modelMessages = input.map((value) => modelMessageSchema.safeParse(value));
+  if (modelMessages.every((result) => result.success)) {
+    return modelMessages.map((result) => result.data);
+  }
+
+  const uiMessages = await safeValidateUIMessages({ messages: input });
+  if (!uiMessages.success) return undefined;
+  try {
+    return await convertToModelMessages(uiMessages.data, { ignoreIncompleteToolCalls: true });
+  } catch {
+    return undefined;
+  }
+}
+
+function modelMessageEvents(
+  messages: ModelMessage[],
+  removeNonUsefulMessages: boolean,
+): ModelContextEvent[] {
+  const useful = removeNonUsefulMessages
+    ? pruneMessages({ messages, reasoning: "all", toolCalls: "all" })
+    : messages;
+  return useful.flatMap((message) => {
+    const role = message.role;
+    if (role === "tool" && removeNonUsefulMessages) return [];
+    const text = modelContentText(message.content, removeNonUsefulMessages);
+    return text ? [{ role, text }] : [];
+  });
+}
+
+function modelContentText(value: unknown, removeNonUsefulMessages: boolean): string {
+  if (typeof value === "string") return value.trim();
+  if (!Array.isArray(value)) return "";
+  return value
+    .flatMap((part) => {
+      if (!isRecord(part)) return [];
+      if (typeof part.text === "string") return [part.text.trim()];
+      if (!removeNonUsefulMessages && typeof part.output === "string") return [part.output.trim()];
+      return [];
+    })
+    .filter(Boolean)
+    .join("\n");
+}
+
+function conversationValue(input: unknown): unknown {
+  if (!isRecord(input)) return input;
+  for (const key of ["messages", "input"] as const) {
+    const value = input[key];
+    if (Array.isArray(value)) return value;
+    if (typeof value === "string" && value.trim()) return value;
+  }
+  return input;
+}
+
+function contextFromEvents(
+  events: readonly ModelContextEvent[],
+  tokenLimit: number,
+  stateEvents: readonly ModelContextEvent[] = events,
+): NormalizedModelContext {
+  const latestUserIndex = lastEventIndex(events, "user");
+  const taskEvent = latestUserIndex >= 0 ? events[latestUserIndex] : events.at(-1);
+  const task = truncateText(taskEvent?.text ?? "", tokenLimit, "start");
+  let remaining = Math.max(0, tokenLimit - 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 stateLatestUserIndex = lastEventIndex(stateEvents, "user");
+  const logicalEvents =
+    stateLatestUserIndex >= 0 ? stateEvents.slice(0, stateLatestUserIndex + 1) : stateEvents;
+  return {
+    task: combined,
+    logicalTurn: logicalEvents.map((event) => `${event.role}:${event.text}`).join("\n"),
+    continuation:
+      stateLatestUserIndex >= 0 &&
+      stateEvents.slice(stateLatestUserIndex + 1).some((event) => event.role !== "system"),
+    events,
+    tokenCount: countTokens(combined),
+  };
+}
+
+function conversationEvents(input: unknown, removeNonUsefulMessages: boolean): ModelContextEvent[] {
+  if (typeof input === "string") return input.trim() ? [{ role: "user", text: input.trim() }] : [];
+  if (Array.isArray(input)) {
+    return input.flatMap((value) => messageEvent(value, removeNonUsefulMessages));
+  }
+  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((item) => messageEvent(item, removeNonUsefulMessages));
+      if (events.length > 0) return events;
+    }
+  }
+  return [];
+}
+
+function messageEvent(value: unknown, removeNonUsefulMessages: boolean): ModelContextEvent[] {
+  if (typeof value === "string") return value.trim() ? [{ role: "user", text: value.trim() }] : [];
+  if (!isRecord(value)) return [];
+  const role = messageRole(value);
+  if (!role || (removeNonUsefulMessages && role === "tool")) return [];
+  if (removeNonUsefulMessages && (value.type === "function_call" || value.type === "reasoning")) {
+    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>>,
+): ModelContextEvent["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",
+  "provider_options",
+  "raw_input",
+  "secret",
+  "token",
+  "tool_calls",
+  "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, tokenLimit: number, edge: "start" | "end"): string {
+  if (tokenLimit <= 0 || !text) return "";
+  const tokens = CONTEXT_ENCODING.encode(text, [], []);
+  if (tokens.length <= tokenLimit) return text;
+  return CONTEXT_ENCODING.decode(
+    edge === "start" ? tokens.slice(0, tokenLimit) : tokens.slice(-tokenLimit),
+  );
+}
+
+function lastEventIndex(
+  events: readonly ModelContextEvent[],
+  role: ModelContextEvent["role"],
+): number {
+  for (let index = events.length - 1; index >= 0; index -= 1) {
+    if (events[index]?.role === role) return index;
+  }
+  return -1;
+}