fix: aggregate user model stats by requested model

This commit is contained in:
Haruka
2026-07-01 19:01:30 +08:00
parent 5eb9da9c93
commit e236bff1ed
2 changed files with 23 additions and 36 deletions
+1 -36
View File
@@ -2786,42 +2786,7 @@ func (r *usageLogRepository) GetUserUsageTrendByUserID(ctx context.Context, user
// GetUserModelStats 获取指定用户的模型统计
func (r *usageLogRepository) GetUserModelStats(ctx context.Context, userID int64, startTime, endTime time.Time) (results []ModelStat, err error) {
query := `
SELECT
model,
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(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as total_tokens,
COALESCE(SUM(total_cost), 0) as cost,
COALESCE(SUM(actual_cost), 0) as actual_cost,
COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) as account_cost
FROM usage_logs
WHERE user_id = $1 AND created_at >= $2 AND created_at < $3
GROUP BY model
ORDER BY total_tokens DESC
`
rows, err := r.sql.QueryContext(ctx, query, userID, startTime, endTime)
if err != nil {
return nil, err
}
defer func() {
// 保持主错误优先;仅在无错误时回传 Close 失败。
// 同时清空返回值,避免误用不完整结果。
if closeErr := rows.Close(); closeErr != nil && err == nil {
err = closeErr
results = nil
}
}()
results, err = scanModelStatsRows(rows)
if err != nil {
return nil, err
}
return results, nil
return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, 0, 0, 0, "", nil, nil, nil, usagestats.ModelSourceRequested, "")
}
// UsageLogFilters represents filters for usage log queries
@@ -434,6 +434,28 @@ func TestUsageLogRepositoryGetModelStatsWithFiltersRequestTypePriority(t *testin
require.NoError(t, mock.ExpectationsWereMet())
}
func TestUsageLogRepositoryGetUserModelStatsUsesRequestedModel(t *testing.T) {
db, mock := newSQLMock(t)
repo := &usageLogRepository{sql: db}
start := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)
end := start.Add(24 * time.Hour)
mock.ExpectQuery("(?s)SELECT\\s+COALESCE\\(NULLIF\\(TRIM\\(requested_model\\), ''\\), model\\) as model,.*WHERE created_at >= \\$1 AND created_at < \\$2\\s+AND user_id = \\$3.*GROUP BY COALESCE\\(NULLIF\\(TRIM\\(requested_model\\), ''\\), model\\) ORDER BY total_tokens DESC").
WithArgs(start, end, int64(7)).
WillReturnRows(sqlmock.NewRows([]string{
"model", "requests", "input_tokens", "output_tokens",
"cache_creation_tokens", "cache_read_tokens", "total_tokens",
"cost", "actual_cost", "account_cost",
}).AddRow("gpt-5.5", int64(2), int64(10), int64(20), int64(0), int64(0), int64(30), 0.1, 0.08, 0.07))
stats, err := repo.GetUserModelStats(context.Background(), 7, start, end)
require.NoError(t, err)
require.Len(t, stats, 1)
require.Equal(t, "gpt-5.5", stats[0].Model)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestUsageLogRepositoryGetStatsWithFiltersRequestedModelSource(t *testing.T) {
db, mock := newSQLMock(t)
repo := &usageLogRepository{sql: db}