mirror of
https://github.com/yangjian102621/geekai.git
synced 2025-09-27 21:56:39 +08:00
329 lines
8.7 KiB
Go
329 lines
8.7 KiB
Go
package chatimpl
|
||
|
||
import (
|
||
"bufio"
|
||
"chatplus/core/types"
|
||
"chatplus/store/model"
|
||
"chatplus/store/vo"
|
||
"chatplus/utils"
|
||
"context"
|
||
"crypto/hmac"
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"fmt"
|
||
"html/template"
|
||
"io"
|
||
"net/http"
|
||
"net/url"
|
||
"strings"
|
||
"time"
|
||
"unicode/utf8"
|
||
)
|
||
|
||
const (
|
||
Region = "cn-beijing"
|
||
Service = "ml_maas"
|
||
)
|
||
|
||
type skylarkResp struct {
|
||
Id string `json:"req_id"`
|
||
Choice struct {
|
||
Message struct {
|
||
Content string `json:"content"`
|
||
} `json:"message"`
|
||
} `json:"choice"`
|
||
}
|
||
|
||
// Skylark 消息发送实现
|
||
func (h *ChatHandler) sendSkylarkMessage(
|
||
chatCtx []interface{},
|
||
req types.ApiRequest,
|
||
userVo vo.User,
|
||
ctx context.Context,
|
||
session *types.ChatSession,
|
||
role model.ChatRole,
|
||
prompt string,
|
||
ws *types.WsClient) error {
|
||
promptCreatedAt := time.Now() // 记录提问时间
|
||
start := time.Now()
|
||
var apiKey model.ApiKey
|
||
response, err := h.doRequest(ctx, req, session.Model.Platform, &apiKey)
|
||
logger.Info("HTTP请求完成，耗时：", time.Since(start))
|
||
if err != nil {
|
||
if strings.Contains(err.Error(), "context canceled") {
|
||
logger.Info("用户取消了请求：", prompt)
|
||
return nil
|
||
} else if strings.Contains(err.Error(), "no available key") {
|
||
utils.ReplyMessage(ws, "抱歉😔😔😔，系统已经没有可用的 API KEY，请联系管理员！")
|
||
return nil
|
||
} else {
|
||
logger.Error(err)
|
||
}
|
||
|
||
utils.ReplyMessage(ws, ErrorMsg)
|
||
utils.ReplyMessage(ws, ErrImg)
|
||
return err
|
||
} else {
|
||
defer response.Body.Close()
|
||
}
|
||
|
||
contentType := response.Header.Get("Content-Type")
|
||
if strings.Contains(contentType, "text/event-stream") {
|
||
replyCreatedAt := time.Now() // 记录回复时间
|
||
// 循环读取 Chunk 消息
|
||
var message = types.Message{}
|
||
var contents = make([]string, 0)
|
||
var content string
|
||
scanner := bufio.NewScanner(response.Body)
|
||
for scanner.Scan() {
|
||
line := scanner.Text()
|
||
if len(line) < 5 || strings.HasPrefix(line, "id:") {
|
||
continue
|
||
}
|
||
|
||
if strings.HasPrefix(line, "data:") {
|
||
content = line[5:]
|
||
}
|
||
|
||
if content == "[DONE]" {
|
||
break
|
||
}
|
||
|
||
// 处理代码换行
|
||
if len(content) == 0 {
|
||
content = "\n"
|
||
}
|
||
|
||
var resp skylarkResp
|
||
err := utils.JsonDecode(content, &resp)
|
||
if err != nil {
|
||
logger.Error("error with parse data line: ", err)
|
||
utils.ReplyMessage(ws, fmt.Sprintf("**解析数据行失败：%s**", err))
|
||
break
|
||
}
|
||
|
||
if len(contents) == 0 {
|
||
utils.ReplyChunkMessage(ws, types.WsMessage{Type: types.WsStart})
|
||
}
|
||
utils.ReplyChunkMessage(ws, types.WsMessage{
|
||
Type: types.WsMiddle,
|
||
Content: utils.InterfaceToString(resp.Choice.Message.Content),
|
||
})
|
||
contents = append(contents, resp.Choice.Message.Content)
|
||
|
||
} // end for
|
||
|
||
if err := scanner.Err(); err != nil {
|
||
if strings.Contains(err.Error(), "context canceled") {
|
||
logger.Info("用户取消了请求：", prompt)
|
||
} else {
|
||
logger.Error("信息读取出错：", err)
|
||
}
|
||
}
|
||
|
||
// 消息发送成功
|
||
if len(contents) > 0 {
|
||
// 更新用户的对话次数
|
||
h.subUserCalls(userVo, session)
|
||
|
||
if message.Role == "" {
|
||
message.Role = "assistant"
|
||
}
|
||
message.Content = strings.Join(contents, "")
|
||
useMsg := types.Message{Role: "user", Content: prompt}
|
||
|
||
// 更新上下文消息，如果是调用函数则不需要更新上下文
|
||
if h.App.ChatConfig.EnableContext {
|
||
chatCtx = append(chatCtx, useMsg) // 提问消息
|
||
chatCtx = append(chatCtx, message) // 回复消息
|
||
h.App.ChatContexts.Put(session.ChatId, chatCtx)
|
||
}
|
||
|
||
// 追加聊天记录
|
||
if h.App.ChatConfig.EnableHistory {
|
||
// for prompt
|
||
promptToken, err := utils.CalcTokens(prompt, req.Model)
|
||
if err != nil {
|
||
logger.Error(err)
|
||
}
|
||
historyUserMsg := model.ChatMessage{
|
||
UserId: userVo.Id,
|
||
ChatId: session.ChatId,
|
||
RoleId: role.Id,
|
||
Type: types.PromptMsg,
|
||
Icon: userVo.Avatar,
|
||
Content: template.HTMLEscapeString(prompt),
|
||
Tokens: promptToken,
|
||
UseContext: true,
|
||
Model: req.Model,
|
||
}
|
||
historyUserMsg.CreatedAt = promptCreatedAt
|
||
historyUserMsg.UpdatedAt = promptCreatedAt
|
||
res := h.db.Save(&historyUserMsg)
|
||
if res.Error != nil {
|
||
logger.Error("failed to save prompt history message: ", res.Error)
|
||
}
|
||
|
||
// for reply
|
||
// 计算本次对话消耗的总 token 数量
|
||
replyToken, _ := utils.CalcTokens(message.Content, req.Model)
|
||
totalTokens := replyToken + getTotalTokens(req)
|
||
historyReplyMsg := model.ChatMessage{
|
||
UserId: userVo.Id,
|
||
ChatId: session.ChatId,
|
||
RoleId: role.Id,
|
||
Type: types.ReplyMsg,
|
||
Icon: role.Icon,
|
||
Content: message.Content,
|
||
Tokens: totalTokens,
|
||
UseContext: true,
|
||
Model: req.Model,
|
||
}
|
||
historyReplyMsg.CreatedAt = replyCreatedAt
|
||
historyReplyMsg.UpdatedAt = replyCreatedAt
|
||
res = h.db.Create(&historyReplyMsg)
|
||
if res.Error != nil {
|
||
logger.Error("failed to save reply history message: ", res.Error)
|
||
}
|
||
// 更新用户信息
|
||
h.incUserTokenFee(userVo.Id, totalTokens)
|
||
}
|
||
|
||
// 保存当前会话
|
||
var chatItem model.ChatItem
|
||
res := h.db.Where("chat_id = ?", session.ChatId).First(&chatItem)
|
||
if res.Error != nil {
|
||
chatItem.ChatId = session.ChatId
|
||
chatItem.UserId = session.UserId
|
||
chatItem.RoleId = role.Id
|
||
chatItem.ModelId = session.Model.Id
|
||
if utf8.RuneCountInString(prompt) > 30 {
|
||
chatItem.Title = string([]rune(prompt)[:30]) + "..."
|
||
} else {
|
||
chatItem.Title = prompt
|
||
}
|
||
chatItem.Model = req.Model
|
||
h.db.Create(&chatItem)
|
||
}
|
||
}
|
||
} else {
|
||
body, err := io.ReadAll(response.Body)
|
||
if err != nil {
|
||
return fmt.Errorf("error with reading response: %v", err)
|
||
}
|
||
|
||
var res struct {
|
||
Code int `json:"error_code"`
|
||
Msg string `json:"error_msg"`
|
||
}
|
||
err = json.Unmarshal(body, &res)
|
||
if err != nil {
|
||
return fmt.Errorf("error with decode response: %v", err)
|
||
}
|
||
utils.ReplyMessage(ws, "请求 Skylark API 失败："+res.Msg)
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
func hmacSHA256(key []byte, content string) []byte {
|
||
mac := hmac.New(sha256.New, key)
|
||
mac.Write([]byte(content))
|
||
return mac.Sum(nil)
|
||
}
|
||
|
||
func getSignedKey(secretKey, date, region, service string) []byte {
|
||
kDate := hmacSHA256([]byte(secretKey), date)
|
||
kRegion := hmacSHA256(kDate, region)
|
||
kService := hmacSHA256(kRegion, service)
|
||
kSigning := hmacSHA256(kService, "request")
|
||
|
||
return kSigning
|
||
}
|
||
|
||
func hashSHA256(data []byte) []byte {
|
||
hash := sha256.New()
|
||
if _, err := hash.Write(data); err != nil {
|
||
logger.Debugf("input hash err:%s", err.Error())
|
||
}
|
||
|
||
return hash.Sum(nil)
|
||
}
|
||
|
||
type SkylarkRequest struct {
|
||
Method string
|
||
Url string
|
||
Body []byte
|
||
AccessKeyID string
|
||
SecretAccessKey string
|
||
}
|
||
|
||
func (h *ChatHandler) makeSkylarkRequestUrl(params SkylarkRequest) (http.Header, error) {
|
||
ul, err := url.Parse(params.Url)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
requestAddr := fmt.Sprintf("%s%s?%s", ul.Host, ul.Path, ul.Query().Encode())
|
||
logger.Debugf("request addr: %s\n", requestAddr)
|
||
|
||
var reqHeaders = http.Header{}
|
||
|
||
now := time.Now()
|
||
date := now.UTC().Format("20060102T150405Z")
|
||
authDate := date[:8]
|
||
reqHeaders.Set("X-Date", date)
|
||
payload := hex.EncodeToString(hashSHA256(params.Body))
|
||
reqHeaders.Set("X-Content-Sha256", payload)
|
||
reqHeaders.Set("Content-Type", "application/json")
|
||
|
||
queryString := strings.Replace(ul.Query().Encode(), "+", "%20", -1)
|
||
signedHeaders := []string{"content-type", "host", "x-content-sha256", "x-date"}
|
||
var headerList []string
|
||
for _, header := range signedHeaders {
|
||
if header == "host" {
|
||
headerList = append(headerList, header+":"+ul.Host)
|
||
} else {
|
||
v := reqHeaders.Get(header)
|
||
headerList = append(headerList, header+":"+strings.TrimSpace(v))
|
||
}
|
||
}
|
||
headerString := strings.Join(headerList, "\n")
|
||
|
||
canonicalString := strings.Join([]string{
|
||
params.Method,
|
||
ul.Path,
|
||
queryString,
|
||
headerString + "\n",
|
||
strings.Join(signedHeaders, ";"),
|
||
payload,
|
||
}, "\n")
|
||
logger.Debugf("canonical string:\n%s\n", canonicalString)
|
||
|
||
hashedCanonicalString := hex.EncodeToString(hashSHA256([]byte(canonicalString)))
|
||
logger.Debugf("hashed canonical string: %s\n", hashedCanonicalString)
|
||
|
||
credentialScope := authDate + "/" + Region + "/" + Service + "/request"
|
||
signString := strings.Join([]string{
|
||
"HMAC-SHA256",
|
||
date,
|
||
credentialScope,
|
||
hashedCanonicalString,
|
||
}, "\n")
|
||
logger.Debugf("sign string:\n%s\n", signString)
|
||
|
||
// 3. 构建认证请求头
|
||
signedKey := getSignedKey(params.SecretAccessKey, authDate, Region, Service)
|
||
signature := hex.EncodeToString(hmacSHA256(signedKey, signString))
|
||
logger.Debugf("signature: %s\n", signature)
|
||
|
||
authorization := "HMAC-SHA256" +
|
||
" Credential=" + params.AccessKeyID + "/" + credentialScope +
|
||
", SignedHeaders=" + strings.Join(signedHeaders, ";") +
|
||
", Signature=" + signature
|
||
reqHeaders.Set("Authorization", authorization)
|
||
logger.Debugf("authorization: %s\n", authorization)
|
||
|
||
return reqHeaders, nil
|
||
}
|