| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099 |
- package repository
- import (
- "context"
- "database/sql"
- "fmt"
- "os"
- "sort"
- "strings"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
- "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
- "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
- "github.com/Wei-Shaw/sub2api/internal/service"
- "github.com/lib/pq"
- )
- // GetUserStatsAggregated returns aggregated usage statistics for a user using database-level aggregation
- func (r *usageLogRepository) GetUserStatsAggregated(ctx context.Context, userID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
- query := `
- SELECT
- COUNT(*) as total_requests,
- COALESCE(SUM(input_tokens), 0) as total_input_tokens,
- COALESCE(SUM(output_tokens), 0) as total_output_tokens,
- COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens,
- COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens,
- COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens,
- COALESCE(SUM(total_cost), 0) as total_cost,
- COALESCE(SUM(actual_cost), 0) as total_actual_cost,
- COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms
- FROM usage_logs
- WHERE user_id = $1 AND created_at >= $2 AND created_at < $3
- `
- var stats usagestats.UsageStats
- if err := scanSingleRow(
- ctx,
- r.sql,
- query,
- []any{userID, startTime, endTime},
- &stats.TotalRequests,
- &stats.TotalInputTokens,
- &stats.TotalOutputTokens,
- &stats.TotalCacheTokens,
- &stats.TotalCacheCreationTokens,
- &stats.TotalCacheReadTokens,
- &stats.TotalCost,
- &stats.TotalActualCost,
- &stats.AverageDurationMs,
- ); err != nil {
- return nil, err
- }
- stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens
- return &stats, nil
- }
- // GetAPIKeyStatsAggregated returns aggregated usage statistics for an API key using database-level aggregation
- func (r *usageLogRepository) GetAPIKeyStatsAggregated(ctx context.Context, apiKeyID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
- query := `
- SELECT
- COUNT(*) as total_requests,
- COALESCE(SUM(input_tokens), 0) as total_input_tokens,
- COALESCE(SUM(output_tokens), 0) as total_output_tokens,
- COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens,
- COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens,
- COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens,
- COALESCE(SUM(total_cost), 0) as total_cost,
- COALESCE(SUM(actual_cost), 0) as total_actual_cost,
- COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms
- FROM usage_logs
- WHERE api_key_id = $1 AND created_at >= $2 AND created_at < $3
- `
- var stats usagestats.UsageStats
- if err := scanSingleRow(
- ctx,
- r.sql,
- query,
- []any{apiKeyID, startTime, endTime},
- &stats.TotalRequests,
- &stats.TotalInputTokens,
- &stats.TotalOutputTokens,
- &stats.TotalCacheTokens,
- &stats.TotalCacheCreationTokens,
- &stats.TotalCacheReadTokens,
- &stats.TotalCost,
- &stats.TotalActualCost,
- &stats.AverageDurationMs,
- ); err != nil {
- return nil, err
- }
- stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens
- return &stats, nil
- }
- // GetAccountStatsAggregated 使用 SQL 聚合统计账号使用数据
- //
- // 性能优化说明:
- // 原实现先查询所有日志记录,再在应用层循环计算统计值:
- // 1. 需要传输大量数据到应用层
- // 2. 应用层循环计算增加 CPU 和内存开销
- //
- // 新实现使用 SQL 聚合函数:
- // 1. 在数据库层完成 COUNT/SUM/AVG 计算
- // 2. 只返回单行聚合结果,大幅减少数据传输量
- // 3. 利用数据库索引优化聚合查询性能
- func (r *usageLogRepository) GetAccountStatsAggregated(ctx context.Context, accountID int64, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
- query := `
- SELECT
- COUNT(*) as total_requests,
- COALESCE(SUM(input_tokens), 0) as total_input_tokens,
- COALESCE(SUM(output_tokens), 0) as total_output_tokens,
- COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens,
- COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens,
- COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens,
- COALESCE(SUM(total_cost), 0) as total_cost,
- COALESCE(SUM(actual_cost), 0) as total_actual_cost,
- COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms
- FROM usage_logs
- WHERE account_id = $1 AND created_at >= $2 AND created_at < $3
- `
- var stats usagestats.UsageStats
- if err := scanSingleRow(
- ctx,
- r.sql,
- query,
- []any{accountID, startTime, endTime},
- &stats.TotalRequests,
- &stats.TotalInputTokens,
- &stats.TotalOutputTokens,
- &stats.TotalCacheTokens,
- &stats.TotalCacheCreationTokens,
- &stats.TotalCacheReadTokens,
- &stats.TotalCost,
- &stats.TotalActualCost,
- &stats.AverageDurationMs,
- ); err != nil {
- return nil, err
- }
- stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens
- return &stats, nil
- }
- // GetModelStatsAggregated 使用 SQL 聚合统计模型使用数据
- // 性能优化:数据库层聚合计算,避免应用层循环统计
- func (r *usageLogRepository) GetModelStatsAggregated(ctx context.Context, modelName string, startTime, endTime time.Time) (*usagestats.UsageStats, error) {
- query := fmt.Sprintf(`
- SELECT
- COUNT(*) as total_requests,
- COALESCE(SUM(input_tokens), 0) as total_input_tokens,
- COALESCE(SUM(output_tokens), 0) as total_output_tokens,
- COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens,
- COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens,
- COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens,
- COALESCE(SUM(total_cost), 0) as total_cost,
- COALESCE(SUM(actual_cost), 0) as total_actual_cost,
- COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms
- FROM usage_logs
- WHERE %s = $1 AND created_at >= $2 AND created_at < $3
- `, rawUsageLogModelColumn)
- var stats usagestats.UsageStats
- if err := scanSingleRow(
- ctx,
- r.sql,
- query,
- []any{modelName, startTime, endTime},
- &stats.TotalRequests,
- &stats.TotalInputTokens,
- &stats.TotalOutputTokens,
- &stats.TotalCacheTokens,
- &stats.TotalCacheCreationTokens,
- &stats.TotalCacheReadTokens,
- &stats.TotalCost,
- &stats.TotalActualCost,
- &stats.AverageDurationMs,
- ); err != nil {
- return nil, err
- }
- stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens
- return &stats, nil
- }
- // GetDailyStatsAggregated 使用 SQL 聚合统计用户的每日使用数据
- // 性能优化:使用 GROUP BY 在数据库层按日期分组聚合,避免应用层循环分组统计
- func (r *usageLogRepository) GetDailyStatsAggregated(ctx context.Context, userID int64, startTime, endTime time.Time) (result []map[string]any, err error) {
- tzName := resolveUsageStatsTimezone()
- query := `
- SELECT
- -- 使用应用时区分组,避免数据库会话时区导致日边界偏移。
- TO_CHAR(created_at AT TIME ZONE $4, 'YYYY-MM-DD') as date,
- COUNT(*) as total_requests,
- COALESCE(SUM(input_tokens), 0) as total_input_tokens,
- COALESCE(SUM(output_tokens), 0) as total_output_tokens,
- COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens,
- COALESCE(SUM(total_cost), 0) as total_cost,
- COALESCE(SUM(actual_cost), 0) as total_actual_cost,
- COALESCE(AVG(COALESCE(duration_ms, 0)), 0) as avg_duration_ms
- FROM usage_logs
- WHERE user_id = $1 AND created_at >= $2 AND created_at < $3
- GROUP BY 1
- ORDER BY 1
- `
- rows, err := r.sql.QueryContext(ctx, query, userID, startTime, endTime, tzName)
- if err != nil {
- return nil, err
- }
- defer func() {
- if closeErr := rows.Close(); closeErr != nil && err == nil {
- err = closeErr
- result = nil
- }
- }()
- result = make([]map[string]any, 0)
- for rows.Next() {
- var (
- date string
- totalRequests int64
- totalInputTokens int64
- totalOutputTokens int64
- totalCacheTokens int64
- totalCost float64
- totalActualCost float64
- avgDurationMs float64
- )
- if err = rows.Scan(
- &date,
- &totalRequests,
- &totalInputTokens,
- &totalOutputTokens,
- &totalCacheTokens,
- &totalCost,
- &totalActualCost,
- &avgDurationMs,
- ); err != nil {
- return nil, err
- }
- result = append(result, map[string]any{
- "date": date,
- "total_requests": totalRequests,
- "total_input_tokens": totalInputTokens,
- "total_output_tokens": totalOutputTokens,
- "total_cache_tokens": totalCacheTokens,
- "total_tokens": totalInputTokens + totalOutputTokens + totalCacheTokens,
- "total_cost": totalCost,
- "total_actual_cost": totalActualCost,
- "average_duration_ms": avgDurationMs,
- })
- }
- if err = rows.Err(); err != nil {
- return nil, err
- }
- return result, nil
- }
- // resolveUsageStatsTimezone 获取用于 SQL 分组的时区名称。
- // 优先使用应用初始化的时区,其次尝试读取 TZ 环境变量,最后回落为 UTC。
- func resolveUsageStatsTimezone() string {
- tzName := timezone.Name()
- if tzName != "" && tzName != "Local" {
- return tzName
- }
- if envTZ := strings.TrimSpace(os.Getenv("TZ")); envTZ != "" {
- return envTZ
- }
- return "UTC"
- }
- // GetAccountTodayStats 获取账号今日统计
- func (r *usageLogRepository) GetAccountTodayStats(ctx context.Context, accountID int64) (*usagestats.AccountStats, error) {
- today := timezone.Today()
- query := `
- SELECT
- COUNT(*) as requests,
- COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens,
- COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as cost,
- COALESCE(SUM(total_cost), 0) as standard_cost,
- COALESCE(SUM(actual_cost), 0) as user_cost
- FROM usage_logs
- WHERE account_id = $1 AND created_at >= $2
- `
- stats := &usagestats.AccountStats{}
- if err := scanSingleRow(
- ctx,
- r.sql,
- query,
- []any{accountID, today},
- &stats.Requests,
- &stats.Tokens,
- &stats.Cost,
- &stats.StandardCost,
- &stats.UserCost,
- ); err != nil {
- return nil, err
- }
- return stats, nil
- }
- // GetAccountWindowStats 获取账号时间窗口内的统计
- func (r *usageLogRepository) GetAccountWindowStats(ctx context.Context, accountID int64, startTime time.Time) (*usagestats.AccountStats, error) {
- query := `
- SELECT
- COUNT(*) as requests,
- COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens,
- COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as cost,
- COALESCE(SUM(total_cost), 0) as standard_cost,
- COALESCE(SUM(actual_cost), 0) as user_cost
- FROM usage_logs
- WHERE account_id = $1 AND created_at >= $2
- `
- stats := &usagestats.AccountStats{}
- if err := scanSingleRow(
- ctx,
- r.sql,
- query,
- []any{accountID, startTime},
- &stats.Requests,
- &stats.Tokens,
- &stats.Cost,
- &stats.StandardCost,
- &stats.UserCost,
- ); err != nil {
- return nil, err
- }
- return stats, nil
- }
- // GetAccountWindowStatsBatch 批量获取同一窗口起点下多个账号的统计数据。
- // 返回 map[accountID]*AccountStats,未命中的账号会返回零值统计,便于上层直接复用。
- func (r *usageLogRepository) GetAccountWindowStatsBatch(ctx context.Context, accountIDs []int64, startTime time.Time) (map[int64]*usagestats.AccountStats, error) {
- result := make(map[int64]*usagestats.AccountStats, len(accountIDs))
- if len(accountIDs) == 0 {
- return result, nil
- }
- query := `
- SELECT
- account_id,
- COUNT(*) as requests,
- COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens,
- COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as cost,
- COALESCE(SUM(total_cost), 0) as standard_cost,
- COALESCE(SUM(actual_cost), 0) as user_cost
- FROM usage_logs
- WHERE account_id = ANY($1) AND created_at >= $2
- GROUP BY account_id
- `
- rows, err := r.sql.QueryContext(ctx, query, pq.Array(accountIDs), startTime)
- if err != nil {
- return nil, err
- }
- defer func() { _ = rows.Close() }()
- for rows.Next() {
- var accountID int64
- stats := &usagestats.AccountStats{}
- if err := rows.Scan(
- &accountID,
- &stats.Requests,
- &stats.Tokens,
- &stats.Cost,
- &stats.StandardCost,
- &stats.UserCost,
- ); err != nil {
- return nil, err
- }
- result[accountID] = stats
- }
- if err := rows.Err(); err != nil {
- return nil, err
- }
- for _, accountID := range accountIDs {
- if _, ok := result[accountID]; !ok {
- result[accountID] = &usagestats.AccountStats{}
- }
- }
- return result, nil
- }
- // GetGeminiUsageTotalsBatch 批量聚合 Gemini 账号在窗口内的 Pro/Flash 请求与用量。
- // 模型分类规则与 service.geminiModelClassFromName 一致:model 包含 flash/lite 视为 flash,其余视为 pro。
- func (r *usageLogRepository) GetGeminiUsageTotalsBatch(ctx context.Context, accountIDs []int64, startTime, endTime time.Time) (map[int64]service.GeminiUsageTotals, error) {
- result := make(map[int64]service.GeminiUsageTotals, len(accountIDs))
- if len(accountIDs) == 0 {
- return result, nil
- }
- query := `
- SELECT
- account_id,
- COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 1 ELSE 0 END), 0) AS flash_requests,
- COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 0 ELSE 1 END), 0) AS pro_requests,
- COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN (input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) ELSE 0 END), 0) AS flash_tokens,
- COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 0 ELSE (input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens) END), 0) AS pro_tokens,
- COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN actual_cost ELSE 0 END), 0) AS flash_cost,
- COALESCE(SUM(CASE WHEN LOWER(COALESCE(model, '')) LIKE '%flash%' OR LOWER(COALESCE(model, '')) LIKE '%lite%' THEN 0 ELSE actual_cost END), 0) AS pro_cost
- FROM usage_logs
- WHERE account_id = ANY($1) AND created_at >= $2 AND created_at < $3
- GROUP BY account_id
- `
- rows, err := r.sql.QueryContext(ctx, query, pq.Array(accountIDs), startTime, endTime)
- if err != nil {
- return nil, err
- }
- defer func() { _ = rows.Close() }()
- for rows.Next() {
- var accountID int64
- var totals service.GeminiUsageTotals
- if err := rows.Scan(
- &accountID,
- &totals.FlashRequests,
- &totals.ProRequests,
- &totals.FlashTokens,
- &totals.ProTokens,
- &totals.FlashCost,
- &totals.ProCost,
- ); err != nil {
- return nil, err
- }
- result[accountID] = totals
- }
- if err := rows.Err(); err != nil {
- return nil, err
- }
- for _, accountID := range accountIDs {
- if _, ok := result[accountID]; !ok {
- result[accountID] = service.GeminiUsageTotals{}
- }
- }
- return result, nil
- }
- // UsageStats represents usage statistics
- type UsageStats = usagestats.UsageStats
- // BatchUserUsageStats represents usage stats for a single user
- type BatchUserUsageStats = usagestats.BatchUserUsageStats
- // PlatformUsage represents per-platform usage breakdown
- type PlatformUsage = usagestats.PlatformUsage
- func normalizePositiveInt64IDs(ids []int64) []int64 {
- if len(ids) == 0 {
- return nil
- }
- seen := make(map[int64]struct{}, len(ids))
- out := make([]int64, 0, len(ids))
- for _, id := range ids {
- if id <= 0 {
- continue
- }
- if _, ok := seen[id]; ok {
- continue
- }
- seen[id] = struct{}{}
- out = append(out, id)
- }
- return out
- }
- // GetBatchUserUsageStats gets today and total actual_cost for multiple users within a time range.
- // If startTime is zero, defaults to 30 days ago.
- func (r *usageLogRepository) GetBatchUserUsageStats(ctx context.Context, userIDs []int64, startTime, endTime time.Time) (map[int64]*BatchUserUsageStats, error) {
- result := make(map[int64]*BatchUserUsageStats)
- normalizedUserIDs := normalizePositiveInt64IDs(userIDs)
- if len(normalizedUserIDs) == 0 {
- return result, nil
- }
- // 默认最近 30 天
- if startTime.IsZero() {
- startTime = time.Now().AddDate(0, 0, -30)
- }
- if endTime.IsZero() {
- endTime = time.Now()
- }
- for _, id := range normalizedUserIDs {
- result[id] = &BatchUserUsageStats{UserID: id}
- }
- // GROUP BY (user_id, effective_platform) 一次查询同时得到总值与按平台拆分。
- // 应用层把同一 user_id 的多行累加为总值,并把非空 platform 行收集到 ByPlatform。
- query := `
- SELECT
- ul.user_id,
- ` + usageLogEffectivePlatformExpr + ` as platform,
- COALESCE(SUM(ul.actual_cost) FILTER (WHERE ul.created_at >= $2 AND ul.created_at < $3), 0) as total_cost,
- COALESCE(SUM(ul.actual_cost) FILTER (WHERE ul.created_at >= $4), 0) as today_cost
- FROM usage_logs ul
- LEFT JOIN groups g ON g.id = ul.group_id
- LEFT JOIN accounts a ON a.id = ul.account_id
- WHERE ul.user_id = ANY($1)
- AND ul.created_at >= LEAST($2, $4)
- AND ` + usageLogSuccessFilterUL + `
- GROUP BY ul.user_id, ` + usageLogEffectivePlatformExpr + `
- `
- today := timezone.Today()
- rows, err := r.sql.QueryContext(ctx, query, pq.Array(normalizedUserIDs), startTime, endTime, today)
- if err != nil {
- return nil, err
- }
- for rows.Next() {
- var userID int64
- var platform sql.NullString
- var total float64
- var todayTotal float64
- if err := rows.Scan(&userID, &platform, &total, &todayTotal); err != nil {
- _ = rows.Close()
- return nil, err
- }
- stats, ok := result[userID]
- if !ok {
- continue
- }
- stats.TotalActualCost += total
- stats.TodayActualCost += todayTotal
- if platform.Valid && platform.String != "" {
- stats.ByPlatform = append(stats.ByPlatform, PlatformUsage{
- Platform: platform.String,
- TotalActualCost: total,
- TodayActualCost: todayTotal,
- })
- }
- }
- if err := rows.Close(); err != nil {
- return nil, err
- }
- if err := rows.Err(); err != nil {
- return nil, err
- }
- return result, nil
- }
- // BatchAPIKeyUsageStats represents usage stats for a single API key
- type BatchAPIKeyUsageStats = usagestats.BatchAPIKeyUsageStats
- // GetBatchAPIKeyUsageStats gets today and total actual_cost for multiple API keys within a time range.
- // If startTime is zero, defaults to 30 days ago.
- func (r *usageLogRepository) GetBatchAPIKeyUsageStats(ctx context.Context, apiKeyIDs []int64, startTime, endTime time.Time) (map[int64]*BatchAPIKeyUsageStats, error) {
- result := make(map[int64]*BatchAPIKeyUsageStats)
- normalizedAPIKeyIDs := normalizePositiveInt64IDs(apiKeyIDs)
- if len(normalizedAPIKeyIDs) == 0 {
- return result, nil
- }
- // 默认最近 30 天
- if startTime.IsZero() {
- startTime = time.Now().AddDate(0, 0, -30)
- }
- if endTime.IsZero() {
- endTime = time.Now()
- }
- for _, id := range normalizedAPIKeyIDs {
- result[id] = &BatchAPIKeyUsageStats{APIKeyID: id}
- }
- query := `
- SELECT
- api_key_id,
- COALESCE(SUM(actual_cost) FILTER (WHERE created_at >= $2 AND created_at < $3), 0) as total_cost,
- COALESCE(SUM(actual_cost) FILTER (WHERE created_at >= $4), 0) as today_cost
- FROM usage_logs
- WHERE api_key_id = ANY($1)
- AND created_at >= LEAST($2, $4)
- GROUP BY api_key_id
- `
- today := timezone.Today()
- rows, err := r.sql.QueryContext(ctx, query, pq.Array(normalizedAPIKeyIDs), startTime, endTime, today)
- if err != nil {
- return nil, err
- }
- for rows.Next() {
- var apiKeyID int64
- var total float64
- var todayTotal float64
- if err := rows.Scan(&apiKeyID, &total, &todayTotal); err != nil {
- _ = rows.Close()
- return nil, err
- }
- if stats, ok := result[apiKeyID]; ok {
- stats.TotalActualCost = total
- stats.TodayActualCost = todayTotal
- }
- }
- if err := rows.Close(); err != nil {
- return nil, err
- }
- if err := rows.Err(); err != nil {
- return nil, err
- }
- return result, nil
- }
- // resolveEndpointColumn maps endpoint type to the corresponding DB column name.
- func resolveEndpointColumn(endpointType string) string {
- switch endpointType {
- case "upstream":
- return "ul.upstream_endpoint"
- case "path":
- return "ul.inbound_endpoint || ' -> ' || ul.upstream_endpoint"
- default:
- return "ul.inbound_endpoint"
- }
- }
- // GetGlobalStats gets usage statistics for all users within a time range
- func (r *usageLogRepository) GetGlobalStats(ctx context.Context, startTime, endTime time.Time) (*UsageStats, error) {
- query := `
- SELECT
- COUNT(*) as total_requests,
- COALESCE(SUM(input_tokens), 0) as total_input_tokens,
- COALESCE(SUM(output_tokens), 0) as total_output_tokens,
- COALESCE(SUM(cache_creation_tokens + cache_read_tokens), 0) as total_cache_tokens,
- COALESCE(SUM(total_cost), 0) as total_cost,
- COALESCE(SUM(actual_cost), 0) as total_actual_cost,
- COALESCE(AVG(duration_ms), 0) as avg_duration_ms
- FROM usage_logs
- WHERE created_at >= $1 AND created_at < $2
- `
- stats := &UsageStats{}
- if err := scanSingleRow(
- ctx,
- r.sql,
- query,
- []any{startTime, endTime},
- &stats.TotalRequests,
- &stats.TotalInputTokens,
- &stats.TotalOutputTokens,
- &stats.TotalCacheTokens,
- &stats.TotalCost,
- &stats.TotalActualCost,
- &stats.AverageDurationMs,
- ); err != nil {
- return nil, err
- }
- stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens
- return stats, nil
- }
- // GetStatsWithFilters gets usage statistics with optional filters
- func (r *usageLogRepository) GetStatsWithFilters(ctx context.Context, filters UsageLogFilters) (*UsageStats, error) {
- conditions := make([]string, 0, 9)
- args := make([]any, 0, 9)
- if filters.UserID > 0 {
- conditions = append(conditions, fmt.Sprintf("user_id = $%d", len(args)+1))
- args = append(args, filters.UserID)
- }
- if filters.APIKeyID > 0 {
- conditions = append(conditions, fmt.Sprintf("api_key_id = $%d", len(args)+1))
- args = append(args, filters.APIKeyID)
- }
- if filters.AccountID > 0 {
- conditions = append(conditions, fmt.Sprintf("account_id = $%d", len(args)+1))
- args = append(args, filters.AccountID)
- }
- if filters.GroupID > 0 {
- conditions = append(conditions, fmt.Sprintf("group_id = $%d", len(args)+1))
- args = append(args, filters.GroupID)
- }
- conditions, args = appendUsageLogModelWhereCondition(conditions, args, filters.Model, filters.ModelFilterSource)
- conditions, args = appendRequestTypeOrStreamWhereCondition(conditions, args, filters.RequestType, filters.Stream)
- conditions, args = appendNativeCompactionV2WhereCondition(conditions, args, filters.NativeCompactionV2, "")
- if filters.BillingType != nil {
- conditions = append(conditions, fmt.Sprintf("billing_type = $%d", len(args)+1))
- args = append(args, int16(*filters.BillingType))
- }
- conditions, args = appendUsageLogBillingModeWhereCondition(conditions, args, filters.BillingMode)
- if filters.UpstreamModelMismatch != nil {
- conditions = append(conditions, upstreamModelMismatchCondition("upstream_model_mismatch", *filters.UpstreamModelMismatch))
- }
- if filters.StartTime != nil {
- conditions = append(conditions, fmt.Sprintf("created_at >= $%d", len(args)+1))
- args = append(args, *filters.StartTime)
- }
- if filters.EndTime != nil {
- conditions = append(conditions, fmt.Sprintf("created_at < $%d", len(args)+1))
- args = append(args, *filters.EndTime)
- }
- query := fmt.Sprintf(`
- WITH scoped AS (
- SELECT
- COALESCE(NULLIF(TRIM(inbound_endpoint), ''), 'unknown') AS inbound_endpoint,
- COALESCE(NULLIF(TRIM(upstream_endpoint), ''), 'unknown') AS upstream_endpoint,
- input_tokens,
- output_tokens,
- cache_creation_tokens,
- cache_read_tokens,
- total_cost,
- actual_cost,
- COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1) AS account_cost,
- duration_ms
- FROM usage_logs
- %s
- )
- SELECT
- GROUPING(inbound_endpoint) AS inbound_grouped,
- GROUPING(upstream_endpoint) AS upstream_grouped,
- inbound_endpoint,
- upstream_endpoint,
- COUNT(*) AS requests,
- COALESCE(SUM(input_tokens), 0) AS input_tokens,
- COALESCE(SUM(output_tokens), 0) AS output_tokens,
- COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens,
- COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
- COALESCE(SUM(total_cost), 0) AS cost,
- COALESCE(SUM(actual_cost), 0) AS actual_cost,
- COALESCE(SUM(account_cost), 0) AS account_cost,
- COALESCE(AVG(duration_ms), 0) AS avg_duration_ms
- FROM scoped
- GROUP BY GROUPING SETS (
- (),
- (inbound_endpoint),
- (upstream_endpoint),
- (inbound_endpoint, upstream_endpoint)
- )
- `, buildWhere(conditions))
- stats := &UsageStats{}
- var totalAccountCost float64
- useAccountCostForEndpoint := filters.AccountID > 0 && filters.UserID == 0 && filters.APIKeyID == 0
- rows, err := r.sql.QueryContext(ctx, query, args...)
- if err != nil {
- return nil, err
- }
- defer func() { _ = rows.Close() }()
- for rows.Next() {
- var (
- inboundGrouped, upstreamGrouped int
- inboundEndpoint, upstreamEndpoint sql.NullString
- requests, inputTokens, outputTokens, cacheCreationTokens, cacheReads int64
- cost, actualCost, accountCost, averageDurationMs float64
- )
- if err := rows.Scan(
- &inboundGrouped,
- &upstreamGrouped,
- &inboundEndpoint,
- &upstreamEndpoint,
- &requests,
- &inputTokens,
- &outputTokens,
- &cacheCreationTokens,
- &cacheReads,
- &cost,
- &actualCost,
- &accountCost,
- &averageDurationMs,
- ); err != nil {
- return nil, err
- }
- totalTokens := inputTokens + outputTokens + cacheCreationTokens + cacheReads
- endpointActualCost := actualCost
- if useAccountCostForEndpoint {
- endpointActualCost = accountCost
- }
- switch {
- case inboundGrouped == 1 && upstreamGrouped == 1:
- stats.TotalRequests = requests
- stats.TotalInputTokens = inputTokens
- stats.TotalOutputTokens = outputTokens
- stats.TotalCacheCreationTokens = cacheCreationTokens
- stats.TotalCacheReadTokens = cacheReads
- stats.TotalCacheTokens = cacheCreationTokens + cacheReads
- stats.TotalCost = cost
- stats.TotalActualCost = actualCost
- totalAccountCost = accountCost
- stats.AverageDurationMs = averageDurationMs
- case inboundGrouped == 0 && upstreamGrouped == 1:
- stats.Endpoints = append(stats.Endpoints, EndpointStat{
- Endpoint: inboundEndpoint.String, Requests: requests, TotalTokens: totalTokens,
- Cost: cost, ActualCost: endpointActualCost,
- })
- case inboundGrouped == 1 && upstreamGrouped == 0:
- stats.UpstreamEndpoints = append(stats.UpstreamEndpoints, EndpointStat{
- Endpoint: upstreamEndpoint.String, Requests: requests, TotalTokens: totalTokens,
- Cost: cost, ActualCost: endpointActualCost,
- })
- case inboundGrouped == 0 && upstreamGrouped == 0:
- stats.EndpointPaths = append(stats.EndpointPaths, EndpointStat{
- Endpoint: inboundEndpoint.String + " -> " + upstreamEndpoint.String,
- Requests: requests, TotalTokens: totalTokens, Cost: cost, ActualCost: endpointActualCost,
- })
- }
- }
- if err := rows.Err(); err != nil {
- return nil, err
- }
- sortEndpointStats := func(values []EndpointStat) {
- sort.Slice(values, func(i, j int) bool {
- if values[i].Requests != values[j].Requests {
- return values[i].Requests > values[j].Requests
- }
- return values[i].Endpoint < values[j].Endpoint
- })
- }
- sortEndpointStats(stats.Endpoints)
- sortEndpointStats(stats.UpstreamEndpoints)
- sortEndpointStats(stats.EndpointPaths)
- stats.TotalAccountCost = &totalAccountCost
- stats.TotalTokens = stats.TotalInputTokens + stats.TotalOutputTokens + stats.TotalCacheTokens
- return stats, nil
- }
- // AccountUsageHistory represents daily usage history for an account
- type AccountUsageHistory = usagestats.AccountUsageHistory
- // AccountUsageSummary represents summary statistics for an account
- type AccountUsageSummary = usagestats.AccountUsageSummary
- // AccountUsageStatsResponse represents the full usage statistics response for an account
- type AccountUsageStatsResponse = usagestats.AccountUsageStatsResponse
- // EndpointStat represents endpoint usage statistics row.
- type EndpointStat = usagestats.EndpointStat
- func (r *usageLogRepository) getEndpointStatsByColumnWithFilters(ctx context.Context, endpointColumn string, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, modelSource string, requestType *int16, stream *bool, billingType *int8, billingMode string) (results []EndpointStat, err error) {
- actualCostExpr := "COALESCE(SUM(actual_cost), 0) as actual_cost"
- if accountID > 0 && userID == 0 && apiKeyID == 0 {
- actualCostExpr = "COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as actual_cost"
- }
- query := fmt.Sprintf(`
- SELECT
- COALESCE(NULLIF(TRIM(%s), ''), 'unknown') AS endpoint,
- COUNT(*) AS requests,
- COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
- COALESCE(SUM(total_cost), 0) as cost,
- %s
- FROM usage_logs
- WHERE created_at >= $1 AND created_at < $2
- `, endpointColumn, actualCostExpr)
- args := []any{startTime, endTime}
- if userID > 0 {
- query += fmt.Sprintf(" AND user_id = $%d", len(args)+1)
- args = append(args, userID)
- }
- if apiKeyID > 0 {
- query += fmt.Sprintf(" AND api_key_id = $%d", len(args)+1)
- args = append(args, apiKeyID)
- }
- if accountID > 0 {
- query += fmt.Sprintf(" AND account_id = $%d", len(args)+1)
- args = append(args, accountID)
- }
- if groupID > 0 {
- query += fmt.Sprintf(" AND group_id = $%d", len(args)+1)
- args = append(args, groupID)
- }
- query, args = appendUsageLogModelQueryFilter(query, args, model, modelSource)
- query, args = appendRequestTypeOrStreamQueryFilter(query, args, requestType, stream)
- if billingType != nil {
- query += fmt.Sprintf(" AND billing_type = $%d", len(args)+1)
- args = append(args, int16(*billingType))
- }
- query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "")
- query += " GROUP BY endpoint ORDER BY requests DESC"
- rows, err := r.sql.QueryContext(ctx, query, args...)
- if err != nil {
- return nil, err
- }
- defer func() {
- if closeErr := rows.Close(); closeErr != nil && err == nil {
- err = closeErr
- results = nil
- }
- }()
- results = make([]EndpointStat, 0)
- for rows.Next() {
- var row EndpointStat
- if err := rows.Scan(&row.Endpoint, &row.Requests, &row.TotalTokens, &row.Cost, &row.ActualCost); err != nil {
- return nil, err
- }
- results = append(results, row)
- }
- if err := rows.Err(); err != nil {
- return nil, err
- }
- return results, nil
- }
- // GetEndpointStatsWithFilters returns inbound endpoint statistics with optional filters.
- func (r *usageLogRepository) GetEndpointStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8) ([]EndpointStat, error) {
- return r.getEndpointStatsByColumnWithFilters(ctx, "inbound_endpoint", startTime, endTime, userID, apiKeyID, accountID, groupID, model, "", requestType, stream, billingType, "")
- }
- // GetUpstreamEndpointStatsWithFilters returns upstream endpoint statistics with optional filters.
- func (r *usageLogRepository) GetUpstreamEndpointStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8) ([]EndpointStat, error) {
- return r.getEndpointStatsByColumnWithFilters(ctx, "upstream_endpoint", startTime, endTime, userID, apiKeyID, accountID, groupID, model, "", requestType, stream, billingType, "")
- }
- // GetAccountUsageStats returns comprehensive usage statistics for an account over a time range
- func (r *usageLogRepository) GetAccountUsageStats(ctx context.Context, accountID int64, startTime, endTime time.Time) (resp *AccountUsageStatsResponse, err error) {
- daysCount := int(endTime.Sub(startTime).Hours()/24) + 1
- if daysCount <= 0 {
- daysCount = 30
- }
- query := `
- SELECT
- TO_CHAR(created_at, 'YYYY-MM-DD') as date,
- COUNT(*) as requests,
- COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as tokens,
- COALESCE(SUM(total_cost), 0) as cost,
- COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as actual_cost,
- COALESCE(SUM(actual_cost), 0) as user_cost
- FROM usage_logs
- WHERE account_id = $1 AND created_at >= $2 AND created_at < $3
- GROUP BY date
- ORDER BY date ASC
- `
- rows, err := r.sql.QueryContext(ctx, query, accountID, startTime, endTime)
- if err != nil {
- return nil, err
- }
- defer func() {
- // 保持主错误优先;仅在无错误时回传 Close 失败。
- // 同时清空返回值,避免误用不完整结果。
- if closeErr := rows.Close(); closeErr != nil && err == nil {
- err = closeErr
- resp = nil
- }
- }()
- history := make([]AccountUsageHistory, 0)
- for rows.Next() {
- var date string
- var requests int64
- var tokens int64
- var cost float64
- var actualCost float64
- var userCost float64
- if err = rows.Scan(&date, &requests, &tokens, &cost, &actualCost, &userCost); err != nil {
- return nil, err
- }
- t, _ := time.Parse("2006-01-02", date)
- history = append(history, AccountUsageHistory{
- Date: date,
- Label: t.Format("01/02"),
- Requests: requests,
- Tokens: tokens,
- Cost: cost,
- ActualCost: actualCost,
- UserCost: userCost,
- })
- }
- if err = rows.Err(); err != nil {
- return nil, err
- }
- var totalAccountCost, totalUserCost, totalStandardCost float64
- var totalRequests, totalTokens int64
- var highestCostDay, highestRequestDay *AccountUsageHistory
- for i := range history {
- h := &history[i]
- totalAccountCost += h.ActualCost
- totalUserCost += h.UserCost
- totalStandardCost += h.Cost
- totalRequests += h.Requests
- totalTokens += h.Tokens
- if highestCostDay == nil || h.ActualCost > highestCostDay.ActualCost {
- highestCostDay = h
- }
- if highestRequestDay == nil || h.Requests > highestRequestDay.Requests {
- highestRequestDay = h
- }
- }
- actualDaysUsed := len(history)
- if actualDaysUsed == 0 {
- actualDaysUsed = 1
- }
- avgQuery := "SELECT COALESCE(AVG(duration_ms), 0) as avg_duration_ms FROM usage_logs WHERE account_id = $1 AND created_at >= $2 AND created_at < $3"
- var avgDuration float64
- if err := scanSingleRow(ctx, r.sql, avgQuery, []any{accountID, startTime, endTime}, &avgDuration); err != nil {
- return nil, err
- }
- summary := AccountUsageSummary{
- Days: daysCount,
- ActualDaysUsed: actualDaysUsed,
- TotalCost: totalAccountCost,
- TotalUserCost: totalUserCost,
- TotalStandardCost: totalStandardCost,
- TotalRequests: totalRequests,
- TotalTokens: totalTokens,
- AvgDailyCost: totalAccountCost / float64(actualDaysUsed),
- AvgDailyUserCost: totalUserCost / float64(actualDaysUsed),
- AvgDailyRequests: float64(totalRequests) / float64(actualDaysUsed),
- AvgDailyTokens: float64(totalTokens) / float64(actualDaysUsed),
- AvgDurationMs: avgDuration,
- }
- todayStr := timezone.Now().Format("2006-01-02")
- for i := range history {
- if history[i].Date == todayStr {
- summary.Today = &struct {
- Date string `json:"date"`
- Cost float64 `json:"cost"`
- UserCost float64 `json:"user_cost"`
- Requests int64 `json:"requests"`
- Tokens int64 `json:"tokens"`
- }{
- Date: history[i].Date,
- Cost: history[i].ActualCost,
- UserCost: history[i].UserCost,
- Requests: history[i].Requests,
- Tokens: history[i].Tokens,
- }
- break
- }
- }
- if highestCostDay != nil {
- summary.HighestCostDay = &struct {
- Date string `json:"date"`
- Label string `json:"label"`
- Cost float64 `json:"cost"`
- UserCost float64 `json:"user_cost"`
- Requests int64 `json:"requests"`
- }{
- Date: highestCostDay.Date,
- Label: highestCostDay.Label,
- Cost: highestCostDay.ActualCost,
- UserCost: highestCostDay.UserCost,
- Requests: highestCostDay.Requests,
- }
- }
- if highestRequestDay != nil {
- summary.HighestRequestDay = &struct {
- Date string `json:"date"`
- Label string `json:"label"`
- Requests int64 `json:"requests"`
- Cost float64 `json:"cost"`
- UserCost float64 `json:"user_cost"`
- }{
- Date: highestRequestDay.Date,
- Label: highestRequestDay.Label,
- Requests: highestRequestDay.Requests,
- Cost: highestRequestDay.ActualCost,
- UserCost: highestRequestDay.UserCost,
- }
- }
- models, err := r.GetModelStatsWithFilters(ctx, startTime, endTime, 0, 0, accountID, 0, nil, nil, nil)
- if err != nil {
- models = []ModelStat{}
- }
- endpoints, endpointErr := r.GetEndpointStatsWithFilters(ctx, startTime, endTime, 0, 0, accountID, 0, "", nil, nil, nil)
- if endpointErr != nil {
- logger.LegacyPrintf("repository.usage_log", "GetEndpointStatsWithFilters failed in GetAccountUsageStats: %v", endpointErr)
- endpoints = []EndpointStat{}
- }
- upstreamEndpoints, upstreamEndpointErr := r.GetUpstreamEndpointStatsWithFilters(ctx, startTime, endTime, 0, 0, accountID, 0, "", nil, nil, nil)
- if upstreamEndpointErr != nil {
- logger.LegacyPrintf("repository.usage_log", "GetUpstreamEndpointStatsWithFilters failed in GetAccountUsageStats: %v", upstreamEndpointErr)
- upstreamEndpoints = []EndpointStat{}
- }
- resp = &AccountUsageStatsResponse{
- History: history,
- Summary: summary,
- Models: models,
- Endpoints: endpoints,
- UpstreamEndpoints: upstreamEndpoints,
- }
- return resp, nil
- }
|