mirror of
https://github.com/yangjian102621/geekai.git
synced 2025-11-09 18:53:43 +08:00
feat: allow bind a chat model for chat role
This commit is contained in:
@@ -8,9 +8,10 @@ import (
|
||||
"chatplus/store/vo"
|
||||
"chatplus/utils"
|
||||
"chatplus/utils/resp"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
"time"
|
||||
)
|
||||
|
||||
type ChatModelHandler struct {
|
||||
|
||||
@@ -8,9 +8,10 @@ import (
|
||||
"chatplus/store/vo"
|
||||
"chatplus/utils"
|
||||
"chatplus/utils/resp"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
"time"
|
||||
)
|
||||
|
||||
type ChatRoleHandler struct {
|
||||
@@ -63,6 +64,25 @@ func (h *ChatRoleHandler) List(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// initialize model mane for role
|
||||
modelIds := make([]int, 0)
|
||||
for _, v := range items {
|
||||
if v.ModelId > 0 {
|
||||
modelIds = append(modelIds, v.ModelId)
|
||||
}
|
||||
}
|
||||
|
||||
modelNameMap := make(map[int]string)
|
||||
if len(modelIds) > 0 {
|
||||
var models []model.ChatModel
|
||||
tx := h.DB.Where("id IN ?", modelIds).Find(&models)
|
||||
if tx.Error == nil {
|
||||
for _, m := range models {
|
||||
modelNameMap[int(m.Id)] = m.Name
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, v := range items {
|
||||
var role vo.ChatRole
|
||||
err := utils.CopyObject(v, &role)
|
||||
@@ -70,6 +90,7 @@ func (h *ChatRoleHandler) List(c *gin.Context) {
|
||||
role.Id = v.Id
|
||||
role.CreatedAt = v.CreatedAt.Unix()
|
||||
role.UpdatedAt = v.UpdatedAt.Unix()
|
||||
role.ModelName = modelNameMap[role.ModelId]
|
||||
roles = append(roles, role)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -68,9 +68,20 @@ func (h *ChatHandler) ChatHandle(c *gin.Context) {
|
||||
modelId := h.GetInt(c, "model_id", 0)
|
||||
|
||||
client := types.NewWsClient(ws)
|
||||
var chatRole model.ChatRole
|
||||
res := h.DB.First(&chatRole, roleId)
|
||||
if res.Error != nil || !chatRole.Enable {
|
||||
utils.ReplyMessage(client, "当前聊天角色不存在或者未启用,连接已关闭!!!")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
// if the role bind a model_id, use role's bind model_id
|
||||
if chatRole.ModelId > 0 {
|
||||
modelId = chatRole.ModelId
|
||||
}
|
||||
// get model info
|
||||
var chatModel model.ChatModel
|
||||
res := h.DB.First(&chatModel, modelId)
|
||||
res = h.DB.First(&chatModel, modelId)
|
||||
if res.Error != nil || chatModel.Enabled == false {
|
||||
utils.ReplyMessage(client, "当前AI模型暂未启用,连接已关闭!!!")
|
||||
c.Abort()
|
||||
@@ -113,13 +124,6 @@ func (h *ChatHandler) ChatHandle(c *gin.Context) {
|
||||
Temperature: chatModel.Temperature,
|
||||
Platform: types.Platform(chatModel.Platform)}
|
||||
logger.Infof("New websocket connected, IP: %s, Username: %s", c.ClientIP(), session.Username)
|
||||
var chatRole model.ChatRole
|
||||
res = h.DB.First(&chatRole, roleId)
|
||||
if res.Error != nil || !chatRole.Enable {
|
||||
utils.ReplyMessage(client, "当前聊天角色不存在或者未启用,连接已关闭!!!")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
h.Init()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user