saw-go/dao/model.go

package dao

import (
	"StuAcaWorksAI/proto"
	"gorm.io/gorm"
)

type Model struct {
	gorm.Model
	Type        string `gorm:"column:type"`        //模型类型
	UserID      uint   `gorm:"column:user_id"`     //用户id
	Url         string `gorm:"column:url"`         //模型地址
	Parameter   string `gorm:"column:parameter"`   //模型参数,存储json
	Description string `gorm:"column:description"` //模型描述
}
type FunctionModel struct {
	gorm.Model
	Name     string `gorm:"column:name"`      //功能名称
	Function string `gorm:"column:function"`  //功能函数,唯一标识
	UserID   uint   `gorm:"column:user_id"`   //用户id
	ModelID  uint   `gorm:"column:model_id"`  //模型id
	ModelIDS string `gorm:"column:model_ids"` //允许模型id,可多个
	Info     string `gorm:"column:info"`      //功能信息，系统功能,对应模型系统参数
	//System   string `gorm:"column:system"`   //系统功能,对应模型系统参数
}

type ModelToken struct {
	gorm.Model
	ModelID   uint `gorm:"column:model_id"`   //模型id
	UserID    uint `gorm:"column:user_id"`    //用户id
	SessionID uint `gorm:"column:session_id"` //会话id
	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}
	var res *gorm.DB
	if proto.Config.SERVER_SQL_LOG {
		res = DB.Debug().Create(&model)
	} else {
		res = DB.Create(&model)
	}
	return res.Error, model.ID
}

// 根据id查找模型
func FindModelByID(id, userID int) []Model {
	var model []Model
	if proto.Config.SERVER_SQL_LOG {
		DB.Debug().Where("id = ? and user_id = ?", id, userID).Find(&model)
	} else {
		DB.Where("id = ? and user_id = ?", id, userID).Find(&model)
	}
	return model
}

// 根据id查找模型
func FindModelByIDV2(id int) []Model {
	var model []Model
	if proto.Config.SERVER_SQL_LOG {
		DB.Debug().Where("id = ?", id).Find(&model)
	} else {
		DB.Where("id = ?", id).Find(&model)
	}
	return model
}

// 根据modelID查找模型
func FindModelByModelID(modelID uint) []Model {
	var model []Model
	if proto.Config.SERVER_SQL_LOG {
		DB.Debug().Where("id = ?", modelID).Find(&model)
	} else {
		DB.Where("id = ?", modelID).Find(&model)
	}
	return model
}

// 根据用户id查找模型
func FindModelByUserID(userID int) []Model {
	var models []Model
	if proto.Config.SERVER_SQL_LOG {
		DB.Debug().Where("user_id = ?", userID).Find(&models)
	} else {
		DB.Where("user_id = ?", userID).Find(&models)
	}
	return models
}

// 根据id删除模型
func DeleteModelByID(id int) error {
	var res *gorm.DB
	if proto.Config.SERVER_SQL_LOG {
		res = DB.Debug().Delete(&Model{}, id)
	} else {
		res = DB.Delete(&Model{}, id)
	}
	return res.Error
}

// 根据id更新模型
func UpdateModelByID(id int, userID uint, modelType, url, parameter, description string) error {
	model := Model{UserID: userID, Type: modelType, Url: url, Parameter: parameter, Description: description}
	var res *gorm.DB
	if proto.Config.SERVER_SQL_LOG {
		res = DB.Debug().Model(&Model{}).Where("id = ?", id).Updates(&model)
	} else {
		res = DB.Model(&Model{}).Where("id = ?", id).Updates(&model)
	}
	return res.Error
}

// 创建功能模型
func CreateFunctionModel(userID uint, modelID uint, name, info, function, modelIDs string) (error, uint) {
	functionModel := FunctionModel{UserID: userID, ModelID: modelID, Name: name, Info: info, Function: function, ModelIDS: modelIDs}
	var res *gorm.DB
	if proto.Config.SERVER_SQL_LOG {
		res = DB.Debug().Create(&functionModel)
	} else {
		res = DB.Create(&functionModel)
	}
	return res.Error, functionModel.ID
}

// 根据id查找功能模型
func FindFunctionModelByID(id, userID int) []FunctionModel {
	var functionModel []FunctionModel
	if proto.Config.SERVER_SQL_LOG {
		DB.Debug().Where("id = ? and user_id = ?", id, userID).Find(&functionModel)
	} else {
		DB.Where("id = ? and user_id = ?", id, userID).Find(&functionModel)
	}
	return functionModel
}

// 根据用户id查找功能模型
func FindFunctionModelByUserID(userID int) []FunctionModel {
	var functionModels []FunctionModel
	if proto.Config.SERVER_SQL_LOG {
		DB.Debug().Where("user_id = ?", userID).Find(&functionModels)
	} else {
		DB.Where("user_id = ?", userID).Find(&functionModels)
	}
	return functionModels
}

// 修改功能模型
func UpdateFunctionModelByID(id int, userID uint, modelID uint, name, info, function, modelIDs string) error {
	functionModel := FunctionModel{UserID: userID, ModelID: modelID, Name: name, Info: info, Function: function, ModelIDS: modelIDs}
	var res *gorm.DB
	if proto.Config.SERVER_SQL_LOG {
		res = DB.Debug().Model(&FunctionModel{}).Where("id = ?", id).Updates(&functionModel)
	} else {
		res = DB.Model(&FunctionModel{}).Where("id = ?", id).Updates(&functionModel)
	}
	return res.Error
}

// 根据id删除功能模型
func DeleteFunctionModelByID(id int) error {
	var res *gorm.DB
	if proto.Config.SERVER_SQL_LOG {
		res = DB.Debug().Delete(&FunctionModel{}, id)
	} else {
		res = DB.Delete(&FunctionModel{}, id)
	}
	return res.Error
}

func FindFunctionModelByFunction(function string) []FunctionModel {
	var models []FunctionModel
	if proto.Config.SERVER_SQL_LOG {
		DB.Debug().Where("function = ?", function).Find(&models)
	} else {
		DB.Where("function = ?", function).Find(&models)
	}
	return models
}

func CreateModelToken(userID uint, modelID uint, sessionID, token uint) (error, uint) {
	modelToken := ModelToken{UserID: userID, ModelID: modelID, Token: token, SessionID: sessionID}
	var res *gorm.DB
	if proto.Config.SERVER_SQL_LOG {
		res = DB.Debug().Create(&modelToken)
	} else {
		res = DB.Create(&modelToken)
	}
	return res.Error, modelToken.ID
}

func FindModelTokenByModelID(modelID uint) []ModelToken {
	var modelToken []ModelToken
	if proto.Config.SERVER_SQL_LOG {
		DB.Debug().Where("model_id = ?", modelID).Find(&modelToken)
	} else {
		DB.Where("model_id = ?", modelID).Find(&modelToken)
	}
	return modelToken
}

func FindModelTokenByUserID(userID uint) []ModelToken {
	var modelToken []ModelToken
	if proto.Config.SERVER_SQL_LOG {
		DB.Debug().Where("user_id = ?", userID).Find(&modelToken)
	} else {
		DB.Where("user_id = ?", userID).Find(&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, sum(token) as token").Where("user_id = ?", userID).Group("session_id").Scan(&modelTokens)
	} else {
		DB.Table("model_tokens").Select("session_id, sum(token) as token").Where("user_id = ?", userID).Group("session_id").Scan(&modelTokens)
	}
	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, sum(token) as token").Where("user_id = ?", userID).Group("model_id").Scan(&modelTokens)
	} else {
		DB.Table("model_tokens").Select("model_id, sum(token) as token").Where("user_id = ?", userID).Group("model_id").Scan(&modelTokens)
	}
	return modelTokens
}