mirror of
https://github.com/Nezumi-2711/9router.git
synced 2026-09-22 13:38:31 +00:00
Fix : usage convert
This commit is contained in:
@@ -27,6 +27,7 @@ export const PROVIDER_MODELS = {
|
||||
{ id: "gemini-2.5-flash-lite", name: "Gemini 2.5 Flash Lite" },
|
||||
],
|
||||
qw: [ // Qwen Code
|
||||
// { id: "qwen3-coder-next", name: "Qwen3 Coder Next" },
|
||||
{ id: "qwen3-coder-plus", name: "Qwen3 Coder Plus" },
|
||||
{ id: "qwen3-coder-flash", name: "Qwen3 Coder Flash" },
|
||||
{ id: "vision-model", name: "Qwen3 Vision Model" },
|
||||
@@ -38,7 +39,7 @@ export const PROVIDER_MODELS = {
|
||||
{ id: "deepseek-r1", name: "DeepSeek R1" },
|
||||
{ id: "deepseek-v3.2-chat", name: "DeepSeek V3.2 Chat" },
|
||||
{ id: "deepseek-v3.2-reasoner", name: "DeepSeek V3.2 Reasoner" },
|
||||
{ id: "minimax-m2", name: "MiniMax M2" },
|
||||
{ id: "minimax-m2.1", name: "MiniMax M2.1" },
|
||||
{ id: "glm-4.7", name: "GLM 4.7" },
|
||||
],
|
||||
ag: [ // Antigravity - special case: models call different backends
|
||||
@@ -59,7 +60,7 @@ export const PROVIDER_MODELS = {
|
||||
{ id: "gpt-5.1-codex", name: "GPT-5.1 Codex" },
|
||||
// { id: "gpt-5.1-codex-mini", name: "GPT-5.1 Codex Mini" },
|
||||
{ id: "gpt-5.1-codex-max", name: "GPT-5.1 Codex Max" },
|
||||
{ id: "gpt-4.1", name: "GPT-4.1" },
|
||||
// { id: "gpt-4.1", name: "GPT-4.1" },
|
||||
{ id: "claude-4.5-sonnet", name: "Claude 4.5 Sonnet" },
|
||||
{ id: "claude-4.5-opus", name: "Claude 4.5 Opus" },
|
||||
{ id: "claude-4.5-haiku", name: "Claude 4.5 Haiku" },
|
||||
@@ -101,6 +102,8 @@ export const PROVIDER_MODELS = {
|
||||
{ id: "glm-4.6v", name: "GLM 4.6V (Vision)" },
|
||||
],
|
||||
kimi: [
|
||||
{ id: "kimi-k2.5", name: "Kimi K2.5" },
|
||||
{ id: "kimi-k2.5-thinking", name: "Kimi K2.5 Thinking" },
|
||||
{ id: "kimi-latest", name: "Kimi Latest" },
|
||||
],
|
||||
minimax: [
|
||||
|
||||
@@ -3,7 +3,7 @@ import { translateRequest, needsTranslation } from "../translator/index.js";
|
||||
import { FORMATS } from "../translator/formats.js";
|
||||
import { createSSETransformStreamWithLogger, createPassthroughStreamWithLogger, COLORS } from "../utils/stream.js";
|
||||
import { createStreamController, pipeWithDisconnect } from "../utils/streamHandler.js";
|
||||
import { addBufferToUsage } from "../utils/usageTracking.js";
|
||||
import { addBufferToUsage, filterUsageForFormat } from "../utils/usageTracking.js";
|
||||
import { refreshWithRetry } from "../services/tokenRefresh.js";
|
||||
import { createRequestLogger } from "../utils/requestLogger.js";
|
||||
import { getModelTargetFormat, PROVIDER_ID_TO_ALIAS } from "../config/providerModels.js";
|
||||
@@ -435,9 +435,10 @@ export async function handleChatCore({ body, modelInfo, credentials, log, onCred
|
||||
? translateNonStreamingResponse(responseBody, targetFormat, sourceFormat)
|
||||
: responseBody;
|
||||
|
||||
// Add buffer to usage for client (to prevent CLI context errors)
|
||||
// Add buffer and filter usage for client (to prevent CLI context errors)
|
||||
if (translatedResponse?.usage) {
|
||||
translatedResponse.usage = addBufferToUsage(translatedResponse.usage);
|
||||
const buffered = addBufferToUsage(translatedResponse.usage);
|
||||
translatedResponse.usage = filterUsageForFormat(buffered, sourceFormat);
|
||||
}
|
||||
|
||||
return {
|
||||
|
||||
@@ -48,6 +48,8 @@ export async function getUsageForProvider(connection) {
|
||||
return await getClaudeUsage(accessToken);
|
||||
case "codex":
|
||||
return await getCodexUsage(accessToken);
|
||||
case "kiro":
|
||||
return await getKiroUsage(accessToken, providerSpecificData);
|
||||
case "qwen":
|
||||
return await getQwenUsage(accessToken, providerSpecificData);
|
||||
case "iflow":
|
||||
@@ -367,6 +369,82 @@ async function getCodexUsage(accessToken) {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Kiro (AWS CodeWhisperer) Usage
|
||||
*/
|
||||
async function getKiroUsage(accessToken, providerSpecificData) {
|
||||
try {
|
||||
const profileArn = providerSpecificData?.profileArn;
|
||||
if (!profileArn) {
|
||||
return { message: "Kiro connected. Profile ARN not available for quota tracking." };
|
||||
}
|
||||
|
||||
// Kiro uses AWS CodeWhisperer GetUsageLimits API
|
||||
const payload = {
|
||||
origin: "AI_EDITOR",
|
||||
profileArn: profileArn,
|
||||
resourceType: "AGENTIC_REQUEST",
|
||||
};
|
||||
|
||||
const response = await fetch("https://codewhisperer.us-east-1.amazonaws.com", {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Authorization": `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/x-amz-json-1.0",
|
||||
"x-amz-target": "AmazonCodeWhispererService.GetUsageLimits",
|
||||
"Accept": "application/json",
|
||||
},
|
||||
body: JSON.stringify(payload),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const errorText = await response.text();
|
||||
throw new Error(`Kiro API error (${response.status}): ${errorText}`);
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
// Parse usage data from usageBreakdownList
|
||||
const usageList = data.usageBreakdownList || [];
|
||||
const quotaInfo = {};
|
||||
|
||||
usageList.forEach((breakdown) => {
|
||||
const resourceType = breakdown.resourceType?.toLowerCase() || "unknown";
|
||||
const used = breakdown.currentUsageWithPrecision || 0;
|
||||
const total = breakdown.usageLimitWithPrecision || 0;
|
||||
|
||||
quotaInfo[resourceType] = {
|
||||
used,
|
||||
total,
|
||||
remaining: total - used,
|
||||
resetTime: data.nextDateReset ? new Date(data.nextDateReset).toISOString() : null,
|
||||
unlimited: false,
|
||||
};
|
||||
|
||||
// Add free trial if available
|
||||
if (breakdown.freeTrialInfo) {
|
||||
const freeUsed = breakdown.freeTrialInfo.currentUsageWithPrecision || 0;
|
||||
const freeTotal = breakdown.freeTrialInfo.usageLimitWithPrecision || 0;
|
||||
|
||||
quotaInfo[`${resourceType}_freetrial`] = {
|
||||
used: freeUsed,
|
||||
total: freeTotal,
|
||||
remaining: freeTotal - freeUsed,
|
||||
resetTime: data.nextDateReset ? new Date(data.nextDateReset).toISOString() : null,
|
||||
unlimited: false,
|
||||
};
|
||||
}
|
||||
});
|
||||
|
||||
return {
|
||||
plan: data.subscriptionInfo?.subscriptionTitle || "Kiro",
|
||||
quotas: quotaInfo,
|
||||
};
|
||||
} catch (error) {
|
||||
throw new Error(`Failed to fetch Kiro usage: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Qwen Usage
|
||||
*/
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import { translateResponse, initState } from "../translator/index.js";
|
||||
import { FORMATS } from "../translator/formats.js";
|
||||
import { trackPendingRequest, appendRequestLog } from "@/lib/usageDb.js";
|
||||
import { extractUsage, hasValidUsage, estimateUsage, logUsage, addBufferToUsage, COLORS } from "./usageTracking.js";
|
||||
import { extractUsage, hasValidUsage, estimateUsage, logUsage, addBufferToUsage, filterUsageForFormat, COLORS } from "./usageTracking.js";
|
||||
|
||||
// Re-export COLORS for backward compatibility
|
||||
export { COLORS };
|
||||
@@ -145,13 +145,14 @@ export function createSSEStream(options = {}) {
|
||||
const isFinishChunk = parsed.choices?.[0]?.finish_reason;
|
||||
if (isFinishChunk && !hasValidUsage(parsed.usage)) {
|
||||
const estimated = estimateUsage(body, totalContentLength, FORMATS.OPENAI);
|
||||
parsed.usage = estimated; // Already has buffer from formatUsage
|
||||
parsed.usage = filterUsageForFormat(estimated, FORMATS.OPENAI); // Filter + already has buffer
|
||||
output = `data: ${JSON.stringify(parsed)}\n`;
|
||||
usage = estimated;
|
||||
injectedUsage = true;
|
||||
} else if (isFinishChunk && usage) {
|
||||
// Add buffer to usage for client (but keep original for logging)
|
||||
parsed.usage = addBufferToUsage(usage);
|
||||
// Add buffer and filter usage for client (but keep original for logging)
|
||||
const buffered = addBufferToUsage(usage);
|
||||
parsed.usage = filterUsageForFormat(buffered, FORMATS.OPENAI);
|
||||
output = `data: ${JSON.stringify(parsed)}\n`;
|
||||
injectedUsage = true;
|
||||
}
|
||||
@@ -234,11 +235,12 @@ export function createSSEStream(options = {}) {
|
||||
const isFinishChunk = item.type === "message_delta" || item.choices?.[0]?.finish_reason;
|
||||
if (state.finishReason && isFinishChunk && !hasValidUsage(item.usage) && totalContentLength > 0) {
|
||||
const estimated = estimateUsage(body, totalContentLength, sourceFormat);
|
||||
item.usage = estimated; // Already has buffer from formatUsage
|
||||
item.usage = filterUsageForFormat(estimated, sourceFormat); // Filter + already has buffer
|
||||
state.usage = estimated;
|
||||
} else if (state.finishReason && isFinishChunk && state.usage) {
|
||||
// Add buffer to usage for client (but keep original in state.usage for logging)
|
||||
item.usage = addBufferToUsage(state.usage);
|
||||
// Add buffer and filter usage for client (but keep original in state.usage for logging)
|
||||
const buffered = addBufferToUsage(state.usage);
|
||||
item.usage = filterUsageForFormat(buffered, sourceFormat);
|
||||
}
|
||||
|
||||
const output = formatSSE(item, sourceFormat);
|
||||
|
||||
@@ -43,14 +43,72 @@ export function addBufferToUsage(usage) {
|
||||
result.prompt_tokens += BUFFER_TOKENS;
|
||||
}
|
||||
|
||||
// Update total_tokens if exists
|
||||
// Calculate or update total_tokens
|
||||
if (result.total_tokens !== undefined) {
|
||||
result.total_tokens += BUFFER_TOKENS;
|
||||
} else if (result.prompt_tokens !== undefined && result.completion_tokens !== undefined) {
|
||||
// Calculate total_tokens if not exists
|
||||
result.total_tokens = result.prompt_tokens + result.completion_tokens;
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
export function filterUsageForFormat(usage, targetFormat) {
|
||||
if (!usage || typeof usage !== "object") return usage;
|
||||
|
||||
// Helper to pick only defined fields from usage
|
||||
const pickFields = (fields) => {
|
||||
const filtered = {};
|
||||
for (const field of fields) {
|
||||
if (usage[field] !== undefined) {
|
||||
filtered[field] = usage[field];
|
||||
}
|
||||
}
|
||||
return filtered;
|
||||
};
|
||||
|
||||
// Define allowed fields for each format
|
||||
const formatFields = {
|
||||
[FORMATS.CLAUDE]: [
|
||||
'input_tokens', 'output_tokens',
|
||||
'cache_read_input_tokens', 'cache_creation_input_tokens',
|
||||
'estimated'
|
||||
],
|
||||
[FORMATS.GEMINI]: [
|
||||
'promptTokenCount', 'candidatesTokenCount', 'totalTokenCount',
|
||||
'cachedContentTokenCount', 'thoughtsTokenCount',
|
||||
'estimated'
|
||||
],
|
||||
[FORMATS.OPENAI_RESPONSES]: [
|
||||
'input_tokens', 'output_tokens',
|
||||
'input_tokens_details', 'output_tokens_details',
|
||||
'estimated'
|
||||
],
|
||||
// OpenAI format (default for OPENAI, CODEX, KIRO, etc.)
|
||||
default: [
|
||||
'prompt_tokens', 'completion_tokens', 'total_tokens',
|
||||
'cached_tokens', 'reasoning_tokens',
|
||||
'prompt_tokens_details', 'completion_tokens_details',
|
||||
'estimated'
|
||||
]
|
||||
};
|
||||
|
||||
// Get fields for target format
|
||||
let fields = formatFields[targetFormat];
|
||||
|
||||
// Use same fields for similar formats
|
||||
if (targetFormat === FORMATS.GEMINI_CLI || targetFormat === FORMATS.ANTIGRAVITY) {
|
||||
fields = formatFields[FORMATS.GEMINI];
|
||||
} else if (targetFormat === FORMATS.OPENAI_RESPONSE) {
|
||||
fields = formatFields[FORMATS.OPENAI_RESPONSES];
|
||||
} else if (!fields) {
|
||||
fields = formatFields.default;
|
||||
}
|
||||
|
||||
return pickFields(fields);
|
||||
}
|
||||
|
||||
/**
|
||||
* Normalize usage object - ensure all values are valid numbers
|
||||
*/
|
||||
@@ -66,6 +124,7 @@ export function normalizeUsage(usage) {
|
||||
|
||||
assignNumber("prompt_tokens", usage?.prompt_tokens);
|
||||
assignNumber("completion_tokens", usage?.completion_tokens);
|
||||
assignNumber("total_tokens", usage?.total_tokens);
|
||||
assignNumber("cache_read_input_tokens", usage?.cache_read_input_tokens);
|
||||
assignNumber("cache_creation_input_tokens", usage?.cache_creation_input_tokens);
|
||||
assignNumber("cached_tokens", usage?.cached_tokens);
|
||||
@@ -141,6 +200,7 @@ export function extractUsage(chunk) {
|
||||
return normalizeUsage({
|
||||
prompt_tokens: chunk.usageMetadata?.promptTokenCount || 0,
|
||||
completion_tokens: chunk.usageMetadata?.candidatesTokenCount || 0,
|
||||
total_tokens: chunk.usageMetadata?.totalTokenCount,
|
||||
cached_tokens: chunk.usageMetadata?.cachedContentTokenCount,
|
||||
reasoning_tokens: chunk.usageMetadata?.thoughtsTokenCount
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user