modelgrid/ts/api/handlers/models.ts

/**
 * Models handler.
 */

import * as http from 'node:http';
import type { IApiError, IListModelsResponse, IModelInfo } from '../../interfaces/api.ts';
import { ClusterCoordinator } from '../../cluster/coordinator.ts';
import { ContainerManager } from '../../containers/container-manager.ts';
import { logger } from '../../logger.ts';
import { ModelRegistry } from '../../models/registry.ts';

export class ModelsHandler {
  private containerManager: ContainerManager;
  private modelRegistry: ModelRegistry;
  private clusterCoordinator: ClusterCoordinator;

  constructor(
    containerManager: ContainerManager,
    modelRegistry: ModelRegistry,
    clusterCoordinator: ClusterCoordinator,
  ) {
    this.containerManager = containerManager;
    this.modelRegistry = modelRegistry;
    this.clusterCoordinator = clusterCoordinator;
  }

  public async handleListModels(res: http.ServerResponse): Promise<void> {
    try {
      const models = await this.getAvailableModels();

      const response: IListModelsResponse = {
        object: 'list',
        data: models,
      };

      res.writeHead(200, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify(response));
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      logger.error(`Failed to list models: ${message}`);
      this.sendError(res, 500, `Failed to list models: ${message}`, 'server_error');
    }
  }

  public async handleGetModel(res: http.ServerResponse, modelId: string): Promise<void> {
    try {
      const models = await this.getAvailableModels();
      const requested = await this.modelRegistry.getModel(modelId);
      const canonicalId = requested?.id || modelId;
      const model = models.find((entry) => entry.id === canonicalId);

      if (!model) {
        this.sendError(res, 404, `Model "${modelId}" not found`, 'model_not_found');
        return;
      }

      res.writeHead(200, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify(model));
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      logger.error(`Failed to get model info: ${message}`);
      this.sendError(res, 500, `Failed to get model info: ${message}`, 'server_error');
    }
  }

  private async getAvailableModels(): Promise<IModelInfo[]> {
    const models: IModelInfo[] = [];
    const seen = new Set<string>();
    const timestamp = Math.floor(Date.now() / 1000);

    const containerModels = await this.containerManager.getAllAvailableModels();
    for (const [modelId, endpoints] of containerModels) {
      if (seen.has(modelId)) {
        continue;
      }

      const primaryEndpoint = endpoints[0];
      seen.add(modelId);
      models.push({
        id: modelId,
        object: 'model',
        created: timestamp,
        owned_by: `modelgrid-${primaryEndpoint?.type || 'vllm'}`,
      });
    }

    const clusterStatus = this.clusterCoordinator.getStatus();
    for (const [modelId, locations] of Object.entries(clusterStatus.models)) {
      if (seen.has(modelId) || locations.length === 0) {
        continue;
      }

      seen.add(modelId);
      models.push({
        id: modelId,
        object: 'model',
        created: timestamp,
        owned_by: `modelgrid-${locations[0].engine}`,
      });
    }

    const catalogModels = await this.modelRegistry.getAllModels();
    for (const model of catalogModels) {
      if (seen.has(model.id)) {
        continue;
      }

      seen.add(model.id);
      models.push({
        id: model.id,
        object: 'model',
        created: timestamp,
        owned_by: `modelgrid-${model.engine}`,
      });
    }

    models.sort((left, right) => left.id.localeCompare(right.id));
    return models;
  }

  private sendError(
    res: http.ServerResponse,
    statusCode: number,
    message: string,
    type: string,
    param?: string,
  ): void {
    const error: IApiError = {
      error: {
        message,
        type,
        param,
      },
    };

    res.writeHead(statusCode, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify(error));
  }
}