feat: different AI model consuming different amounts of use_calls

This commit is contained in:
RockYang
2023-10-26 14:38:06 +08:00
parent 7ccb4c5f06
commit 7677ae254f
9 changed files with 27 additions and 16 deletions

View File

@@ -47,6 +47,7 @@ type ChatModel struct {
Id uint `json:"id"`
Platform Platform `json:"platform"`
Value string `json:"value"`
Weight int `json:"weight"`
}
type ApiError struct {

View File

@@ -168,9 +168,7 @@ func (h *ChatHandler) sendAzureMessage(
// 消息发送成功
if len(contents) > 0 {
// 更新用户的对话次数
if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
}
h.subUserCalls(userVo, session)
if message.Role == "" {
message.Role = "assistant"

View File

@@ -124,9 +124,7 @@ func (h *ChatHandler) sendBaiduMessage(
// 消息发送成功
if len(contents) > 0 {
// 更新用户的对话次数
if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
}
h.subUserCalls(userVo, session)
if message.Role == "" {
message.Role = "assistant"

View File

@@ -103,6 +103,7 @@ func (h *ChatHandler) ChatHandle(c *gin.Context) {
session.Model = types.ChatModel{
Id: chatModel.Id,
Value: chatModel.Value,
Weight: chatModel.Weight,
Platform: types.Platform(chatModel.Platform)}
logger.Infof("New websocket connected, IP: %s, Username: %s", c.ClientIP(), session.Username)
var chatRole model.ChatRole
@@ -463,3 +464,15 @@ func (h *ChatHandler) doRequest(ctx context.Context, req types.ApiRequest, platf
}
return client.Do(request)
}
// 扣减用户的对话次数
func (h *ChatHandler) subUserCalls(userVo vo.User, session *types.ChatSession) {
// 仅当用户没有导入自己的 API KEY 时才进行扣减
if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
num := 1
if session.Model.Weight > 0 {
num = session.Model.Weight
}
h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", num))
}
}

View File

@@ -104,9 +104,7 @@ func (h *ChatHandler) sendChatGLMMessage(
// 消息发送成功
if len(contents) > 0 {
// 更新用户的对话次数
if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
}
h.subUserCalls(userVo, session)
if message.Role == "" {
message.Role = "assistant"

View File

@@ -167,9 +167,7 @@ func (h *ChatHandler) sendOpenAiMessage(
// 消息发送成功
if len(contents) > 0 {
// 更新用户的对话次数
if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
}
h.subUserCalls(userVo, session)
if message.Role == "" {
message.Role = "assistant"

View File

@@ -167,9 +167,7 @@ func (h *ChatHandler) sendXunFeiMessage(
// 消息发送成功
if len(contents) > 0 {
// 更新用户的对话次数
if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
}
h.subUserCalls(userVo, session)
if message.Role == "" {
message.Role = "assistant"