添加会话tokens使用量返回

dao/model.go

@@ -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
+}

proto/im.go

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

service/imService.go

@@ -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