267 lines
11 KiB
JavaScript
267 lines
11 KiB
JavaScript
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 };
|
||
|
||
/**
|
||
* 從回應 header 擷取速率配額剩餘量／上限。
|
||
* 支援 OpenAI 相容（x-ratelimit-*-tokens）與 Anthropic（anthropic-ratelimit-tokens-*），
|
||
* 兩者皆缺時退而採用 requests 維度。記錄「最近一次」的數值（即最新的視窗狀態）。
|
||
*/
|
||
export function recordRateLimit(headers) {
|
||
if (!headers || typeof headers !== 'object') return;
|
||
const h = {};
|
||
for (const k of Object.keys(headers)) h[k.toLowerCase()] = headers[k];
|
||
|
||
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，避免被偽造的 baseURL
|
||
* （如 `openrouter.ai.evil.com` 或 `evil.com/openrouter.ai`）矇騙而把 API key 送往惡意主機。
|
||
*/
|
||
function isOpenRouterBaseURL(baseURL) {
|
||
try {
|
||
const host = new URL(baseURL).hostname.toLowerCase();
|
||
return host === 'openrouter.ai' || host.endsWith('.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: '次數' };
|
||
|
||
/**
|
||
* 計算「剩餘可用百分比」，依優先序擇一：
|
||
* 1. 帳號額度（quota 有上限）→ 剩餘 credits / 上限；
|
||
* 2. 速率配額（rate limit header）→ 當前視窗剩餘 / 上限；
|
||
* 皆無法取得時回傳 { percent: null, reason }。
|
||
*/
|
||
export function resolveRemainingPercent(quota, rate) {
|
||
if (quota?.available && quota.limit != null && Number(quota.limit) > 0) {
|
||
const limit = Number(quota.limit);
|
||
const remaining = quota.remaining == null ? limit - num(quota.used) : Number(quota.remaining);
|
||
return { percent: round1((remaining / limit) * 100), basis: '帳號額度', remaining, limit, unit: quota.currency || '' };
|
||
}
|
||
if (rate?.hasData && Number(rate.limit) > 0) {
|
||
const limit = Number(rate.limit);
|
||
const remaining = Number(rate.remaining);
|
||
const kindLabel = RATE_KIND_LABEL[rate.kind] || rate.kind;
|
||
return { percent: round1((remaining / limit) * 100), basis: `速率配額（當前視窗，${kindLabel}）`, remaining, 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;
|
||
}
|