fix: update token usage calculation

This commit is contained in:
CaIon
2025-08-16 19:11:15 +08:00
parent aa1f5c6e4e
commit c834694992
3 changed files with 13 additions and 20 deletions

View File

@@ -5,6 +5,7 @@ import (
"fmt" "fmt"
"one-api/common" "one-api/common"
"one-api/logger" "one-api/logger"
"one-api/types"
"os" "os"
"strings" "strings"
"time" "time"
@@ -150,10 +151,10 @@ type RecordConsumeLogParams struct {
} }
func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) { func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) {
logger.LogInfo(c, fmt.Sprintf("record consume log: userId=%d, params=%s", userId, common.GetJsonString(params)))
if !common.LogConsumeEnabled { if !common.LogConsumeEnabled {
return return
} }
logger.LogInfo(c, fmt.Sprintf("record consume log: userId=%d, params=%s", userId, common.GetJsonString(params)))
username := c.GetString("username") username := c.GetString("username")
otherStr := common.MapToJsonStr(params.Other) otherStr := common.MapToJsonStr(params.Other)
// 判断是否需要记录 IP // 判断是否需要记录 IP
@@ -236,26 +237,22 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName
return nil, 0, err return nil, 0, err
} }
channelIdsMap := make(map[int]struct{}) channelIds := types.NewSet[int]()
channelMap := make(map[int]string)
for _, log := range logs { for _, log := range logs {
if log.ChannelId != 0 { if log.ChannelId != 0 {
channelIdsMap[log.ChannelId] = struct{}{} channelIds.Add(log.ChannelId)
} }
} }
channelIds := make([]int, 0, len(channelIdsMap)) if channelIds.Len() > 0 {
for channelId := range channelIdsMap {
channelIds = append(channelIds, channelId)
}
if len(channelIds) > 0 {
var channels []struct { var channels []struct {
Id int `gorm:"column:id"` Id int `gorm:"column:id"`
Name string `gorm:"column:name"` Name string `gorm:"column:name"`
} }
if err = DB.Table("channels").Select("id, name").Where("id IN ?", channelIds).Find(&channels).Error; err != nil { if err = DB.Table("channels").Select("id, name").Where("id IN ?", channelIds.Items()).Find(&channels).Error; err != nil {
return logs, total, err return logs, total, err
} }
channelMap := make(map[int]string, len(channels))
for _, channel := range channels { for _, channel := range channels {
channelMap[channel.Id] = channel.Name channelMap[channel.Id] = channel.Name
} }

View File

@@ -21,12 +21,6 @@ type QuotaData struct {
} }
func UpdateQuotaData() { func UpdateQuotaData() {
// recover
defer func() {
if r := recover(); r != nil {
common.SysLog(fmt.Sprintf("UpdateQuotaData panic: %s", r))
}
}()
for { for {
if common.DataExportEnabled { if common.DataExportEnabled {
common.SysLog("正在更新数据看板数据...") common.SysLog("正在更新数据看板数据...")

View File

@@ -103,12 +103,14 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
// 非正常结束,使用输出文本的 token 数量 // 非正常结束,使用输出文本的 token 数量
completionTokens := service.CountTextToken(tempStr, info.UpstreamModelName) completionTokens := service.CountTextToken(tempStr, info.UpstreamModelName)
usage.CompletionTokens = completionTokens usage.CompletionTokens = completionTokens
if usage.PromptTokens == 0 {
usage.PromptTokens = info.PromptTokens
}
} }
} }
if usage.PromptTokens == 0 && usage.CompletionTokens != 0 {
usage.PromptTokens = usage.CompletionTokens
} else {
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
}
return usage, nil return usage, nil
} }