diff --git a/controller/token.go b/controller/token.go index 399ccb4f..8ed8b957 100644 --- a/controller/token.go +++ b/controller/token.go @@ -5,6 +5,7 @@ import ( "one-api/common" "one-api/model" "strconv" + "strings" "github.com/gin-gonic/gin" ) @@ -82,6 +83,57 @@ func GetTokenStatus(c *gin.Context) { }) } +func GetTokenUsage(c *gin.Context) { + authHeader := c.GetHeader("Authorization") + if authHeader == "" { + c.JSON(http.StatusUnauthorized, gin.H{ + "success": false, + "message": "No Authorization header", + }) + return + } + + parts := strings.Split(authHeader, " ") + if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" { + c.JSON(http.StatusUnauthorized, gin.H{ + "success": false, + "message": "Invalid Bearer token", + }) + return + } + tokenKey := parts[1] + + token, err := model.GetTokenByKey(strings.TrimPrefix(tokenKey, "sk-"), false) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + + expiredAt := token.ExpiredTime + if expiredAt == -1 { + expiredAt = 0 + } + + c.JSON(http.StatusOK, gin.H{ + "code": true, + "message": "ok", + "data": gin.H{ + "object": "token_usage", + "name": token.Name, + "total_granted": token.RemainQuota + token.UsedQuota, + "total_used": token.UsedQuota, + "total_available": token.RemainQuota, + "unlimited_quota": token.UnlimitedQuota, + "model_limits": token.GetModelLimitsMap(), + "model_limits_enabled": token.ModelLimitsEnabled, + "expires_at": expiredAt, + }, + }) +} + func AddToken(c *gin.Context) { token := model.Token{} err := c.ShouldBindJSON(&token) diff --git a/router/api-router.go b/router/api-router.go index aa3cba6d..be721b05 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -145,6 +145,17 @@ func SetApiRouter(router *gin.Engine) { tokenRoute.DELETE("/:id", controller.DeleteToken) tokenRoute.POST("/batch", controller.DeleteTokenBatch) } + + usageRoute := apiRouter.Group("/usage") + usageRoute.Use(middleware.CriticalRateLimit()) + { + tokenUsageRoute := usageRoute.Group("/token") + tokenUsageRoute.Use(middleware.TokenAuth()) + { + tokenUsageRoute.GET("/", controller.GetTokenUsage) + } + } + redemptionRoute := apiRouter.Group("/redemption") redemptionRoute.Use(middleware.AdminAuth()) { @@ -172,7 +183,6 @@ func SetApiRouter(router *gin.Engine) { logRoute.Use(middleware.CORS()) { logRoute.GET("/token", controller.GetLogByKey) - } groupRoute := apiRouter.Group("/group") groupRoute.Use(middleware.AdminAuth())