329 lines
8.8 KiB
Go
329 lines
8.8 KiB
Go
package workers
|
||
|
||
import (
|
||
"context"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"fmt"
|
||
"go-bot/tools"
|
||
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"os"
|
||
"regexp"
|
||
"strconv"
|
||
"strings"
|
||
|
||
openai "github.com/sashabaranov/go-openai"
|
||
)
|
||
|
||
func init() {
|
||
RegisterWorkerFactory("ai", func(parms []string, uid, gid, role, mid, rawMsg string) Worker {
|
||
return &AI{
|
||
StdAns: NewStdAns(parms, uid, gid, role, mid, rawMsg),
|
||
}
|
||
})
|
||
}
|
||
|
||
type AI struct {
|
||
*StdAns
|
||
}
|
||
|
||
func (a *AI) GetMsg() string {
|
||
if len(a.Parms) < 2 {
|
||
return "使用!ai xxx 向我提问吧"
|
||
}
|
||
|
||
ask := a.Parms[1]
|
||
if ask == "" || strings.HasPrefix(ask, "[CQ:reply,id=") {
|
||
return "不问问题你说个屁！"
|
||
}
|
||
|
||
OPENAI_API_KEY, OPENAI_BaseURL, MODEL, PROMPT, CONTEXT := tools.GetOAIConfig()
|
||
if OPENAI_API_KEY == "" {
|
||
return "OPENAI_API_KEY 未配置"
|
||
}
|
||
oaiconfig := openai.DefaultConfig(OPENAI_API_KEY)
|
||
oaiconfig.BaseURL = OPENAI_BaseURL
|
||
|
||
client := openai.NewClientWithConfig(oaiconfig)
|
||
msg := fmt.Sprintf("[CQ:at,qq=%s] ", a.UID)
|
||
if strings.ToLower(a.Parms[1]) != "models" {
|
||
|
||
// 如果非回复消息
|
||
if !strings.HasPrefix(a.Parms[len(a.Parms)-1], "[CQ:reply,id=") {
|
||
messages := make([]openai.ChatCompletionMessage, 0)
|
||
messages = append(messages, openai.ChatCompletionMessage{
|
||
Role: openai.ChatMessageRoleSystem,
|
||
Content: PROMPT,
|
||
})
|
||
// println("messages0:", messages)
|
||
var key string
|
||
redisClient := tools.GetRedisClient()
|
||
// 如果redisClient不为空，则获取上下文
|
||
if redisClient != nil {
|
||
// 限制请求频率
|
||
if !tools.CheckRequestFrequency(a.GID, a.UID, 10) {
|
||
log.Printf("请求过于频繁。\n")
|
||
return "请求过于频繁。"
|
||
}
|
||
|
||
key = fmt.Sprintf("context:%s:%s:%s", "ai", a.GID, a.UID)
|
||
// 获取上下文
|
||
message := openai.ChatCompletionMessage{
|
||
Role: openai.ChatMessageRoleUser,
|
||
Content: a.RawMsg[strings.Index(a.RawMsg, " ")+1:],
|
||
}
|
||
// 序列化为 JSON 字符串
|
||
jsonMessage, err := json.Marshal(message)
|
||
if err != nil {
|
||
log.Println("序列化错误:", err)
|
||
log.Println("存储的 JSON 字符串:", string(jsonMessage))
|
||
|
||
return "序列化错误"
|
||
}
|
||
tools.AddToContext(key, jsonMessage, CONTEXT)
|
||
// }
|
||
// println("RawMsg:", a.RawMsg[strings.Index(a.RawMsg, " ")+1:])
|
||
length := tools.GetListLength(key)
|
||
// if length > 0 {
|
||
for i := 0; i < int(length); i++ {
|
||
message, err := tools.GetListValue(key, int64(i))
|
||
if err != nil {
|
||
log.Println("获取上下文失败:", err)
|
||
return "获取上下文失败"
|
||
}
|
||
// log.Println("读取的 JSON 字符串:", message)
|
||
|
||
var msg openai.ChatCompletionMessage
|
||
err = json.Unmarshal([]byte(message), &msg)
|
||
if err != nil {
|
||
log.Println("反序列化错误:", err)
|
||
|
||
return "反序列化错误"
|
||
}
|
||
messages = append(messages, msg)
|
||
}
|
||
} else {
|
||
|
||
messages = append(messages, openai.ChatCompletionMessage{
|
||
Role: openai.ChatMessageRoleUser,
|
||
Content: a.RawMsg[strings.Index(a.RawMsg, " ")+1:],
|
||
})
|
||
|
||
}
|
||
// println("messages:", messages)
|
||
// for i, msg := range messages {
|
||
// fmt.Printf("消息 %d:\n", i+1)
|
||
// fmt.Printf(" 角色: %s\n", msg.Role)
|
||
// fmt.Printf(" 内容: %s\n", msg.Content)
|
||
// fmt.Println()
|
||
// }
|
||
resp, err := client.CreateChatCompletion(
|
||
context.Background(),
|
||
openai.ChatCompletionRequest{
|
||
Model: MODEL,
|
||
Stream: false,
|
||
Messages: messages,
|
||
},
|
||
)
|
||
if err != nil {
|
||
log.Println("ChatCompletion error: ", err)
|
||
tools.RemoveLastValueFromList(key)
|
||
return "请求失败"
|
||
}
|
||
// println(resp.Choices[0].Message.Content)
|
||
content := resp.Choices[0].Message.Content
|
||
if redisClient != nil {
|
||
|
||
message := openai.ChatCompletionMessage{
|
||
Role: openai.ChatMessageRoleAssistant,
|
||
Content: stripMarkdown(content)}
|
||
// 序列化为 JSON 字符串
|
||
jsonMessage, err := json.Marshal(message)
|
||
if err != nil {
|
||
log.Println("Assistant消息序列化错误:", err)
|
||
return "Assistant消息序列化错误"
|
||
}
|
||
tools.AddToContext(key, jsonMessage, CONTEXT)
|
||
|
||
}
|
||
return msg + stripMarkdown(content)
|
||
} else {
|
||
// 匹配回复消息
|
||
pattern := `^\[CQ:reply,id=(-?\d+)\]`
|
||
re := regexp.MustCompile(pattern)
|
||
matches := re.FindStringSubmatch(a.Parms[len(a.Parms)-1])
|
||
var msgId string
|
||
if len(matches) > 0 {
|
||
|
||
msgId = matches[1]
|
||
println("msgId:", msgId)
|
||
} else {
|
||
msgId = ""
|
||
log.Println("未找到回复消息")
|
||
return "未找到回复消息"
|
||
}
|
||
file, picUrl, fileSizeStr := a.GetHisMsg(msgId)
|
||
// println("file:", file, "picUrl:", picUrl, "fileSizeStr:", fileSizeStr)
|
||
|
||
if picUrl == "" {
|
||
log.Println("未找到文件信息")
|
||
return "未找到文件信息"
|
||
}
|
||
fileSize, err := strconv.ParseFloat(fileSizeStr, 64)
|
||
if err != nil {
|
||
fmt.Println("获取图片大小失败:", err)
|
||
return "获取图片大小失败"
|
||
}
|
||
if fileSize/1024/1024 > 5.0 {
|
||
log.Println("文件大小超过5M")
|
||
return "文件大小超过5M"
|
||
}
|
||
filePath := a.GetImage(file)
|
||
// println("filePath:", filePath)
|
||
if filePath == "" {
|
||
log.Println("获取图片失败")
|
||
return "获取图片失败"
|
||
}
|
||
base64Img := Image2Base64(filePath, picUrl)
|
||
if base64Img == "" {
|
||
log.Println("图片转换base64失败")
|
||
return "图片转换base64失败"
|
||
}
|
||
// 找到第一个空格的位置
|
||
firstSpaceIndex := strings.Index(a.RawMsg, " ")
|
||
|
||
// 找到最后一个空格的位置
|
||
lastSpaceIndex := strings.LastIndex(a.RawMsg, " ")
|
||
// 调用图片分析
|
||
resp, err := client.CreateChatCompletion(
|
||
context.Background(),
|
||
openai.ChatCompletionRequest{
|
||
Model: MODEL,
|
||
Messages: []openai.ChatCompletionMessage{
|
||
{
|
||
Role: "user",
|
||
MultiContent: []openai.ChatMessagePart{
|
||
{
|
||
Type: openai.ChatMessagePartTypeText,
|
||
Text: a.RawMsg[firstSpaceIndex+1 : lastSpaceIndex],
|
||
}, {
|
||
Type: openai.ChatMessagePartTypeImageURL,
|
||
ImageURL: &openai.ChatMessageImageURL{
|
||
URL: base64Img,
|
||
Detail: openai.ImageURLDetailAuto,
|
||
},
|
||
},
|
||
},
|
||
},
|
||
},
|
||
},
|
||
)
|
||
if err != nil {
|
||
log.Println("ChatCompletion error: ", err)
|
||
return "请求失败"
|
||
|
||
}
|
||
msg += resp.Choices[0].Message.Content
|
||
|
||
}
|
||
return stripMarkdown(msg)
|
||
}
|
||
models, err := client.ListModels(context.Background())
|
||
if err != nil {
|
||
log.Println("ListModels error: ", err)
|
||
return "查询支持模型失败！"
|
||
|
||
}
|
||
var modelList string
|
||
for _, model := range models.Models {
|
||
if MODEL == model.ID {
|
||
model.ID = model.ID + "\t(当前使用)"
|
||
}
|
||
modelList += model.ID + "\n"
|
||
}
|
||
return modelList
|
||
// return handleModelRequest(OPENAI_API_KEY, OPENAI_BaseURL)
|
||
}
|
||
|
||
func stripMarkdown(text string) string {
|
||
|
||
// println("before:", text)
|
||
// 移除代码块，但保留代码内容
|
||
re := regexp.MustCompile("(?s)```(?:\\w+\\s*\n)?(.*?)```")
|
||
text = re.ReplaceAllStringFunc(text, func(match string) string {
|
||
submatches := re.FindStringSubmatch(match)
|
||
if len(submatches) > 1 {
|
||
return strings.TrimSpace(submatches[1])
|
||
}
|
||
return match
|
||
})
|
||
|
||
// 移除行内代码
|
||
re = regexp.MustCompile("`([^`]+)`")
|
||
text = re.ReplaceAllString(text, "$1")
|
||
|
||
// 移除标题
|
||
re = regexp.MustCompile(`(?m)^#{1,6}\s+(.+)$`)
|
||
text = re.ReplaceAllString(text, "$1")
|
||
|
||
// 移除粗体和斜体
|
||
re = regexp.MustCompile(`(\*\*|__)(.+?)(\*\*|__)`)
|
||
text = re.ReplaceAllString(text, "$2")
|
||
re = regexp.MustCompile(`(\*|_)(.+?)(\*|_)`)
|
||
text = re.ReplaceAllString(text, "$2")
|
||
|
||
// 移除链接，保留链接文本
|
||
re = regexp.MustCompile(`\[([^\]]+)\]\([^\)]+\)`)
|
||
text = re.ReplaceAllString(text, "$1")
|
||
|
||
// 移除图片
|
||
re = regexp.MustCompile(`!\[([^\]]*)\]\([^\)]+\)`)
|
||
text = re.ReplaceAllString(text, "$1")
|
||
|
||
// 移除水平线
|
||
re = regexp.MustCompile(`(?m)^-{3,}|_{3,}|\*{3,}$`)
|
||
text = re.ReplaceAllString(text, "")
|
||
|
||
// 移除块引用
|
||
re = regexp.MustCompile(`(?m)^>\s+(.+)`)
|
||
text = re.ReplaceAllString(text, "$1")
|
||
|
||
// 移除列表标记
|
||
re = regexp.MustCompile(`(?m)^[ \t]*[\*\-+][ \t]+`)
|
||
text = re.ReplaceAllString(text, "")
|
||
re = regexp.MustCompile(`(?m)^[ \t]*\d+\.[ \t]+`)
|
||
text = re.ReplaceAllString(text, "")
|
||
|
||
// 移除多余的空行
|
||
re = regexp.MustCompile(`\n{3,}`)
|
||
text = re.ReplaceAllString(text, "\n\n")
|
||
// println("after:", text)
|
||
return strings.TrimSpace(text)
|
||
}
|
||
|
||
func Image2Base64(path string, picUrl string) string {
|
||
// 尝试从文件路径读取图片
|
||
file, err := os.Open(path)
|
||
if err == nil {
|
||
defer file.Close()
|
||
if data, err := io.ReadAll(file); err == nil {
|
||
return "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(data)
|
||
}
|
||
}
|
||
|
||
// 如果文件路径不可用，则尝试从 URL 下载图片
|
||
resp, err := http.Get(picUrl)
|
||
if err != nil {
|
||
return ""
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if data, err := io.ReadAll(resp.Body); err == nil {
|
||
return "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(data)
|
||
}
|
||
return ""
|
||
|
||
}
|