From eb4c0f940433cf377fe903c5663d88ee7664a32d Mon Sep 17 00:00:00 2001
From: junleea <354425203@qq.com>
Date: Sun, 30 Mar 2025 14:38:49 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=A8=A1=E5=9E=8Btokens?=
 =?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            | 20 ++++++++++++++++++++
 service/modelService.go | 34 ++++++++++++++++++++++++++++++----
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/dao/model.go b/dao/model.go
index 44e26f5..3fae14b 100644
--- a/dao/model.go
+++ b/dao/model.go
@@ -31,6 +31,11 @@ type ModelToken struct {
 	Token     uint `gorm:"column:token"`      //模型使用token
 }
 
+type ModelResponse struct {
+	Model
+	Token uint `json:"token"`
+}
+
 // 创建模型
 func CreateModel(userID uint, modelType, url, parameter, description string) (error, uint) {
 	model := Model{UserID: userID, Type: modelType, Url: url, Parameter: parameter, Description: description}
@@ -211,3 +216,18 @@ func FindModelTotalTokensBySessionID(userID int) []ModelTokensSession {
 	}
 	return modelTokens
 }
+
+type ModelTokensModel struct {
+	ModelID uint `gorm:"column:model_id"`
+	Token   uint `gorm:"column:token"`
+}
+
+func FindModelTotalTokensByModelID(userID int) []ModelTokensModel {
+	var modelTokens []ModelTokensModel
+	if proto.Config.SERVER_SQL_LOG {
+		DB.Debug().Table("model_tokens").Select("model_id, count(*) as token").Where("user_id = ?", userID).Group("model_id").Scan(&modelTokens)
+	} else {
+		DB.Table("model_tokens").Select("model_id, count(*) as token").Where("user_id = ?", userID).Group("model_id").Scan(&modelTokens)
+	}
+	return modelTokens
+}
diff --git a/service/modelService.go b/service/modelService.go
index 66fd137..0a8364b 100644
--- a/service/modelService.go
+++ b/service/modelService.go
@@ -19,13 +19,39 @@ func CreateModel(userID uint, modelType, url, parameter, description string) (er
 }
 
 // 根据id查找模型
-func FindModelByID(id, userID int) []dao.Model {
-	return dao.FindModelByID(id, userID)
+func FindModelByID(id, userID int) []dao.ModelResponse {
+	models := dao.FindModelByID(id, userID)
+	modelTokens := dao.FindModelTotalTokensByModelID(id)
+	var res []dao.ModelResponse
+	mt := make(map[uint]uint)
+	for _, v := range modelTokens {
+		mt[v.ModelID] = v.Token
+	}
+	for _, v := range models {
+		var model dao.ModelResponse
+		model.Model = v
+		model.Token = mt[v.ID]
+		res = append(res, model)
+	}
+	return res
 }
 
 // 根据用户id查找模型
-func FindModelByUserID(userID int) []dao.Model {
-	return dao.FindModelByUserID(userID)
+func FindModelByUserID(userID int) []dao.ModelResponse {
+	models := dao.FindModelByUserID(userID)
+	modelTokens := dao.FindModelTotalTokensByModelID(userID)
+	var res []dao.ModelResponse
+	mt := make(map[uint]uint)
+	for _, v := range modelTokens {
+		mt[v.ModelID] = v.Token
+	}
+	for _, v := range models {
+		var model dao.ModelResponse
+		model.Model = v
+		model.Token = mt[v.ID]
+		res = append(res, model)
+	}
+	return res
 }
 
 // 根据id删除模型