flynn/src/models/router.ts

import type { ChatRequest, ChatResponse, ModelClient } from './types.js';

export type ModelTier = 'fast' | 'default' | 'complex' | 'local';

export interface ModelRouterConfig {
  default: ModelClient;
  fast?: ModelClient;
  complex?: ModelClient;
  local?: ModelClient;
  fallbackChain: ModelClient[];
}

export class ModelRouter implements ModelClient {
  private clients: Map<ModelTier, ModelClient>;
  private defaultClient: ModelClient;
  private fallbackChain: ModelClient[];

  constructor(config: ModelRouterConfig) {
    this.clients = new Map();
    this.defaultClient = config.default;
    this.fallbackChain = config.fallbackChain;

    this.clients.set('default', config.default);
    if (config.fast) this.clients.set('fast', config.fast);
    if (config.complex) this.clients.set('complex', config.complex);
    if (config.local) this.clients.set('local', config.local);
  }

  async chat(request: ChatRequest, tier?: ModelTier): Promise<ChatResponse> {
    const primaryClient = tier ? this.clients.get(tier) ?? this.defaultClient : this.defaultClient;
    const errors: Error[] = [];

    // Try primary client
    try {
      return await primaryClient.chat(request);
    } catch (error) {
      errors.push(error instanceof Error ? error : new Error(String(error)));
      console.warn(`Primary model failed: ${errors[0].message}`);
    }

    // Try fallback chain
    for (const fallbackClient of this.fallbackChain) {
      try {
        console.log('Trying fallback model...');
        return await fallbackClient.chat(request);
      } catch (error) {
        errors.push(error instanceof Error ? error : new Error(String(error)));
        console.warn(`Fallback model failed: ${errors[errors.length - 1].message}`);
      }
    }

    throw new Error(`All model providers failed: ${errors.map(e => e.message).join(', ')}`);
  }

  getClient(tier: ModelTier): ModelClient | undefined {
    return this.clients.get(tier);
  }
}