go-wxhelper/handler/at_message.go

package handler

import (
	"context"
	"fmt"
	"github.com/sashabaranov/go-openai"
	"go-wechat/client"
	"go-wechat/config"
	"go-wechat/entity"
	"go-wechat/model"
	"go-wechat/utils"
	"log"
	"regexp"
	"strings"
)

// handleAtMessage
// @description: 处理At机器人的消息
// @param m
func handleAtMessage(m model.Message) {
	if !config.Conf.Ai.Enable {
		return
	}

	// 取出所有启用了AI的好友或群组
	var count int64
	client.MySQL.Model(&entity.Friend{}).Where("enable_ai IS TRUE").Where("wxid = ?", m.FromUser).Count(&count)
	if count < 1 {
		return
	}

	// 预处理一下发送的消息，用正则去掉@机器人的内容
	re := regexp.MustCompile(`@([^ ]+)`)
	matches := re.FindStringSubmatch(m.Content)
	if len(matches) > 0 {
		// 过滤掉第一个匹配到的
		m.Content = strings.Replace(m.Content, matches[0], "", 1)
	}

	// 组装消息体
	messages := make([]openai.ChatCompletionMessage, 0)
	if config.Conf.Ai.Personality != "" {
		// 填充人设
		messages = append(messages, openai.ChatCompletionMessage{
			Role:    openai.ChatMessageRoleSystem,
			Content: config.Conf.Ai.Personality,
		})
	}
	// 填充用户消息
	messages = append(messages, openai.ChatCompletionMessage{
		Role:    openai.ChatMessageRoleUser,
		Content: m.Content,
	})

	// 配置模型
	chatModel := openai.GPT3Dot5Turbo0613
	if config.Conf.Ai.Model != "" {
		chatModel = config.Conf.Ai.Model
	}

	// 默认使用AI回复
	conf := openai.DefaultConfig(config.Conf.Ai.ApiKey)
	if config.Conf.Ai.BaseUrl != "" {
		conf.BaseURL = fmt.Sprintf("%s/v1", config.Conf.Ai.BaseUrl)
	}
	client := openai.NewClientWithConfig(conf)
	resp, err := client.CreateChatCompletion(
		context.Background(),
		openai.ChatCompletionRequest{
			Model:    chatModel,
			Messages: messages,
		},
	)

	if err != nil {
		log.Printf("OpenAI聊天发起失败: %v", err.Error())
		utils.SendMessage(m.FromUser, m.GroupUser, "AI炸啦~", 0)
		return
	}

	// 发送消息
	utils.SendMessage(m.FromUser, m.GroupUser, "\n"+resp.Choices[0].Message.Content, 0)
}