Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

314 строки
7.8 KiB

  1. package controller
  2. import (
  3. "fmt"
  4. "net/http"
  5. "strconv"
  6. "strings"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/QuantumNous/new-api/i18n"
  9. "github.com/QuantumNous/new-api/model"
  10. "github.com/QuantumNous/new-api/setting/operation_setting"
  11. "github.com/gin-gonic/gin"
  12. )
  13. func GetAllTokens(c *gin.Context) {
  14. userId := c.GetInt("id")
  15. pageInfo := common.GetPageQuery(c)
  16. tokens, err := model.GetAllUserTokens(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize())
  17. if err != nil {
  18. common.ApiError(c, err)
  19. return
  20. }
  21. total, _ := model.CountUserTokens(userId)
  22. pageInfo.SetTotal(int(total))
  23. pageInfo.SetItems(tokens)
  24. common.ApiSuccess(c, pageInfo)
  25. return
  26. }
  27. func SearchTokens(c *gin.Context) {
  28. userId := c.GetInt("id")
  29. keyword := c.Query("keyword")
  30. token := c.Query("token")
  31. pageInfo := common.GetPageQuery(c)
  32. tokens, total, err := model.SearchUserTokens(userId, keyword, token, pageInfo.GetStartIdx(), pageInfo.GetPageSize())
  33. if err != nil {
  34. common.ApiError(c, err)
  35. return
  36. }
  37. pageInfo.SetTotal(int(total))
  38. pageInfo.SetItems(tokens)
  39. common.ApiSuccess(c, pageInfo)
  40. return
  41. }
  42. func GetToken(c *gin.Context) {
  43. id, err := strconv.Atoi(c.Param("id"))
  44. userId := c.GetInt("id")
  45. if err != nil {
  46. common.ApiError(c, err)
  47. return
  48. }
  49. token, err := model.GetTokenByIds(id, userId)
  50. if err != nil {
  51. common.ApiError(c, err)
  52. return
  53. }
  54. c.JSON(http.StatusOK, gin.H{
  55. "success": true,
  56. "message": "",
  57. "data": token,
  58. })
  59. return
  60. }
  61. func GetTokenStatus(c *gin.Context) {
  62. tokenId := c.GetInt("token_id")
  63. userId := c.GetInt("id")
  64. token, err := model.GetTokenByIds(tokenId, userId)
  65. if err != nil {
  66. common.ApiError(c, err)
  67. return
  68. }
  69. expiredAt := token.ExpiredTime
  70. if expiredAt == -1 {
  71. expiredAt = 0
  72. }
  73. c.JSON(http.StatusOK, gin.H{
  74. "object": "credit_summary",
  75. "total_granted": token.RemainQuota,
  76. "total_used": 0, // not supported currently
  77. "total_available": token.RemainQuota,
  78. "expires_at": expiredAt * 1000,
  79. })
  80. }
  81. func GetTokenUsage(c *gin.Context) {
  82. authHeader := c.GetHeader("Authorization")
  83. if authHeader == "" {
  84. c.JSON(http.StatusUnauthorized, gin.H{
  85. "success": false,
  86. "message": "No Authorization header",
  87. })
  88. return
  89. }
  90. parts := strings.Split(authHeader, " ")
  91. if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" {
  92. c.JSON(http.StatusUnauthorized, gin.H{
  93. "success": false,
  94. "message": "Invalid Bearer token",
  95. })
  96. return
  97. }
  98. tokenKey := parts[1]
  99. token, err := model.GetTokenByKey(strings.TrimPrefix(tokenKey, "sk-"), false)
  100. if err != nil {
  101. common.SysError("failed to get token by key: " + err.Error())
  102. common.ApiErrorI18n(c, i18n.MsgTokenGetInfoFailed)
  103. return
  104. }
  105. expiredAt := token.ExpiredTime
  106. if expiredAt == -1 {
  107. expiredAt = 0
  108. }
  109. c.JSON(http.StatusOK, gin.H{
  110. "code": true,
  111. "message": "ok",
  112. "data": gin.H{
  113. "object": "token_usage",
  114. "name": token.Name,
  115. "total_granted": token.RemainQuota + token.UsedQuota,
  116. "total_used": token.UsedQuota,
  117. "total_available": token.RemainQuota,
  118. "unlimited_quota": token.UnlimitedQuota,
  119. "model_limits": token.GetModelLimitsMap(),
  120. "model_limits_enabled": token.ModelLimitsEnabled,
  121. "expires_at": expiredAt,
  122. },
  123. })
  124. }
  125. func AddToken(c *gin.Context) {
  126. token := model.Token{}
  127. err := c.ShouldBindJSON(&token)
  128. if err != nil {
  129. common.ApiError(c, err)
  130. return
  131. }
  132. if len(token.Name) > 50 {
  133. common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong)
  134. return
  135. }
  136. // 非无限额度时,检查额度值是否超出有效范围
  137. if !token.UnlimitedQuota {
  138. if token.RemainQuota < 0 {
  139. common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative)
  140. return
  141. }
  142. maxQuotaValue := int((1000000000 * common.QuotaPerUnit))
  143. if token.RemainQuota > maxQuotaValue {
  144. common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue})
  145. return
  146. }
  147. }
  148. // 检查用户令牌数量是否已达上限
  149. maxTokens := operation_setting.GetMaxUserTokens()
  150. count, err := model.CountUserTokens(c.GetInt("id"))
  151. if err != nil {
  152. common.ApiError(c, err)
  153. return
  154. }
  155. if int(count) >= maxTokens {
  156. c.JSON(http.StatusOK, gin.H{
  157. "success": false,
  158. "message": fmt.Sprintf("已达到最大令牌数量限制 (%d)", maxTokens),
  159. })
  160. return
  161. }
  162. key, err := common.GenerateKey()
  163. if err != nil {
  164. common.ApiErrorI18n(c, i18n.MsgTokenGenerateFailed)
  165. common.SysLog("failed to generate token key: " + err.Error())
  166. return
  167. }
  168. cleanToken := model.Token{
  169. UserId: c.GetInt("id"),
  170. Name: token.Name,
  171. Key: key,
  172. CreatedTime: common.GetTimestamp(),
  173. AccessedTime: common.GetTimestamp(),
  174. ExpiredTime: token.ExpiredTime,
  175. RemainQuota: token.RemainQuota,
  176. UnlimitedQuota: token.UnlimitedQuota,
  177. ModelLimitsEnabled: token.ModelLimitsEnabled,
  178. ModelLimits: token.ModelLimits,
  179. AllowIps: token.AllowIps,
  180. Group: token.Group,
  181. CrossGroupRetry: token.CrossGroupRetry,
  182. BoundChannelId: token.BoundChannelId,
  183. }
  184. err = cleanToken.Insert()
  185. if err != nil {
  186. common.ApiError(c, err)
  187. return
  188. }
  189. c.JSON(http.StatusOK, gin.H{
  190. "success": true,
  191. "message": "",
  192. })
  193. return
  194. }
  195. func DeleteToken(c *gin.Context) {
  196. id, _ := strconv.Atoi(c.Param("id"))
  197. userId := c.GetInt("id")
  198. err := model.DeleteTokenById(id, userId)
  199. if err != nil {
  200. common.ApiError(c, err)
  201. return
  202. }
  203. c.JSON(http.StatusOK, gin.H{
  204. "success": true,
  205. "message": "",
  206. })
  207. return
  208. }
  209. func UpdateToken(c *gin.Context) {
  210. userId := c.GetInt("id")
  211. statusOnly := c.Query("status_only")
  212. token := model.Token{}
  213. err := c.ShouldBindJSON(&token)
  214. if err != nil {
  215. common.ApiError(c, err)
  216. return
  217. }
  218. if len(token.Name) > 50 {
  219. common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong)
  220. return
  221. }
  222. if !token.UnlimitedQuota {
  223. if token.RemainQuota < 0 {
  224. common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative)
  225. return
  226. }
  227. maxQuotaValue := int((1000000000 * common.QuotaPerUnit))
  228. if token.RemainQuota > maxQuotaValue {
  229. common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue})
  230. return
  231. }
  232. }
  233. cleanToken, err := model.GetTokenByIds(token.Id, userId)
  234. if err != nil {
  235. common.ApiError(c, err)
  236. return
  237. }
  238. if token.Status == common.TokenStatusEnabled {
  239. if cleanToken.Status == common.TokenStatusExpired && cleanToken.ExpiredTime <= common.GetTimestamp() && cleanToken.ExpiredTime != -1 {
  240. common.ApiErrorI18n(c, i18n.MsgTokenExpiredCannotEnable)
  241. return
  242. }
  243. if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota {
  244. common.ApiErrorI18n(c, i18n.MsgTokenExhaustedCannotEable)
  245. return
  246. }
  247. }
  248. if statusOnly != "" {
  249. cleanToken.Status = token.Status
  250. } else {
  251. // If you add more fields, please also update token.Update()
  252. cleanToken.Name = token.Name
  253. cleanToken.ExpiredTime = token.ExpiredTime
  254. cleanToken.RemainQuota = token.RemainQuota
  255. cleanToken.UnlimitedQuota = token.UnlimitedQuota
  256. cleanToken.ModelLimitsEnabled = token.ModelLimitsEnabled
  257. cleanToken.ModelLimits = token.ModelLimits
  258. cleanToken.AllowIps = token.AllowIps
  259. cleanToken.Group = token.Group
  260. cleanToken.CrossGroupRetry = token.CrossGroupRetry
  261. cleanToken.BoundChannelId = token.BoundChannelId
  262. }
  263. err = cleanToken.Update()
  264. if err != nil {
  265. common.ApiError(c, err)
  266. return
  267. }
  268. c.JSON(http.StatusOK, gin.H{
  269. "success": true,
  270. "message": "",
  271. "data": cleanToken,
  272. })
  273. }
  274. type TokenBatch struct {
  275. Ids []int `json:"ids"`
  276. }
  277. func DeleteTokenBatch(c *gin.Context) {
  278. tokenBatch := TokenBatch{}
  279. if err := c.ShouldBindJSON(&tokenBatch); err != nil || len(tokenBatch.Ids) == 0 {
  280. common.ApiErrorI18n(c, i18n.MsgInvalidParams)
  281. return
  282. }
  283. userId := c.GetInt("id")
  284. count, err := model.BatchDeleteTokens(tokenBatch.Ids, userId)
  285. if err != nil {
  286. common.ApiError(c, err)
  287. return
  288. }
  289. c.JSON(http.StatusOK, gin.H{
  290. "success": true,
  291. "message": "",
  292. "data": count,
  293. })
  294. }