artifacts/api-server/src/lib/llm-roles/model-selection.ts

/**
 * Role → model 选择策略 + role-specific timeout。
 *
 * 设计原则(A0 第 5 章):
 *   - judge / diagnoser:**便宜**为先(deepseek / qwen),反正只是评分 + 诊断
 *   - planner / code-mutator:**强**为先(anthropic / openai),负责真改东西
 *   - 但任意一个角色的"理想 model"不可用时,**降级到任何可用 model**
 *     而非默认抛错 — 因为 LLM disabled 已经有 LlmRoleDisabled 单独处理,
 *     这里不该再误伤 caller。
 *
 * timeout 表:
 *   - judge       8s   (短 prompt 短 output)
 *   - diagnoser  15s   (中等 prompt + 中等推理)
 *   - planner    20s   (中等 prompt + 多 candidate)
 *   - code-mutator 45s (大 prompt + 大 output)
 */
import { listModels, type ModelInfo } from "../models";
import type { LlmRole } from "./types";

export const ROLE_TIMEOUT_MS: Record<LlmRole, number> = {
  judge: 8_000,
  diagnoser: 15_000,
  planner: 20_000,
  "code-mutator": 45_000,
};

/**
 * 角色到首选 model id 列表(按优先级);第一个 listModels() 中 available
 * 的赢。**这些 id 必须与 lib/models.ts 完全一致**,拼错了会静默降级到任意可用
 * model(可能更贵或更慢),违反"判官/诊断便宜、规划/代码强"的设计意图。
 */
const ROLE_PREFERENCE: Record<LlmRole, string[]> = {
  judge: ["mdl_deepseek-v3", "mdl_qwen-max", "mdl_glm-4_6", "mdl_deepseek-r1"],
  diagnoser: [
    "mdl_deepseek-r1",
    "mdl_qwen-max",
    "mdl_glm-4_6",
    "mdl_deepseek-v3",
  ],
  planner: ["mdl_claude-sonnet-4-6", "mdl_gpt-5_2", "mdl_deepseek-r1"],
  "code-mutator": [
    "mdl_claude-sonnet-4-6",
    "mdl_gpt-5_2",
    "mdl_deepseek-r1",
  ],
};

/**
 * 已知 per-1k token 美元价(units: USD per 1000 tokens)。当 ModelInfo.pricing
 * 为 0(目前所有 model 都是)时,我们用这表估算 cost,以便 budget guard
 * 能真生效。这是粗估,不是 billing 用。
 */
const KNOWN_MODEL_PRICING_USD_PER_1K: Record<
  string,
  { input: number; output: number }
> = {
  "mdl_claude-sonnet-4-6": { input: 0.003, output: 0.015 },
  "mdl_gpt-5_2": { input: 0.0025, output: 0.01 },
  "mdl_gemini-2_5-pro": { input: 0.00125, output: 0.005 },
  "mdl_deepseek-r1": { input: 0.00055, output: 0.00219 },
  "mdl_deepseek-v3": { input: 0.00027, output: 0.0011 },
  "mdl_qwen-max": { input: 0.0004, output: 0.0012 },
  "mdl_minimax-m2": { input: 0.0003, output: 0.0012 },
  "mdl_kimi-k2": { input: 0.0006, output: 0.0025 },
  "mdl_glm-4_6": { input: 0.0005, output: 0.002 },
};

export interface SelectedModel {
  modelId: string;
  modelInfo: ModelInfo;
  /** 估算每 1k tokens 的 USD 价(input / output)。 */
  pricing: { input_per_1k: number; output_per_1k: number };
}

/** 选 role 对应的首个可用 model;无可用则 null。 */
export function pickModelForRole(role: LlmRole): SelectedModel | null {
  const all = listModels();
  const preferred = ROLE_PREFERENCE[role];
  for (const id of preferred) {
    const m = all.find((mm) => mm.id === id && mm.available);
    if (m) return { modelId: m.id, modelInfo: m, pricing: pricingFor(m) };
  }
  // 全降级:任意可用 model。
  const any = all.find((mm) => mm.available);
  if (any)
    return { modelId: any.id, modelInfo: any, pricing: pricingFor(any) };
  return null;
}

function pricingFor(m: ModelInfo): {
  input_per_1k: number;
  output_per_1k: number;
} {
  // ModelInfo.pricing 优先,有值用之;否则查内置 known table;再否则 0。
  if (
    m.pricing &&
    (m.pricing.input_per_1k > 0 || m.pricing.output_per_1k > 0)
  ) {
    return {
      input_per_1k: m.pricing.input_per_1k,
      output_per_1k: m.pricing.output_per_1k,
    };
  }
  const known = KNOWN_MODEL_PRICING_USD_PER_1K[m.id];
  if (known) {
    return { input_per_1k: known.input, output_per_1k: known.output };
  }
  return { input_per_1k: 0, output_per_1k: 0 };
}

export function estimateCostUsd(
  pricing: { input_per_1k: number; output_per_1k: number },
  inputTokens: number,
  outputTokens: number,
): number {
  return (
    (inputTokens / 1000) * pricing.input_per_1k +
    (outputTokens / 1000) * pricing.output_per_1k
  );
}