basant307/AI_Governance_Project

1/**2 * @license3 * Copyright 2025 Qwen Team4 * SPDX-License-Identifier: Apache-2.05 */6 7import { AuthType } from '../core/contentGenerator.js';8import { defaultModalities } from '../core/modalityDefaults.js';9import { tokenLimit } from '../core/tokenLimits.js';10import { DEFAULT_OPENAI_BASE_URL } from '../core/openaiContentGenerator/constants.js';11import {12  type ModelConfig,13  type ModelProvidersConfig,14  type ProviderProtocolConfig,15  type ResolvedModelConfig,16  type AvailableModel,17} from './types.js';18import { DEFAULT_QWEN_MODEL } from '../config/models.js';19import { QWEN_OAUTH_MODELS } from './constants.js';20import { createDebugLogger } from '../utils/debugLogger.js';21 22const debugLogger = createDebugLogger('MODEL_REGISTRY');23 24export { QWEN_OAUTH_MODELS } from './constants.js';25 26/**27 * Validates if a string key is a valid AuthType enum value.28 * @param key - The key to validate29 * @returns The validated AuthType or undefined if invalid30 */31function validateAuthTypeKey(key: string): AuthType | undefined {32  // Check if the key is a valid AuthType enum value33  if (Object.values(AuthType).includes(key as AuthType)) {34    return key as AuthType;35  }36 37  // Invalid key38  return undefined;39}40 41/**42 * Resolve the SDK protocol (an {@link AuthType}) that should route a43 * `modelProviders` provider id.44 *45 * Precedence:46 *  1. An explicit {@link ProviderProtocolConfig} entry for the provider id.47 *  2. The provider id itself when it is already a built-in protocol48 *     (e.g. `openai`, `gemini`) — preserves the pre-existing behavior.49 *50 * Returns `undefined` for an unknown provider id with no mapping, or an explicit51 * mapping whose value is not a known protocol, so the caller skips it (keeping52 * the typo guard for hand-edited settings). Pure: callers decide how loudly to53 * report a skip. Additive — configs without `providerProtocol` behave as before.54 */55export function resolveProviderProtocol(56  providerId: string,57  providerProtocol?: ProviderProtocolConfig,58): AuthType | undefined {59  const explicit =60    providerProtocol && Object.hasOwn(providerProtocol, providerId)61      ? providerProtocol[providerId]62      : undefined;63  if (explicit !== undefined) {64    return validateAuthTypeKey(explicit);65  }66  return validateAuthTypeKey(providerId);67}68 69function shouldUseCanonicalModalities(modelId: string): boolean {70  return /^minimax-m3/i.test(modelId.trim().toLowerCase());71}72 73/**74 * Build a composite registry key from model id and optional baseUrl.75 * Two models with the same id but different baseUrls are distinct entries.76 * When baseUrl is omitted/empty the key is just the id (backward compatible).77 */78export function modelRegistryKey(id: string, baseUrl?: string): string {79  return baseUrl ? `${id}\0${baseUrl}` : id;80}81 82/**83 * Central registry for managing model configurations.84 * Models are organized by authType.85 */86export class ModelRegistry {87  private modelsByAuthType: Map<AuthType, Map<string, ResolvedModelConfig>>;88 89  /** providerId -> SDK protocol mapping; persists across reloads. */90  private providerProtocolConfig: ProviderProtocolConfig;91 92  private getDefaultBaseUrl(authType: AuthType): string {93    switch (authType) {94      case AuthType.QWEN_OAUTH:95        return 'DYNAMIC_QWEN_OAUTH_BASE_URL';96      case AuthType.USE_OPENAI:97        return DEFAULT_OPENAI_BASE_URL;98      default:99        return '';100    }101  }102 103  constructor(104    modelProvidersConfig?: ModelProvidersConfig,105    providerProtocolConfig?: ProviderProtocolConfig,106  ) {107    this.modelsByAuthType = new Map();108    this.providerProtocolConfig = providerProtocolConfig ?? {};109 110    // Always register qwen-oauth models (hard-coded, cannot be overridden)111    this.registerAuthTypeModels(AuthType.QWEN_OAUTH, QWEN_OAUTH_MODELS);112 113    // Register user-configured models for other providers114    this.registerProvidersConfig(modelProvidersConfig);115  }116 117  /**118   * Register every user-configured provider under its resolved SDK protocol.119   * A provider id maps to a protocol via {@link resolveProviderProtocol}120   * (explicit `providerProtocol` entry, or the id itself when it is a built-in121   * protocol). Unmapped unknown ids are skipped with a warning.122   */123  private registerProvidersConfig(124    modelProvidersConfig?: ModelProvidersConfig,125  ): void {126    if (!modelProvidersConfig) return;127 128    for (const [providerId, models] of Object.entries(modelProvidersConfig)) {129      const protocol = resolveProviderProtocol(130        providerId,131        this.providerProtocolConfig,132      );133 134      if (!protocol) {135        const knownProtocols = Object.values(AuthType).join(', ');136        const mapped = Object.hasOwn(this.providerProtocolConfig, providerId)137          ? this.providerProtocolConfig[providerId]138          : undefined;139        const message =140          mapped !== undefined141            ? `Provider "${providerId}" maps to "${mapped}" via providerProtocol, ` +142              `which is not a known protocol (${knownProtocols}); skipping.`143            : `Provider "${providerId}" in modelProviders is not a built-in protocol ` +144              `(${knownProtocols}) and has no providerProtocol mapping; skipping. ` +145              `Add providerProtocol["${providerId}"] to route it to an SDK protocol.`;146        debugLogger.warn(message);147        continue;148      }149 150      // qwen-oauth uses hard-coded models and cannot be overridden151      if (protocol === AuthType.QWEN_OAUTH) {152        continue;153      }154 155      this.registerAuthTypeModels(protocol, models, providerId);156    }157  }158 159  /**160   * Register models for an authType.161   * Uniqueness is determined by the composite key (id + baseUrl).162   * Two models with the same id but different baseUrls are treated as distinct.163   * If multiple models share both id and baseUrl, the first one takes precedence.164   */165  private registerAuthTypeModels(166    authType: AuthType,167    models: ModelConfig[],168    providerId?: string,169  ): void {170    // Defensive: runtime data from settings.json can violate the static type —171    // e.g. a hand-edited file, or one still in the reverted #5089 V5 shape172    // ({ protocol, models }) that the CLI v5->v4 migration has not yet173    // rewritten. Skip such entries with a clear warning instead of throwing an174    // opaque "models is not iterable" from the loop below.175    if (!Array.isArray(models)) {176      debugLogger.warn(177        `modelProviders for provider "${providerId ?? authType}" is not an array; ` +178          `skipping. Expected ModelConfig[]; legacy { protocol, models } entries ` +179          `are normally rewritten by the v5->v4 settings migration.`,180      );181      return;182    }183 184    // Merge into any existing map for this protocol: multiple provider ids can185    // resolve to the same protocol (e.g. `openai` and a custom `idealab` both186    // routing to the openai protocol). First registration of a composite187    // (id + baseUrl) key wins.188    const modelMap =189      this.modelsByAuthType.get(authType) ??190      new Map<string, ResolvedModelConfig>();191    const providerLabel =192      providerId && providerId !== authType193        ? ` (provider "${providerId}")`194        : '';195 196    for (const config of models) {197      const key = modelRegistryKey(config.id, config.baseUrl);198      if (modelMap.has(key)) {199        debugLogger.warn(200          `Duplicate model id "${config.id}"${config.baseUrl ? ` with baseUrl "${config.baseUrl}"` : ''} for protocol "${authType}"${providerLabel}. Using the first registered config.`,201        );202        continue;203      }204      const resolved = this.resolveModelConfig(config, authType);205      modelMap.set(key, resolved);206    }207 208    this.modelsByAuthType.set(authType, modelMap);209  }210 211  /**212   * Get all models for a specific authType.213   * This is used by /model command to show only relevant models.214   */215  getModelsForAuthType(authType: AuthType): AvailableModel[] {216    const models = this.modelsByAuthType.get(authType);217    if (!models) return [];218 219    return Array.from(models.values()).map((model) => ({220      id: model.id,221      label: model.name,222      description: model.description,223      capabilities: model.capabilities,224      authType: model.authType,225      isVision: model.capabilities?.vision ?? false,226      contextWindowSize:227        model.generationConfig.contextWindowSize ?? tokenLimit(model.id),228      // `modalities` is auto-filled in `resolveModelConfig`, so it is229      // always defined on `ResolvedModelConfig` — no fallback needed here.230      modalities: model.generationConfig.modalities,231      baseUrl: model.baseUrl,232      envKey: model.envKey,233      fastOnly: model.fastOnly,234      voiceOnly: model.voiceOnly,235    }));236  }237 238  /**239   * Get model configuration by authType and modelId.240   * When baseUrl is provided, looks up by the exact composite key (id+baseUrl).241   * When baseUrl is omitted, tries the plain id first (backward compatible),242   * then scans all entries for the first match by model id.243   */244  getModel(245    authType: AuthType,246    modelId: string,247    baseUrl?: string,248  ): ResolvedModelConfig | undefined {249    const models = this.modelsByAuthType.get(authType);250    if (!models) return undefined;251 252    if (baseUrl) {253      return models.get(modelRegistryKey(modelId, baseUrl));254    }255 256    // Try plain id key first (models registered without explicit baseUrl)257    const plain = models.get(modelId);258    if (plain) return plain;259 260    // Scan for the first entry with matching model id261    for (const model of models.values()) {262      if (model.id === modelId) return model;263    }264    return undefined;265  }266 267  /**268   * Check if model exists for given authType.269   * When baseUrl is provided, checks the exact composite key.270   * When baseUrl is omitted, checks plain id and scans by model id.271   */272  hasModel(authType: AuthType, modelId: string, baseUrl?: string): boolean {273    return this.getModel(authType, modelId, baseUrl) !== undefined;274  }275 276  /**277   * Get default model for an authType.278   * For qwen-oauth, returns the coder model.279   * For others, returns the first configured model.280   */281  getDefaultModelForAuthType(282    authType: AuthType,283  ): ResolvedModelConfig | undefined {284    if (authType === AuthType.QWEN_OAUTH) {285      return this.getModel(authType, DEFAULT_QWEN_MODEL);286    }287    const models = this.modelsByAuthType.get(authType);288    if (!models || models.size === 0) return undefined;289    return Array.from(models.values())[0];290  }291 292  /**293   * Resolve model config by applying defaults294   */295  private resolveModelConfig(296    config: ModelConfig,297    authType: AuthType,298  ): ResolvedModelConfig {299    this.validateModelConfig(config, authType);300 301    const generationConfig = { ...(config.generationConfig ?? {}) };302    // Auto-fill modalities from the model name when the provider didn't set303    // them explicitly. Without this, downstream consumers that read straight304    // from the registry (e.g. sub-agents via getResolvedModel) would inherit305    // the parent session's modalities instead of the agent's own.306    if (307      generationConfig.modalities === undefined ||308      shouldUseCanonicalModalities(config.id)309    ) {310      generationConfig.modalities = defaultModalities(config.id);311    }312 313    return {314      ...config,315      authType,316      name: config.name || config.id,317      baseUrl: config.baseUrl || this.getDefaultBaseUrl(authType),318      generationConfig,319      capabilities: config.capabilities || {},320    };321  }322 323  /**324   * Validate model configuration325   */326  private validateModelConfig(config: ModelConfig, authType: AuthType): void {327    if (!config.id) {328      throw new Error(329        `Model config in authType '${authType}' missing required field: id`,330      );331    }332    if (config.fastOnly && config.voiceOnly) {333      debugLogger.warn(334        `Model "${config.id}" in authType "${authType}" has both fastOnly and voiceOnly set. It will be unreachable in all model selectors.`,335      );336    }337  }338 339  /**340   * Reload models from updated configuration.341   * Clears existing user-configured models and re-registers from new config.342   * Preserves hard-coded qwen-oauth models.343   *344   * @param providerProtocolConfig - Updated provider->protocol map. `undefined`345   *   PRESERVES the existing map (so a reload carrying only modelProviders does346   *   not lose the mapping); any object value REPLACES it, so passing `{}`347   *   clears the mapping. Callers that want to preserve must omit the argument,348   *   not pass `settings.providerProtocol ?? {}`.349   */350  reloadModels(351    modelProvidersConfig?: ModelProvidersConfig,352    providerProtocolConfig?: ProviderProtocolConfig,353  ): void {354    if (providerProtocolConfig !== undefined) {355      this.providerProtocolConfig = providerProtocolConfig;356    }357 358    // Clear existing user-configured models (preserve qwen-oauth)359    for (const authType of this.modelsByAuthType.keys()) {360      if (authType !== AuthType.QWEN_OAUTH) {361        this.modelsByAuthType.delete(authType);362      }363    }364 365    // Re-register user-configured models under their resolved protocol366    this.registerProvidersConfig(modelProvidersConfig);367  }368}369