添加会话tokens使用量返回

This commit is contained in:
junleea 2025-03-30 14:33:52 +08:00
parent 81e342b154
commit b471b32405
3 changed files with 26 additions and 4 deletions

View File

@ -196,3 +196,18 @@ func FindModelTokenByUserID(userID uint) []ModelToken {
}
return modelToken
}
type ModelTokensSession struct {
SessionID uint `gorm:"column:session_id"`
Token uint `gorm:"column:token"`
}
func FindModelTotalTokensBySessionID(userID int) []ModelTokensSession {
var modelTokens []ModelTokensSession
if proto.Config.SERVER_SQL_LOG {
DB.Debug().Table("model_tokens").Select("session_id, count(*) as token").Where("user_id = ?", userID).Group("session_id").Scan(&modelTokens)
} else {
DB.Table("model_tokens").Select("session_id, count(*) as token").Where("user_id = ?", userID).Group("session_id").Scan(&modelTokens)
}
return modelTokens
}

View File

@ -58,10 +58,11 @@ func (m *ModelParam) SetDefaultParams() {
type SessionResponse struct {
gorm.Model
UserID int `gorm:"column:user_id"` //只能由用户创建
Name string `gorm:"column:name"`
Context []int `gorm:"column:context;type:json"` //会话上下文
MsgCount int `gorm:"column:msg_count"` //消息数量
UserID int `gorm:"column:user_id"` //只能由用户创建
Name string `gorm:"column:name"`
Context []int `gorm:"column:context;type:json"` //会话上下文
MsgCount int `gorm:"column:msg_count"` //消息数量
TokenUsage uint `gorm:"column:token_usage"` //token使用数量
}
type IMParamContext struct {

View File

@ -22,10 +22,15 @@ func CreateSession(userID int, name string) (error, uint) {
func FindSessionByUserID(userID int) []proto.SessionResponse {
sessions := dao.FindSessionByUserID(userID)
sessionsMsgCounts := dao.FindSessionMessageCountByUserID(userID)
sessionsTokens := dao.FindModelTotalTokensBySessionID(userID)
smap := make(map[int]int)
st := make(map[int]uint)
for _, v := range sessionsMsgCounts {
smap[v.SessionID] = v.Count
}
for _, v := range sessionsTokens {
st[int(v.SessionID)] = v.Token
}
var res []proto.SessionResponse
for _, v := range sessions {
var session proto.SessionResponse
@ -35,6 +40,7 @@ func FindSessionByUserID(userID int) []proto.SessionResponse {
session.Name = v.Name
session.UserID = v.UserID
session.MsgCount = smap[int(v.ID)]
session.TokenUsage = st[int(v.ID)]
res = append(res, session)
}
return res