modelgrid/ts/containers/base-container.ts

/**
 * Base Container
 *
 * Abstract base class for AI model containers.
 */

import type {
  IContainerConfig,
  IContainerStatus,
  ILoadedModel,
  TContainerType,
} from '../interfaces/container.ts';
import type { IChatCompletionRequest, IChatCompletionResponse } from '../interfaces/api.ts';
import { ContainerRuntime } from '../docker/container-runtime.ts';
import { logger } from '../logger.ts';

/**
 * Model pull progress callback
 */
export type TModelPullProgress = (progress: {
  model: string;
  status: string;
  percent?: number;
}) => void;

/**
 * Abstract base class for AI model containers
 */
export abstract class BaseContainer {
  /** Container type */
  public abstract readonly type: TContainerType;

  /** Display name */
  public abstract readonly displayName: string;

  /** Default Docker image */
  public abstract readonly defaultImage: string;

  /** Default internal port */
  public abstract readonly defaultPort: number;

  /** Container configuration */
  protected config: IContainerConfig;

  /** Container runtime */
  protected runtime: ContainerRuntime;

  constructor(config: IContainerConfig) {
    this.config = config;
    this.runtime = new ContainerRuntime();
  }

  /**
   * Get the container configuration
   */
  public getConfig(): IContainerConfig {
    return this.config;
  }

  /**
   * Get the endpoint URL for this container
   */
  public getEndpoint(): string {
    const port = this.config.externalPort || this.config.port;
    return `http://localhost:${port}`;
  }

  /**
   * Start the container
   */
  public async start(): Promise<boolean> {
    logger.info(`Starting ${this.displayName} container: ${this.config.name}`);
    return this.runtime.startContainer(this.config);
  }

  /**
   * Stop the container
   */
  public async stop(): Promise<boolean> {
    logger.info(`Stopping ${this.displayName} container: ${this.config.name}`);
    return this.runtime.stopContainer(this.config.id);
  }

  /**
   * Restart the container
   */
  public async restart(): Promise<boolean> {
    logger.info(`Restarting ${this.displayName} container: ${this.config.name}`);
    return this.runtime.restartContainer(this.config.id);
  }

  /**
   * Remove the container
   */
  public async remove(): Promise<boolean> {
    logger.info(`Removing ${this.displayName} container: ${this.config.name}`);
    return this.runtime.removeContainer(this.config.id);
  }

  /**
   * Get container status
   */
  public async getStatus(): Promise<IContainerStatus> {
    return this.runtime.getContainerStatus(this.config);
  }

  /**
   * Get container logs
   */
  public async getLogs(lines: number = 100): Promise<string> {
    return this.runtime.getLogs(this.config.id, { lines });
  }

  /**
   * Check if the container is healthy
   */
  public abstract isHealthy(): Promise<boolean>;

  /**
   * Get list of available models
   */
  public abstract listModels(): Promise<string[]>;

  /**
   * Get list of loaded models with details
   */
  public abstract getLoadedModels(): Promise<ILoadedModel[]>;

  /**
   * Pull a model
   */
  public abstract pullModel(modelName: string, onProgress?: TModelPullProgress): Promise<boolean>;

  /**
   * Remove a model
   */
  public abstract removeModel(modelName: string): Promise<boolean>;

  /**
   * Send a chat completion request
   */
  public abstract chatCompletion(request: IChatCompletionRequest): Promise<IChatCompletionResponse>;

  /**
   * Stream a chat completion request
   */
  public abstract chatCompletionStream(
    request: IChatCompletionRequest,
    onChunk: (chunk: string) => void,
  ): Promise<void>;

  /**
   * Make HTTP request to container
   */
  protected async fetch(
    path: string,
    options: {
      method?: string;
      headers?: Record<string, string>;
      body?: unknown;
      timeout?: number;
    } = {},
  ): Promise<Response> {
    const endpoint = this.getEndpoint();
    const url = `${endpoint}${path}`;

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

    try {
      const response = await fetch(url, {
        method: options.method || 'GET',
        headers: {
          'Content-Type': 'application/json',
          ...options.headers,
        },
        body: options.body ? JSON.stringify(options.body) : undefined,
        signal: controller.signal,
      });

      return response;
    } finally {
      clearTimeout(timeoutId);
    }
  }

  /**
   * Make HTTP request and parse JSON response
   */
  protected async fetchJson<T>(
    path: string,
    options: {
      method?: string;
      headers?: Record<string, string>;
      body?: unknown;
      timeout?: number;
    } = {},
  ): Promise<T> {
    const response = await this.fetch(path, options);

    if (!response.ok) {
      const errorText = await response.text();
      throw new Error(`HTTP ${response.status}: ${errorText}`);
    }

    return response.json();
  }

  /**
   * Generate a unique request ID
   */
  protected generateRequestId(): string {
    return `chatcmpl-${Date.now().toString(36)}-${Math.random().toString(36).substring(2, 8)}`;
  }
}