feat: implement usage tracking for AI requests
Adds local token usage tracking for all AI providers. Usage data is captured during stream processing and stored in a local database. Includes a new Usage tab in the Providers dashboard to visualize historical token consumption. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -8,6 +8,46 @@ import { createRequestLogger } from "../utils/requestLogger.js";
|
||||
import { getModelTargetFormat, PROVIDER_ID_TO_ALIAS } from "../config/providerModels.js";
|
||||
import { createErrorResult, parseUpstreamError, formatProviderError } from "../utils/error.js";
|
||||
import { handleBypassRequest } from "../utils/bypassHandler.js";
|
||||
import { saveRequestUsage } from "@/lib/usageDb.js";
|
||||
|
||||
/**
|
||||
* Extract usage from non-streaming response body
|
||||
* Handles different provider response formats
|
||||
*/
|
||||
function extractUsageFromResponse(responseBody, provider) {
|
||||
if (!responseBody) return null;
|
||||
|
||||
// OpenAI format
|
||||
if (responseBody.usage) {
|
||||
return {
|
||||
prompt_tokens: responseBody.usage.prompt_tokens || 0,
|
||||
completion_tokens: responseBody.usage.completion_tokens || 0,
|
||||
cached_tokens: responseBody.usage.prompt_tokens_details?.cached_tokens,
|
||||
reasoning_tokens: responseBody.usage.completion_tokens_details?.reasoning_tokens
|
||||
};
|
||||
}
|
||||
|
||||
// Claude format
|
||||
if (responseBody.usage?.input_tokens !== undefined || responseBody.usage?.output_tokens !== undefined) {
|
||||
return {
|
||||
prompt_tokens: responseBody.usage.input_tokens || 0,
|
||||
completion_tokens: responseBody.usage.output_tokens || 0,
|
||||
cache_read_input_tokens: responseBody.usage.cache_read_input_tokens,
|
||||
cache_creation_input_tokens: responseBody.usage.cache_creation_input_tokens
|
||||
};
|
||||
}
|
||||
|
||||
// Gemini format
|
||||
if (responseBody.usageMetadata) {
|
||||
return {
|
||||
prompt_tokens: responseBody.usageMetadata.promptTokenCount || 0,
|
||||
completion_tokens: responseBody.usageMetadata.candidatesTokenCount || 0,
|
||||
reasoning_tokens: responseBody.usageMetadata.thoughtsTokenCount
|
||||
};
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Core chat handler - shared between SSE and Worker
|
||||
@@ -20,8 +60,9 @@ import { handleBypassRequest } from "../utils/bypassHandler.js";
|
||||
* @param {function} options.onCredentialsRefreshed - Callback when credentials are refreshed
|
||||
* @param {function} options.onRequestSuccess - Callback when request succeeds (to clear error status)
|
||||
* @param {function} options.onDisconnect - Callback when client disconnects
|
||||
* @param {string} options.connectionId - Connection ID for usage tracking
|
||||
*/
|
||||
export async function handleChatCore({ body, modelInfo, credentials, log, onCredentialsRefreshed, onRequestSuccess, onDisconnect, clientRawRequest }) {
|
||||
export async function handleChatCore({ body, modelInfo, credentials, log, onCredentialsRefreshed, onRequestSuccess, onDisconnect, clientRawRequest, connectionId }) {
|
||||
const { provider, model } = modelInfo;
|
||||
|
||||
const sourceFormat = detectFormat(body);
|
||||
@@ -220,12 +261,29 @@ export async function handleChatCore({ body, modelInfo, credentials, log, onCred
|
||||
// Non-streaming response
|
||||
if (!stream) {
|
||||
const responseBody = await providerResponse.json();
|
||||
|
||||
|
||||
// Notify success - caller can clear error status if needed
|
||||
if (onRequestSuccess) {
|
||||
await onRequestSuccess();
|
||||
}
|
||||
|
||||
|
||||
// Log usage for non-streaming responses
|
||||
const usage = extractUsageFromResponse(responseBody, provider);
|
||||
if (usage) {
|
||||
const msg = `[${new Date().toLocaleTimeString("en-US", { hour12: false, hour: "2-digit", minute: "2-digit" })}] 📊 [USAGE] ${provider.toUpperCase()} | in=${usage.prompt_tokens || 0} | out=${usage.completion_tokens || 0}${connectionId ? ` | account=${connectionId.slice(0, 8)}...` : ""}`;
|
||||
console.log(`${COLORS.green}${msg}${COLORS.reset}`);
|
||||
|
||||
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);
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
success: true,
|
||||
response: new Response(JSON.stringify(responseBody), {
|
||||
@@ -254,9 +312,9 @@ export async function handleChatCore({ body, modelInfo, credentials, log, onCred
|
||||
// Create transform stream with logger for streaming response
|
||||
let transformStream;
|
||||
if (needsTranslation(targetFormat, sourceFormat)) {
|
||||
transformStream = createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider, reqLogger, toolNameMap);
|
||||
transformStream = createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider, reqLogger, toolNameMap, model, connectionId);
|
||||
} else {
|
||||
transformStream = createPassthroughStreamWithLogger(provider, reqLogger);
|
||||
transformStream = createPassthroughStreamWithLogger(provider, reqLogger, model, connectionId);
|
||||
}
|
||||
|
||||
// Pipe response through transform with disconnect detection
|
||||
|
||||
Reference in New Issue
Block a user