All files / src/router index.ts

74.5% Statements 76/102
57.69% Branches 15/26
50% Functions 4/8
74.5% Lines 76/102

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171                        1x 1x                                                               1x 1x 1x       1x 1x 1x 1x       1x 5x 5x 5x   5x 7x 21x 21x 21x 21x 21x   21x 21x   21x 7x     7x 2x 2x 7x       7x   1x 1x 1x   6x 6x   6x 7x 2x 2x 2x 7x 1x 1x 1x 7x       7x 3x 3x 3x 3x 7x       7x     7x   6x 7x   5x   5x 5x             5x               5x 7x 7x 7x   5x       5x     5x   5x 5x       1x 1x 1x 1x 1x 1x 1x 1x 1x  
/**
 * AgentKits — Smart Routing Module
 *
 * Auto-select the optimal provider based on cost, latency, capability.
 * Routes requests to the cheapest/fastest provider that meets requirements.
 *
 * Usage:
 *   import { createRouter } from 'agentkits/router';
 *   const router = createRouter({ providers: [...], strategy: 'cheapest' });
 *   const result = await router.complete('Hello');
 */
 
import { createChat, type ChatConfig, type ChatMessage, type ChatResponse } from '../llm/index.js';
import { getPricing } from '../cost/index.js';
 
// ── Types ──────────────────────────────────────────────────────────
 
export type RoutingStrategy = 'cheapest' | 'fastest' | 'quality' | 'round-robin' | 'random';
 
export interface RouterConfig {
  providers: Array<ChatConfig & { priority?: number; maxTokens?: number }>;
  strategy?: RoutingStrategy;
  /** For quality strategy: prefer these providers */
  qualityOrder?: string[];
  /** Max latency in ms — skip providers slower than this */
  maxLatencyMs?: number;
  /** Budget per request in USD — skip providers more expensive */
  maxCostUsd?: number;
}
 
export interface RouterClient {
  /** Route and complete a message */
  complete(prompt: string, options?: { system?: string }): Promise<string & { provider?: string }>;
  /** Route and chat with message history */
  chat(messages: ChatMessage[], options?: {}): Promise<ChatResponse & { provider: string }>;
  /** Get the selected provider without calling it */
  select(estimatedTokens?: number): { provider: string; reason: string };
  /** Last used provider */
  lastProvider?: string;
  /** Current config */
  readonly config: Readonly<RouterConfig>;
}
 
// ── Quality Rankings ───────────────────────────────────────────────
 
const DEFAULT_QUALITY_ORDER = [
  'openai', 'gemini', 'deepseek', 'dashscope', 'zhipu', 'moonshot', 'minimax', 'ollama',
];
 
// ── Latency Estimates (ms, approximate) ────────────────────────────
 
const LATENCY_ESTIMATES: Record<string, number> = {
  ollama: 200, openai: 400, gemini: 500, deepseek: 600,
  dashscope: 500, zhipu: 600, moonshot: 700, minimax: 700,
};
 
// ── Factory ────────────────────────────────────────────────────────
 
export function createRouter(config: RouterConfig): RouterClient {
  const strategy = config.strategy ?? 'cheapest';
  const qualityOrder = config.qualityOrder ?? DEFAULT_QUALITY_ORDER;
  let roundRobinIndex = 0;
 
  function selectProvider(estimatedTokens = 1000): { provider: string; reason: string; chatConfig: ChatConfig } {
    let candidates = config.providers.map(p => {
      const provider = p.provider ?? 'openai';
      const pricing = getPricing(provider);
      const p0 = Array.isArray(pricing) ? pricing[0] : pricing;
      const estCost = p0
        ? (p0.inputPer1M * estimatedTokens / 1000000) + (p0.outputPer1M * estimatedTokens / 1000000)
        : Infinity;
      const latency = LATENCY_ESTIMATES[provider] ?? 1000;
      const quality = qualityOrder.indexOf(provider);
 
      return { ...p, provider, estCost, latency, quality: quality === -1 ? 999 : quality };
    });
 
    // Filter by constraints
    if (config.maxLatencyMs) {
      candidates = candidates.filter(c => c.latency <= config.maxLatencyMs!);
    }
    if (config.maxCostUsd) {
      candidates = candidates.filter(c => c.estCost <= config.maxCostUsd!);
    }
 
    if (candidates.length === 0) {
      // Fallback to first provider if all filtered out
      const fallback = config.providers[0];
      return { provider: fallback.provider ?? 'openai', reason: 'fallback (all filtered)', chatConfig: fallback };
    }
 
    let selected: typeof candidates[0];
    let reason: string;
 
    switch (strategy) {
      case 'cheapest':
        selected = candidates.sort((a, b) => a.estCost - b.estCost)[0];
        reason = `cheapest ($${selected.estCost.toFixed(6)}/req est)`;
        break;
      case 'fastest':
        selected = candidates.sort((a, b) => a.latency - b.latency)[0];
        reason = `fastest (~${selected.latency}ms est)`;
        break;
      case 'quality':
        selected = candidates.sort((a, b) => a.quality - b.quality)[0];
        reason = `highest quality (#${selected.quality + 1})`;
        break;
      case 'round-robin':
        selected = candidates[roundRobinIndex % candidates.length];
        roundRobinIndex++;
        reason = `round-robin (index ${roundRobinIndex - 1})`;
        break;
      case 'random':
        selected = candidates[Math.floor(Math.random() * candidates.length)];
        reason = 'random';
        break;
      default:
        selected = candidates[0];
        reason = 'default';
    }
 
    return { provider: selected.provider, reason, chatConfig: selected };
  }
 
  let lastUsedProvider: string | undefined;
 
  const client: RouterClient = {
    async complete(prompt, options = {}) {
      const { provider, chatConfig } = selectProvider();
      lastUsedProvider = provider;
      const chat = createChat(chatConfig);
      return chat.complete(prompt);
    },
 
    async chat(messages, options = {}) {
      const { provider, chatConfig } = selectProvider();
      lastUsedProvider = provider;
      const chat = createChat(chatConfig);
      const result = await chat.chat(messages);
      return Object.assign({}, result, { provider });
    },
 
    select(estimatedTokens) {
      const { provider, reason } = selectProvider(estimatedTokens);
      return { provider, reason };
    },
 
    get config() {
      return config;
    },
 
    get lastProvider() {
      return lastUsedProvider;
    },
  };
 
  return client;
}
 
// ── Convenience ────────────────────────────────────────────────────
 
export function listStrategies(): Array<{ id: RoutingStrategy; description: string }> {
  return [
    { id: 'cheapest',    description: 'Select lowest cost provider' },
    { id: 'fastest',     description: 'Select lowest latency provider' },
    { id: 'quality',     description: 'Select highest quality provider' },
    { id: 'round-robin', description: 'Rotate through providers' },
    { id: 'random',      description: 'Random selection' },
  ];
}