mirror of
https://github.com/Wei-Shaw/claude-relay-service.git
synced 2026-01-23 18:14:51 +00:00
feat: 增加每日费用限制
This commit is contained in:
KevinLiao
@@ -26,7 +26,8 @@ class ApiKeyService {
|
||||
enableModelRestriction = false,
|
||||
restrictedModels = [],
|
||||
enableClientRestriction = false,
|
||||
allowedClients = []
|
||||
allowedClients = [],
|
||||
dailyCostLimit = 0
|
||||
} = options;
|
||||
|
||||
// 生成简单的API Key (64字符十六进制)
|
||||
@@ -51,6 +52,7 @@ class ApiKeyService {
|
||||
restrictedModels: JSON.stringify(restrictedModels || []),
|
||||
enableClientRestriction: String(enableClientRestriction || false),
|
||||
allowedClients: JSON.stringify(allowedClients || []),
|
||||
dailyCostLimit: String(dailyCostLimit || 0),
|
||||
createdAt: new Date().toISOString(),
|
||||
lastUsedAt: '',
|
||||
expiresAt: expiresAt || '',
|
||||
@@ -79,6 +81,7 @@ class ApiKeyService {
|
||||
restrictedModels: JSON.parse(keyData.restrictedModels),
|
||||
enableClientRestriction: keyData.enableClientRestriction === 'true',
|
||||
allowedClients: JSON.parse(keyData.allowedClients || '[]'),
|
||||
dailyCostLimit: parseFloat(keyData.dailyCostLimit || 0),
|
||||
createdAt: keyData.createdAt,
|
||||
expiresAt: keyData.expiresAt,
|
||||
createdBy: keyData.createdBy
|
||||
@@ -114,6 +117,9 @@ class ApiKeyService {
|
||||
|
||||
// 获取使用统计（供返回数据使用）
|
||||
const usage = await redis.getUsageStats(keyData.id);
|
||||
|
||||
// 获取当日费用统计
|
||||
const dailyCost = await redis.getDailyCost(keyData.id);
|
||||
|
||||
// 更新最后使用时间（优化：只在实际API调用时更新，而不是验证时）
|
||||
// 注意：lastUsedAt的更新已移至recordUsage方法中
|
||||
@@ -152,6 +158,8 @@ class ApiKeyService {
|
||||
restrictedModels: restrictedModels,
|
||||
enableClientRestriction: keyData.enableClientRestriction === 'true',
|
||||
allowedClients: allowedClients,
|
||||
dailyCostLimit: parseFloat(keyData.dailyCostLimit || 0),
|
||||
dailyCost: dailyCost || 0,
|
||||
usage
|
||||
}
|
||||
};
|
||||
@@ -178,6 +186,8 @@ class ApiKeyService {
|
||||
key.enableModelRestriction = key.enableModelRestriction === 'true';
|
||||
key.enableClientRestriction = key.enableClientRestriction === 'true';
|
||||
key.permissions = key.permissions || 'all'; // 兼容旧数据
|
||||
key.dailyCostLimit = parseFloat(key.dailyCostLimit || 0);
|
||||
key.dailyCost = await redis.getDailyCost(key.id) || 0;
|
||||
try {
|
||||
key.restrictedModels = key.restrictedModels ? JSON.parse(key.restrictedModels) : [];
|
||||
} catch (e) {
|
||||
@@ -207,7 +217,7 @@ class ApiKeyService {
|
||||
}
|
||||
|
||||
// 允许更新的字段
|
||||
const allowedUpdates = ['name', 'description', 'tokenLimit', 'concurrencyLimit', 'rateLimitWindow', 'rateLimitRequests', 'isActive', 'claudeAccountId', 'geminiAccountId', 'permissions', 'expiresAt', 'enableModelRestriction', 'restrictedModels', 'enableClientRestriction', 'allowedClients'];
|
||||
const allowedUpdates = ['name', 'description', 'tokenLimit', 'concurrencyLimit', 'rateLimitWindow', 'rateLimitRequests', 'isActive', 'claudeAccountId', 'geminiAccountId', 'permissions', 'expiresAt', 'enableModelRestriction', 'restrictedModels', 'enableClientRestriction', 'allowedClients', 'dailyCostLimit'];
|
||||
const updatedData = { ...keyData };
|
||||
|
||||
for (const [field, value] of Object.entries(updates)) {
|
||||
@@ -261,9 +271,26 @@ class ApiKeyService {
|
||||
try {
|
||||
const totalTokens = inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens;
|
||||
|
||||
// 计算费用
|
||||
const CostCalculator = require('../utils/costCalculator');
|
||||
const costInfo = CostCalculator.calculateCost({
|
||||
input_tokens: inputTokens,
|
||||
output_tokens: outputTokens,
|
||||
cache_creation_input_tokens: cacheCreateTokens,
|
||||
cache_read_input_tokens: cacheReadTokens
|
||||
}, model);
|
||||
|
||||
// 记录API Key级别的使用统计
|
||||
await redis.incrementTokenUsage(keyId, totalTokens, inputTokens, outputTokens, cacheCreateTokens, cacheReadTokens, model);
|
||||
|
||||
// 记录费用统计
|
||||
if (costInfo.costs.total > 0) {
|
||||
await redis.incrementDailyCost(keyId, costInfo.costs.total);
|
||||
logger.database(`💰 Recorded cost for ${keyId}: $${costInfo.costs.total.toFixed(6)}, model: ${model}`);
|
||||
} else {
|
||||
logger.debug(`💰 No cost recorded for ${keyId} - zero cost for model: ${model}`);
|
||||
}
|
||||
|
||||
// 获取API Key数据以确定关联的账户
|
||||
const keyData = await redis.getApiKey(keyId);
|
||||
if (keyData && Object.keys(keyData).length > 0) {
|
||||
@@ -276,7 +303,7 @@ class ApiKeyService {
|
||||
await redis.incrementAccountUsage(accountId, totalTokens, inputTokens, outputTokens, cacheCreateTokens, cacheReadTokens, model);
|
||||
logger.database(`📊 Recorded account usage: ${accountId} - ${totalTokens} tokens (API Key: ${keyId})`);
|
||||
} else {
|
||||
logger.debug(`⚠️ No accountId provided for usage recording, skipping account-level statistics`);
|
||||
logger.debug('⚠️ No accountId provided for usage recording, skipping account-level statistics');
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
182
src/services/costInitService.js
Normal file
182
src/services/costInitService.js
Normal file
@@ -0,0 +1,182 @@
|
||||
const redis = require('../models/redis');
|
||||
const apiKeyService = require('./apiKeyService');
|
||||
const CostCalculator = require('../utils/costCalculator');
|
||||
const logger = require('../utils/logger');
|
||||
|
||||
class CostInitService {
|
||||
/**
|
||||
* 初始化所有API Key的费用数据
|
||||
* 扫描历史使用记录并计算费用
|
||||
*/
|
||||
async initializeAllCosts() {
|
||||
try {
|
||||
logger.info('💰 Starting cost initialization for all API Keys...');
|
||||
|
||||
const apiKeys = await apiKeyService.getAllApiKeys();
|
||||
const client = redis.getClientSafe();
|
||||
|
||||
let processedCount = 0;
|
||||
let errorCount = 0;
|
||||
|
||||
for (const apiKey of apiKeys) {
|
||||
try {
|
||||
await this.initializeApiKeyCosts(apiKey.id, client);
|
||||
processedCount++;
|
||||
|
||||
if (processedCount % 10 === 0) {
|
||||
logger.info(`💰 Processed ${processedCount} API Keys...`);
|
||||
}
|
||||
} catch (error) {
|
||||
errorCount++;
|
||||
logger.error(`❌ Failed to initialize costs for API Key ${apiKey.id}:`, error);
|
||||
}
|
||||
}
|
||||
|
||||
logger.success(`💰 Cost initialization completed! Processed: ${processedCount}, Errors: ${errorCount}`);
|
||||
return { processed: processedCount, errors: errorCount };
|
||||
} catch (error) {
|
||||
logger.error('❌ Failed to initialize costs:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 初始化单个API Key的费用数据
|
||||
*/
|
||||
async initializeApiKeyCosts(apiKeyId, client) {
|
||||
// 获取所有时间的模型使用统计
|
||||
const modelKeys = await client.keys(`usage:${apiKeyId}:model:*:*:*`);
|
||||
|
||||
// 按日期分组统计
|
||||
const dailyCosts = new Map(); // date -> cost
|
||||
const monthlyCosts = new Map(); // month -> cost
|
||||
const hourlyCosts = new Map(); // hour -> cost
|
||||
|
||||
for (const key of modelKeys) {
|
||||
// 解析key格式: usage:{keyId}:model:{period}:{model}:{date}
|
||||
const match = key.match(/usage:(.+):model:(daily|monthly|hourly):(.+):(\d{4}-\d{2}(?:-\d{2})?(?::\d{2})?)$/);
|
||||
if (!match) continue;
|
||||
|
||||
const [, , period, model, dateStr] = match;
|
||||
|
||||
// 获取使用数据
|
||||
const data = await client.hgetall(key);
|
||||
if (!data || Object.keys(data).length === 0) continue;
|
||||
|
||||
// 计算费用
|
||||
const usage = {
|
||||
input_tokens: parseInt(data.totalInputTokens) || parseInt(data.inputTokens) || 0,
|
||||
output_tokens: parseInt(data.totalOutputTokens) || parseInt(data.outputTokens) || 0,
|
||||
cache_creation_input_tokens: parseInt(data.totalCacheCreateTokens) || parseInt(data.cacheCreateTokens) || 0,
|
||||
cache_read_input_tokens: parseInt(data.totalCacheReadTokens) || parseInt(data.cacheReadTokens) || 0
|
||||
};
|
||||
|
||||
const costResult = CostCalculator.calculateCost(usage, model);
|
||||
const cost = costResult.costs.total;
|
||||
|
||||
// 根据period分组累加费用
|
||||
if (period === 'daily') {
|
||||
const currentCost = dailyCosts.get(dateStr) || 0;
|
||||
dailyCosts.set(dateStr, currentCost + cost);
|
||||
} else if (period === 'monthly') {
|
||||
const currentCost = monthlyCosts.get(dateStr) || 0;
|
||||
monthlyCosts.set(dateStr, currentCost + cost);
|
||||
} else if (period === 'hourly') {
|
||||
const currentCost = hourlyCosts.get(dateStr) || 0;
|
||||
hourlyCosts.set(dateStr, currentCost + cost);
|
||||
}
|
||||
}
|
||||
|
||||
// 将计算出的费用写入Redis
|
||||
const promises = [];
|
||||
|
||||
// 写入每日费用
|
||||
for (const [date, cost] of dailyCosts) {
|
||||
const key = `usage:cost:daily:${apiKeyId}:${date}`;
|
||||
promises.push(
|
||||
client.set(key, cost.toString()),
|
||||
client.expire(key, 86400 * 30) // 30天过期
|
||||
);
|
||||
}
|
||||
|
||||
// 写入每月费用
|
||||
for (const [month, cost] of monthlyCosts) {
|
||||
const key = `usage:cost:monthly:${apiKeyId}:${month}`;
|
||||
promises.push(
|
||||
client.set(key, cost.toString()),
|
||||
client.expire(key, 86400 * 90) // 90天过期
|
||||
);
|
||||
}
|
||||
|
||||
// 写入每小时费用
|
||||
for (const [hour, cost] of hourlyCosts) {
|
||||
const key = `usage:cost:hourly:${apiKeyId}:${hour}`;
|
||||
promises.push(
|
||||
client.set(key, cost.toString()),
|
||||
client.expire(key, 86400 * 7) // 7天过期
|
||||
);
|
||||
}
|
||||
|
||||
// 计算总费用
|
||||
let totalCost = 0;
|
||||
for (const cost of dailyCosts.values()) {
|
||||
totalCost += cost;
|
||||
}
|
||||
|
||||
// 写入总费用
|
||||
if (totalCost > 0) {
|
||||
const totalKey = `usage:cost:total:${apiKeyId}`;
|
||||
promises.push(client.set(totalKey, totalCost.toString()));
|
||||
}
|
||||
|
||||
await Promise.all(promises);
|
||||
|
||||
logger.debug(`💰 Initialized costs for API Key ${apiKeyId}: Daily entries: ${dailyCosts.size}, Total cost: $${totalCost.toFixed(2)}`);
|
||||
}
|
||||
|
||||
/**
|
||||
* 检查是否需要初始化费用数据
|
||||
*/
|
||||
async needsInitialization() {
|
||||
try {
|
||||
const client = redis.getClientSafe();
|
||||
|
||||
// 检查是否有任何费用数据
|
||||
const costKeys = await client.keys('usage:cost:*');
|
||||
|
||||
// 如果没有费用数据，需要初始化
|
||||
if (costKeys.length === 0) {
|
||||
logger.info('💰 No cost data found, initialization needed');
|
||||
return true;
|
||||
}
|
||||
|
||||
// 检查是否有使用数据但没有对应的费用数据
|
||||
const sampleKeys = await client.keys('usage:*:model:daily:*:*');
|
||||
if (sampleKeys.length > 10) {
|
||||
// 抽样检查
|
||||
const sampleSize = Math.min(10, sampleKeys.length);
|
||||
for (let i = 0; i < sampleSize; i++) {
|
||||
const usageKey = sampleKeys[Math.floor(Math.random() * sampleKeys.length)];
|
||||
const match = usageKey.match(/usage:(.+):model:daily:(.+):(\d{4}-\d{2}-\d{2})$/);
|
||||
if (match) {
|
||||
const [, keyId, , date] = match;
|
||||
const costKey = `usage:cost:daily:${keyId}:${date}`;
|
||||
const hasCost = await client.exists(costKey);
|
||||
if (!hasCost) {
|
||||
logger.info(`💰 Found usage without cost data for key ${keyId} on ${date}, initialization needed`);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.info('💰 Cost data appears to be up to date');
|
||||
return false;
|
||||
} catch (error) {
|
||||
logger.error('❌ Failed to check initialization status:', error);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = new CostInitService();
|
||||
Reference in New Issue
Block a user