mirror of
https://github.com/Nezumi-2711/9router.git
synced 2026-09-22 13:38:31 +00:00
feat: Update response handling and logging for improved usage tracking
This commit is contained in:
decolua
+79
-136
@@ -1,130 +1,15 @@
|
||||
import { translateResponse, initState } from "../translator/index.js";
|
||||
import { FORMATS } from "../translator/formats.js";
|
||||
import { saveRequestUsage, trackPendingRequest, appendRequestLog } from "@/lib/usageDb.js";
|
||||
import { trackPendingRequest, appendRequestLog } from "@/lib/usageDb.js";
|
||||
import { extractUsage, hasValidUsage, estimateUsage, logUsage, COLORS } from "./usageTracking.js";
|
||||
|
||||
// Re-export COLORS for backward compatibility
|
||||
export { COLORS };
|
||||
|
||||
// Singleton TextEncoder/Decoder for performance (reuse across all streams)
|
||||
const sharedDecoder = new TextDecoder();
|
||||
const sharedEncoder = new TextEncoder();
|
||||
|
||||
// Get HH:MM:SS timestamp
|
||||
function getTimeString() {
|
||||
return new Date().toLocaleTimeString("en-US", { hour12: false, hour: "2-digit", minute: "2-digit", second: "2-digit" });
|
||||
}
|
||||
|
||||
// Extract usage from any format (Claude, OpenAI, Gemini, Responses API)
|
||||
function extractUsage(chunk) {
|
||||
if (!chunk || typeof chunk !== "object") return null;
|
||||
|
||||
// Claude format (message_delta event)
|
||||
if (chunk.type === "message_delta" && chunk.usage && typeof chunk.usage === 'object') {
|
||||
return normalizeUsage({
|
||||
prompt_tokens: chunk.usage.input_tokens || 0,
|
||||
completion_tokens: chunk.usage.output_tokens || 0,
|
||||
cache_read_input_tokens: chunk.usage.cache_read_input_tokens,
|
||||
cache_creation_input_tokens: chunk.usage.cache_creation_input_tokens
|
||||
});
|
||||
}
|
||||
// OpenAI Responses API format (response.completed or response.done)
|
||||
if ((chunk.type === "response.completed" || chunk.type === "response.done") && chunk.response?.usage && typeof chunk.response.usage === 'object') {
|
||||
const usage = chunk.response.usage;
|
||||
return normalizeUsage({
|
||||
prompt_tokens: usage.input_tokens || usage.prompt_tokens || 0,
|
||||
completion_tokens: usage.output_tokens || usage.completion_tokens || 0,
|
||||
cached_tokens: usage.input_tokens_details?.cached_tokens,
|
||||
reasoning_tokens: usage.output_tokens_details?.reasoning_tokens
|
||||
});
|
||||
}
|
||||
// OpenAI format
|
||||
if (chunk.usage && typeof chunk.usage === 'object' && chunk.usage.prompt_tokens !== undefined) {
|
||||
return normalizeUsage({
|
||||
prompt_tokens: chunk.usage.prompt_tokens,
|
||||
completion_tokens: chunk.usage.completion_tokens || 0,
|
||||
cached_tokens: chunk.usage.prompt_tokens_details?.cached_tokens,
|
||||
reasoning_tokens: chunk.usage.completion_tokens_details?.reasoning_tokens
|
||||
});
|
||||
}
|
||||
// Gemini format (Antigravity)
|
||||
if (chunk.usageMetadata && typeof chunk.usageMetadata === 'object') {
|
||||
return normalizeUsage({
|
||||
prompt_tokens: chunk.usageMetadata?.promptTokenCount || 0,
|
||||
completion_tokens: chunk.usageMetadata?.candidatesTokenCount || 0,
|
||||
cached_tokens: chunk.usageMetadata?.cachedContentTokenCount,
|
||||
reasoning_tokens: chunk.usageMetadata?.thoughtsTokenCount
|
||||
});
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
function normalizeUsage(usage) {
|
||||
if (!usage || typeof usage !== "object" || Array.isArray(usage)) return null;
|
||||
|
||||
const normalized = {};
|
||||
const assignNumber = (key, value) => {
|
||||
if (value === undefined || value === null) return;
|
||||
const numeric = Number(value);
|
||||
if (Number.isFinite(numeric)) normalized[key] = numeric;
|
||||
};
|
||||
|
||||
assignNumber("prompt_tokens", usage?.prompt_tokens);
|
||||
assignNumber("completion_tokens", usage?.completion_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);
|
||||
assignNumber("reasoning_tokens", usage?.reasoning_tokens);
|
||||
|
||||
if (Object.keys(normalized).length === 0) return null;
|
||||
return normalized;
|
||||
}
|
||||
|
||||
// ANSI color codes
|
||||
export const COLORS = {
|
||||
reset: "\x1b[0m",
|
||||
red: "\x1b[31m",
|
||||
green: "\x1b[32m",
|
||||
yellow: "\x1b[33m",
|
||||
blue: "\x1b[34m",
|
||||
cyan: "\x1b[36m"
|
||||
};
|
||||
|
||||
// Log usage with cache info (green color)
|
||||
function logUsage(provider, usage, model = null, connectionId = null) {
|
||||
if (!usage || typeof usage !== 'object') return;
|
||||
|
||||
const p = provider?.toUpperCase() || "UNKNOWN";
|
||||
|
||||
// Support both formats:
|
||||
// - OpenAI: prompt_tokens, completion_tokens
|
||||
// - Claude: input_tokens, output_tokens
|
||||
const inTokens = usage?.prompt_tokens || usage?.input_tokens || 0;
|
||||
const outTokens = usage?.completion_tokens || usage?.output_tokens || 0;
|
||||
|
||||
let msg = `[${getTimeString()}] 📊 [USAGE] ${p} | in=${inTokens} | out=${outTokens}`;
|
||||
if (connectionId) msg += ` | account=${connectionId.slice(0, 8)}...`;
|
||||
|
||||
// Support both formats: cache_read_input_tokens (Claude) and cached_tokens (OpenAI/Gemini)
|
||||
const cacheRead = usage.cache_read_input_tokens || usage.cached_tokens;
|
||||
if (cacheRead) msg += ` | cache_read=${cacheRead}`;
|
||||
|
||||
if (usage.cache_creation_input_tokens) msg += ` | cache_write=${usage.cache_creation_input_tokens}`;
|
||||
if (usage.reasoning_tokens) msg += ` | reasoning=${usage.reasoning_tokens}`;
|
||||
|
||||
console.log(`${COLORS.green}${msg}${COLORS.reset}`);
|
||||
|
||||
// Log to log.txt
|
||||
appendRequestLog({ model, provider, connectionId, tokens: usage, status: "200 OK" }).catch(() => { });
|
||||
|
||||
// Save to DB
|
||||
saveRequestUsage({
|
||||
provider: provider || "unknown",
|
||||
model: model || "unknown",
|
||||
tokens: usage,
|
||||
timestamp: new Date().toISOString(),
|
||||
connectionId: connectionId || undefined
|
||||
}).catch(err => {
|
||||
console.error("Failed to save usage stats:", err.message);
|
||||
});
|
||||
}
|
||||
|
||||
// Parse SSE data line (optimized - reduce string operations)
|
||||
function parseSSELine(line) {
|
||||
if (!line || line.charCodeAt(0) !== 100) return null; // 'd' = 100
|
||||
@@ -199,6 +84,7 @@ const STREAM_MODE = {
|
||||
* @param {object} options.reqLogger - Request logger instance
|
||||
* @param {string} options.model - Model name
|
||||
* @param {string} options.connectionId - Connection ID for usage tracking
|
||||
* @param {object} options.body - Request body (for input token estimation)
|
||||
*/
|
||||
export function createSSEStream(options = {}) {
|
||||
const {
|
||||
@@ -209,7 +95,8 @@ export function createSSEStream(options = {}) {
|
||||
reqLogger = null,
|
||||
toolNameMap = null,
|
||||
model = null,
|
||||
connectionId = null
|
||||
connectionId = null,
|
||||
body = null
|
||||
} = options;
|
||||
|
||||
let buffer = "";
|
||||
@@ -218,6 +105,9 @@ export function createSSEStream(options = {}) {
|
||||
// State for translate mode
|
||||
const state = mode === STREAM_MODE.TRANSLATE ? { ...initState(sourceFormat), provider, toolNameMap } : null;
|
||||
|
||||
// Track content length for usage estimation (both modes)
|
||||
let totalContentLength = 0;
|
||||
|
||||
return new TransformStream({
|
||||
transform(chunk, controller) {
|
||||
const text = sharedDecoder.decode(chunk, { stream: true });
|
||||
@@ -232,20 +122,46 @@ export function createSSEStream(options = {}) {
|
||||
|
||||
// Passthrough mode: normalize and forward
|
||||
if (mode === STREAM_MODE.PASSTHROUGH) {
|
||||
let output;
|
||||
let injectedUsage = false;
|
||||
|
||||
if (trimmed.startsWith("data:") && trimmed.slice(5).trim() !== "[DONE]") {
|
||||
try {
|
||||
const parsed = JSON.parse(trimmed.slice(5).trim());
|
||||
|
||||
// Track content length for estimation
|
||||
const content = parsed.choices?.[0]?.delta?.content || parsed.choices?.[0]?.delta?.reasoning_content;
|
||||
if (content && typeof content === "string") {
|
||||
totalContentLength += content.length;
|
||||
}
|
||||
|
||||
// Extract usage from chunk
|
||||
const extracted = extractUsage(parsed);
|
||||
if (extracted) usage = extracted;
|
||||
if (extracted) {
|
||||
usage = extracted;
|
||||
}
|
||||
|
||||
// Inject estimated usage into final chunk (has finish_reason but no valid usage)
|
||||
const isFinishChunk = parsed.choices?.[0]?.finish_reason;
|
||||
if (isFinishChunk && !hasValidUsage(parsed.usage)) {
|
||||
const estimated = estimateUsage(body, totalContentLength, FORMATS.OPENAI);
|
||||
parsed.usage = estimated;
|
||||
output = `data: ${JSON.stringify(parsed)}\n`;
|
||||
usage = estimated;
|
||||
injectedUsage = true;
|
||||
}
|
||||
} catch { }
|
||||
}
|
||||
// Normalize: ensure "data: " has space
|
||||
let output;
|
||||
if (line.startsWith("data:") && !line.startsWith("data: ")) {
|
||||
output = "data: " + line.slice(5) + "\n";
|
||||
} else {
|
||||
output = line + "\n";
|
||||
|
||||
// Normalize if not already injected
|
||||
if (!injectedUsage) {
|
||||
if (line.startsWith("data:") && !line.startsWith("data: ")) {
|
||||
output = "data: " + line.slice(5) + "\n";
|
||||
} else {
|
||||
output = line + "\n";
|
||||
}
|
||||
}
|
||||
|
||||
reqLogger?.appendConvertedChunk?.(output);
|
||||
controller.enqueue(sharedEncoder.encode(output));
|
||||
continue;
|
||||
@@ -264,6 +180,14 @@ export function createSSEStream(options = {}) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Track content length for estimation (from various formats)
|
||||
const content = parsed.delta?.text || // Claude
|
||||
parsed.choices?.[0]?.delta?.content || // OpenAI
|
||||
parsed.candidates?.[0]?.content?.parts?.[0]?.text; // Gemini
|
||||
if (content && typeof content === "string") {
|
||||
totalContentLength += content.length;
|
||||
}
|
||||
|
||||
// Extract usage
|
||||
const extracted = extractUsage(parsed);
|
||||
if (extracted) state.usage = extracted;
|
||||
@@ -281,6 +205,14 @@ export function createSSEStream(options = {}) {
|
||||
|
||||
if (translated?.length > 0) {
|
||||
for (const item of translated) {
|
||||
// Inject estimated usage if finish chunk has no valid usage
|
||||
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;
|
||||
state.usage = estimated;
|
||||
}
|
||||
|
||||
const output = formatSSE(item, sourceFormat);
|
||||
reqLogger?.appendConvertedChunk?.(output);
|
||||
controller.enqueue(sharedEncoder.encode(output));
|
||||
@@ -304,10 +236,15 @@ export function createSSEStream(options = {}) {
|
||||
reqLogger?.appendConvertedChunk?.(output);
|
||||
controller.enqueue(sharedEncoder.encode(output));
|
||||
}
|
||||
if (usage && typeof usage === 'object') {
|
||||
|
||||
// Estimate usage if provider didn't return valid usage (PASSTHROUGH is always OpenAI format)
|
||||
if (!hasValidUsage(usage) && totalContentLength > 0) {
|
||||
usage = estimateUsage(body, totalContentLength, FORMATS.OPENAI);
|
||||
}
|
||||
|
||||
if (hasValidUsage(usage)) {
|
||||
logUsage(provider, usage, model, connectionId);
|
||||
} else {
|
||||
// No usage data available - still mark request as completed
|
||||
appendRequestLog({ model, provider, connectionId, tokens: null, status: "200 OK" }).catch(() => { });
|
||||
}
|
||||
return;
|
||||
@@ -361,10 +298,14 @@ export function createSSEStream(options = {}) {
|
||||
reqLogger?.appendConvertedChunk?.(doneOutput);
|
||||
controller.enqueue(sharedEncoder.encode(doneOutput));
|
||||
|
||||
if (state?.usage && typeof state.usage === 'object') {
|
||||
// Estimate usage if provider didn't return valid usage (for translate mode)
|
||||
if (!hasValidUsage(state?.usage) && totalContentLength > 0) {
|
||||
state.usage = estimateUsage(body, totalContentLength, sourceFormat);
|
||||
}
|
||||
|
||||
if (hasValidUsage(state?.usage)) {
|
||||
logUsage(state.provider || targetFormat, state.usage, model, connectionId);
|
||||
} else {
|
||||
// No usage data available - still mark request as completed
|
||||
appendRequestLog({ model, provider, connectionId, tokens: null, status: "200 OK" }).catch(() => { });
|
||||
}
|
||||
} catch (error) {
|
||||
@@ -375,7 +316,7 @@ export function createSSEStream(options = {}) {
|
||||
}
|
||||
|
||||
// Convenience functions for backward compatibility
|
||||
export function createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider = null, reqLogger = null, toolNameMap = null, model = null, connectionId = null) {
|
||||
export function createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider = null, reqLogger = null, toolNameMap = null, model = null, connectionId = null, body = null) {
|
||||
return createSSEStream({
|
||||
mode: STREAM_MODE.TRANSLATE,
|
||||
targetFormat,
|
||||
@@ -384,16 +325,18 @@ export function createSSETransformStreamWithLogger(targetFormat, sourceFormat, p
|
||||
reqLogger,
|
||||
toolNameMap,
|
||||
model,
|
||||
connectionId
|
||||
connectionId,
|
||||
body
|
||||
});
|
||||
}
|
||||
|
||||
export function createPassthroughStreamWithLogger(provider = null, reqLogger = null, model = null, connectionId = null) {
|
||||
export function createPassthroughStreamWithLogger(provider = null, reqLogger = null, model = null, connectionId = null, body = null) {
|
||||
return createSSEStream({
|
||||
mode: STREAM_MODE.PASSTHROUGH,
|
||||
provider,
|
||||
reqLogger,
|
||||
model,
|
||||
connectionId
|
||||
connectionId,
|
||||
body
|
||||
});
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user