From b471b32405c434d25ca46a900f92ada271bfd44d Mon Sep 17 00:00:00 2001
From: junleea <354425203@qq.com>
Date: Sun, 30 Mar 2025 14:33:52 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E4=BC=9A=E8=AF=9Dtokens?=
 =?UTF-8?q?=E4=BD=BF=E7=94=A8=E9=87=8F=E8=BF=94=E5=9B=9E?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 dao/model.go         | 15 +++++++++++++++
 proto/im.go          |  9 +++++----
 service/imService.go |  6 ++++++
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/dao/model.go b/dao/model.go
index 8568f58..44e26f5 100644
--- a/dao/model.go
+++ b/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
+}
diff --git a/proto/im.go b/proto/im.go
index 6ec1d39..0e942aa 100644
--- a/proto/im.go
+++ b/proto/im.go
@@ -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 {
diff --git a/service/imService.go b/service/imService.go
index 29b9f6b..1ecb889 100644
--- a/service/imService.go
+++ b/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