feat: 画像内コンテンツの検索

cmd/config.go

@@ -97,6 +97,16 @@ type Config struct {
 		Password string `mapstructure:"password" yaml:"password"`
 	} `mapstructure:"es" yaml:"es"`
 
+	// ImageSearch 画像検索用外部サーバー設定
+	ImageSearch struct {
+		// URL 外部サーバーのURL (default: "" = 無効)
+		URL string `mapstructure:"url" yaml:"url"`
+		// Timeout リクエストタイムアウト秒数 (default: 10)
+		Timeout int `mapstructure:"timeout" yaml:"timeout"`
+		// VectorDimension Embeddingベクトルの次元数 (default: 768)
+		VectorDimension int `mapstructure:"vectorDimension" yaml:"vectorDimension"`
+	} `mapstructure:"imageSearch" yaml:"imageSearch"`
+
 	// Storage ファイルストレージ設定
 	Storage struct {
 		// Type ストレージタイプ (default: local)
@@ -288,6 +298,9 @@ func init() {
 	viper.SetDefault("es.url", "")
 	viper.SetDefault("es.username", "elastic")
 	viper.SetDefault("es.password", "password")
+	viper.SetDefault("imageSearch.url", "")
+	viper.SetDefault("imageSearch.timeout", 10)
+	viper.SetDefault("imageSearch.vectorDimension", 768)
 	viper.SetDefault("storage.type", "local")
 	viper.SetDefault("storage.local.dir", "./storage")
 	viper.SetDefault("storage.swift.username", "")
@@ -458,9 +471,9 @@ func newFCMClientIfAvailable(repo repository.Repository, logger *zap.Logger, unr
 	return fcm.NewNullClient(), nil
 }
 
-func initSearchServiceIfAvailable(mm message.Manager, cm channel.Manager, repo repository.Repository, logger *zap.Logger, config search.ESEngineConfig) (search.Engine, error) {
+func initSearchServiceIfAvailable(mm message.Manager, cm channel.Manager, repo repository.Repository, fs storage.FileStorage, logger *zap.Logger, config search.ESEngineConfig) (search.Engine, error) {
 	if len(config.URL) > 0 {
-		return search.NewESEngine(mm, cm, repo, logger, config)
+		return search.NewESEngine(mm, cm, repo, fs, logger, config)
 	}
 	return search.NewNullEngine(), nil
 }
@@ -478,6 +491,11 @@ func provideESEngineConfig(c *Config) search.ESEngineConfig {
 		URL:      c.ES.URL,
 		Username: c.ES.Username,
 		Password: c.ES.Password,
+		ImageSearch: search.ImageSearchConfig{
+			URL:             c.ImageSearch.URL,
+			Timeout:         time.Duration(c.ImageSearch.Timeout) * time.Second,
+			VectorDimension: c.ImageSearch.VectorDimension,
+		},
 	}
 }
 

cmd/image_index.go

@@ -0,0 +1,201 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/spf13/cobra"
+	"go.uber.org/zap"
+
+	"github.com/traPtitech/traQ/repository"
+	"github.com/traPtitech/traQ/repository/gorm"
+	"github.com/traPtitech/traQ/service/channel"
+	"github.com/traPtitech/traQ/service/message"
+	"github.com/traPtitech/traQ/service/search"
+	"github.com/traPtitech/traQ/utils/optional"
+)
+
+const imageIndexBatchSize = 100
+
+// imageIndexCommand 画像インデックス再作成コマンド
+func imageIndexCommand() *cobra.Command {
+	var clearIndex bool
+
+	cmd := cobra.Command{
+		Use:   "image-index",
+		Short: "Reindex image OCR/embedding data for Elasticsearch",
+		Long: `Process messages with image attachments and update Elasticsearch index
+with OCR text and embedding vectors.
+
+Only messages that have not yet been processed are targeted.
+Use --clear to remove existing image index data before processing,
+which is useful when the embedding model has been changed.`,
+		Run: func(_ *cobra.Command, _ []string) {
+			logger, gormLogger := getCLILoggers()
+			defer logger.Sync()
+
+			logger.Info("starting image index batch")
+
+			// Database
+			engine, err := c.getDatabase()
+			if err != nil {
+				logger.Fatal("failed to connect database", zap.Error(err))
+			}
+			engine.Logger = gormLogger
+			db, err := engine.DB()
+			if err != nil {
+				logger.Fatal("failed to get *sql.DB", zap.Error(err))
+			}
+			defer db.Close()
+
+			// FileStorage
+			fs, err := c.getFileStorage()
+			if err != nil {
+				logger.Fatal("failed to setup file storage", zap.Error(err))
+			}
+
+			// Repository
+			repo, _, err := gorm.NewGormRepository(engine, nil, logger, false)
+			if err != nil {
+				logger.Fatal("failed to initialize repository", zap.Error(err))
+			}
+
+			// Channel Manager
+			cm, err := channel.InitChannelManager(repo, logger)
+			if err != nil {
+				logger.Fatal("failed to initialize channel manager", zap.Error(err))
+			}
+
+			// Message Manager
+			mm, err := message.NewMessageManager(repo, cm, logger)
+			if err != nil {
+				logger.Fatal("failed to initialize message manager", zap.Error(err))
+			}
+
+			// Search Engine
+			esConfig := provideESEngineConfig(&c)
+			searchEngine, err := search.NewESEngine(mm, cm, repo, fs, logger, esConfig)
+			if err != nil {
+				logger.Fatal("failed to initialize search engine", zap.Error(err))
+			}
+			defer searchEngine.Close()
+
+			ctx := context.Background()
+
+			// --clear: 既存の画像インデックスデータをクリア
+			if clearIndex {
+				logger.Info("clearing existing image index data")
+				if err := searchEngine.ClearImageIndex(ctx); err != nil {
+					logger.Fatal("failed to clear image index", zap.Error(err))
+				}
+			}
+
+			// ESから未処理のメッセージIDを取得
+			unprocessedIDs, err := searchEngine.GetUnprocessedImageMessageIDs(ctx)
+			if err != nil {
+				logger.Fatal("failed to get unprocessed message IDs", zap.Error(err))
+			}
+			logger.Info(fmt.Sprintf("found %d unprocessed messages with images", len(unprocessedIDs)))
+
+			// バッチごとにDBからメッセージを取得して処理
+			totalProcessed := 0
+			for i := 0; i < len(unprocessedIDs); i += imageIndexBatchSize {
+				end := i + imageIndexBatchSize
+				if end > len(unprocessedIDs) {
+					end = len(unprocessedIDs)
+				}
+				batch := unprocessedIDs[i:end]
+
+				messages, _, err := repo.GetMessages(ctx, repository.MessagesQuery{
+					IDIn:           optional.From(batch),
+					Limit:          len(batch),
+					DisablePreload: true,
+				})
+				if err != nil {
+					logger.Fatal("failed to get messages from DB", zap.Error(err))
+				}
+
+				if err := searchEngine.ProcessImagesForMessages(ctx, messages); err != nil {
+					logger.Error("failed to process image batch, progress is saved in ES",
+						zap.Error(err),
+						zap.Int("processed", totalProcessed))
+					break
+				}
+
+				totalProcessed += len(messages)
+				logger.Info(fmt.Sprintf("processed %d/%d messages", totalProcessed, len(unprocessedIDs)))
+			}
+
+			logger.Info(fmt.Sprintf("image index batch completed. processed: %d", totalProcessed))
+		},
+	}
+
+	cmd.Flags().BoolVar(&clearIndex, "clear", false, "Clear existing image index data before processing (use when embedding model has changed)")
+
+	return &cmd
+}
+
+// imageIndexClearCommand 画像インデックスデータのみをクリアするコマンド
+func imageIndexClearCommand() *cobra.Command {
+	cmd := cobra.Command{
+		Use:   "image-index-clear",
+		Short: "Clear image OCR/embedding data from Elasticsearch index",
+		Long:  "Remove imageText and imageVector fields from all documents. Run image-index afterwards to re-process.",
+		Run: func(_ *cobra.Command, _ []string) {
+			logger, gormLogger := getCLILoggers()
+			defer logger.Sync()
+
+			// Database
+			engine, err := c.getDatabase()
+			if err != nil {
+				logger.Fatal("failed to connect database", zap.Error(err))
+			}
+			engine.Logger = gormLogger
+			db, err := engine.DB()
+			if err != nil {
+				logger.Fatal("failed to get *sql.DB", zap.Error(err))
+			}
+			defer db.Close()
+
+			// FileStorage
+			fs, err := c.getFileStorage()
+			if err != nil {
+				logger.Fatal("failed to setup file storage", zap.Error(err))
+			}
+
+			// Repository
+			repo, _, err := gorm.NewGormRepository(engine, nil, logger, false)
+			if err != nil {
+				logger.Fatal("failed to initialize repository", zap.Error(err))
+			}
+
+			// Channel Manager
+			cm, err := channel.InitChannelManager(repo, logger)
+			if err != nil {
+				logger.Fatal("failed to initialize channel manager", zap.Error(err))
+			}
+
+			// Message Manager
+			mm, err := message.NewMessageManager(repo, cm, logger)
+			if err != nil {
+				logger.Fatal("failed to initialize message manager", zap.Error(err))
+			}
+
+			// Search Engine
+			esConfig := provideESEngineConfig(&c)
+			searchEngine, err := search.NewESEngine(mm, cm, repo, fs, logger, esConfig)
+			if err != nil {
+				logger.Fatal("failed to initialize search engine", zap.Error(err))
+			}
+			defer searchEngine.Close()
+
+			ctx := context.Background()
+			if err := searchEngine.ClearImageIndex(ctx); err != nil {
+				logger.Fatal("failed to clear image index", zap.Error(err))
+			}
+			logger.Info("image index data cleared successfully")
+		},
+	}
+
+	return &cmd
+}

cmd/root.go

@@ -75,6 +75,8 @@ func init() {
 		stampCommand(),
 		versionCommand(),
 		healthcheckCommand(),
+		imageIndexCommand(),
+		imageIndexClearCommand(),
 	)
 
 	flags := rootCommand.PersistentFlags()

service/search/engine.go

@@ -1,6 +1,7 @@
 package search
 
 import (
+	"context"
 	"errors"
 	"regexp"
 	"strings"
@@ -9,13 +10,17 @@ import (
 	vd "github.com/go-ozzo/ozzo-validation/v4"
 	"github.com/gofrs/uuid"
 
+	"github.com/traPtitech/traQ/model"
 	"github.com/traPtitech/traQ/service/message"
 	"github.com/traPtitech/traQ/utils/optional"
 )
 
 // ErrServiceUnavailable エラー 現在検索サービスが利用できません
 var ErrServiceUnavailable = errors.New("search service is unavailable")
 
+// ErrImageSearchUnavailable エラー 画像検索サービスが利用できません
+var ErrImageSearchUnavailable = errors.New("image search service is unavailable")
+
 // Engine 検索エンジンインターフェイス
 type Engine interface {
 	// Do 与えられたクエリで検索を実行します
@@ -24,6 +29,12 @@ type Engine interface {
 	Available() bool
 	// Close 検索サービスを終了します
 	Close() error
+	// ProcessImagesForMessages バッチ用: 指定メッセージの画像を処理してインデックスを更新します
+	ProcessImagesForMessages(ctx context.Context, messages []*model.Message) error
+	// GetUnprocessedImageMessageIDs 画像付きだがインデックス未処理のメッセージIDを取得します
+	GetUnprocessedImageMessageIDs(ctx context.Context) ([]uuid.UUID, error)
+	// ClearImageIndex 全メッセージの画像インデックスデータ(imageText, imageVector)をクリアします
+	ClearImageIndex(ctx context.Context) error
 }
 
 // Query 検索クエリ