modelgrid/ts/models/registry.ts

/**
 * Model Registry
 *
 * Manages the greenlit model list and model availability.
 */

import type { IGreenlitModel, IGreenlitModelsList } from '../interfaces/config.ts';
import type { TContainerType } from '../interfaces/container.ts';
import { MODEL_REGISTRY, TIMING } from '../constants.ts';
import { logger } from '../logger.ts';

/**
 * Model registry for managing greenlit models
 */
export class ModelRegistry {
  private greenlistUrl: string;
  private cachedGreenlist: IGreenlitModelsList | null = null;
  private cacheTime: number = 0;

  constructor(greenlistUrl: string = MODEL_REGISTRY.DEFAULT_GREENLIST_URL) {
    this.greenlistUrl = greenlistUrl;
  }

  /**
   * Set the greenlist URL
   */
  public setGreenlistUrl(url: string): void {
    this.greenlistUrl = url;
    this.cachedGreenlist = null;
    this.cacheTime = 0;
  }

  /**
   * Fetch the greenlit model list from remote URL
   */
  public async fetchGreenlist(forceRefresh: boolean = false): Promise<IGreenlitModelsList> {
    // Return cached data if still valid
    if (
      !forceRefresh &&
      this.cachedGreenlist &&
      Date.now() - this.cacheTime < TIMING.GREENLIST_CACHE_DURATION_MS
    ) {
      return this.cachedGreenlist;
    }

    try {
      logger.dim(`Fetching greenlit models from: ${this.greenlistUrl}`);

      const controller = new AbortController();
      const timeout = setTimeout(() => controller.abort(), 30000);

      const response = await fetch(this.greenlistUrl, {
        signal: controller.signal,
        headers: {
          'Accept': 'application/json',
          'User-Agent': 'ModelGrid/1.0',
        },
      });

      clearTimeout(timeout);

      if (!response.ok) {
        throw new Error(`HTTP ${response.status}: ${response.statusText}`);
      }

      const greenlist = await response.json() as IGreenlitModelsList;

      // Validate structure
      if (!greenlist.models || !Array.isArray(greenlist.models)) {
        throw new Error('Invalid greenlist format: missing models array');
      }

      // Cache the result
      this.cachedGreenlist = greenlist;
      this.cacheTime = Date.now();

      logger.dim(`Loaded ${greenlist.models.length} greenlit models`);
      return greenlist;
    } catch (error) {
      logger.warn(`Failed to fetch greenlist: ${error instanceof Error ? error.message : String(error)}`);

      // Return fallback if we have no cache
      if (!this.cachedGreenlist) {
        logger.dim('Using fallback greenlist');
        return this.getFallbackGreenlist();
      }

      // Return stale cache
      return this.cachedGreenlist;
    }
  }

  /**
   * Get fallback greenlist
   */
  private getFallbackGreenlist(): IGreenlitModelsList {
    return {
      version: '1.0',
      lastUpdated: new Date().toISOString(),
      models: MODEL_REGISTRY.FALLBACK_GREENLIST as unknown as IGreenlitModel[],
    };
  }

  /**
   * Check if a model is greenlit
   */
  public async isModelGreenlit(modelName: string): Promise<boolean> {
    const greenlist = await this.fetchGreenlist();
    return greenlist.models.some((m) => this.normalizeModelName(m.name) === this.normalizeModelName(modelName));
  }

  /**
   * Get greenlit model info
   */
  public async getGreenlitModel(modelName: string): Promise<IGreenlitModel | null> {
    const greenlist = await this.fetchGreenlist();
    const normalized = this.normalizeModelName(modelName);
    return greenlist.models.find((m) => this.normalizeModelName(m.name) === normalized) || null;
  }

  /**
   * Get all greenlit models
   */
  public async getAllGreenlitModels(): Promise<IGreenlitModel[]> {
    const greenlist = await this.fetchGreenlist();
    return greenlist.models;
  }

  /**
   * Get greenlit models by container type
   */
  public async getModelsByContainer(containerType: TContainerType): Promise<IGreenlitModel[]> {
    const greenlist = await this.fetchGreenlist();
    return greenlist.models.filter((m) => m.container === containerType);
  }

  /**
   * Get greenlit models that fit within VRAM limit
   */
  public async getModelsWithinVram(maxVramGb: number): Promise<IGreenlitModel[]> {
    const greenlist = await this.fetchGreenlist();
    return greenlist.models.filter((m) => m.minVram <= maxVramGb);
  }

  /**
   * Get recommended container type for a model
   */
  public async getRecommendedContainer(modelName: string): Promise<TContainerType | null> {
    const model = await this.getGreenlitModel(modelName);
    return model ? model.container : null;
  }

  /**
   * Get minimum VRAM required for a model
   */
  public async getMinVram(modelName: string): Promise<number | null> {
    const model = await this.getGreenlitModel(modelName);
    return model ? model.minVram : null;
  }

  /**
   * Check if model fits in available VRAM
   */
  public async modelFitsInVram(modelName: string, availableVramGb: number): Promise<boolean> {
    const minVram = await this.getMinVram(modelName);
    if (minVram === null) {
      // Model not in greenlist, assume it might fit
      return true;
    }
    return availableVramGb >= minVram;
  }

  /**
   * Normalize model name for comparison
   * Handles variations like "llama3:8b" vs "llama3:8B" vs "meta-llama/llama-3-8b"
   */
  private normalizeModelName(name: string): string {
    return name
      .toLowerCase()
      .replace(/[^a-z0-9:.-]/g, '')
      .trim();
  }

  /**
   * Search models by name pattern
   */
  public async searchModels(pattern: string): Promise<IGreenlitModel[]> {
    const greenlist = await this.fetchGreenlist();
    const normalizedPattern = pattern.toLowerCase();

    return greenlist.models.filter((m) =>
      m.name.toLowerCase().includes(normalizedPattern) ||
      m.description?.toLowerCase().includes(normalizedPattern) ||
      m.tags?.some((t) => t.toLowerCase().includes(normalizedPattern))
    );
  }

  /**
   * Get models by tags
   */
  public async getModelsByTags(tags: string[]): Promise<IGreenlitModel[]> {
    const greenlist = await this.fetchGreenlist();
    const normalizedTags = tags.map((t) => t.toLowerCase());

    return greenlist.models.filter((m) =>
      m.tags?.some((t) => normalizedTags.includes(t.toLowerCase()))
    );
  }

  /**
   * Clear the cached greenlist
   */
  public clearCache(): void {
    this.cachedGreenlist = null;
    this.cacheTime = 0;
  }

  /**
   * Print greenlist summary
   */
  public async printSummary(): Promise<void> {
    const greenlist = await this.fetchGreenlist();

    // Group by container type
    const byContainer = new Map<string, IGreenlitModel[]>();
    for (const model of greenlist.models) {
      if (!byContainer.has(model.container)) {
        byContainer.set(model.container, []);
      }
      byContainer.get(model.container)!.push(model);
    }

    logger.logBoxTitle('Greenlit Models', 60, 'info');
    logger.logBoxLine(`Version: ${greenlist.version}`);
    logger.logBoxLine(`Last Updated: ${greenlist.lastUpdated}`);
    logger.logBoxLine(`Total Models: ${greenlist.models.length}`);
    logger.logBoxLine('');

    for (const [container, models] of byContainer) {
      logger.logBoxLine(`${container.toUpperCase()} (${models.length}):`);
      for (const model of models.slice(0, 5)) {
        logger.logBoxLine(`  - ${model.name} (${model.minVram}GB VRAM)`);
      }
      if (models.length > 5) {
        logger.logBoxLine(`  ... and ${models.length - 5} more`);
      }
      logger.logBoxLine('');
    }

    logger.logBoxEnd();
  }
}