🎨 优化ai画图代码

Author: guohuiyuan

README.md

@@ -418,7 +418,7 @@ print("run[CQ:image,file="+j["img"]+"]")
 
   `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/aiimage"`
 
-  - [x] 设置AI画图密钥	
+  - [x] 设置AI画图密钥xxx	
   - [x] 设置AI画图接口地址https://api.siliconflow.cn/v1/images/generations
   - [x] 设置AI画图模型名Kwai-Kolors/Kolors
   - [x] 查看AI画图配置

plugin/aiimage/config.go

@@ -3,60 +3,40 @@ package aiimage
 
 import (
 	"fmt"
-	"os"
 	"strings"
 	"sync"
-	"time"
 
 	sql "github.com/FloatTech/sqlite"
 )
 
-// Storage 管理画图配置存储
-type Storage struct {
+// storage 管理画图配置存储
+type storage struct {
 	sync.RWMutex
 	db sql.Sqlite
 }
 
-var (
-	sdb = &Storage{
-		db: sql.New("data/aiimage/config.db"),
-	}
-)
-
-func init() {
-	if err := os.MkdirAll("data/aiimage", 0755); err != nil {
-		panic(err)
-	}
-	if err := sdb.db.Open(time.Hour * 24); err != nil {
-		panic(err)
-	}
-	if err := sdb.db.Create("config", &ImageConfig{}); err != nil {
-		panic(err)
-	}
-}
-
-// ImageConfig 存储AI画图配置信息
-type ImageConfig struct {
+// imageConfig 存储AI画图配置信息
+type imageConfig struct {
 	ID        int64  `db:"id"`        // 主键ID
 	APIKey    string `db:"apiKey"`    // API密钥
 	APIURL    string `db:"apiUrl"`    // API地址
 	ModelName string `db:"modelName"` // 画图模型名称
 }
 
-// GetConfig 获取当前配置
-func GetConfig() ImageConfig {
+// getConfig 获取当前配置
+func (sdb *storage) getConfig() imageConfig {
 	sdb.RLock()
 	defer sdb.RUnlock()
-	cfg := ImageConfig{}
+	cfg := imageConfig{}
 	_ = sdb.db.Find("config", &cfg, "WHERE id = 1")
 	return cfg
 }
 
-// SetConfig 设置AI画图配置
-func SetConfig(apiKey, apiURL, modelName string) error {
+// setConfig 设置AI画图配置
+func (sdb *storage) setConfig(apiKey, apiURL, modelName string) error {
 	sdb.Lock()
 	defer sdb.Unlock()
-	return sdb.db.Insert("config", &ImageConfig{
+	return sdb.db.Insert("config", &imageConfig{
 		ID:        1,
 		APIKey:    apiKey,
 		APIURL:    apiURL,
@@ -65,8 +45,8 @@ func SetConfig(apiKey, apiURL, modelName string) error {
 }
 
 // PrintConfig 返回格式化后的配置信息
-func PrintConfig() string {
-	cfg := GetConfig()
+func (sdb *storage) PrintConfig() string {
+	cfg := sdb.getConfig()
 	var builder strings.Builder
 	builder.WriteString("当前AI画图配置:\n")
 	builder.WriteString(fmt.Sprintf("• 密钥: %s\n", cfg.APIKey))

plugin/aiimage/main.go

@@ -6,8 +6,11 @@ import (
 	"encoding/json"
 	"net/http"
 	"strings"
+	"time"
 
+	fcext "github.com/FloatTech/floatbox/ctxext"
 	"github.com/FloatTech/floatbox/web"
+	sql "github.com/FloatTech/sqlite"
 	"github.com/tidwall/gjson"
 	zero "github.com/wdvxdr1123/ZeroBot"
 	"github.com/wdvxdr1123/ZeroBot/message"
@@ -18,60 +21,78 @@ import (
 )
 
 func init() {
+	var sdb = &storage{}
+
 	en := control.AutoRegister(&ctrl.Options[*zero.Ctx]{
 		DisableOnDefault: false,
 		Extra:            control.ExtraFromString("aiimage"),
 		Brief:            "AI画图",
-		Help: "- 设置AI画图密钥\n" +
+		Help: "- 设置AI画图密钥xxx\n" +
 			"- 设置AI画图接口地址https://api.siliconflow.cn/v1/images/generations\n" +
 			"- 设置AI画图模型名Kwai-Kolors/Kolors\n" +
 			"- 查看AI画图配置\n" +
 			"- AI画图 [描述]",
 		PrivateDataFolder: "aiimage",
 	})
 
-	en.OnPrefix("设置AI画图密钥", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
+	getdb := fcext.DoOnceOnSuccess(func(ctx *zero.Ctx) bool {
+		sdb.db = sql.New(en.DataFolder() + "aiimage.db")
+		err := sdb.db.Open(time.Hour)
+		if err == nil {
+			// 创建配置表
+			err = sdb.db.Create("config", &imageConfig{})
+			if err != nil {
+				ctx.SendChain(message.Text("[ERROR]:", err))
+				return false
+			}
+			return true
+		}
+		ctx.SendChain(message.Text("[ERROR]:", err))
+		return false
+	})
+
+	en.OnPrefix("设置AI画图密钥", getdb, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
 			apiKey := strings.TrimSpace(ctx.State["args"].(string))
-			cfg := GetConfig()
-			err := SetConfig(apiKey, cfg.APIURL, cfg.ModelName)
+			cfg := sdb.getConfig()
+			err := sdb.setConfig(apiKey, cfg.APIURL, cfg.ModelName)
 			if err != nil {
 				ctx.SendChain(message.Text("ERROR: 设置API密钥失败: ", err))
 				return
 			}
 			ctx.SendChain(message.Text("成功设置API密钥"))
 		})
 
-	en.OnPrefix("设置AI画图接口地址", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
+	en.OnPrefix("设置AI画图接口地址", getdb, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
 			apiURL := strings.TrimSpace(ctx.State["args"].(string))
-			cfg := GetConfig()
-			err := SetConfig(cfg.APIKey, apiURL, cfg.ModelName)
+			cfg := sdb.getConfig()
+			err := sdb.setConfig(cfg.APIKey, apiURL, cfg.ModelName)
 			if err != nil {
 				ctx.SendChain(message.Text("ERROR: 设置API地址失败: ", err))
 				return
 			}
 			ctx.SendChain(message.Text("成功设置API地址"))
 		})
 
-	en.OnPrefix("设置AI画图模型名", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
+	en.OnPrefix("设置AI画图模型名", getdb, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
 			modelName := strings.TrimSpace(ctx.State["args"].(string))
-			cfg := GetConfig()
-			err := SetConfig(cfg.APIKey, cfg.APIURL, modelName)
+			cfg := sdb.getConfig()
+			err := sdb.setConfig(cfg.APIKey, cfg.APIURL, modelName)
 			if err != nil {
 				ctx.SendChain(message.Text("ERROR: 设置模型失败: ", err))
 				return
 			}
 			ctx.SendChain(message.Text("成功设置模型: ", modelName))
 		})
 
-	en.OnFullMatch("查看AI画图配置", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
+	en.OnFullMatch("查看AI画图配置", getdb, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
-			ctx.SendChain(message.Text(PrintConfig()))
+			ctx.SendChain(message.Text(sdb.PrintConfig()))
 		})
 
-	en.OnPrefix("AI画图").SetBlock(true).
+	en.OnPrefix("AI画图", getdb).SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
 			ctx.SendChain(message.Text("少女思考中..."))
 			prompt := strings.TrimSpace(ctx.State["args"].(string))
@@ -80,22 +101,21 @@ func init() {
 				return
 			}
 
-			cfg := GetConfig()
+			cfg := sdb.getConfig()
 			if cfg.APIKey == "" || cfg.APIURL == "" || cfg.ModelName == "" {
 				ctx.SendChain(message.Text("请先配置API密钥、地址和模型"))
 				return
 			}
 
 			// 准备请求数据
-			reqData := map[string]interface{}{
+			reqBytes, _ := json.Marshal(map[string]interface{}{
 				"model":               cfg.ModelName,
 				"prompt":              prompt,
 				"image_size":          "1024x1024",
 				"batch_size":          4,
 				"num_inference_steps": 20,
 				"guidance_scale":      7.5,
-			}
-			reqBytes, _ := json.Marshal(reqData)
+			})
 
 			// 发送API请求
 			data, err := web.RequestDataWithHeaders(