diff --git a/cmd/config.go b/cmd/config.go index cf3c448b2..3f67a6543 100644 --- a/cmd/config.go +++ b/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, + }, } } diff --git a/cmd/image_index.go b/cmd/image_index.go new file mode 100644 index 000000000..62db5ff29 --- /dev/null +++ b/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 +} diff --git a/cmd/root.go b/cmd/root.go index f42646388..32a07e965 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -75,6 +75,8 @@ func init() { stampCommand(), versionCommand(), healthcheckCommand(), + imageIndexCommand(), + imageIndexClearCommand(), ) flags := rootCommand.PersistentFlags() diff --git a/cmd/wire_gen.go b/cmd/wire_gen.go index 5fe906a35..d3a217a30 100644 --- a/cmd/wire_gen.go +++ b/cmd/wire_gen.go @@ -88,7 +88,7 @@ func newServer(hub2 *hub.Hub, db *gorm.DB, repo repository.Repository, fs storag } oidcService := provideOIDCService(c2, repo, rbacRBAC) esEngineConfig := provideESEngineConfig(c2) - engine, err := initSearchServiceIfAvailable(messageManager, manager, repo, logger, esEngineConfig) + engine, err := initSearchServiceIfAvailable(messageManager, manager, repo, fs, logger, esEngineConfig) if err != nil { return nil, err } diff --git a/service/search/engine.go b/service/search/engine.go index cd22dc2a7..ce0dfbeb5 100644 --- a/service/search/engine.go +++ b/service/search/engine.go @@ -1,6 +1,7 @@ package search import ( + "context" "errors" "regexp" "strings" @@ -9,6 +10,7 @@ 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" ) @@ -16,6 +18,9 @@ import ( // 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 検索クエリ diff --git a/service/search/es.go b/service/search/es.go index 9152cedd1..833dbd5cd 100644 --- a/service/search/es.go +++ b/service/search/es.go @@ -14,9 +14,11 @@ import ( "github.com/gofrs/uuid" "go.uber.org/zap" + "github.com/traPtitech/traQ/model" "github.com/traPtitech/traQ/repository" "github.com/traPtitech/traQ/service/channel" "github.com/traPtitech/traQ/service/message" + "github.com/traPtitech/traQ/utils/storage" ) const ( @@ -38,35 +40,47 @@ type ESEngineConfig struct { Username string // Password ESのパスワード Password string + // ImageSearch 画像検索設定 + ImageSearch ImageSearchConfig } // esEngine search.Engine 実装 type esEngine struct { - client *elasticsearch.Client - mm message.Manager - cm channel.Manager - repo repository.Repository - l *zap.Logger - done chan<- struct{} + client *elasticsearch.Client + mm message.Manager + cm channel.Manager + repo repository.Repository + fs storage.FileStorage + imageClient ImageSearchClient + imageConfig ImageSearchConfig + l *zap.Logger + done chan<- struct{} +} + +// esImageVector 画像ごとのEmbeddingベクトル +type esImageVector struct { + Vector []float64 `json:"vector"` } // esMessageDoc Elasticsearchに入るメッセージの情報 type esMessageDoc struct { - ID uuid.UUID `json:"-"` - UserID uuid.UUID `json:"userId"` - ChannelID uuid.UUID `json:"channelId"` - IsPublic bool `json:"isPublic"` - Bot bool `json:"bot"` - Text string `json:"text"` - CreatedAt time.Time `json:"createdAt"` - UpdatedAt time.Time `json:"updatedAt"` - To []uuid.UUID `json:"to"` - Citation []uuid.UUID `json:"citation"` - HasURL bool `json:"hasURL"` - HasAttachments bool `json:"hasAttachments"` - HasImage bool `json:"hasImage"` - HasVideo bool `json:"hasVideo"` - HasAudio bool `json:"hasAudio"` + ID uuid.UUID `json:"-"` + UserID uuid.UUID `json:"userId"` + ChannelID uuid.UUID `json:"channelId"` + IsPublic bool `json:"isPublic"` + Bot bool `json:"bot"` + Text string `json:"text"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` + To []uuid.UUID `json:"to"` + Citation []uuid.UUID `json:"citation"` + HasURL bool `json:"hasURL"` + HasAttachments bool `json:"hasAttachments"` + HasImage bool `json:"hasImage"` + HasVideo bool `json:"hasVideo"` + HasAudio bool `json:"hasAudio"` + ImageText string `json:"imageText,omitempty"` + ImageVectors []esImageVector `json:"imageVector,omitempty"` } // esMessageDocUpdate Update用 Elasticsearchに入るメッセージの部分的な情報 @@ -81,6 +95,12 @@ type esMessageDocUpdate struct { HasAudio bool `json:"hasAudio"` } +// esImageDocUpdate 画像処理結果によるインデックスの部分更新 +type esImageDocUpdate struct { + ImageText string `json:"imageText,omitempty"` + ImageVectors []esImageVector `json:"imageVector,omitempty"` +} + type esCreateIndexBody struct { Mappings m `json:"mappings"` Settings m `json:"settings"` @@ -88,10 +108,10 @@ type esCreateIndexBody struct { type m map[string]any -// esMapping Elasticsearchに入るメッセージの情報 -// esMessageDoc と同じにする -var esMapping = m{ - "properties": m{ +// esMessageMapping Elasticsearchに入るメッセージのマッピングを生成する +// vectorDimension が 0 の場合は画像ベクトル関連フィールドを含めない +func esMessageMapping(vectorDimension int) m { + properties := m{ "userId": m{ "type": "keyword", }, @@ -137,7 +157,23 @@ var esMapping = m{ "hasAudio": m{ "type": "boolean", }, - }, + "imageText": m{ + "type": "text", + "analyzer": "sudachi_analyzer", + }, + } + if vectorDimension > 0 { + properties["imageVector"] = m{ + "type": "nested", + "properties": m{ + "vector": m{ + "type": "dense_vector", + "dims": vectorDimension, + }, + }, + } + } + return m{"properties": properties} } // esSetting Indexに追加するsetting情報 @@ -173,7 +209,7 @@ var esSetting = m{ } // NewESEngine Elasticsearch検索エンジンを生成します -func NewESEngine(mm message.Manager, cm channel.Manager, repo repository.Repository, logger *zap.Logger, config ESEngineConfig) (Engine, error) { +func NewESEngine(mm message.Manager, cm channel.Manager, repo repository.Repository, fs storage.FileStorage, logger *zap.Logger, config ESEngineConfig) (Engine, error) { // esクライアント作成 client, err := elasticsearch.NewClient(elasticsearch.Config{ Addresses: []string{config.URL}, @@ -222,7 +258,7 @@ func NewESEngine(mm message.Manager, cm channel.Manager, repo repository.Reposit // indexが存在しなかったら作成 if existsRes.StatusCode == http.StatusNotFound { reqBody, err := json.Marshal(esCreateIndexBody{ - Mappings: esMapping, + Mappings: esMessageMapping(config.ImageSearch.VectorDimension), Settings: esSetting, }) if err != nil { @@ -242,14 +278,19 @@ func NewESEngine(mm message.Manager, cm channel.Manager, repo repository.Reposit defer createIndexRes.Body.Close() } + imageClient := NewImageSearchClient(config.ImageSearch, logger) + done := make(chan struct{}) engine := &esEngine{ - client: client, - mm: mm, - cm: cm, - repo: repo, - l: logger.Named("search"), - done: done, + client: client, + mm: mm, + cm: cm, + repo: repo, + fs: fs, + imageClient: imageClient, + imageConfig: config.ImageSearch, + l: logger.Named("search"), + done: done, } go engine.syncLoop(done) @@ -305,7 +346,7 @@ func (e *esEngine) Do(q *Query) (Result, error) { if q.Word.Valid { body := simpleQueryString{ Query: q.Word.V, - Fields: []string{"text"}, + Fields: []string{"text", "imageText"}, DefaultOperator: "AND", } @@ -403,6 +444,23 @@ func (e *esEngine) Do(q *Query) (Result, error) { // NOTE: 現状`sort.Key`はそのままesのソートキーとして使える前提 sort := q.GetSortKey() + // ベクトル検索が可能かチェックし、ハイブリッド検索を試みる + if q.Word.Valid && e.imageClient.Available() && e.imageConfig.VectorDimension > 0 { + ctx, cancel := context.WithTimeout(context.Background(), e.imageConfig.Timeout) + defer cancel() + + queryVector, err := e.imageClient.EmbedText(ctx, q.Word.V) + if err == nil && queryVector != nil { + return e.doHybridSearch(musts, queryVector, sort, limit, offset) + } + e.l.Debug("falling back to text-only search", zap.Error(err)) + } + + return e.doTextSearch(musts, sort, limit, offset) +} + +// doTextSearch 通常のテキスト検索を実行する +func (e *esEngine) doTextSearch(musts []searchQuery, sort string, limit, offset int) (Result, error) { b, err := json.Marshal(newSearchBody(musts)) if err != nil { return nil, fmt.Errorf("failed to marshal search query: %w", err) @@ -433,6 +491,69 @@ func (e *esEngine) Do(q *Query) (Result, error) { return e.parseResultFromResponse(res) } +// doHybridSearch RRFによるハイブリッド検索（テキスト + ベクトル）を実行する +func (e *esEngine) doHybridSearch(musts []searchQuery, queryVector []float64, sort string, limit, offset int) (Result, error) { + // RRF (Reciprocal Rank Fusion) を使ったハイブリッド検索 + // テキスト検索クエリ + textQuery := newSearchBody(musts) + + // knnクエリ（nested field用） + knnQuery := m{ + "field": "imageVector.vector", + "query_vector": queryVector, + "k": limit + offset, + "num_candidates": 100, + "nested": m{ + "path": "imageVector", + }, + } + + hybridBody := m{ + "query": textQuery.Query, + "knn": knnQuery, + "rank": m{ + "rrf": m{}, + }, + "size": limit, + "from": offset, + } + + b, err := json.Marshal(hybridBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal hybrid search query: %w", err) + } + + sr, err := e.client.Search( + e.client.Search.WithIndex(getIndexName(esMessageIndex)), + e.client.Search.WithBody(bytes.NewBuffer(b)), + ) + if err != nil { + // ハイブリッド検索に失敗した場合はテキスト検索にフォールバック + e.l.Warn("hybrid search failed, falling back to text search", zap.Error(err)) + return e.doTextSearch(musts, sort, limit, offset) + } + if sr.IsError() { + e.l.Warn("hybrid search returned error, falling back to text search") + sr.Body.Close() + return e.doTextSearch(musts, sort, limit, offset) + } + defer sr.Body.Close() + + var res esSearchResponse + err = json.NewDecoder(sr.Body).Decode(&res) + if err != nil { + return nil, err + } + + e.l.Debug("hybrid search result", zap.Reflect("hits", res.Hits)) + return e.parseResultFromResponse(res) +} + +// generateImageURL ファイルの署名付きURLを生成する +func (e *esEngine) generateImageURL(fileID uuid.UUID) (string, error) { + return e.fs.GenerateAccessURL(fileID.String(), model.FileTypeUserFile) +} + func (e *esEngine) Available() bool { // このクライアントにはライフサイクルが無いので、常にtrueを返す。 return true diff --git a/service/search/es_sync.go b/service/search/es_sync.go index c09103217..cbc46b7b8 100644 --- a/service/search/es_sync.go +++ b/service/search/es_sync.go @@ -31,6 +31,8 @@ type attributes struct { HasImage bool HasVideo bool HasAudio bool + // ImageFileIDs アニメーションでない画像の添付ファイルIDリスト（非同期画像処理用） + ImageFileIDs []uuid.UUID } // ユーザーがbotかどうかのcache @@ -100,6 +102,10 @@ func (e *esEngine) getAttributes(m *model.Message, parseResult *message.ParseRes } if strings.HasPrefix(meta.Mime, "image/") { attr.HasImage = true + // アニメーション画像でない画像のファイルIDを収集 + if !meta.IsAnimatedImage { + attr.ImageFileIDs = append(attr.ImageFileIDs, attachmentID) + } } else if strings.HasPrefix(meta.Mime, "video/") { attr.HasVideo = true } else if strings.HasPrefix(meta.Mime, "audio/") { @@ -210,6 +216,12 @@ func (e *esEngine) sync() error { return nil } +// imageProcessTask 非同期画像処理のタスク +type imageProcessTask struct { + MessageID uuid.UUID + ImageFileIDs []uuid.UUID +} + func syncNewMessages(e *esEngine, messages []*model.Message, lastInsert time.Time, lastSynced time.Time, userCache userCache) (err error) { bulkIndexer, err := esutil.NewBulkIndexer(esutil.BulkIndexerConfig{ Client: e.client, @@ -219,6 +231,8 @@ func syncNewMessages(e *esEngine, messages []*model.Message, lastInsert time.Tim return err } + var imageTasks []imageProcessTask + defer func() { closeErr := bulkIndexer.Close(context.Background()) if err != nil && closeErr != nil { // エラーが発生してからdeferに来た時、エラーの上書きを防ぐ。 @@ -232,15 +246,30 @@ func syncNewMessages(e *esEngine, messages []*model.Message, lastInsert time.Tim e.l.Info(fmt.Sprintf("indexed %v message(s) to index, updated %v message(s) on index, failed %v message(s), last insert %v", bulkIndexer.Stats().NumIndexed, bulkIndexer.Stats().NumUpdated, bulkIndexer.Stats().NumFailed, lastInsert)) + + // 非同期で画像処理を開始 + if len(imageTasks) > 0 && e.imageClient.Available() { + go e.processImageTasks(imageTasks) + } }() for _, v := range messages { if v.CreatedAt.After(lastSynced) { - doc, err := e.convertMessageCreated(v, message.Parse(v.Text), userCache) + parseResult := message.Parse(v.Text) + doc, err := e.convertMessageCreated(v, parseResult, userCache) if err != nil { return err } + // 画像が含まれている場合、非同期処理タスクを作成 + attr := e.getAttributes(v, parseResult) + if len(attr.ImageFileIDs) > 0 { + imageTasks = append(imageTasks, imageProcessTask{ + MessageID: v.ID, + ImageFileIDs: attr.ImageFileIDs, + }) + } + data, err := json.Marshal(*doc) if err != nil { return err @@ -276,6 +305,120 @@ func syncNewMessages(e *esEngine, messages []*model.Message, lastInsert time.Tim return nil } +// processImageTasks 画像処理タスクを非同期で実行し、ESインデックスを更新する +func (e *esEngine) processImageTasks(tasks []imageProcessTask) { + for _, task := range tasks { + if err := e.processImageTask(task); err != nil { + e.l.Warn("failed to process image task", + zap.Stringer("messageID", task.MessageID), + zap.Error(err)) + } + } +} + +// processImageTask 1メッセージ分の画像処理を実行する +func (e *esEngine) processImageTask(task imageProcessTask) error { + // 署名付きURLを生成 + imageURLs := make([]string, 0, len(task.ImageFileIDs)) + for _, fileID := range task.ImageFileIDs { + url, err := e.generateImageURL(fileID) + if err != nil { + e.l.Warn("failed to generate image URL", + zap.Stringer("fileID", fileID), + zap.Error(err)) + continue + } + if url == "" { + continue + } + imageURLs = append(imageURLs, url) + } + if len(imageURLs) == 0 { + return nil + } + + ctx, cancel := context.WithTimeout(context.Background(), e.imageConfig.Timeout) + defer cancel() + + results, err := e.imageClient.ProcessImages(ctx, imageURLs) + if err != nil { + return fmt.Errorf("failed to process images: %w", err) + } + + // OCR結果を結合 + var texts []string + var vectors []esImageVector + for _, r := range results { + if r.Text != "" { + texts = append(texts, r.Text) + } + if r.Vector != nil { + vectors = append(vectors, esImageVector{Vector: r.Vector}) + } + } + + // ESインデックスを部分更新 + update := esImageDocUpdate{ + ImageText: strings.Join(texts, "\n"), + ImageVectors: vectors, + } + + data, err := json.Marshal(map[string]any{"doc": update}) + if err != nil { + return fmt.Errorf("failed to marshal image update: %w", err) + } + + res, err := e.client.Update( + getIndexName(esMessageIndex), + task.MessageID.String(), + bytes.NewReader(data), + ) + if err != nil { + return fmt.Errorf("failed to update image data in ES: %w", err) + } + defer res.Body.Close() + + if res.IsError() { + return fmt.Errorf("ES update error for message %s: %s", task.MessageID, res.String()) + } + + e.l.Info("updated image data for message", + zap.Stringer("messageID", task.MessageID), + zap.Int("imageCount", len(results))) + return nil +} + +// ProcessImagesForMessages バッチ用: 指定メッセージの画像を処理してインデックスを更新する +func (e *esEngine) ProcessImagesForMessages(ctx context.Context, messages []*model.Message) error { + for _, msg := range messages { + parseResult := message.Parse(msg.Text) + attr := e.getAttributes(msg, parseResult) + if len(attr.ImageFileIDs) == 0 { + continue + } + + task := imageProcessTask{ + MessageID: msg.ID, + ImageFileIDs: attr.ImageFileIDs, + } + if err := e.processImageTask(task); err != nil { + e.l.Warn("batch: failed to process image task", + zap.Stringer("messageID", msg.ID), + zap.Error(err)) + // バッチ処理ではエラーをスキップして続行 + continue + } + + // コンテキストキャンセルを確認 + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + } + return nil +} + func syncDeletedMessages(e *esEngine, messages []*model.Message, lastDelete time.Time, lastSynced time.Time) (err error) { bulkIndexer, err := esutil.NewBulkIndexer(esutil.BulkIndexerConfig{ Client: e.client, @@ -316,6 +459,120 @@ func syncDeletedMessages(e *esEngine, messages []*model.Message, lastDelete time return nil } +// GetUnprocessedImageMessageIDs hasImage=trueだがimageVectorが未設定のメッセージIDをESから取得する +func (e *esEngine) GetUnprocessedImageMessageIDs(ctx context.Context) ([]uuid.UUID, error) { + const scrollBatchSize = 1000 + var allIDs []uuid.UUID + var searchAfter []any + + for { + body := m{ + "query": m{ + "bool": m{ + "must": []m{ + {"term": m{"hasImage": m{"value": true}}}, + }, + "must_not": []m{ + {"nested": m{ + "path": "imageVector", + "query": m{"exists": m{"field": "imageVector.vector"}}, + }}, + }, + }, + }, + "size": scrollBatchSize, + "sort": []m{{"_doc": "asc"}}, + "_source": false, + } + if searchAfter != nil { + body["search_after"] = searchAfter + } + + b, err := json.Marshal(body) + if err != nil { + return nil, fmt.Errorf("failed to marshal query: %w", err) + } + + sr, err := e.client.Search( + e.client.Search.WithContext(ctx), + e.client.Search.WithIndex(getIndexName(esMessageIndex)), + e.client.Search.WithBody(bytes.NewReader(b)), + ) + if err != nil { + return nil, fmt.Errorf("failed to search unprocessed images: %w", err) + } + + var res esSearchResponse + err = json.NewDecoder(sr.Body).Decode(&res) + sr.Body.Close() + if err != nil { + return nil, fmt.Errorf("failed to decode search response: %w", err) + } + + hits := res.Hits.Hits + if len(hits) == 0 { + break + } + + for _, hit := range hits { + id, err := uuid.FromString(hit.ID) + if err != nil { + e.l.Warn("invalid message ID in ES", zap.String("id", hit.ID)) + continue + } + allIDs = append(allIDs, id) + } + + lastHit := hits[len(hits)-1] + searchAfter = make([]any, len(lastHit.Sort)) + for i, v := range lastHit.Sort { + searchAfter[i] = v + } + + if len(hits) < scrollBatchSize { + break + } + } + + return allIDs, nil +} + +// ClearImageIndex hasImage=trueの全ドキュメントからimageTextとimageVectorを削除する +func (e *esEngine) ClearImageIndex(ctx context.Context) error { + body := m{ + "query": m{ + "term": m{"hasImage": m{"value": true}}, + }, + "script": m{ + "source": "ctx._source.remove('imageText'); ctx._source.remove('imageVector');", + "lang": "painless", + }, + } + + b, err := json.Marshal(body) + if err != nil { + return fmt.Errorf("failed to marshal update_by_query body: %w", err) + } + + res, err := e.client.UpdateByQuery( + []string{getIndexName(esMessageIndex)}, + e.client.UpdateByQuery.WithContext(ctx), + e.client.UpdateByQuery.WithBody(bytes.NewReader(b)), + e.client.UpdateByQuery.WithRefresh(true), + ) + if err != nil { + return fmt.Errorf("failed to clear image index: %w", err) + } + defer res.Body.Close() + + if res.IsError() { + return fmt.Errorf("ES update_by_query error: %s", res.String()) + } + + e.l.Info("cleared image index data from all messages") + return nil +} + // lastInsertedUpdated esに存在している、updatedAtが一番新しいメッセージの値を取得します func (e *esEngine) lastInsertedUpdated() (time.Time, error) { sr, err := e.client.Search( diff --git a/service/search/image_client.go b/service/search/image_client.go new file mode 100644 index 000000000..e800c9c53 --- /dev/null +++ b/service/search/image_client.go @@ -0,0 +1,99 @@ +package search + +import ( + "context" + "time" + + "github.com/gofrs/uuid" + "go.uber.org/zap" +) + +// ImageSearchConfig 画像検索用外部サーバー設定 +type ImageSearchConfig struct { + // URL 外部サーバーのURL (空の場合は画像検索無効) + URL string + // Timeout リクエストタイムアウト + Timeout time.Duration + // VectorDimension Embeddingベクトルの次元数 + VectorDimension int +} + +// ImageProcessResult 外部サーバーによる画像処理結果 +type ImageProcessResult struct { + // Text OCRで抽出されたテキスト + Text string + // Vector Embeddingベクトル + Vector []float64 +} + +// ImageSearchClient 画像検索用外部サーバークライアントインターフェース +type ImageSearchClient interface { + // EmbedText テキストをベクトルに変換する（検索クエリ用） + EmbedText(ctx context.Context, text string) ([]float64, error) + // ProcessImages 画像URLリストからOCR/Embeddingを実行する + ProcessImages(ctx context.Context, imageURLs []string) ([]ImageProcessResult, error) + // Available 外部サーバーが利用可能かどうか + Available() bool +} + +// stubImageSearchClient スタブ実装 +type stubImageSearchClient struct { + config ImageSearchConfig + l *zap.Logger +} + +// NewImageSearchClient 画像検索クライアントを生成する +// 現在はスタブ実装を返す +func NewImageSearchClient(config ImageSearchConfig, logger *zap.Logger) ImageSearchClient { + if config.URL == "" { + return &nullImageSearchClient{} + } + return &stubImageSearchClient{ + config: config, + l: logger.Named("image-search-client"), + } +} + +func (c *stubImageSearchClient) EmbedText(ctx context.Context, text string) ([]float64, error) { + // TODO: 外部サーバーにテキストを送信してベクトルを取得する + // POST {config.URL}/embed/text + // Request: { "text": "..." } + // Response: { "vector": [0.1, 0.2, ...] } + c.l.Debug("stub: EmbedText called", zap.String("text", text)) + return nil, ErrImageSearchUnavailable +} + +func (c *stubImageSearchClient) ProcessImages(ctx context.Context, imageURLs []string) ([]ImageProcessResult, error) { + // TODO: 外部サーバーに画像の署名付きURLを送信してOCR/Embeddingを取得する + // POST {config.URL}/process/images + // Request: { "imageUrls": ["https://...signed-url..."] } + // Response: { "results": [{ "text": "OCR結果", "vector": [0.1, ...] }] } + c.l.Debug("stub: ProcessImages called", zap.Int("imageCount", len(imageURLs))) + return nil, ErrImageSearchUnavailable +} + +func (c *stubImageSearchClient) Available() bool { + // TODO: 外部サーバーのヘルスチェック + return false +} + +// nullImageSearchClient 画像検索無効時のnull実装 +type nullImageSearchClient struct{} + +func (c *nullImageSearchClient) EmbedText(context.Context, string) ([]float64, error) { + return nil, ErrImageSearchUnavailable +} + +func (c *nullImageSearchClient) ProcessImages(context.Context, []string) ([]ImageProcessResult, error) { + return nil, ErrImageSearchUnavailable +} + +func (c *nullImageSearchClient) Available() bool { + return false +} + +// ImageAttachment 画像添付ファイル情報 +type ImageAttachment struct { + FileID uuid.UUID + Key string +} diff --git a/service/search/null.go b/service/search/null.go index 16f6c9284..51d83a033 100644 --- a/service/search/null.go +++ b/service/search/null.go @@ -1,5 +1,13 @@ package search +import ( + "context" + + "github.com/gofrs/uuid" + + "github.com/traPtitech/traQ/model" +) + var nullE = &nullEngine{} type nullEngine struct{} @@ -20,3 +28,15 @@ func (n *nullEngine) Available() bool { func (n *nullEngine) Close() error { return nil } + +func (n *nullEngine) ProcessImagesForMessages(context.Context, []*model.Message) error { + return ErrServiceUnavailable +} + +func (n *nullEngine) GetUnprocessedImageMessageIDs(context.Context) ([]uuid.UUID, error) { + return nil, ErrServiceUnavailable +} + +func (n *nullEngine) ClearImageIndex(context.Context) error { + return ErrServiceUnavailable +}