From 735c3cdde1416f5b4d12d6bcc7994de43d0c10ef Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
<41315874+fumiama@users.noreply.github.com>
Date: Fri, 14 Feb 2025 17:23:00 +0900
Subject: [PATCH] feat: add plugin aichat
---
README.md | 11 +++
go.mod | 1 +
go.sum | 2 +
main.go | 2 +
plugin/aichat/list.go | 51 +++++++++++
plugin/aichat/main.go | 191 ++++++++++++++++++++++++++++++++++++++++++
6 files changed, 258 insertions(+)
create mode 100644 plugin/aichat/list.go
create mode 100644 plugin/aichat/main.go
diff --git a/README.md b/README.md
index 93e9f1ff..0777ad20 100644
--- a/README.md
+++ b/README.md
@@ -1537,6 +1537,17 @@ print("run[CQ:image,file="+j["img"]+"]")
### *低优先级*
+
+ OpenAI聊天
+
+ `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/aichat"`
+
+ - [x] 设置AI聊天触发概率10
+ - [x] 设置AI聊天密钥xxx
+ - [x] 设置AI聊天模型名xxx
+ - [x] 设置AI聊天系统提示词xxx
+
+
骂人
diff --git a/go.mod b/go.mod
index ce2691b7..9bd81def 100644
--- a/go.mod
+++ b/go.mod
@@ -21,6 +21,7 @@ require (
github.com/disintegration/imaging v1.6.2
github.com/fumiama/ahsai v0.1.0
github.com/fumiama/cron v1.3.0
+ github.com/fumiama/deepinfra v0.0.0-20250214072937-12ba46058885
github.com/fumiama/go-base16384 v1.7.0
github.com/fumiama/go-registry v0.2.7
github.com/fumiama/gotracemoe v0.0.3
diff --git a/go.sum b/go.sum
index ec376e14..2121beb1 100644
--- a/go.sum
+++ b/go.sum
@@ -57,6 +57,8 @@ github.com/fumiama/ahsai v0.1.0 h1:LXD61Kaj6kJHa3AEGsLIfKNzcgaVxg7JB72OR4yNNZ4=
github.com/fumiama/ahsai v0.1.0/go.mod h1:fFeNnqgo44i8FIaguK659aQryuZeFy+4klYLQu/rfdk=
github.com/fumiama/cron v1.3.0 h1:ZWlwuexF+HQHl3cYytEE5HNwD99q+3vNZF1GrEiXCFo=
github.com/fumiama/cron v1.3.0/go.mod h1:bz5Izvgi/xEUI8tlBN8BI2jr9Moo8N4or0KV8xXuPDY=
+github.com/fumiama/deepinfra v0.0.0-20250214072937-12ba46058885 h1:AHuorF/H+9q/+A3CclMbr5W+kbpaMw1r5E4UUC7ETUQ=
+github.com/fumiama/deepinfra v0.0.0-20250214072937-12ba46058885/go.mod h1:pNn32xTo/u72cTCIq3EejJQPTZqg420Xb3XI+Ou7ZmU=
github.com/fumiama/go-base16384 v1.7.0 h1:6fep7XPQWxRlh4Hu+KsdH+6+YdUp+w6CwRXtMWSsXCA=
github.com/fumiama/go-base16384 v1.7.0/go.mod h1:OEn+947GV5gsbTAnyuUW/SrfxJYUdYupSIQXOuGOcXM=
github.com/fumiama/go-registry v0.2.7 h1:tLEqgEpsiybQMqBv0dLHm5leia/z1DhajMupwnOHeNs=
diff --git a/main.go b/main.go
index 444f5d8a..987f782c 100644
--- a/main.go
+++ b/main.go
@@ -167,6 +167,8 @@ import (
// vvvvvvvvvvvvvv //
// vvvv //
+ _ "github.com/FloatTech/ZeroBot-Plugin/plugin/aichat" // AI聊天
+
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/curse" // 骂人
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/thesaurus" // 词典匹配回复
diff --git a/plugin/aichat/list.go b/plugin/aichat/list.go
new file mode 100644
index 00000000..d681914f
--- /dev/null
+++ b/plugin/aichat/list.go
@@ -0,0 +1,51 @@
+package aichat
+
+import (
+ "sync"
+
+ "github.com/fumiama/deepinfra"
+ "github.com/fumiama/deepinfra/model"
+)
+
+const listcap = 6
+
+type list struct {
+ mu sync.RWMutex
+ m map[int64][]string
+}
+
+func newlist() list {
+ return list{
+ m: make(map[int64][]string, 64),
+ }
+}
+
+func (l *list) add(grp int64, txt string) {
+ l.mu.Lock()
+ defer l.mu.Unlock()
+ msgs, ok := l.m[grp]
+ if !ok {
+ msgs = make([]string, 1, listcap)
+ msgs[0] = txt
+ l.m[grp] = msgs
+ return
+ }
+ if len(msgs) < cap(msgs) {
+ msgs = append(msgs, txt)
+ l.m[grp] = msgs
+ return
+ }
+ copy(msgs[:], msgs[1:])
+ msgs[len(msgs)-1] = txt
+ l.m[grp] = msgs
+}
+
+func (l *list) body(mn, sysp string, grp int64) deepinfra.Model {
+ m := model.NewCustom(mn, "", 0.7, 0.9, 1024).System(sysp)
+ l.mu.RLock()
+ defer l.mu.RUnlock()
+ for _, msg := range l.m[grp] {
+ _ = m.User(msg)
+ }
+ return m
+}
diff --git a/plugin/aichat/main.go b/plugin/aichat/main.go
new file mode 100644
index 00000000..bb75dc7f
--- /dev/null
+++ b/plugin/aichat/main.go
@@ -0,0 +1,191 @@
+// Package aichat OpenAI聊天
+package aichat
+
+import (
+ "math/rand"
+ "os"
+ "strconv"
+ "strings"
+ "sync/atomic"
+ "unsafe"
+
+ "github.com/fumiama/deepinfra"
+ "github.com/sirupsen/logrus"
+
+ zero "github.com/wdvxdr1123/ZeroBot"
+ "github.com/wdvxdr1123/ZeroBot/message"
+
+ "github.com/FloatTech/floatbox/file"
+ "github.com/FloatTech/floatbox/process"
+ ctrl "github.com/FloatTech/zbpctrl"
+ "github.com/FloatTech/zbputils/control"
+)
+
+var (
+ api *deepinfra.API
+ en = control.AutoRegister(&ctrl.Options[*zero.Ctx]{
+ DisableOnDefault: false,
+ Extra: control.ExtraFromString("aichat"),
+ Brief: "OpenAI聊天",
+ Help: "- 设置AI聊天触发概率10\n- 设置AI聊天密钥xxx\n- 设置AI聊天模型名xxx\n- 设置AI聊天系统提示词xxx",
+ PrivateDataFolder: "aichat",
+ })
+ lst = newlist()
+)
+
+var (
+ modelname = "deepseek-ai/DeepSeek-R1"
+ systemprompt = "你正在QQ群与用户聊天，用户发送了消息。按自己的心情简短思考后，条理清晰地回应**一句话**，禁止回应多句。"
+)
+
+func init() {
+ mf := en.DataFolder() + "model.txt"
+ sf := en.DataFolder() + "system.txt"
+ if file.IsExist(mf) {
+ data, err := os.ReadFile(mf)
+ if err != nil {
+ logrus.Warnln("read model", err)
+ } else {
+ modelname = string(data)
+ }
+ }
+ if file.IsExist(sf) {
+ data, err := os.ReadFile(sf)
+ if err != nil {
+ logrus.Warnln("read system", err)
+ } else {
+ systemprompt = string(data)
+ }
+ }
+
+ en.OnMessage(func(ctx *zero.Ctx) bool {
+ txt := ctx.ExtractPlainText()
+ ctx.State["aichat_txt"] = txt
+ return txt != ""
+ }).SetBlock(false).Handle(func(ctx *zero.Ctx) {
+ lst.add(ctx.Event.GroupID, ctx.State["aichat_txt"].(string))
+ gid := ctx.Event.GroupID
+ if gid == 0 {
+ gid = -ctx.Event.UserID
+ }
+ c, ok := ctx.State["manager"].(*ctrl.Control[*zero.Ctx])
+ if !ok {
+ return
+ }
+ rate := c.GetData(gid)
+ if !ctx.Event.IsToMe && rand.Intn(100) >= int(rate) {
+ return
+ }
+ key := ""
+ err := c.GetExtra(&key)
+ if err != nil {
+ logrus.Warnln("ERROR: get extra err:", err)
+ return
+ }
+ if key == "" {
+ logrus.Warnln("ERROR: get extra err: empty key")
+ return
+ }
+ var x deepinfra.API
+ y := &x
+ if api == nil {
+ x = deepinfra.NewAPI(deepinfra.APIDeepInfra, key)
+ atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&api)), unsafe.Pointer(&x))
+ } else {
+ y = api
+ }
+ data, err := y.Request(lst.body(modelname, systemprompt, gid))
+ if err != nil {
+ logrus.Warnln("[niniqun] post err:", err)
+ return
+ }
+ txt := strings.Trim(data, "\n ")
+ if len(txt) > 0 {
+ lst.add(ctx.Event.GroupID, txt)
+ nick := zero.BotConfig.NickName[rand.Intn(len(zero.BotConfig.NickName))]
+ txt = strings.ReplaceAll(txt, "{name}", ctx.CardOrNickName(ctx.Event.UserID))
+ txt = strings.ReplaceAll(txt, "{me}", nick)
+ id := any(nil)
+ if ctx.Event.IsToMe {
+ id = ctx.Event.MessageID
+ }
+ for _, t := range strings.Split(txt, "{segment}") {
+ if t == "" {
+ continue
+ }
+ if id != nil {
+ id = ctx.SendChain(message.Reply(id), message.Text(t))
+ } else {
+ id = ctx.SendChain(message.Text(t))
+ }
+ process.SleepAbout1sTo2s()
+ }
+ }
+ })
+ en.OnPrefix("设置AI聊天触发概率", zero.AdminPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
+ args := strings.TrimSpace(ctx.State["args"].(string))
+ if args == "" {
+ ctx.SendChain(message.Text("ERROR: empty args"))
+ return
+ }
+ c, ok := ctx.State["manager"].(*ctrl.Control[*zero.Ctx])
+ if !ok {
+ ctx.SendChain(message.Text("ERROR: no such plugin"))
+ return
+ }
+ r, err := strconv.Atoi(args)
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: parse rate err: ", err))
+ return
+ }
+ gid := ctx.Event.GroupID
+ if gid == 0 {
+ gid = -ctx.Event.UserID
+ }
+ c.SetData(gid, int64(r&0xff))
+ ctx.SendChain(message.Text("成功"))
+ })
+ en.OnPrefix("设置AI聊天密钥", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
+ args := strings.TrimSpace(ctx.State["args"].(string))
+ if args == "" {
+ ctx.SendChain(message.Text("ERROR: empty args"))
+ return
+ }
+ c, ok := ctx.State["manager"].(*ctrl.Control[*zero.Ctx])
+ if !ok {
+ ctx.SendChain(message.Text("ERROR: no such plugin"))
+ return
+ }
+ err := c.SetExtra(&args)
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ })
+ en.OnPrefix("设置AI聊天模型名", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
+ args := strings.TrimSpace(ctx.State["args"].(string))
+ if args == "" {
+ ctx.SendChain(message.Text("ERROR: empty args"))
+ return
+ }
+ modelname = args
+ err := os.WriteFile(mf, []byte(args), 0644)
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ })
+ en.OnPrefix("设置AI聊天系统提示词", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
+ args := strings.TrimSpace(ctx.State["args"].(string))
+ if args == "" {
+ ctx.SendChain(message.Text("ERROR: empty args"))
+ return
+ }
+ systemprompt = args
+ err := os.WriteFile(sf, []byte(args), 0644)
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ })
+}