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删除模型