feat: add role management and usage tracking for AI code review

- Implemented role parsing and loading from markdown files, including attributes like name, side, focus, badge, color, and personality. - Created functions to build prompts for analysis, line location, and verdicts based on roles. - Added tests for role management functionalities to ensure correct parsing and loading of roles. - Developed usage tracking for AI assistant interactions, including token usage and rate limits. - Implemented functions to extract and record usage data from various LLM providers. - Added tests for usage tracking functionalities to validate correct accumulation and reporting of usage statistics.
2026-06-25 09:34:59 +00:00
parent 120b83c904
commit 525f6f9350
37 changed files with 6405 additions and 23 deletions
@@ -0,0 +1,289 @@
+import axios from 'axios';
+import { warn } from './log.js';
+
+/** 本次執行的 token 累計（跨所有 LLM 呼叫）。 */
+const runUsage = { calls: 0, promptTokens: 0, completionTokens: 0, totalTokens: 0 };
+
+function num(x) {
+  const n = Number(x);
+  return Number.isFinite(n) ? n : 0;
+}
+
+/**
+ * 把各平台回應中的 token usage 正規化成 { promptTokens, completionTokens, totalTokens }。
+ * 支援：OpenAI 相容 usage、OpenAI Responses（input/output_tokens）、
+ * Gemini usageMetadata、Ollama 原生 eval_count、OpenCode tokens。
+ * 回應中沒有任何可辨識的 usage 時回傳 null。
+ */
+export function extractUsage(data) {
+  if (!data || typeof data !== 'object') return null;
+
+  // OpenAI 相容 / OpenAI Responses
+  const u = data.usage;
+  if (u && typeof u === 'object') {
+    const prompt = num(u.prompt_tokens ?? u.input_tokens);
+    const completion = num(u.completion_tokens ?? u.output_tokens);
+    const total = u.total_tokens != null ? num(u.total_tokens) : prompt + completion;
+    if (prompt || completion || total) return { promptTokens: prompt, completionTokens: completion, totalTokens: total };
+  }
+
+  // Gemini 原生 usageMetadata
+  const g = data.usageMetadata;
+  if (g && typeof g === 'object') {
+    const prompt = num(g.promptTokenCount);
+    const completion = num(g.candidatesTokenCount);
+    const total = g.totalTokenCount != null ? num(g.totalTokenCount) : prompt + completion;
+    if (prompt || completion || total) return { promptTokens: prompt, completionTokens: completion, totalTokens: total };
+  }
+
+  // Ollama 原生回應
+  if (data.prompt_eval_count != null || data.eval_count != null) {
+    const prompt = num(data.prompt_eval_count);
+    const completion = num(data.eval_count);
+    return { promptTokens: prompt, completionTokens: completion, totalTokens: prompt + completion };
+  }
+
+  // OpenCode（tokens 可能位於 data.tokens 或 data.info.tokens）
+  const t = data.tokens || data.info?.tokens || data.data?.info?.tokens;
+  if (t && typeof t === 'object') {
+    const prompt = num(t.input ?? t.prompt);
+    const completion = num(t.output ?? t.completion);
+    const total = t.total != null ? num(t.total) : prompt + completion;
+    if (prompt || completion || total) return { promptTokens: prompt, completionTokens: completion, totalTokens: total };
+  }
+
+  return null;
+}
+
+/** 記錄一次 LLM 呼叫的 usage（無法解析時仍計一次呼叫，但 token 計 0）。 */
+export function recordUsage(data) {
+  runUsage.calls += 1;
+  const u = extractUsage(data);
+  if (u) {
+    runUsage.promptTokens += u.promptTokens;
+    runUsage.completionTokens += u.completionTokens;
+    runUsage.totalTokens += u.totalTokens;
+  }
+  return u;
+}
+
+/** 取得本次執行至今的 token 累計（複本）。 */
+export function getRunUsage() {
+  return { ...runUsage };
+}
+
+/** 重置累計（測試用）。 */
+export function resetRunUsage() {
+  runUsage.calls = 0;
+  runUsage.promptTokens = 0;
+  runUsage.completionTokens = 0;
+  runUsage.totalTokens = 0;
+}
+
+/** 最近一次回應的速率配額（rate limit）快照，用來計算「當前視窗剩餘百分比」。 */
+const rateLimit = { hasData: false, remaining: null, limit: null, kind: null };
+
+/** 將物件的 key 全部轉小寫，方便對大小寫不敏感的 HTTP header 取值。 */
+function lowerCaseKeys(obj) {
+  const out = {};
+  for (const k of Object.keys(obj)) out[k.toLowerCase()] = obj[k];
+  return out;
+}
+
+/**
+ * 從回應 header 擷取速率配額剩餘量／上限。
+ * 支援 OpenAI 相容（x-ratelimit-*-tokens）與 Anthropic（anthropic-ratelimit-tokens-*），
+ * 兩者皆缺時退而採用 requests 維度。記錄「最近一次」的數值（即最新的視窗狀態）。
+ */
+export function recordRateLimit(headers) {
+  if (!headers || typeof headers !== 'object') return;
+  const h = lowerCaseKeys(headers);
+
+  let remaining = h['x-ratelimit-remaining-tokens'] ?? h['anthropic-ratelimit-tokens-remaining'];
+  let limit = h['x-ratelimit-limit-tokens'] ?? h['anthropic-ratelimit-tokens-limit'];
+  let kind = 'tokens';
+  if (remaining == null || limit == null) {
+    remaining = h['x-ratelimit-remaining-requests'] ?? h['anthropic-ratelimit-requests-remaining'];
+    limit = h['x-ratelimit-limit-requests'] ?? h['anthropic-ratelimit-requests-limit'];
+    kind = 'requests';
+  }
+  if (remaining == null || limit == null) return;
+
+  rateLimit.hasData = true;
+  rateLimit.remaining = num(remaining);
+  rateLimit.limit = num(limit);
+  rateLimit.kind = kind;
+}
+
+/** 取得最近一次的速率配額快照（複本）。 */
+export function getRateLimit() {
+  return { ...rateLimit };
+}
+
+/** 重置速率配額快照（測試用）。 */
+export function resetRateLimit() {
+  rateLimit.hasData = false;
+  rateLimit.remaining = null;
+  rateLimit.limit = null;
+  rateLimit.kind = null;
+}
+
+const stripSlash = (s) => String(s || '').replace(/\/$/, '');
+
+/**
+ * 以實際 hostname 精確比對是否為 OpenRouter（僅接受 apex 域名 `openrouter.ai`），
+ * 避免被偽造的 baseURL（如 `openrouter.ai.evil.com`、`evil.com/openrouter.ai` 或任何子網域）
+ * 矇騙而把 API key 送往非 OpenRouter 主機。
+ */
+function isOpenRouterBaseURL(baseURL) {
+  try {
+    return new URL(baseURL).hostname.toLowerCase() === 'openrouter.ai';
+  } catch {
+    return false;
+  }
+}
+
+/**
+ * OpenRouter：以 API key 呼叫 GET /auth/key 取得額度（可靠）。
+ * 回傳金額單位為 USD credits。
+ */
+async function fetchOpenRouterQuota({ apiKey, baseURL }, get) {
+  const resp = await get(`${stripSlash(baseURL)}/auth/key`, {
+    headers: { Authorization: `Bearer ${apiKey}` },
+    timeout: 30000,
+  });
+  const d = resp.data?.data || {};
+  const used = num(d.usage);
+  const limit = d.limit == null ? null : num(d.limit);
+  const remaining = d.limit_remaining == null ? (limit == null ? null : limit - used) : num(d.limit_remaining);
+  return { available: true, used, limit, remaining, currency: 'USD', source: 'openrouter' };
+}
+
+/**
+ * 各平台帳號額度查詢策略。
+ * 多數官方平台無法僅憑 API key 取得帳號額度（需 org/admin 權限），故誠實回報「無法取得」並附原因；
+ * 本地/自架服務（ollama/opencode）則回報「不適用」。
+ */
+const QUOTA_STRATEGIES = {
+  openai: async (cfg, get) => {
+    if (isOpenRouterBaseURL(cfg.baseURL)) return fetchOpenRouterQuota(cfg, get);
+    return { available: false, reason: 'OpenAI 帳號額度需 dashboard session 權限，API key 無法取得' };
+  },
+  claude: async () => ({ available: false, reason: 'Anthropic 額度需 Admin API 權限，一般 API key 無法取得' }),
+  gemini: async () => ({ available: false, reason: 'Gemini 額度由 Google Cloud quota 管理，API key 無法直接查詢' }),
+  amazonq: async () => ({ available: false, reason: 'Amazon Q 額度由 AWS 帳務管理，需 AWS 憑證查詢' }),
+  ollama: async () => ({ available: false, reason: '本地服務，無帳號額度概念' }),
+  opencode: async () => ({ available: false, reason: '自架服務，無帳號額度概念' }),
+};
+
+/**
+ * 取得指定平台的帳號額度。任何失敗都降級為 { available: false, reason }，不丟例外。
+ * deps.get 可注入以利測試（預設 axios.get）。
+ */
+export async function fetchAccountQuota(provider, config = {}, deps = {}) {
+  const get = deps.get || axios.get;
+  const strategy = QUOTA_STRATEGIES[provider];
+  if (!strategy) return { available: false, reason: `未支援 ${provider} 額度查詢` };
+  const apiKey = Array.isArray(config.apiKeys) ? config.apiKeys[0] : config.apiKey;
+  try {
+    return await strategy({ apiKey, baseURL: config.baseURL }, get);
+  } catch (e) {
+    warn(`取得 ${provider} 帳號額度失敗（視為無法取得）: ${e.message}`);
+    return { available: false, reason: e.message };
+  }
+}
+
+/** 千分位整數/小數格式。 */
+function fmt(n) {
+  if (n == null || Number.isNaN(Number(n))) return '0';
+  const [int, frac] = String(Number(n)).split('.');
+  const withCommas = int.replace(/\B(?=(\d{3})+(?!\d))/g, ',');
+  return frac ? `${withCommas}.${frac}` : withCommas;
+}
+
+function money(currency, n) {
+  return currency ? `${currency} ${fmt(n)}` : fmt(n);
+}
+
+function round1(n) {
+  return Math.round(Number(n) * 10) / 10;
+}
+
+const RATE_KIND_LABEL = { tokens: 'token', requests: '次數' };
+
+/**
+ * 計算「剩餘百分比」= remaining / limit × 100。
+ * limit 或 remaining 為 null／undefined／NaN／Infinity，或 limit ≤ 0 時回 null，
+ * 避免算出 Infinity%／NaN%／負百分比或除以零。
+ */
+function calculatePercent(remaining, limit) {
+  if (remaining == null || limit == null) return null;
+  const rem = Number(remaining);
+  const lim = Number(limit);
+  if (!Number.isFinite(rem) || !Number.isFinite(lim) || lim <= 0) return null;
+  return round1((rem / lim) * 100);
+}
+
+/**
+ * 計算「剩餘可用百分比」，依優先序擇一：
+ * 1. 帳號額度（quota 有有效上限）→ 剩餘 credits / 上限；
+ * 2. 速率配額（rate limit header，有有效上限）→ 當前視窗剩餘 / 上限；
+ * 上限或剩餘為無效值（null／0／負數／NaN／Infinity）時跳過計算，落到 { percent: null, reason }。
+ */
+export function resolveRemainingPercent(quota, rate) {
+  if (quota?.available && quota.limit != null) {
+    const limit = Number(quota.limit);
+    const remaining = quota.remaining == null ? limit - num(quota.used) : Number(quota.remaining);
+    const percent = calculatePercent(remaining, limit);
+    if (percent != null) {
+      return { percent, basis: '帳號額度', remaining, limit, unit: quota.currency || '' };
+    }
+  }
+  if (rate?.hasData) {
+    const percent = calculatePercent(rate.remaining, rate.limit);
+    if (percent != null) {
+      const kindLabel = RATE_KIND_LABEL[rate.kind] || rate.kind;
+      return { percent, basis: `速率配額（當前視窗，${kindLabel}）`, remaining: Number(rate.remaining), limit: Number(rate.limit), unit: '' };
+    }
+  }
+  let reason;
+  if (quota?.available && quota.limit == null) reason = '帳號額度無上限，無法計算百分比';
+  else if (quota && !quota.available) reason = quota.reason || '平台未提供額度';
+  else reason = '平台未提供額度或速率配額資訊';
+  return { percent: null, reason };
+}
+
+function remainingLine(pct) {
+  if (pct.percent == null) return `剩餘可用：無法計算百分比（${pct.reason}）`;
+  const detail = `${pct.basis}：${money(pct.unit, pct.remaining)} / ${money(pct.unit, pct.limit)}`;
+  return `剩餘可用 **${pct.percent}%**（${detail}）`;
+}
+
+/** 產生 PR Review 本文用的「AI 助理使用量」Markdown 區塊。 */
+export function formatUsageStats(provider, model, usage, quota, rate) {
+  const pct = resolveRemainingPercent(quota, rate);
+  const lines = [
+    '## 🤖 AI 助理使用量',
+    '',
+    `**本次審查**（${provider} / ${model}，共 ${usage.calls} 次呼叫）`,
+    '',
+    '| 提示 token | 回應 token | 合計 |',
+    '| --- | --- | --- |',
+    `| ${fmt(usage.promptTokens)} | ${fmt(usage.completionTokens)} | ${fmt(usage.totalTokens)} |`,
+    '',
+    '**剩餘可用**',
+    '',
+    remainingLine(pct),
+  ];
+  return lines.join('\n');
+}
+
+/** 產生單行 log 用的使用量摘要。 */
+export function formatUsageStatsLine(provider, model, usage, quota, rate) {
+  const pct = resolveRemainingPercent(quota, rate);
+  const tokenPart = `本次 ${provider}/${model}: 提示${usage.promptTokens} + 回應${usage.completionTokens} = ${usage.totalTokens} token（${usage.calls} 次呼叫）`;
+  const pctPart = pct.percent == null
+    ? `；剩餘可用: 無法計算（${pct.reason}）`
+    : `；剩餘可用: ${pct.percent}%（${pct.basis} ${money(pct.unit, pct.remaining)}/${money(pct.unit, pct.limit)}）`;
+  return tokenPart + pctPart;
+}