feat: add hybrid search command and implement related functionality

Author: krau

errs/errors.go

@@ -11,12 +11,12 @@ var (
 	ErrFileTooLarge     = errors.New("file too large (>20MB)")
 	ErrNoAvailableFile  = errors.New("no available file")
 
-	ErrNilArtwork          = errors.New("artwork is nil")
-	ErrNotFoundArtworks    = errors.New("artworks not found")
-	ErrArtworkAlreadyExist = errors.New("artwork already exists")
-	ErrSourceNotSupported  = errors.New("source not supported")
-	ErrIndexOOB            = errors.New("index out of bounds")
-	ErrFailedToFetchArtwork  = errors.New("failed to fetch artwork")
+	ErrNilArtwork           = errors.New("artwork is nil")
+	ErrNotFoundArtworks     = errors.New("artworks not found")
+	ErrArtworkAlreadyExist  = errors.New("artwork already exists")
+	ErrSourceNotSupported   = errors.New("source not supported")
+	ErrIndexOOB             = errors.New("index out of bounds")
+	ErrFailedToFetchArtwork = errors.New("failed to fetch artwork")
 
 	ErrStorageUnkown = errors.New("unknown storage")
 
@@ -31,4 +31,6 @@ var (
 	ErrChatIDNotSet = errors.New("chat id not set")
 
 	ErrAliasAlreadyUsed = errors.New("alias already used")
+
+	ErrNotEnabledHybridSearch = errors.New("hybrid search not enabled")
 )

service/artwork_search.go

@@ -0,0 +1,54 @@
+package service
+
+import (
+	"context"
+
+	"github.com/bytedance/sonic"
+	"github.com/krau/ManyACG/adapter"
+	"github.com/krau/ManyACG/common"
+	"github.com/krau/ManyACG/config"
+	"github.com/krau/ManyACG/dao"
+	"github.com/krau/ManyACG/errs"
+	"github.com/krau/ManyACG/types"
+	"github.com/meilisearch/meilisearch-go"
+	"go.mongodb.org/mongo-driver/bson/primitive"
+)
+
+func HybridSearchArtworks(ctx context.Context, queryText string, hybridSemanticRatio float64, limit int64, options ...*types.AdapterOption) ([]*types.Artwork, error) {
+	if common.MeilisearchClient == nil {
+		return nil, errs.ErrNotEnabledHybridSearch
+	}
+	index := common.MeilisearchClient.Index(config.Cfg.Search.MeiliSearch.Index)
+	resp, err := index.SearchWithContext(ctx, queryText, &meilisearch.SearchRequest{
+		Limit: limit,
+		Hybrid: &meilisearch.SearchRequestHybrid{
+			SemanticRatio: hybridSemanticRatio,
+		},
+	})
+	if err != nil {
+		return nil, err
+	}
+	hits := resp.Hits
+	artworkSearchDocs := make([]*types.ArtworkSearchDocument, 0, len(hits))
+	hitsBytes, err := sonic.Marshal(hits)
+	if err != nil {
+		return nil, err
+	}
+	err = sonic.Unmarshal(hitsBytes, &artworkSearchDocs)
+	if err != nil {
+		return nil, err
+	}
+	artworkModels := make([]*types.ArtworkModel, 0, len(artworkSearchDocs))
+	for _, doc := range artworkSearchDocs {
+		objectID, err := primitive.ObjectIDFromHex(doc.ID)
+		if err != nil {
+			return nil, err
+		}
+		artworkModel, err := dao.GetArtworkByID(ctx, objectID)
+		if err != nil {
+			return nil, err
+		}
+		artworkModels = append(artworkModels, artworkModel)
+	}
+	return adapter.ConvertToArtworks(ctx, artworkModels, options...)
+}

service/change_stream.go

@@ -68,12 +68,14 @@ func (m *artworkSyncManager) ProcessArtworkChangeEvent(event bson.M) {
 	}()
 	operationType := event["operationType"].(string)
 	switch operationType {
-	case "insert", "update":
+	case "update":
 		m.ProcessArtworkUpdateEvent(event)
 	case "delete":
 		m.ProcessArtworkDeleteEvent(event)
 	case "replace":
 		m.ProcessArtworkReplaceEvent(event)
+	case "insert":
+		// do nothing
 	default:
 		common.Logger.Debugf("unknown operation type: %s", operationType)
 	}

telegram/bot.go

@@ -63,6 +63,10 @@ var (
 			Command:     "help",
 			Description: "食用指南",
 		},
+		{
+			Command:     "query",
+			Description: "混合搜索作品",
+		},
 	}
 
 	AdminCommands = []telego.BotCommand{

telegram/handlers/handlers.go

@@ -30,13 +30,13 @@ func RegisterHandlers(hg *telegohandler.HandlerGroup) {
 	mg.HandleMessageCtx(GetArtworkInfoCommand, telegohandler.CommandEqual("info"))
 	mg.HandleMessageCtx(CalculatePicture, telegohandler.CommandEqual("hash"))
 	mg.HandleMessageCtx(GetStats, telegohandler.CommandEqual("stats"))
+	mg.HandleMessageCtx(HybridSearchArtworks, telegohandler.CommandEqual("query"))
 
 	mg.HandleMessageCtx(SetAdmin, telegohandler.CommandEqual("set_admin"))
 	mg.HandleMessageCtx(DeleteArtwork, telegohandler.Or(telegohandler.CommandEqual("delete"), telegohandler.CommandEqual("del")))
 	mg.HandleMessageCtx(ProcessPicturesHashAndSize, telegohandler.CommandEqual("process_pictures_hashsize"))
 	mg.HandleMessageCtx(ProcessPicturesStorage, telegohandler.CommandEqual("process_pictures_storage"))
 	mg.HandleMessageCtx(FixTwitterArtists, telegohandler.CommandEqual("fix_twitter_artists"))
-
 	mg.HandleMessageCtx(ToggleArtworkR18, telegohandler.CommandEqual("r18"))
 	mg.HandleMessageCtx(SetArtworkTags, telegohandler.Or(telegohandler.CommandEqual("tags"), telegohandler.CommandEqual("addtags"), telegohandler.CommandEqual("deltags")))
 	mg.HandleMessageCtx(EditArtworkTitle, telegohandler.CommandEqual("title"))

telegram/handlers/help.go

@@ -20,6 +20,7 @@ func Help(ctx context.Context, bot *telego.Bot, message telego.Message) {
 /hash - 计算图片信息
 /stats - 获取统计数据
 /files - 获取作品原图
+/query - 混合搜索作品
 `
 	helpText += `
 随机图片相关功能中支持使用以下格式的参数:

telegram/handlers/random_picture.go telegram/handlers/query_artwork.go

@@ -5,10 +5,12 @@ import (
 	"errors"
 	"fmt"
 	"math/rand"
+	"strconv"
 	"strings"
 
 	"github.com/krau/ManyACG/adapter"
 	"github.com/krau/ManyACG/common"
+	"github.com/krau/ManyACG/config"
 
 	"github.com/krau/ManyACG/service"
 	"github.com/krau/ManyACG/sources"
@@ -76,3 +78,62 @@ func RandomPicture(ctx context.Context, bot *telego.Bot, message telego.Message)
 		}
 	}
 }
+
+func HybridSearchArtworks(ctx context.Context, bot *telego.Bot, message telego.Message) {
+	if common.MeilisearchClient == nil {
+		utils.ReplyMessage(bot, message, "未启用混合搜索功能")
+		return
+	}
+	_, _, args := telegoutil.ParseCommand(message.Text)
+	if len(args) == 0 {
+		utils.ReplyMessage(bot, message, "使用方法: /query <搜索内容> [语义比例]\n语义比例为0-1的浮点数, 应位于参数列表最后, 越大越趋向于基于语义搜索, 若不提供, 使用默认值0.8")
+		return
+	}
+	var hybridSemanticRatio float64
+	var queryText string
+	hybridSemanticRatio, err := strconv.ParseFloat(args[len(args)-1], 64)
+	if err != nil {
+		hybridSemanticRatio = 0.8
+		queryText = strings.Join(args, " ")
+	} else {
+		if hybridSemanticRatio < 0 || hybridSemanticRatio > 1 {
+			utils.ReplyMessage(bot, message, "参数错误: 语义比例应为0-1的浮点数")
+			return
+		}
+		queryText = strings.Join(args[:len(args)-1], " ")
+	}
+	artworks, err := service.HybridSearchArtworks(ctx, queryText, hybridSemanticRatio, 10)
+	if err != nil {
+		common.Logger.Errorf("搜索失败: %s", err)
+		utils.ReplyMessage(bot, message, "搜索失败, 请联系管理员检查搜索引擎设置与状态")
+		return
+	}
+	if len(artworks) == 0 {
+		utils.ReplyMessage(bot, message, "未找到相关图片")
+		return
+	}
+
+	if len(artworks) > 10 {
+		artworks = artworks[:10]
+	}
+
+	inputMedias := make([]telego.InputMedia, 0, len(artworks))
+	for _, artwork := range artworks {
+		picture := artwork.Pictures[0]
+		var file telego.InputFile
+		if picture.TelegramInfo != nil && picture.TelegramInfo.PhotoFileID != "" {
+			file = telegoutil.FileFromID(picture.TelegramInfo.PhotoFileID)
+		} else {
+			photoURL := fmt.Sprintf("%s/?url=%s&w=2560&h=2560&we&output=jpg", config.Cfg.WSRVURL, picture.Original)
+			file = telegoutil.FileFromURL(photoURL)
+		}
+		caption := fmt.Sprintf("<a href=\"%s\">%s</a>", artwork.SourceURL, common.EscapeHTML(artwork.Title))
+		inputMedias = append(inputMedias, telegoutil.MediaPhoto(file).WithCaption(caption).WithParseMode(telego.ModeHTML))
+	}
+	mediaGroup := telegoutil.MediaGroup(message.Chat.ChatID(), inputMedias...)
+	_, err = bot.SendMediaGroup(mediaGroup)
+	if err != nil {
+		common.Logger.Errorf("发送图片失败: %s", err)
+	}
+
+}