From d5227f1159fe2f3683c0fc2b853d1ab5faf9d3dd 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: Sun, 30 Mar 2025 23:28:55 +0900
Subject: [PATCH] optimize(aichat): use config struct

---
 plugin/aichat/cfg.go  |  96 ++++++++++++++++++
 plugin/aichat/main.go | 231 +++++++-----------------------------------
 2 files changed, 132 insertions(+), 195 deletions(-)
 create mode 100644 plugin/aichat/cfg.go

diff --git a/plugin/aichat/cfg.go b/plugin/aichat/cfg.go
new file mode 100644
index 00000000..545e3e2c
--- /dev/null
+++ b/plugin/aichat/cfg.go
@@ -0,0 +1,96 @@
+package aichat
+
+import (
+	"strings"
+
+	ctrl "github.com/FloatTech/zbpctrl"
+	"github.com/FloatTech/zbputils/chat"
+	"github.com/fumiama/deepinfra"
+	"github.com/fumiama/deepinfra/model"
+	"github.com/sirupsen/logrus"
+	zero "github.com/wdvxdr1123/ZeroBot"
+	"github.com/wdvxdr1123/ZeroBot/message"
+)
+
+var cfg = newconfig()
+
+type config struct {
+	ModelName string
+	Type      int
+	SystemP   string
+	API       string
+	Key       string
+	Separator string
+	NoReplyAT bool
+	NoSystemP bool
+}
+
+func newconfig() config {
+	return config{
+		ModelName: model.ModelDeepDeek,
+		SystemP:   chat.SystemPrompt,
+		API:       deepinfra.OpenAIDeepInfra,
+	}
+}
+
+func (c *config) isvalid() bool {
+	return c.ModelName != "" && c.API != "" && c.Key != ""
+}
+
+func ensureconfig(ctx *zero.Ctx) bool {
+	c, ok := ctx.State["manager"].(*ctrl.Control[*zero.Ctx])
+	if !ok {
+		return false
+	}
+	if !cfg.isvalid() {
+		err := c.GetExtra(&cfg)
+		if err != nil {
+			logrus.Warnln("ERROR: get extra err:", err)
+		}
+		if !cfg.isvalid() {
+			cfg = newconfig()
+		}
+	}
+	return true
+}
+
+func newextrasetstr(ptr *string) func(ctx *zero.Ctx) {
+	return 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
+		}
+		*ptr = args
+		err := c.SetExtra(&cfg)
+		if err != nil {
+			ctx.SendChain(message.Text("ERROR: set extra err: ", err))
+			return
+		}
+		ctx.SendChain(message.Text("成功"))
+	}
+}
+
+func newextrasetbool(ptr *bool) func(ctx *zero.Ctx) {
+	return func(ctx *zero.Ctx) {
+		args := ctx.State["regex_matched"].([]string)
+		isno := args[1] == "不"
+		c, ok := ctx.State["manager"].(*ctrl.Control[*zero.Ctx])
+		if !ok {
+			ctx.SendChain(message.Text("ERROR: no such plugin"))
+			return
+		}
+		*ptr = isno
+		err := c.SetExtra(&cfg)
+		if err != nil {
+			ctx.SendChain(message.Text("ERROR: set extra err: ", err))
+			return
+		}
+		ctx.SendChain(message.Text("成功"))
+	}
+}
diff --git a/plugin/aichat/main.go b/plugin/aichat/main.go
index 9c6ad71e..82e8d152 100644
--- a/plugin/aichat/main.go
+++ b/plugin/aichat/main.go
@@ -3,7 +3,6 @@ package aichat
 
 import (
 	"math/rand"
-	"os"
 	"strconv"
 	"strings"
 
@@ -14,7 +13,6 @@ import (
 	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/chat"
@@ -22,7 +20,7 @@ import (
 )
 
 var (
-	// en data [4 cfg] [4 type] [8 temp] [8 rate] LSB
+	// en data [8 temp] [8 rate] LSB
 	en = control.AutoRegister(&ctrl.Options[*zero.Ctx]{
 		DisableOnDefault: false,
 		Extra:            control.ExtraFromString("aichat"),
@@ -43,15 +41,6 @@ var (
 	})
 )
 
-var (
-	modelname      = model.ModelDeepDeek
-	systemprompt   = chat.SystemPrompt
-	api            = deepinfra.OpenAIDeepInfra
-	sepstr         = ""
-	noreplyat      = false
-	nosystemprompt = false
-)
-
 var apitypes = map[string]uint8{
 	"OpenAI": 0,
 	"OLLaMA": 1,
@@ -59,49 +48,9 @@ var apitypes = map[string]uint8{
 }
 
 func init() {
-	mf := en.DataFolder() + "model.txt"
-	sf := en.DataFolder() + "system.txt"
-	pf := en.DataFolder() + "sep.txt"
-	af := en.DataFolder() + "api.txt"
-	nf := en.DataFolder() + "NoReplyAT"
-	syspf := en.DataFolder() + "NoSystemPrompt"
-	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)
-		}
-	}
-	if file.IsExist(pf) {
-		data, err := os.ReadFile(pf)
-		if err != nil {
-			logrus.Warnln("read sep", err)
-		} else {
-			sepstr = string(data)
-		}
-	}
-	if file.IsExist(af) {
-		data, err := os.ReadFile(af)
-		if err != nil {
-			logrus.Warnln("read api", err)
-		} else {
-			api = string(data)
-		}
-	}
-	noreplyat = file.IsExist(nf)
-	nosystemprompt = file.IsExist(syspf)
-
-	en.OnMessage(func(ctx *zero.Ctx) bool {
-		return ctx.ExtractPlainText() != "" && (!noreplyat || (noreplyat && !ctx.Event.IsToMe))
+	en.OnMessage(ensureconfig, func(ctx *zero.Ctx) bool {
+		return ctx.ExtractPlainText() != "" &&
+			(!cfg.NoReplyAT || (cfg.NoReplyAT && !ctx.Event.IsToMe))
 	}).SetBlock(false).Handle(func(ctx *zero.Ctx) {
 		gid := ctx.Event.GroupID
 		if gid == 0 {
@@ -113,7 +62,6 @@ func init() {
 		}
 		rate := c.GetData(gid)
 		temp := (rate >> 8) & 0xff
-		typ := (rate >> 16) & 0x0f
 		rate &= 0xff
 		if !ctx.Event.IsToMe && rand.Intn(100) >= int(rate) {
 			return
@@ -121,13 +69,7 @@ func init() {
 		if ctx.Event.IsToMe {
 			ctx.Block()
 		}
-		key := ""
-		err := c.GetExtra(&key)
-		if err != nil {
-			logrus.Warnln("ERROR: get extra err:", err)
-			return
-		}
-		if key == "" {
+		if cfg.Key == "" {
 			logrus.Warnln("ERROR: get extra err: empty key")
 			return
 		}
@@ -139,31 +81,31 @@ func init() {
 			temp = 100
 		}
 
-		x := deepinfra.NewAPI(api, key)
+		x := deepinfra.NewAPI(cfg.API, cfg.Key)
 		var mod model.Protocol
 
-		switch typ {
+		switch cfg.Type {
 		case 0:
 			mod = model.NewOpenAI(
-				modelname, sepstr,
+				cfg.ModelName, cfg.Separator,
 				float32(temp)/100, 0.9, 4096,
 			)
 		case 1:
 			mod = model.NewOLLaMA(
-				modelname, sepstr,
+				cfg.ModelName, cfg.Separator,
 				float32(temp)/100, 0.9, 4096,
 			)
 		case 2:
 			mod = model.NewGenAI(
-				modelname,
+				cfg.ModelName,
 				float32(temp)/100, 0.9, 4096,
 			)
 		default:
-			logrus.Warnln("[aichat] unsupported AI type", typ)
+			logrus.Warnln("[aichat] unsupported AI type", cfg.Type)
 			return
 		}
 
-		data, err := x.Request(chat.Ask(mod, gid, systemprompt, nosystemprompt))
+		data, err := x.Request(chat.Ask(mod, gid, cfg.SystemP, cfg.NoSystemP))
 		if err != nil {
 			logrus.Warnln("[aichat] post err:", err)
 			return
@@ -258,7 +200,7 @@ func init() {
 		}
 		ctx.SendChain(message.Text("成功"))
 	})
-	en.OnPrefix("设置AI聊天接口类型", zero.SuperUserPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
+	en.OnPrefix("设置AI聊天接口类型", ensureconfig, 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"))
@@ -274,144 +216,43 @@ func init() {
 			ctx.SendChain(message.Text("ERROR: 未知类型 ", args))
 			return
 		}
-		gid := ctx.Event.GroupID
-		if gid == 0 {
-			gid = -ctx.Event.UserID
-		}
-		val := c.GetData(gid) & (^0x0f0000)
-		err := c.SetData(gid, val|(int64(typ&0x0f)<<16))
+		cfg.Type = int(typ)
+		err := c.SetExtra(&cfg)
 		if err != nil {
-			ctx.SendChain(message.Text("ERROR: set data err: ", err))
+			ctx.SendChain(message.Text("ERROR: set extra err: ", err))
 			return
 		}
 		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
-		}
-		api = args
-		err := os.WriteFile(af, []byte(args), 0644)
-		if err != nil {
-			ctx.SendChain(message.Text("ERROR: ", err))
-			return
-		}
-		ctx.SendChain(message.Text("成功"))
+	en.OnPrefix("设置AI聊天接口地址", ensureconfig, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
+		Handle(newextrasetstr(&cfg.API))
+	en.OnPrefix("设置AI聊天密钥", ensureconfig, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
+		Handle(newextrasetstr(&cfg.Key))
+	en.OnPrefix("设置AI聊天模型名", ensureconfig, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
+		Handle(newextrasetstr(&cfg.ModelName))
+	en.OnPrefix("设置AI聊天系统提示词", ensureconfig, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
+		Handle(newextrasetstr(&cfg.SystemP))
+	en.OnFullMatch("查看AI聊天系统提示词", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
+		ctx.SendChain(message.Text(cfg.SystemP))
 	})
-	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
-		}
+	en.OnFullMatch("重置AI聊天系统提示词", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
 		c, ok := ctx.State["manager"].(*ctrl.Control[*zero.Ctx])
 		if !ok {
 			ctx.SendChain(message.Text("ERROR: no such plugin"))
 			return
 		}
-		err := c.SetExtra(&args)
+		cfg.SystemP = chat.SystemPrompt
+		err := c.SetExtra(&cfg)
 		if err != nil {
-			ctx.SendChain(message.Text("ERROR: ", err))
+			ctx.SendChain(message.Text("ERROR: set extra err: ", err))
 			return
 		}
 		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
-		}
-		modelname = args
-		err := os.WriteFile(mf, []byte(args), 0644)
-		if err != nil {
-			ctx.SendChain(message.Text("ERROR: ", err))
-			return
-		}
-		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
-		}
-		systemprompt = args
-		err := os.WriteFile(sf, []byte(args), 0644)
-		if err != nil {
-			ctx.SendChain(message.Text("ERROR: ", err))
-			return
-		}
-		ctx.SendChain(message.Text("成功"))
-	})
-	en.OnFullMatch("查看AI聊天系统提示词", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
-		ctx.SendChain(message.Text(systemprompt))
-	})
-	en.OnFullMatch("重置AI聊天系统提示词", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
-		systemprompt = chat.SystemPrompt
-		_ = os.Remove(sf)
-		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 == "" {
-			sepstr = ""
-			_ = os.Remove(pf)
-			ctx.SendChain(message.Text("清除成功"))
-			return
-		}
-		sepstr = args
-		err := os.WriteFile(pf, []byte(args), 0644)
-		if err != nil {
-			ctx.SendChain(message.Text("ERROR: ", err))
-			return
-		}
-		ctx.SendChain(message.Text("设置成功"))
-	})
-	en.OnRegex("^设置AI聊天(不)?响应AT$", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
-		args := ctx.State["regex_matched"].([]string)
-		isno := args[1] == "不"
-		if isno {
-			f, err := os.Create(nf)
-			if err != nil {
-				ctx.SendChain(message.Text("ERROR: ", err))
-				return
-			}
-			defer f.Close()
-			_, err = f.WriteString("PLACEHOLDER")
-			if err != nil {
-				ctx.SendChain(message.Text("ERROR: ", err))
-				return
-			}
-			noreplyat = true
-		} else {
-			_ = os.Remove(nf)
-			noreplyat = false
-		}
-		ctx.SendChain(message.Text("成功"))
-	})
-	en.OnRegex("^设置AI聊天(不)?支持系统提示词$", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).Handle(func(ctx *zero.Ctx) {
-		args := ctx.State["regex_matched"].([]string)
-		isno := args[1] == "不"
-		if isno {
-			f, err := os.Create(syspf)
-			if err != nil {
-				ctx.SendChain(message.Text("ERROR: ", err))
-				return
-			}
-			defer f.Close()
-			_, err = f.WriteString("PLACEHOLDER")
-			if err != nil {
-				ctx.SendChain(message.Text("ERROR: ", err))
-				return
-			}
-			nosystemprompt = true
-		} else {
-			_ = os.Remove(syspf)
-			nosystemprompt = false
-		}
-		ctx.SendChain(message.Text("成功"))
-	})
+	en.OnPrefix("设置AI聊天分隔符", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
+		Handle(newextrasetstr(&cfg.Separator))
+	en.OnRegex("^设置AI聊天(不)?响应AT$", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
+		Handle(newextrasetbool(&cfg.NoReplyAT))
+	en.OnRegex("^设置AI聊天(不)?支持系统提示词$", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
+		Handle(newextrasetbool(&cfg.NoSystemP))
 }