claude-relay-service/src/services/costInitService.js

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()