Browse Source

feat: restructure token usage routes and enhance token retrieval logic

CaIon 6 months ago
parent
commit
93ce48aca8
2 changed files with 22 additions and 11 deletions
  1. 11 10
      controller/token.go
  2. 11 1
      router/api-router.go

+ 11 - 10
controller/token.go

@@ -103,7 +103,7 @@ func GetTokenUsage(c *gin.Context) {
 	}
 	tokenKey := parts[1]
 
-	token, err := model.GetTokenByKey(tokenKey, true)
+	token, err := model.GetTokenByKey(strings.TrimPrefix(tokenKey, "sk-"), false)
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{
 			"success": false,
@@ -118,17 +118,18 @@ func GetTokenUsage(c *gin.Context) {
 	}
 
 	c.JSON(http.StatusOK, gin.H{
-		"code": true,
+		"code":    true,
 		"message": "ok",
 		"data": gin.H{
-			"object":          "token_usage",
-			"id":              token.Id,
-			"name":            token.Name,
-			"total_granted":   token.RemainQuota + token.UsedQuota,
-			"total_used":      token.UsedQuota,
-			"total_available": token.RemainQuota,
-			"unlimited_quota": token.UnlimitedQuota,
-			"expires_at":      expiredAt,
+			"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,
 		},
 	})
 }

+ 11 - 1
router/api-router.go

@@ -139,13 +139,23 @@ func SetApiRouter(router *gin.Engine) {
 		{
 			tokenRoute.GET("/", controller.GetAllTokens)
 			tokenRoute.GET("/search", controller.SearchTokens)
-			tokenRoute.GET("/usage", controller.GetTokenUsage)
 			tokenRoute.GET("/:id", controller.GetToken)
 			tokenRoute.POST("/", controller.AddToken)
 			tokenRoute.PUT("/", controller.UpdateToken)
 			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())
 		{