modelgrid/ts/api/router.ts

/**
 * API Router
 *
 * Routes incoming requests to appropriate handlers.
 */

import * as http from 'node:http';
import type { IApiError } from '../interfaces/api.ts';
import { logger } from '../logger.ts';
import { ContainerManager } from '../containers/container-manager.ts';
import { ModelRegistry } from '../models/registry.ts';
import { ModelLoader } from '../models/loader.ts';
import { ChatHandler } from './handlers/chat.ts';
import { ModelsHandler } from './handlers/models.ts';
import { EmbeddingsHandler } from './handlers/embeddings.ts';
import { AuthMiddleware } from './middleware/auth.ts';
import { SanityMiddleware } from './middleware/sanity.ts';

/**
 * API Router - routes requests to handlers
 */
export class ApiRouter {
  private containerManager: ContainerManager;
  private modelRegistry: ModelRegistry;
  private modelLoader: ModelLoader;
  private chatHandler: ChatHandler;
  private modelsHandler: ModelsHandler;
  private embeddingsHandler: EmbeddingsHandler;
  private authMiddleware: AuthMiddleware;
  private sanityMiddleware: SanityMiddleware;

  constructor(
    containerManager: ContainerManager,
    modelRegistry: ModelRegistry,
    modelLoader: ModelLoader,
    apiKeys: string[],
  ) {
    this.containerManager = containerManager;
    this.modelRegistry = modelRegistry;
    this.modelLoader = modelLoader;

    // Initialize handlers
    this.chatHandler = new ChatHandler(containerManager, modelLoader);
    this.modelsHandler = new ModelsHandler(containerManager, modelRegistry);
    this.embeddingsHandler = new EmbeddingsHandler(containerManager);

    // Initialize middleware
    this.authMiddleware = new AuthMiddleware(apiKeys);
    this.sanityMiddleware = new SanityMiddleware(modelRegistry);
  }

  /**
   * Route a request to the appropriate handler
   */
  public async route(
    req: http.IncomingMessage,
    res: http.ServerResponse,
    path: string,
  ): Promise<void> {
    // OpenAI API endpoints
    if (path === '/v1/chat/completions') {
      await this.handleChatCompletions(req, res);
      return;
    }

    if (path === '/v1/completions') {
      await this.handleCompletions(req, res);
      return;
    }

    if (path === '/v1/models' || path === '/v1/models/') {
      await this.handleModels(req, res);
      return;
    }

    if (path.startsWith('/v1/models/')) {
      await this.handleModelInfo(req, res, path);
      return;
    }

    if (path === '/v1/embeddings') {
      await this.handleEmbeddings(req, res);
      return;
    }

    // Not found
    this.sendError(res, 404, `Endpoint not found: ${path}`, 'invalid_request_error');
  }

  /**
   * Handle POST /v1/chat/completions
   */
  private async handleChatCompletions(
    req: http.IncomingMessage,
    res: http.ServerResponse,
  ): Promise<void> {
    if (req.method !== 'POST') {
      this.sendError(res, 405, 'Method not allowed', 'invalid_request_error');
      return;
    }

    // Authenticate
    if (!this.authMiddleware.authenticate(req)) {
      this.sendError(res, 401, 'Invalid API key', 'authentication_error');
      return;
    }

    // Parse body
    const body = await this.parseRequestBody(req);
    if (!body) {
      this.sendError(res, 400, 'Invalid JSON body', 'invalid_request_error');
      return;
    }

    // Validate request
    const validation = this.sanityMiddleware.validateChatRequest(body);
    if (!validation.valid) {
      this.sendError(res, 400, validation.error || 'Invalid request', 'invalid_request_error');
      return;
    }

    // Handle request
    await this.chatHandler.handleChatCompletion(req, res, body);
  }

  /**
   * Handle POST /v1/completions (legacy endpoint)
   */
  private async handleCompletions(
    req: http.IncomingMessage,
    res: http.ServerResponse,
  ): Promise<void> {
    if (req.method !== 'POST') {
      this.sendError(res, 405, 'Method not allowed', 'invalid_request_error');
      return;
    }

    // Authenticate
    if (!this.authMiddleware.authenticate(req)) {
      this.sendError(res, 401, 'Invalid API key', 'authentication_error');
      return;
    }

    // Parse body
    const body = await this.parseRequestBody(req);
    if (!body) {
      this.sendError(res, 400, 'Invalid JSON body', 'invalid_request_error');
      return;
    }

    // Convert to chat format and handle
    const chatBody = this.convertCompletionToChat(body);
    await this.chatHandler.handleChatCompletion(req, res, chatBody);
  }

  /**
   * Handle GET /v1/models
   */
  private async handleModels(
    req: http.IncomingMessage,
    res: http.ServerResponse,
  ): Promise<void> {
    if (req.method !== 'GET') {
      this.sendError(res, 405, 'Method not allowed', 'invalid_request_error');
      return;
    }

    // Authenticate
    if (!this.authMiddleware.authenticate(req)) {
      this.sendError(res, 401, 'Invalid API key', 'authentication_error');
      return;
    }

    await this.modelsHandler.handleListModels(res);
  }

  /**
   * Handle GET /v1/models/:model
   */
  private async handleModelInfo(
    req: http.IncomingMessage,
    res: http.ServerResponse,
    path: string,
  ): Promise<void> {
    if (req.method !== 'GET') {
      this.sendError(res, 405, 'Method not allowed', 'invalid_request_error');
      return;
    }

    // Authenticate
    if (!this.authMiddleware.authenticate(req)) {
      this.sendError(res, 401, 'Invalid API key', 'authentication_error');
      return;
    }

    const modelId = path.replace('/v1/models/', '');
    await this.modelsHandler.handleGetModel(res, modelId);
  }

  /**
   * Handle POST /v1/embeddings
   */
  private async handleEmbeddings(
    req: http.IncomingMessage,
    res: http.ServerResponse,
  ): Promise<void> {
    if (req.method !== 'POST') {
      this.sendError(res, 405, 'Method not allowed', 'invalid_request_error');
      return;
    }

    // Authenticate
    if (!this.authMiddleware.authenticate(req)) {
      this.sendError(res, 401, 'Invalid API key', 'authentication_error');
      return;
    }

    // Parse body
    const body = await this.parseRequestBody(req);
    if (!body) {
      this.sendError(res, 400, 'Invalid JSON body', 'invalid_request_error');
      return;
    }

    await this.embeddingsHandler.handleEmbeddings(res, body);
  }

  /**
   * Parse request body
   */
  private async parseRequestBody(req: http.IncomingMessage): Promise<unknown | null> {
    return new Promise((resolve) => {
      let body = '';

      req.on('data', (chunk) => {
        body += chunk.toString();
        // Limit body size
        if (body.length > 10 * 1024 * 1024) {
          resolve(null);
        }
      });

      req.on('end', () => {
        try {
          resolve(JSON.parse(body));
        } catch {
          resolve(null);
        }
      });

      req.on('error', () => {
        resolve(null);
      });
    });
  }

  /**
   * Convert legacy completion request to chat format
   */
  private convertCompletionToChat(body: Record<string, unknown>): Record<string, unknown> {
    const prompt = body.prompt as string | string[];
    const promptText = Array.isArray(prompt) ? prompt.join('\n') : prompt;

    return {
      model: body.model,
      messages: [
        { role: 'user', content: promptText },
      ],
      max_tokens: body.max_tokens,
      temperature: body.temperature,
      top_p: body.top_p,
      n: body.n,
      stream: body.stream,
      stop: body.stop,
    };
  }

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

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