qq_bot/handler/headmaster/headmaster.go

103 lines
3.1 KiB
Go

package headmaster
import (
"context"
"fmt"
"strings"
"time"
"git.lxtend.com/qqbot/config"
"git.lxtend.com/qqbot/constants"
"git.lxtend.com/qqbot/handler"
"git.lxtend.com/qqbot/model"
"git.lxtend.com/qqbot/util"
"github.com/sashabaranov/go-openai"
)
var (
histories = make(map[string][]openai.ChatCompletionMessage)
histories_time = make(map[string]time.Time)
)
func init() {
handler.RegisterHandler("校长", headmasterHandler, constants.LEVEL_USER)
handler.RegisterPrivateHandler(headmasterHandler)
}
func headmasterHandler(msg model.Message) (reply model.Reply) {
var from string
if msg.GroupInfo.IsGroupMsg {
from = util.From(msg.GroupInfo.GroupId, 0)
} else {
from = util.From(0, msg.UserId)
}
if len(msg.RawMsg) > 7 && msg.RawMsg[0:7] == "校长 " {
return model.Reply{
ReplyMsg: ask(from, fmt.Sprintf("{\"qqid\":%d,\"nickname\":\"%s\",\"referid\":%d,\"msg\":\"%s\"}", msg.UserId, msg.UserNickName, msg.OriginMsgId, msg.RawMsg[7:])),
ReferOriginMsg: false,
FromMsg: msg,
}
}
return model.Reply{
ReplyMsg: ask(from, fmt.Sprintf("{\"qqid\":%d,\"nickname\":\"%s\",\"referid\":%d,\"msg\":\"%s\"}", msg.UserId, msg.UserNickName, msg.OriginMsgId, msg.RawMsg)),
ReferOriginMsg: false,
FromMsg: msg,
}
}
func ask(from string, question string) (reply string) {
llmConfig := openai.DefaultAzureConfig(config.ConfigManager.GetProperty("openai_api_key"), config.ConfigManager.GetProperty("openai_api_base_url"))
llmConfig.APIType = openai.APITypeOpenAI
llmConfig.APIVersion = ""
client := openai.NewClientWithConfig(llmConfig)
resp, err := client.CreateChatCompletion(
context.Background(),
openai.ChatCompletionRequest{
Model: config.ConfigManager.GetProperty("openai_model_name"),
Messages: GenRequestFromUsr(from, question),
},
)
if err != nil {
fmt.Printf("ChatCompletion error: %v\n", err)
return
}
AppendReplyToHistory(from, resp.Choices[0].Message.Content)
return enterFormatter(resp.Choices[0].Message.Content)
}
func enterFormatter(msgIn string) string {
return strings.ReplaceAll(msgIn, "\\n", "\n")
}
func GenRequestFromUsr(from string, question string) []openai.ChatCompletionMessage {
if _, ok := histories[from]; !ok || histories_time[from].Add(2*time.Minute).Before(time.Now()) {
histories[from] = make([]openai.ChatCompletionMessage, 0)
histories[from] = append(histories[from], openai.ChatCompletionMessage{
Role: openai.ChatMessageRoleSystem,
Content: config.ConfigManager.GetProperty("openai_prompt"),
},
)
}
histories[from] = append(histories[from], openai.ChatCompletionMessage{
Role: openai.ChatMessageRoleUser,
Content: question,
})
return histories[from]
}
func AppendReplyToHistory(from string, reply string) {
histories[from] = append(histories[from], openai.ChatCompletionMessage{
Role: openai.ChatMessageRoleAssistant,
Content: reply,
})
histories_time[from] = time.Now()
for len(histories[from]) > 10 {
histories[from] = histories[from][1:]
histories[from][0] = openai.ChatCompletionMessage{
Role: openai.ChatMessageRoleSystem,
Content: config.ConfigManager.GetProperty("openai_prompt"),
}
}
}