feat: add tagger integration and enable automatic tagging for artworks

krau · krau · commit bfd1591be4a3 · 2025-01-25T23:53:48.000+08:00
diff --git a/common/common.go b/common/common.go
@@ -8,8 +8,10 @@ import (
 
 func Init() {
 	initHttpClient()
-	initResendClient()
 	initLogger()
+	if config.Cfg.Auth.Resend.APIKey != "" {
+		initResendClient()
+	}
 	if searchCfg := config.Cfg.Search; searchCfg.Enable {
 		switch searchCfg.Engine {
 		case "meilisearch":
@@ -19,4 +21,7 @@ func Init() {
 			os.Exit(1)
 		}
 	}
+	if taggerCfg := config.Cfg.Tagger; taggerCfg.Enable {
+		initTaggerClient()
+	}
 }
diff --git a/common/tagger_client.go b/common/tagger_client.go
@@ -0,0 +1,79 @@
+package common
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/imroc/req/v3"
+	"github.com/krau/ManyACG/config"
+)
+
+type taggerClient struct {
+	Client  *req.Client
+	host    string
+	token   string
+	timeout time.Duration
+}
+
+func (c *taggerClient) Health() (string, error) {
+	var health struct {
+		Status string `json:"status"`
+	}
+	resp, err := c.Client.R().Get("/health")
+	if err != nil {
+		return "", err
+	}
+	if err := json.Unmarshal(resp.Bytes(), &health); err != nil {
+		return "", err
+	}
+	return health.Status, nil
+}
+
+type taggerPredictResponse struct {
+	PredictedTags []string           `json:"predicted_tags"`
+	Scores        map[string]float64 `json:"scores"`
+}
+
+func (c *taggerClient) Predict(ctx context.Context, file []byte) (*taggerPredictResponse, error) {
+	resp, err := c.Client.R().SetContext(ctx).SetFileBytes("file", "image", file).Post("/predict")
+	if err != nil {
+		return nil, err
+	}
+	if resp.IsErrorState() {
+		return nil, fmt.Errorf("tagger predict failed: %s", resp.Status)
+	}
+	var predict taggerPredictResponse
+	if err := json.Unmarshal(resp.Bytes(), &predict); err != nil {
+		return nil, err
+	}
+	return &predict, nil
+}
+
+var TaggerClient *taggerClient
+
+func initTaggerClient() {
+	if config.Cfg.Tagger.Host == "" || config.Cfg.Tagger.Token == "" {
+		Logger.Fatalf("Tagger configuration is incomplete")
+		os.Exit(1)
+	}
+	client := req.C().
+		SetCommonBearerAuthToken(config.Cfg.Tagger.Token).
+		SetBaseURL(config.Cfg.Tagger.Host).
+		SetTimeout(time.Duration(config.Cfg.Tagger.Timeout) * time.Second).
+		SetUserAgent("ManyACG/" + Version)
+	TaggerClient = &taggerClient{
+		Client:  client,
+		host:    config.Cfg.Tagger.Host,
+		token:   config.Cfg.Tagger.Token,
+		timeout: time.Duration(config.Cfg.Tagger.Timeout) * time.Second,
+	}
+	if status, err := TaggerClient.Health(); err != nil {
+		Logger.Fatalf("Tagger health check failed: %s", err)
+		os.Exit(1)
+	} else {
+		Logger.Infof("Tagger health check: %s", status)
+	}
+}
diff --git a/config/tagger.go b/config/tagger.go
@@ -0,0 +1,9 @@
+package config
+
+type taggerConfig struct {
+	Enable  bool   `toml:"enable" mapstructure:"enable" json:"enable" yaml:"enable"`
+	Host    string `toml:"host" mapstructure:"host" json:"host" yaml:"host"`
+	Token   string `toml:"token" mapstructure:"token" json:"token" yaml:"token"`
+	Timeout int    `toml:"timeout" mapstructure:"timeout" json:"timeout" yaml:"timeout"`
+	TagNew  bool   `toml:"tagnew" mapstructure:"tagnew" json:"tagnew" yaml:"tagnew"`
+}
diff --git a/config/viper.go b/config/viper.go
@@ -12,7 +12,6 @@ import (
 type Config struct {
 	Debug    bool           `toml:"debug" mapstructure:"debug" json:"debug" yaml:"debug"`
 	WSRVURL  string         `toml:"wsrv_url" mapstructure:"wsrv_url" json:"wsrv_url" yaml:"wsrv_url"`
-	Search   searchConfig   `toml:"search" mapstructure:"search" json:"search" yaml:"search"`
 	API      apiConfig      `toml:"api" mapstructure:"api" json:"api" yaml:"api"`
 	Auth     authConfig     `toml:"auth" mapstructure:"auth" json:"auth" yaml:"auth"`
 	Fetcher  fetcherConfig  `toml:"fetcher" mapstructure:"fetcher" json:"fetcher" yaml:"fetcher"`
@@ -21,6 +20,8 @@ type Config struct {
 	Storage  storageConfigs `toml:"storage" mapstructure:"storage" json:"storage" yaml:"storage"`
 	Telegram telegramConfig `toml:"telegram" mapstructure:"telegram" json:"telegram" yaml:"telegram"`
 	Database databaseConfig `toml:"database" mapstructure:"database" json:"database" yaml:"database"`
+	Search   searchConfig   `toml:"search" mapstructure:"search" json:"search" yaml:"search"`
+	Tagger   taggerConfig   `toml:"tagger" mapstructure:"tagger" json:"tagger" yaml:"tagger"`
 }
 
 type fetcherConfig struct {
diff --git a/service/change_stream.go b/service/change_stream.go
@@ -112,10 +112,10 @@ func (m *artworkSyncManager) ProcessArtworkUpdateEvent(event bson.M) {
 	}
 	task, err := common.MeilisearchClient.Index(config.Cfg.Search.MeiliSearch.Index).UpdateDocuments(artworkJSON)
 	if err != nil {
-		common.Logger.Errorf("add artwork to meilisearch error: %s", err)
+		common.Logger.Errorf("update artwork to meilisearch error: %s", err)
 		return
 	}
-	common.Logger.Debugf("commited add artwork task to meilisearch: %d", task.TaskUID)
+	common.Logger.Debugf("commited update artwork task to meilisearch: %d", task.TaskUID)
 }
 
 func (m *artworkSyncManager) ProcessArtworkDeleteEvent(event bson.M) {
diff --git a/service/service.go b/service/service.go
@@ -12,6 +12,9 @@ func InitService() {
 	if config.Cfg.Search.Enable {
 		go syncArtworkToSearchEngine()
 	}
+	if config.Cfg.Tagger.Enable {
+		go listenPredictArtworkTagsTask()
+	}
 }
 
 type Service struct{}
diff --git a/service/tags.go b/service/tags.go
@@ -6,8 +6,10 @@ import (
 	"fmt"
 
 	"github.com/duke-git/lancet/v2/slice"
+	"github.com/krau/ManyACG/common"
 	"github.com/krau/ManyACG/dao"
 	"github.com/krau/ManyACG/errs"
+	"github.com/krau/ManyACG/storage"
 	"github.com/krau/ManyACG/types"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/bson/primitive"
@@ -167,3 +169,64 @@ func AddTagAliasByID(ctx context.Context, tagID primitive.ObjectID, alias ...str
 	}
 	return result.(*types.TagModel), nil
 }
+
+func PredictArtworkTagsByIDAndUpdate(ctx context.Context, artworkID primitive.ObjectID) error {
+	if common.TaggerClient == nil {
+		return errors.New("tagger not available")
+	}
+	artwork, err := GetArtworkByID(ctx, artworkID)
+	if err != nil {
+		return err
+	}
+	predictedTags := make([]string, 0)
+	for _, picture := range artwork.Pictures {
+		var pictureFile []byte
+		if picture.StorageInfo.Regular != nil {
+			pictureFile, err = storage.GetFile(ctx, picture.StorageInfo.Regular)
+		} else if picture.StorageInfo.Original != nil {
+			pictureFile, err = storage.GetFile(ctx, picture.StorageInfo.Original)
+		} else {
+			pictureFile, err = common.DownloadWithCache(ctx, picture.Original, nil)
+		}
+		if err != nil {
+			return err
+		}
+		common.Logger.Debugf("predict picture %s", picture.Original)
+		result, err := common.TaggerClient.Predict(ctx, pictureFile)
+		if err != nil {
+			common.Logger.Errorf("predict picture %s error: %s", picture.Original, err)
+			continue
+		}
+		if len(result.PredictedTags) == 0 {
+			continue
+		}
+		predictedTags = slice.Union(predictedTags, result.PredictedTags)
+	}
+	newTags := slice.Compact(slice.Union(artwork.Tags, predictedTags))
+	if err := UpdateArtworkTagsByURL(ctx, artwork.SourceURL, newTags); err != nil {
+		return err
+	}
+	return nil
+}
+
+type predictArtworkTagsTask struct {
+	ArtworkID primitive.ObjectID
+	Ctx       context.Context
+}
+
+var predictArtworkTagsTaskChan = make(chan *predictArtworkTagsTask)
+
+func AddPredictArtworkTagTask(ctx context.Context, artworkID primitive.ObjectID) {
+	predictArtworkTagsTaskChan <- &predictArtworkTagsTask{
+		ArtworkID: artworkID,
+		Ctx:       ctx,
+	}
+}
+
+func listenPredictArtworkTagsTask() {
+	for task := range predictArtworkTagsTaskChan {
+		if err := PredictArtworkTagsByIDAndUpdate(task.Ctx, task.ArtworkID); err != nil {
+			common.Logger.Errorf("predict artwork %s tags error: %s", task.ArtworkID.Hex(), err)
+		}
+	}
+}
diff --git a/telegram/handlers/admin_edit_artwork.go b/telegram/handlers/admin_edit_artwork.go
@@ -3,12 +3,15 @@ package handlers
 import (
 	"context"
 	"errors"
+	"fmt"
 	"strconv"
 	"strings"
 
+	"github.com/duke-git/lancet/v2/slice"
 	"github.com/krau/ManyACG/common"
 	"github.com/krau/ManyACG/service"
 	"github.com/krau/ManyACG/sources"
+	"github.com/krau/ManyACG/storage"
 	"github.com/krau/ManyACG/telegram/utils"
 	"github.com/krau/ManyACG/types"
 
@@ -324,3 +327,119 @@ func ReCaptionArtwork(ctx context.Context, bot *telego.Bot, message telego.Messa
 	})
 	utils.ReplyMessage(bot, message, "已重新生成作品描述")
 }
+
+func AutoTaggingArtwork(ctx context.Context, bot *telego.Bot, message telego.Message) {
+	if !CheckPermissionInGroup(ctx, message, types.PermissionEditArtwork) {
+		utils.ReplyMessage(bot, message, "你没有编辑作品的权限")
+		return
+	}
+	if common.TaggerClient == nil {
+		utils.ReplyMessage(bot, message, "Tagger is not available")
+	}
+	var sourceURL string
+	var findUrlInArgs bool
+	if message.ReplyToMessage != nil {
+		sourceURL = utils.FindSourceURLForMessage(message.ReplyToMessage)
+	} else {
+		sourceURL = sources.FindSourceURL(message.Text)
+		findUrlInArgs = true
+	}
+	if sourceURL == "" {
+		utils.ReplyMessage(bot, message, "请回复一条消息, 或者指定作品链接")
+		return
+	}
+
+	artwork, err := service.GetArtworkByURL(ctx, sourceURL)
+	if err != nil {
+		utils.ReplyMessage(bot, message, "获取作品信息失败: "+err.Error())
+		return
+	}
+	selectAllPictures := true
+	pictureIndex := 1
+	_, _, args := telegoutil.ParseCommand(message.Text)
+	if len(args) > func() int {
+		if findUrlInArgs {
+			return 1
+		}
+		return 0
+	}() {
+		selectAllPictures = false
+		pictureIndex, err := strconv.Atoi(args[len(args)-1])
+		if err != nil {
+			utils.ReplyMessage(bot, message, "图片序号错误")
+			return
+		}
+		if pictureIndex < 1 || pictureIndex > len(artwork.Pictures) {
+			utils.ReplyMessage(bot, message, "图片序号超出范围")
+			return
+		}
+	}
+	pictures := make([]*types.Picture, 0)
+	if selectAllPictures {
+		pictures = artwork.Pictures
+	} else {
+		picture := artwork.Pictures[pictureIndex-1]
+		pictures[0] = picture
+	}
+	msg, err := utils.ReplyMessage(bot, message, "正在请求...")
+	if err != nil {
+		common.Logger.Errorf("Reply message failed: %s", err)
+		return
+	}
+	for i, picture := range pictures {
+		var file []byte
+		file, err = storage.GetFile(ctx, func() *types.StorageDetail {
+			if picture.StorageInfo.Regular != nil {
+				return picture.StorageInfo.Regular
+			} else {
+				return picture.StorageInfo.Original
+			}
+		}())
+		if err != nil {
+			file, err = common.DownloadWithCache(ctx, picture.Original, nil)
+		}
+		if err != nil {
+			common.Logger.Errorf("Download picture %s failed: %s", picture.Original, err)
+			continue
+		}
+		common.Logger.Debugf("Predicting tags for %s", picture.Original)
+		predict, err := common.TaggerClient.Predict(ctx, file)
+		if err != nil {
+			common.Logger.Errorf("Predict tags failed: %s", err)
+			utils.ReplyMessage(bot, message, "Predict tags failed")
+			return
+		}
+		if len(predict.PredictedTags) == 0 {
+			utils.ReplyMessage(bot, message, "No tags predicted")
+			return
+		}
+		newTags := slice.Union(artwork.Tags, predict.PredictedTags)
+		if err := service.UpdateArtworkTagsByURL(ctx, artwork.SourceURL, newTags); err != nil {
+			utils.ReplyMessage(bot, message, "更新作品标签失败: "+err.Error())
+			return
+		}
+		artwork, err = service.GetArtworkByURL(ctx, artwork.SourceURL)
+		if err != nil {
+			utils.ReplyMessage(bot, message, "获取更新后的作品信息失败: "+err.Error())
+			return
+		}
+		if artwork.Pictures[0].TelegramInfo.MessageID != 0 {
+			bot.EditMessageCaption(&telego.EditMessageCaptionParams{
+				ChatID:    ChannelChatID,
+				MessageID: artwork.Pictures[0].TelegramInfo.MessageID,
+				Caption:   utils.GetArtworkHTMLCaption(artwork),
+				ParseMode: telego.ModeHTML,
+			})
+		}
+		bot.EditMessageText(&telego.EditMessageTextParams{
+			ChatID:    msg.Chat.ChatID(),
+			MessageID: msg.MessageID,
+			Text:      fmt.Sprintf("选择的第 %d 张图片预测标签成功", i+1),
+		})
+	}
+	bot.EditMessageText(&telego.EditMessageTextParams{
+		ChatID:    msg.Chat.ChatID(),
+		MessageID: msg.MessageID,
+		Text:      "更新作品标签成功",
+	})
+}
diff --git a/telegram/handlers/handlers.go b/telegram/handlers/handlers.go
@@ -46,6 +46,7 @@ func RegisterHandlers(hg *telegohandler.HandlerGroup) {
 	mg.HandleMessageCtx(AddTagAlias, telegohandler.CommandEqual("tagalias"))
 	mg.HandleMessageCtx(DumpArtworkInfo, telegohandler.CommandEqual("dump"))
 	mg.HandleMessageCtx(ReCaptionArtwork, telegohandler.CommandEqual("recaption"))
+	mg.HandleMessageCtx(AutoTaggingArtwork, telegohandler.CommandEqual("autotag"))
 
 	hg.HandleCallbackQueryCtx(PostArtworkCallbackQuery, telegohandler.CallbackDataContains("post_artwork"))
 	hg.HandleCallbackQueryCtx(SearchPictureCallbackQuery, telegohandler.CallbackDataPrefix("search_picture"))
diff --git a/telegram/utils/post_artwork.go b/telegram/utils/post_artwork.go
@@ -248,8 +248,18 @@ func PostAndCreateArtwork(ctx context.Context, artwork *types.Artwork, bot *tele
 }
 
 func afterCreate(ctx context.Context, artwork *types.Artwork, bot *telego.Bot, fromID int64) {
-	for _, picture := range artwork.Pictures {
-		service.AddProcessPictureTask(ctx, picture)
+	go func() {
+		for _, picture := range artwork.Pictures {
+			service.AddProcessPictureTask(ctx, picture)
+		}
+	}()
+	if config.Cfg.Tagger.TagNew {
+		objectID, err := primitive.ObjectIDFromHex(artwork.ID)
+		if err != nil {
+			common.Logger.Fatalf("invalid ObjectID: %s", artwork.ID)
+			return
+		}
+		go service.AddPredictArtworkTagTask(ctx, objectID)
 	}
 	go checkDuplicate(ctx, artwork, bot, fromID)
 	go prettyPostedArtworkTagCaption(ctx, artwork, bot)

Original file line number	Diff line number	Diff line change
`@@ -8,8 +8,10 @@ import (`
`8`	`8`
`9`	`9`	`func Init() {`
`10`	`10`	`initHttpClient()`
`11`		`- initResendClient()`
`12`	`11`	`initLogger()`
	`12`	`+ if config.Cfg.Auth.Resend.APIKey != "" {`
	`13`	`+ initResendClient()`
	`14`	`+ }`
`13`	`15`	`if searchCfg := config.Cfg.Search; searchCfg.Enable {`
`14`	`16`	`switch searchCfg.Engine {`
`15`	`17`	`case "meilisearch":`
`@@ -19,4 +21,7 @@ func Init() {`
`19`	`21`	`os.Exit(1)`
`20`	`22`	`}`
`21`	`23`	`}`
	`24`	`+ if taggerCfg := config.Cfg.Tagger; taggerCfg.Enable {`
	`25`	`+ initTaggerClient()`
	`26`	`+ }`
`22`	`27`	`}`
Original file line number	Diff line number	Diff line change
`@@ -112,10 +112,10 @@ func (m *artworkSyncManager) ProcessArtworkUpdateEvent(event bson.M) {`
`112`	`112`	`}`
`113`	`113`	`task, err := common.MeilisearchClient.Index(config.Cfg.Search.MeiliSearch.Index).UpdateDocuments(artworkJSON)`
`114`	`114`	`if err != nil {`
`115`		`- common.Logger.Errorf("add artwork to meilisearch error: %s", err)`
	`115`	`+ common.Logger.Errorf("update artwork to meilisearch error: %s", err)`
`116`	`116`	`return`
`117`	`117`	`}`
`118`		`- common.Logger.Debugf("commited add artwork task to meilisearch: %d", task.TaskUID)`
	`118`	`+ common.Logger.Debugf("commited update artwork task to meilisearch: %d", task.TaskUID)`
`119`	`119`	`}`
`120`	`120`
`121`	`121`	`func (m *artworkSyncManager) ProcessArtworkDeleteEvent(event bson.M) {`
Original file line number	Diff line number	Diff line change
`@@ -12,6 +12,9 @@ func InitService() {`
`12`	`12`	`if config.Cfg.Search.Enable {`
`13`	`13`	`go syncArtworkToSearchEngine()`
`14`	`14`	`}`
	`15`	`+ if config.Cfg.Tagger.Enable {`
	`16`	`+ go listenPredictArtworkTagsTask()`
	`17`	`+ }`
`15`	`18`	`}`
`16`	`19`
`17`	`20`	`type Service struct{}`