Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.
 
 
 

504 Zeilen
15 KiB

  1. package model
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "time"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/QuantumNous/new-api/logger"
  9. relaycommon "github.com/QuantumNous/new-api/relay/common"
  10. "github.com/QuantumNous/new-api/types"
  11. "github.com/gin-gonic/gin"
  12. "github.com/bytedance/gopkg/util/gopool"
  13. "gorm.io/gorm"
  14. )
  15. type Log struct {
  16. Id int `json:"id" gorm:"index:idx_created_at_id,priority:1;index:idx_user_id_id,priority:2"`
  17. UserId int `json:"user_id" gorm:"index;index:idx_user_id_id,priority:1"`
  18. CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_created_at_id,priority:2;index:idx_created_at_type"`
  19. Type int `json:"type" gorm:"index:idx_created_at_type"`
  20. Content string `json:"content"`
  21. Username string `json:"username" gorm:"index;index:index_username_model_name,priority:2;default:''"`
  22. TokenName string `json:"token_name" gorm:"index;default:''"`
  23. ModelName string `json:"model_name" gorm:"index;index:index_username_model_name,priority:1;default:''"`
  24. Quota int `json:"quota" gorm:"default:0"`
  25. PromptTokens int `json:"prompt_tokens" gorm:"default:0"`
  26. CompletionTokens int `json:"completion_tokens" gorm:"default:0"`
  27. UseTime int `json:"use_time" gorm:"default:0"`
  28. IsStream bool `json:"is_stream"`
  29. ChannelId int `json:"channel" gorm:"index"`
  30. ChannelName string `json:"channel_name" gorm:"->"`
  31. TokenId int `json:"token_id" gorm:"default:0;index"`
  32. Group string `json:"group" gorm:"index"`
  33. Ip string `json:"ip" gorm:"index;default:''"`
  34. RequestId string `json:"request_id,omitempty" gorm:"type:varchar(64);index:idx_logs_request_id;default:''"`
  35. ChatId string `json:"chat_id,omitempty" gorm:"type:varchar(128);index:idx_logs_chat_id;default:''"`
  36. UpstreamId string `json:"upstream_id,omitempty" gorm:"type:varchar(128);index:idx_logs_upstream_id;default:''"`
  37. Other string `json:"other"`
  38. }
  39. // don't use iota, avoid change log type value
  40. const (
  41. LogTypeUnknown = 0
  42. LogTypeTopup = 1
  43. LogTypeConsume = 2
  44. LogTypeManage = 3
  45. LogTypeSystem = 4
  46. LogTypeError = 5
  47. LogTypeRefund = 6
  48. )
  49. func formatUserLogs(logs []*Log, startIdx int) {
  50. for i := range logs {
  51. logs[i].ChannelName = ""
  52. var otherMap map[string]interface{}
  53. otherMap, _ = common.StrToMap(logs[i].Other)
  54. if otherMap != nil {
  55. // Remove admin-only debug fields.
  56. delete(otherMap, "admin_info")
  57. delete(otherMap, "reject_reason")
  58. }
  59. logs[i].Other = common.MapToJsonStr(otherMap)
  60. logs[i].Id = startIdx + i + 1
  61. }
  62. }
  63. func GetLogByTokenId(tokenId int) (logs []*Log, err error) {
  64. err = LOG_DB.Model(&Log{}).Where("token_id = ?", tokenId).Order("id desc").Limit(common.MaxRecentItems).Find(&logs).Error
  65. formatUserLogs(logs, 0)
  66. return logs, err
  67. }
  68. func RecordLog(userId int, logType int, content string) {
  69. if logType == LogTypeConsume && !common.LogConsumeEnabled {
  70. return
  71. }
  72. username, _ := GetUsernameById(userId, false)
  73. log := &Log{
  74. UserId: userId,
  75. Username: username,
  76. CreatedAt: common.GetTimestamp(),
  77. Type: logType,
  78. Content: content,
  79. }
  80. err := LOG_DB.Create(log).Error
  81. if err != nil {
  82. common.SysLog("failed to record log: " + err.Error())
  83. }
  84. }
  85. func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string, tokenName string, content string, tokenId int, useTimeSeconds int,
  86. isStream bool, group string, other map[string]interface{}) {
  87. logger.LogInfo(c, fmt.Sprintf("record error log: userId=%d, channelId=%d, modelName=%s, tokenName=%s, content=%s", userId, channelId, modelName, tokenName, content))
  88. username := c.GetString("username")
  89. requestId := c.GetString(common.RequestIdKey)
  90. chatId := relaycommon.GetRelayChatID(c)
  91. upstreamId := relaycommon.GetRelayUpstreamID(c)
  92. otherStr := common.MapToJsonStr(other)
  93. // 判断是否需要记录 IP
  94. needRecordIp := false
  95. if settingMap, err := GetUserSetting(userId, false); err == nil {
  96. if settingMap.RecordIpLog {
  97. needRecordIp = true
  98. }
  99. }
  100. log := &Log{
  101. UserId: userId,
  102. Username: username,
  103. CreatedAt: common.GetTimestamp(),
  104. Type: LogTypeError,
  105. Content: content,
  106. PromptTokens: 0,
  107. CompletionTokens: 0,
  108. TokenName: tokenName,
  109. ModelName: modelName,
  110. Quota: 0,
  111. ChannelId: channelId,
  112. TokenId: tokenId,
  113. UseTime: useTimeSeconds,
  114. IsStream: isStream,
  115. Group: group,
  116. Ip: func() string {
  117. if needRecordIp {
  118. return c.ClientIP()
  119. }
  120. return ""
  121. }(),
  122. RequestId: requestId,
  123. ChatId: chatId,
  124. UpstreamId: upstreamId,
  125. Other: otherStr,
  126. }
  127. err := LOG_DB.Create(log).Error
  128. if err != nil {
  129. logger.LogError(c, "failed to record log: "+err.Error())
  130. }
  131. }
  132. type RecordConsumeLogParams struct {
  133. ChannelId int `json:"channel_id"`
  134. PromptTokens int `json:"prompt_tokens"`
  135. CompletionTokens int `json:"completion_tokens"`
  136. ModelName string `json:"model_name"`
  137. TokenName string `json:"token_name"`
  138. Quota int `json:"quota"`
  139. Content string `json:"content"`
  140. TokenId int `json:"token_id"`
  141. UseTimeSeconds int `json:"use_time_seconds"`
  142. IsStream bool `json:"is_stream"`
  143. Group string `json:"group"`
  144. Other map[string]interface{} `json:"other"`
  145. }
  146. func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) {
  147. if !common.LogConsumeEnabled {
  148. return
  149. }
  150. logger.LogInfo(c, fmt.Sprintf("record consume log: userId=%d, params=%s", userId, common.GetJsonString(params)))
  151. username := c.GetString("username")
  152. requestId := c.GetString(common.RequestIdKey)
  153. chatId := relaycommon.GetRelayChatID(c)
  154. upstreamId := relaycommon.GetRelayUpstreamID(c)
  155. otherStr := common.MapToJsonStr(params.Other)
  156. // 判断是否需要记录 IP
  157. needRecordIp := false
  158. if settingMap, err := GetUserSetting(userId, false); err == nil {
  159. if settingMap.RecordIpLog {
  160. needRecordIp = true
  161. }
  162. }
  163. log := &Log{
  164. UserId: userId,
  165. Username: username,
  166. CreatedAt: common.GetTimestamp(),
  167. Type: LogTypeConsume,
  168. Content: params.Content,
  169. PromptTokens: params.PromptTokens,
  170. CompletionTokens: params.CompletionTokens,
  171. TokenName: params.TokenName,
  172. ModelName: params.ModelName,
  173. Quota: params.Quota,
  174. ChannelId: params.ChannelId,
  175. TokenId: params.TokenId,
  176. UseTime: params.UseTimeSeconds,
  177. IsStream: params.IsStream,
  178. Group: params.Group,
  179. Ip: func() string {
  180. if needRecordIp {
  181. return c.ClientIP()
  182. }
  183. return ""
  184. }(),
  185. RequestId: requestId,
  186. ChatId: chatId,
  187. UpstreamId: upstreamId,
  188. Other: otherStr,
  189. }
  190. err := LOG_DB.Create(log).Error
  191. if err != nil {
  192. logger.LogError(c, "failed to record log: "+err.Error())
  193. }
  194. if common.DataExportEnabled {
  195. gopool.Go(func() {
  196. LogQuotaData(userId, username, params.ModelName, params.Quota, common.GetTimestamp(), params.PromptTokens+params.CompletionTokens)
  197. })
  198. }
  199. }
  200. type RecordTaskBillingLogParams struct {
  201. UserId int
  202. LogType int
  203. Content string
  204. ChannelId int
  205. ModelName string
  206. Quota int
  207. TokenId int
  208. Group string
  209. Other map[string]interface{}
  210. }
  211. func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
  212. if params.LogType == LogTypeConsume && !common.LogConsumeEnabled {
  213. return
  214. }
  215. username, _ := GetUsernameById(params.UserId, false)
  216. tokenName := ""
  217. if params.TokenId > 0 {
  218. if token, err := GetTokenById(params.TokenId); err == nil {
  219. tokenName = token.Name
  220. }
  221. }
  222. log := &Log{
  223. UserId: params.UserId,
  224. Username: username,
  225. CreatedAt: common.GetTimestamp(),
  226. Type: params.LogType,
  227. Content: params.Content,
  228. TokenName: tokenName,
  229. ModelName: params.ModelName,
  230. Quota: params.Quota,
  231. ChannelId: params.ChannelId,
  232. TokenId: params.TokenId,
  233. Group: params.Group,
  234. Other: common.MapToJsonStr(params.Other),
  235. }
  236. err := LOG_DB.Create(log).Error
  237. if err != nil {
  238. common.SysLog("failed to record task billing log: " + err.Error())
  239. }
  240. }
  241. func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string, requestId string, chatId string, upstreamId string) (logs []*Log, total int64, err error) {
  242. var tx *gorm.DB
  243. if logType == LogTypeUnknown {
  244. tx = LOG_DB
  245. } else {
  246. tx = LOG_DB.Where("logs.type = ?", logType)
  247. }
  248. if modelName != "" {
  249. tx = tx.Where("logs.model_name like ?", modelName)
  250. }
  251. if username != "" {
  252. tx = tx.Where("logs.username = ?", username)
  253. }
  254. if tokenName != "" {
  255. tx = tx.Where("logs.token_name = ?", tokenName)
  256. }
  257. if requestId != "" {
  258. tx = tx.Where("logs.request_id = ?", requestId)
  259. }
  260. if chatId != "" {
  261. tx = tx.Where("logs.chat_id = ?", chatId)
  262. }
  263. if upstreamId != "" {
  264. tx = tx.Where("logs.upstream_id = ?", upstreamId)
  265. }
  266. if startTimestamp != 0 {
  267. tx = tx.Where("logs.created_at >= ?", startTimestamp)
  268. }
  269. if endTimestamp != 0 {
  270. tx = tx.Where("logs.created_at <= ?", endTimestamp)
  271. }
  272. if channel != 0 {
  273. tx = tx.Where("logs.channel_id = ?", channel)
  274. }
  275. if group != "" {
  276. tx = tx.Where("logs."+logGroupCol+" = ?", group)
  277. }
  278. err = tx.Model(&Log{}).Count(&total).Error
  279. if err != nil {
  280. return nil, 0, err
  281. }
  282. err = tx.Order("logs.id desc").Limit(num).Offset(startIdx).Find(&logs).Error
  283. if err != nil {
  284. return nil, 0, err
  285. }
  286. channelIds := types.NewSet[int]()
  287. for _, log := range logs {
  288. if log.ChannelId != 0 {
  289. channelIds.Add(log.ChannelId)
  290. }
  291. }
  292. if channelIds.Len() > 0 {
  293. var channels []struct {
  294. Id int `gorm:"column:id"`
  295. Name string `gorm:"column:name"`
  296. }
  297. if common.MemoryCacheEnabled {
  298. // Cache get channel
  299. for _, channelId := range channelIds.Items() {
  300. if cacheChannel, err := CacheGetChannel(channelId); err == nil {
  301. channels = append(channels, struct {
  302. Id int `gorm:"column:id"`
  303. Name string `gorm:"column:name"`
  304. }{
  305. Id: channelId,
  306. Name: cacheChannel.Name,
  307. })
  308. }
  309. }
  310. } else {
  311. // Bulk query channels from DB
  312. if err = DB.Table("channels").Select("id, name").Where("id IN ?", channelIds.Items()).Find(&channels).Error; err != nil {
  313. return logs, total, err
  314. }
  315. }
  316. channelMap := make(map[int]string, len(channels))
  317. for _, channel := range channels {
  318. channelMap[channel.Id] = channel.Name
  319. }
  320. for i := range logs {
  321. logs[i].ChannelName = channelMap[logs[i].ChannelId]
  322. }
  323. }
  324. return logs, total, err
  325. }
  326. const logSearchCountLimit = 10000
  327. func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string, requestId string, chatId string, upstreamId string) (logs []*Log, total int64, err error) {
  328. var tx *gorm.DB
  329. if logType == LogTypeUnknown {
  330. tx = LOG_DB.Where("logs.user_id = ?", userId)
  331. } else {
  332. tx = LOG_DB.Where("logs.user_id = ? and logs.type = ?", userId, logType)
  333. }
  334. if modelName != "" {
  335. modelNamePattern, err := sanitizeLikePattern(modelName)
  336. if err != nil {
  337. return nil, 0, err
  338. }
  339. tx = tx.Where("logs.model_name LIKE ? ESCAPE '!'", modelNamePattern)
  340. }
  341. if tokenName != "" {
  342. tx = tx.Where("logs.token_name = ?", tokenName)
  343. }
  344. if requestId != "" {
  345. tx = tx.Where("logs.request_id = ?", requestId)
  346. }
  347. if chatId != "" {
  348. tx = tx.Where("logs.chat_id = ?", chatId)
  349. }
  350. if upstreamId != "" {
  351. tx = tx.Where("logs.upstream_id = ?", upstreamId)
  352. }
  353. if startTimestamp != 0 {
  354. tx = tx.Where("logs.created_at >= ?", startTimestamp)
  355. }
  356. if endTimestamp != 0 {
  357. tx = tx.Where("logs.created_at <= ?", endTimestamp)
  358. }
  359. if group != "" {
  360. tx = tx.Where("logs."+logGroupCol+" = ?", group)
  361. }
  362. err = tx.Model(&Log{}).Limit(logSearchCountLimit).Count(&total).Error
  363. if err != nil {
  364. common.SysError("failed to count user logs: " + err.Error())
  365. return nil, 0, errors.New("查询日志失败")
  366. }
  367. err = tx.Order("logs.id desc").Limit(num).Offset(startIdx).Find(&logs).Error
  368. if err != nil {
  369. common.SysError("failed to search user logs: " + err.Error())
  370. return nil, 0, errors.New("查询日志失败")
  371. }
  372. formatUserLogs(logs, startIdx)
  373. return logs, total, err
  374. }
  375. type Stat struct {
  376. Quota int `json:"quota"`
  377. Rpm int `json:"rpm"`
  378. Tpm int `json:"tpm"`
  379. }
  380. func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int, group string) (stat Stat, err error) {
  381. tx := LOG_DB.Table("logs").Select("sum(quota) quota")
  382. // 为rpm和tpm创建单独的查询
  383. rpmTpmQuery := LOG_DB.Table("logs").Select("count(*) rpm, sum(prompt_tokens) + sum(completion_tokens) tpm")
  384. if username != "" {
  385. tx = tx.Where("username = ?", username)
  386. rpmTpmQuery = rpmTpmQuery.Where("username = ?", username)
  387. }
  388. if tokenName != "" {
  389. tx = tx.Where("token_name = ?", tokenName)
  390. rpmTpmQuery = rpmTpmQuery.Where("token_name = ?", tokenName)
  391. }
  392. if startTimestamp != 0 {
  393. tx = tx.Where("created_at >= ?", startTimestamp)
  394. }
  395. if endTimestamp != 0 {
  396. tx = tx.Where("created_at <= ?", endTimestamp)
  397. }
  398. if modelName != "" {
  399. modelNamePattern, err := sanitizeLikePattern(modelName)
  400. if err != nil {
  401. return stat, err
  402. }
  403. tx = tx.Where("model_name LIKE ? ESCAPE '!'", modelNamePattern)
  404. rpmTpmQuery = rpmTpmQuery.Where("model_name LIKE ? ESCAPE '!'", modelNamePattern)
  405. }
  406. if channel != 0 {
  407. tx = tx.Where("channel_id = ?", channel)
  408. rpmTpmQuery = rpmTpmQuery.Where("channel_id = ?", channel)
  409. }
  410. if group != "" {
  411. tx = tx.Where(logGroupCol+" = ?", group)
  412. rpmTpmQuery = rpmTpmQuery.Where(logGroupCol+" = ?", group)
  413. }
  414. tx = tx.Where("type = ?", LogTypeConsume)
  415. rpmTpmQuery = rpmTpmQuery.Where("type = ?", LogTypeConsume)
  416. // 只统计最近60秒的rpm和tpm
  417. rpmTpmQuery = rpmTpmQuery.Where("created_at >= ?", time.Now().Add(-60*time.Second).Unix())
  418. // 执行查询
  419. if err := tx.Scan(&stat).Error; err != nil {
  420. common.SysError("failed to query log stat: " + err.Error())
  421. return stat, errors.New("查询统计数据失败")
  422. }
  423. if err := rpmTpmQuery.Scan(&stat).Error; err != nil {
  424. common.SysError("failed to query rpm/tpm stat: " + err.Error())
  425. return stat, errors.New("查询统计数据失败")
  426. }
  427. return stat, nil
  428. }
  429. func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string) (token int) {
  430. tx := LOG_DB.Table("logs").Select("ifnull(sum(prompt_tokens),0) + ifnull(sum(completion_tokens),0)")
  431. if username != "" {
  432. tx = tx.Where("username = ?", username)
  433. }
  434. if tokenName != "" {
  435. tx = tx.Where("token_name = ?", tokenName)
  436. }
  437. if startTimestamp != 0 {
  438. tx = tx.Where("created_at >= ?", startTimestamp)
  439. }
  440. if endTimestamp != 0 {
  441. tx = tx.Where("created_at <= ?", endTimestamp)
  442. }
  443. if modelName != "" {
  444. tx = tx.Where("model_name = ?", modelName)
  445. }
  446. tx.Where("type = ?", LogTypeConsume).Scan(&token)
  447. return token
  448. }
  449. func DeleteOldLog(ctx context.Context, targetTimestamp int64, limit int) (int64, error) {
  450. var total int64 = 0
  451. for {
  452. if nil != ctx.Err() {
  453. return total, ctx.Err()
  454. }
  455. result := LOG_DB.Where("created_at < ?", targetTimestamp).Limit(limit).Delete(&Log{})
  456. if nil != result.Error {
  457. return total, result.Error
  458. }
  459. total += result.RowsAffected
  460. if result.RowsAffected < int64(limit) {
  461. break
  462. }
  463. }
  464. return total, nil
  465. }