2024-11-06 01:06:38 +08:00
|
|
|
|
package glm
|
|
|
|
|
|
|
|
|
|
import (
|
|
|
|
|
"catface/app/global/variable"
|
|
|
|
|
"context"
|
|
|
|
|
"errors"
|
2024-11-19 02:22:39 +08:00
|
|
|
|
"fmt"
|
2024-11-16 18:18:07 +08:00
|
|
|
|
"strings"
|
|
|
|
|
"time"
|
2024-11-06 01:06:38 +08:00
|
|
|
|
|
|
|
|
|
"github.com/yankeguo/zhipu"
|
2024-11-19 02:22:39 +08:00
|
|
|
|
"go.uber.org/zap"
|
2024-11-06 01:06:38 +08:00
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
// ChatWithGLM 封装了与GLM模型进行对话的逻辑
|
2024-11-19 02:22:39 +08:00
|
|
|
|
func Chat(message string, client *zhipu.ChatCompletionService) (string, error) {
|
|
|
|
|
service := client.AddMessage(zhipu.ChatCompletionMessage{
|
|
|
|
|
Role: "user",
|
|
|
|
|
Content: message,
|
|
|
|
|
})
|
2024-11-06 01:06:38 +08:00
|
|
|
|
|
|
|
|
|
res, err := service.Do(context.Background())
|
|
|
|
|
if err != nil {
|
|
|
|
|
apiErrorCode := zhipu.GetAPIErrorCode(err)
|
|
|
|
|
return "", errors.New(apiErrorCode) // 将字符串包装成 error 类型
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
return res.Choices[0].Message.Content, nil
|
|
|
|
|
}
|
2024-11-16 14:00:57 +08:00
|
|
|
|
|
|
|
|
|
// ChatStream 接收一个消息和一个通道,将流式响应发送到通道中
|
2024-11-19 02:22:39 +08:00
|
|
|
|
func ChatStream(message string, ch chan<- string, client *zhipu.ChatCompletionService) error {
|
|
|
|
|
service := client.AddMessage(zhipu.ChatCompletionMessage{Role: "user", Content: message}).
|
2024-11-16 14:00:57 +08:00
|
|
|
|
SetStreamHandler(func(chunk zhipu.ChatCompletionResponse) error {
|
|
|
|
|
content := chunk.Choices[0].Delta.Content
|
|
|
|
|
if content != "" {
|
|
|
|
|
ch <- content // 将内容发送到通道
|
|
|
|
|
}
|
|
|
|
|
return nil
|
|
|
|
|
})
|
|
|
|
|
|
2024-11-19 02:22:39 +08:00
|
|
|
|
// Test
|
|
|
|
|
messages := client.GetMessages()
|
|
|
|
|
for id, message := range messages {
|
|
|
|
|
variable.ZapLog.Info(fmt.Sprintf("message-%d", id+1), zap.String("message", message.(zhipu.ChatCompletionMessage).Role), zap.String("content", message.(zhipu.ChatCompletionMessage).Content))
|
|
|
|
|
}
|
|
|
|
|
|
2024-11-16 14:00:57 +08:00
|
|
|
|
// 执行服务调用
|
2024-11-19 02:22:39 +08:00
|
|
|
|
res, err := service.Do(context.Background())
|
2024-11-16 14:00:57 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
return err
|
|
|
|
|
}
|
2024-11-19 02:22:39 +08:00
|
|
|
|
// 增加 AI 回答的消息记录。
|
|
|
|
|
client.AddMessage(zhipu.ChatCompletionMessage{Role: "assistant", Content: res.Choices[0].Message.Content})
|
2024-11-16 14:00:57 +08:00
|
|
|
|
|
|
|
|
|
return nil
|
|
|
|
|
}
|
2024-11-16 18:18:07 +08:00
|
|
|
|
|
|
|
|
|
// 带缓冲机制的 ChatStream;计数 & 计时 双判定。
|
2024-11-19 02:22:39 +08:00
|
|
|
|
func BufferedChatStream(message string, ch chan<- string, client *zhipu.ChatCompletionService) error {
|
2024-11-16 18:18:07 +08:00
|
|
|
|
bufferedCh := make(chan string) // 带缓冲的通道,缓冲大小为10
|
|
|
|
|
timer := time.NewTimer(500 * time.Millisecond) // 定时器,500毫秒
|
|
|
|
|
|
|
|
|
|
go func() {
|
2024-11-19 02:22:39 +08:00
|
|
|
|
err := ChatStream(message, bufferedCh, client)
|
2024-11-16 18:18:07 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
return
|
|
|
|
|
}
|
|
|
|
|
close(bufferedCh)
|
|
|
|
|
}()
|
|
|
|
|
|
|
|
|
|
var buffer strings.Builder
|
|
|
|
|
for {
|
|
|
|
|
select {
|
|
|
|
|
case c, ok := <-bufferedCh:
|
|
|
|
|
if !ok {
|
|
|
|
|
if buffer.Len() > 0 {
|
|
|
|
|
ch <- buffer.String()
|
|
|
|
|
}
|
|
|
|
|
return nil // 依靠这里停止函数。
|
|
|
|
|
}
|
|
|
|
|
buffer.WriteString(c)
|
|
|
|
|
if buffer.Len() >= 10 {
|
|
|
|
|
ch <- buffer.String()
|
|
|
|
|
buffer.Reset()
|
|
|
|
|
timer.Reset(500 * time.Millisecond)
|
|
|
|
|
}
|
|
|
|
|
case <-timer.C:
|
|
|
|
|
if buffer.Len() > 0 {
|
|
|
|
|
ch <- buffer.String()
|
|
|
|
|
buffer.Reset()
|
|
|
|
|
}
|
|
|
|
|
timer.Reset(500 * time.Millisecond)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|