feat: allow user to set custom api keys for different platforms

This commit is contained in:
RockYang
2023-09-04 17:34:29 +08:00
parent f7a427d2c0
commit 2820adad53
11 changed files with 76 additions and 68 deletions

View File

@@ -29,7 +29,7 @@ func (h *ChatHandler) sendChatGLMMessage(
ws *types.WsClient) error {
promptCreatedAt := time.Now() // 记录提问时间
start := time.Now()
var apiKey string
var apiKey = userVo.ChatConfig.ApiKeys[session.Model.Platform]
response, err := h.doRequest(ctx, req, session.Model.Platform, &apiKey)
logger.Info("HTTP请求完成耗时", time.Now().Sub(start))
if err != nil {
@@ -103,7 +103,9 @@ func (h *ChatHandler) sendChatGLMMessage(
// 消息发送成功
if len(contents) > 0 {
// 更新用户的对话次数
h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
}
if message.Role == "" {
message.Role = "assistant"
@@ -112,14 +114,14 @@ func (h *ChatHandler) sendChatGLMMessage(
useMsg := types.Message{Role: "user", Content: prompt}
// 更新上下文消息,如果是调用函数则不需要更新上下文
if userVo.ChatConfig.EnableContext {
if h.App.ChatConfig.EnableContext {
chatCtx = append(chatCtx, useMsg) // 提问消息
chatCtx = append(chatCtx, message) // 回复消息
h.App.ChatContexts.Put(session.ChatId, chatCtx)
}
// 追加聊天记录
if userVo.ChatConfig.EnableHistory {
if h.App.ChatConfig.EnableHistory {
// for prompt
promptToken, err := utils.CalcTokens(prompt, req.Model)
if err != nil {
@@ -167,8 +169,6 @@ func (h *ChatHandler) sendChatGLMMessage(
// 计算本次对话消耗的总 token 数量
var totalTokens = 0
totalTokens = replyToken + getTotalTokens(req)
h.db.Model(&model.User{}).Where("id = ?", userVo.Id).
UpdateColumn("tokens", gorm.Expr("tokens + ?", totalTokens))
h.db.Model(&model.User{}).Where("id = ?", userVo.Id).
UpdateColumn("total_tokens", gorm.Expr("total_tokens + ?", totalTokens))
}
@@ -205,7 +205,7 @@ func (h *ChatHandler) sendChatGLMMessage(
return fmt.Errorf("error with decode response: %v", err)
}
if !res.Success {
utils.ReplyMessage(ws, "请求 ChatGML 失败:"+res.Msg)
utils.ReplyMessage(ws, "请求 ChatGLM 失败:"+res.Msg)
}
}