diff --git a/common/constants.go b/common/constants.go index c7d5637c8e9..126e69a549f 100644 --- a/common/constants.go +++ b/common/constants.go @@ -68,6 +68,7 @@ var TaskEnabled = true var DataExportEnabled = true var DataExportInterval = 5 // unit: minute var DataExportDefaultTime = "hour" // unit: minute +var ApiKeyStatsEnabled = false var DefaultCollapseSidebar = false // default value of collapse sidebar // Any options with "Secret", "Token" in its key won't be return by GetOptions diff --git a/controller/log.go b/controller/log.go index a43f7b75227..f0284bf9f94 100644 --- a/controller/log.go +++ b/controller/log.go @@ -6,6 +6,7 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/gin-gonic/gin" ) @@ -150,6 +151,58 @@ func GetLogsSelfStat(c *gin.Context) { return } +// GetLogStatsByToken returns API key usage statistics for admins. +func GetLogStatsByToken(c *gin.Context) { + if !isApiKeyStatsEnabled(c) { + return + } + + startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) + endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) + data, err := model.GetLogStatsByToken(0, startTimestamp, endTimestamp) + if err != nil { + common.ApiError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": data, + }) +} + +// GetUserLogStatsByToken returns API key usage statistics for the current user. +func GetUserLogStatsByToken(c *gin.Context) { + if !isApiKeyStatsEnabled(c) { + return + } + + userId := c.GetInt("id") + startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) + endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) + data, err := model.GetLogStatsByToken(userId, startTimestamp, endTimestamp) + if err != nil { + common.ApiError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": data, + }) +} + +func isApiKeyStatsEnabled(c *gin.Context) bool { + if operation_setting.SelfUseModeEnabled && common.ApiKeyStatsEnabled { + return true + } + c.JSON(http.StatusForbidden, gin.H{ + "success": false, + "message": "api key statistics is disabled", + }) + return false +} + func DeleteHistoryLogs(c *gin.Context) { targetTimestamp, _ := strconv.ParseInt(c.Query("target_timestamp"), 10, 64) if targetTimestamp == 0 { diff --git a/controller/log_token_stat_test.go b/controller/log_token_stat_test.go new file mode 100644 index 00000000000..75c660db405 --- /dev/null +++ b/controller/log_token_stat_test.go @@ -0,0 +1,67 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func withApiKeyStatsSettings(t *testing.T, selfUseModeEnabled, apiKeyStatsEnabled bool) { + t.Helper() + + originalSelfUseModeEnabled := operation_setting.SelfUseModeEnabled + originalApiKeyStatsEnabled := common.ApiKeyStatsEnabled + operation_setting.SelfUseModeEnabled = selfUseModeEnabled + common.ApiKeyStatsEnabled = apiKeyStatsEnabled + t.Cleanup(func() { + operation_setting.SelfUseModeEnabled = originalSelfUseModeEnabled + common.ApiKeyStatsEnabled = originalApiKeyStatsEnabled + }) +} + +func performTokenStatsRequest(handler gin.HandlerFunc) *httptest.ResponseRecorder { + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + req := httptest.NewRequest(http.MethodGet, "/api/log/stat/tokens", nil) + ctx.Request = req + handler(ctx) + return recorder +} + +func TestGetLogStatsByTokenRejectsWhenApiKeyStatsDisabled(t *testing.T) { + gin.SetMode(gin.TestMode) + withApiKeyStatsSettings(t, true, false) + + recorder := performTokenStatsRequest(GetLogStatsByToken) + + require.Equal(t, http.StatusForbidden, recorder.Code) + require.Contains(t, recorder.Body.String(), "api key statistics is disabled") +} + +func TestGetLogStatsByTokenRejectsWhenSelfUseModeDisabled(t *testing.T) { + gin.SetMode(gin.TestMode) + withApiKeyStatsSettings(t, false, true) + + recorder := performTokenStatsRequest(GetLogStatsByToken) + + require.Equal(t, http.StatusForbidden, recorder.Code) + require.Contains(t, recorder.Body.String(), "api key statistics is disabled") +} + +func TestGetLogStatsByTokenAllowsWhenEnabled(t *testing.T) { + gin.SetMode(gin.TestMode) + withApiKeyStatsSettings(t, true, true) + db := setupModelListControllerTestDB(t) + require.NoError(t, db.AutoMigrate(&model.Log{})) + + recorder := performTokenStatsRequest(GetLogStatsByToken) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"success":true`) +} diff --git a/controller/misc.go b/controller/misc.go index 344cda7715d..3b21f024367 100644 --- a/controller/misc.go +++ b/controller/misc.go @@ -87,6 +87,7 @@ func GetStatus(c *gin.Context) { "chats": setting.Chats, "demo_site_enabled": operation_setting.DemoSiteEnabled, "self_use_mode_enabled": operation_setting.SelfUseModeEnabled, + "api_key_stats_enabled": common.ApiKeyStatsEnabled && operation_setting.SelfUseModeEnabled, "register_enabled": common.RegisterEnabled, "password_register_enabled": common.PasswordRegisterEnabled, "default_use_auto_group": setting.DefaultUseAutoGroup, diff --git a/model/log_token_stat.go b/model/log_token_stat.go new file mode 100644 index 00000000000..c6e2aff4baa --- /dev/null +++ b/model/log_token_stat.go @@ -0,0 +1,82 @@ +package model + +import ( + "fmt" + "time" + + "github.com/QuantumNous/new-api/common" +) + +const ( + tokenStatsDefaultRangeSeconds = 24 * 60 * 60 + tokenStatsMaxRangeSeconds = 30 * 24 * 60 * 60 +) + +// TokenQuotaData is hourly aggregated usage data by API key. +type TokenQuotaData struct { + TokenId int `json:"token_id"` + TokenName string `json:"token_name"` + CreatedAt int64 `json:"created_at"` + Count int `json:"count"` + Quota int `json:"quota"` + TokenUsed int `json:"token_used"` +} + +// GetLogStatsByToken aggregates consume logs by API key and hour. +// userId = 0 returns all users' data; userId > 0 limits data to that user. +func GetLogStatsByToken(userId int, startTime, endTime int64) ([]*TokenQuotaData, error) { + var results []*TokenQuotaData + + startTime, endTime, err := normalizeTokenStatsTimeRange(startTime, endTime, time.Now().Unix()) + if err != nil { + return nil, err + } + + query := LOG_DB.Table("logs"). + Select(`token_id, + token_name, + (created_at - created_at % 3600) AS created_at, + COUNT(*) AS count, + COALESCE(SUM(quota), 0) AS quota, + COALESCE(SUM(prompt_tokens), 0) + COALESCE(SUM(completion_tokens), 0) AS token_used`). + Where("type = ?", LogTypeConsume). + Where("token_id != 0"). + Where("token_name != ''"). + Group("token_id, token_name, (created_at - created_at % 3600)"). + Order("(created_at - created_at % 3600)") + + if userId > 0 { + query = query.Where("user_id = ?", userId) + } + if startTime != 0 { + query = query.Where("created_at >= ?", startTime) + } + if endTime != 0 { + query = query.Where("created_at <= ?", endTime) + } + + if err := query.Scan(&results).Error; err != nil { + common.SysError("failed to query token quota data: " + err.Error()) + return nil, err + } + return results, nil +} + +func normalizeTokenStatsTimeRange(startTime, endTime, now int64) (int64, int64, error) { + if startTime == 0 && endTime == 0 { + endTime = now + startTime = endTime - tokenStatsDefaultRangeSeconds + } else if endTime == 0 { + endTime = now + } else if startTime == 0 { + startTime = endTime - tokenStatsMaxRangeSeconds + } + + if endTime < startTime { + return 0, 0, fmt.Errorf("invalid time range: end_timestamp < start_timestamp") + } + if endTime-startTime > tokenStatsMaxRangeSeconds { + startTime = endTime - tokenStatsMaxRangeSeconds + } + return startTime, endTime, nil +} diff --git a/model/log_token_stat_test.go b/model/log_token_stat_test.go new file mode 100644 index 00000000000..bd3debbb575 --- /dev/null +++ b/model/log_token_stat_test.go @@ -0,0 +1,38 @@ +package model + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestNormalizeTokenStatsTimeRangeDefaultsToRecentDay(t *testing.T) { + startTime, endTime, err := normalizeTokenStatsTimeRange(0, 0, 1_700_000_000) + require.NoError(t, err) + require.Equal(t, int64(1_700_000_000), endTime) + require.Equal(t, int64(1_700_000_000-tokenStatsDefaultRangeSeconds), startTime) +} + +func TestNormalizeTokenStatsTimeRangeFillsMissingBoundary(t *testing.T) { + startTime, endTime, err := normalizeTokenStatsTimeRange(1_699_999_000, 0, 1_700_000_000) + require.NoError(t, err) + require.Equal(t, int64(1_699_999_000), startTime) + require.Equal(t, int64(1_700_000_000), endTime) + + startTime, endTime, err = normalizeTokenStatsTimeRange(0, 1_700_000_000, 1_700_000_100) + require.NoError(t, err) + require.Equal(t, int64(1_700_000_000-tokenStatsMaxRangeSeconds), startTime) + require.Equal(t, int64(1_700_000_000), endTime) +} + +func TestNormalizeTokenStatsTimeRangeCapsRange(t *testing.T) { + startTime, endTime, err := normalizeTokenStatsTimeRange(1_600_000_000, 1_700_000_000, 1_700_000_000) + require.NoError(t, err) + require.Equal(t, int64(1_700_000_000-tokenStatsMaxRangeSeconds), startTime) + require.Equal(t, int64(1_700_000_000), endTime) +} + +func TestNormalizeTokenStatsTimeRangeRejectsInvalidRange(t *testing.T) { + _, _, err := normalizeTokenStatsTimeRange(1_700_000_001, 1_700_000_000, 1_700_000_000) + require.Error(t, err) +} diff --git a/model/option.go b/model/option.go index e0a3048d34f..43b9d10c6c3 100644 --- a/model/option.go +++ b/model/option.go @@ -52,6 +52,7 @@ func InitOptionMap() { common.OptionMap["DrawingEnabled"] = strconv.FormatBool(common.DrawingEnabled) common.OptionMap["TaskEnabled"] = strconv.FormatBool(common.TaskEnabled) common.OptionMap["DataExportEnabled"] = strconv.FormatBool(common.DataExportEnabled) + common.OptionMap["ApiKeyStatsEnabled"] = strconv.FormatBool(common.ApiKeyStatsEnabled) common.OptionMap["ChannelDisableThreshold"] = strconv.FormatFloat(common.ChannelDisableThreshold, 'f', -1, 64) common.OptionMap["EmailDomainRestrictionEnabled"] = strconv.FormatBool(common.EmailDomainRestrictionEnabled) common.OptionMap["EmailAliasRestrictionEnabled"] = strconv.FormatBool(common.EmailAliasRestrictionEnabled) @@ -295,6 +296,8 @@ func updateOptionMap(key string, value string) (err error) { common.TaskEnabled = boolValue case "DataExportEnabled": common.DataExportEnabled = boolValue + case "ApiKeyStatsEnabled": + common.ApiKeyStatsEnabled = boolValue case "DefaultCollapseSidebar": common.DefaultCollapseSidebar = boolValue case "MjNotifyEnabled": diff --git a/router/api-router.go b/router/api-router.go index da026ed92f4..9ad4f9a683b 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -297,7 +297,9 @@ func SetApiRouter(router *gin.Engine) { logRoute.GET("/", middleware.AdminAuth(), controller.GetAllLogs) logRoute.DELETE("/", middleware.AdminAuth(), controller.DeleteHistoryLogs) logRoute.GET("/stat", middleware.AdminAuth(), controller.GetLogsStat) + logRoute.GET("/stat/tokens", middleware.AdminAuth(), controller.GetLogStatsByToken) logRoute.GET("/self/stat", middleware.UserAuth(), controller.GetLogsSelfStat) + logRoute.GET("/self/stat/tokens", middleware.UserAuth(), controller.GetUserLogStatsByToken) logRoute.GET("/channel_affinity_usage_cache", middleware.AdminAuth(), controller.GetChannelAffinityUsageCacheStats) logRoute.GET("/search", middleware.AdminAuth(), controller.SearchAllLogs) logRoute.GET("/self", middleware.UserAuth(), controller.GetUserLogs) diff --git a/web/default/src/components/ui/sidebar.tsx b/web/default/src/components/ui/sidebar.tsx index 5c16eece5d8..129443fb238 100644 --- a/web/default/src/components/ui/sidebar.tsx +++ b/web/default/src/components/ui/sidebar.tsx @@ -475,7 +475,7 @@ function SidebarMenu({ className, ...props }: React.ComponentProps<'ul'>) {