Skip to content

Commit c0fb8f4

Browse files
kraukrau
committed
feat!: enhance artwork search functionality with R18 type filtering and context management
1 parent 99137e3 commit c0fb8f4

File tree

3 files changed

+42
-11
lines changed

3 files changed

+42
-11
lines changed

‎api/restful/routers/artwork/list.go

+4-4
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ func GetArtworkList(ctx *gin.Context) {
132132
}
133133
if request.Keyword != "" {
134134
if request.Hybrid {
135-
getArtworkListHybrid(ctx, request.Keyword, 0.8, request.Page*request.PageSize, request.PageSize, adapterOption)
135+
getArtworkListHybrid(ctx, request.Keyword, 0.8, request.Page*request.PageSize, request.PageSize, r18Type, adapterOption)
136136
return
137137
}
138138
keywordSlice := common.ParseStringTo2DArray(request.Keyword, ",", ";")
@@ -158,7 +158,7 @@ func GetArtworkList(ctx *gin.Context) {
158158
common.GinErrorResponse(ctx, err, http.StatusInternalServerError, "Failed to get target artwork")
159159
return
160160
}
161-
artworks, err := service.SearchSimilarArtworks(ctx, artworkID.Hex(), request.Page*request.PageSize, request.PageSize, adapterOption)
161+
artworks, err := service.SearchSimilarArtworks(ctx, artworkID.Hex(), request.Page*request.PageSize, request.PageSize, r18Type, adapterOption)
162162
if err != nil {
163163
common.GinErrorResponse(ctx, err, http.StatusInternalServerError, "Failed to get similar artworks")
164164
return
@@ -238,8 +238,8 @@ func getArtworkListByKeyword(ctx *gin.Context, keywordSlice [][]string, r18Type
238238
ctx.JSON(http.StatusOK, ResponseFromArtworks(artworks, ctx.GetBool("auth")))
239239
}
240240

241-
func getArtworkListHybrid(ctx *gin.Context, queryText string, hybridSemanticRatio float64, offset, limit int64, adapterOption ...*types.AdapterOption) {
242-
artworks, err := service.HybridSearchArtworks(ctx, queryText, hybridSemanticRatio, offset, limit, adapterOption...)
241+
func getArtworkListHybrid(ctx *gin.Context, queryText string, hybridSemanticRatio float64, offset, limit int64, r18Type types.R18Type, adapterOption ...*types.AdapterOption) {
242+
artworks, err := service.HybridSearchArtworks(ctx, queryText, hybridSemanticRatio, offset, limit, r18Type, adapterOption...)
243243
if err != nil {
244244
if errors.Is(err, mongo.ErrNoDocuments) {
245245
common.GinErrorResponse(ctx, err, http.StatusNotFound, "Artworks not found")

‎service/artwork_search.go

+26-2
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,27 @@ import (
1414
"go.mongodb.org/mongo-driver/bson/primitive"
1515
)
1616

17-
func HybridSearchArtworks(ctx context.Context, queryText string, hybridSemanticRatio float64, offset, limit int64, options ...*types.AdapterOption) ([]*types.Artwork, error) {
17+
func HybridSearchArtworks(ctx context.Context, queryText string, hybridSemanticRatio float64, offset, limit int64, r18 types.R18Type, options ...*types.AdapterOption) ([]*types.Artwork, error) {
1818
if common.MeilisearchClient == nil {
1919
return nil, errs.ErrSearchEngineUnavailable
2020
}
21+
22+
var filter string
23+
switch r18 {
24+
case types.R18TypeAll:
25+
filter = ""
26+
case types.R18TypeNone:
27+
filter = "r18 = false"
28+
case types.R18TypeOnly:
29+
filter = "r18 = true"
30+
}
31+
2132
index := common.MeilisearchClient.Index(config.Cfg.Search.MeiliSearch.Index)
2233
resp, err := index.SearchWithContext(ctx, queryText, &meilisearch.SearchRequest{
2334
Offset: offset,
2435
Limit: limit,
2536
AttributesToRetrieve: []string{"id"},
37+
Filter: filter,
2638
Hybrid: &meilisearch.SearchRequestHybrid{
2739
Embedder: config.Cfg.Search.MeiliSearch.Embedder,
2840
SemanticRatio: hybridSemanticRatio,
@@ -56,10 +68,21 @@ func HybridSearchArtworks(ctx context.Context, queryText string, hybridSemanticR
5668
return adapter.ConvertToArtworks(ctx, artworkModels, options...)
5769
}
5870

59-
func SearchSimilarArtworks(ctx context.Context, artworkIdStr string, offset, limit int64, options ...*types.AdapterOption) ([]*types.Artwork, error) {
71+
func SearchSimilarArtworks(ctx context.Context, artworkIdStr string, offset, limit int64, r18 types.R18Type, options ...*types.AdapterOption) ([]*types.Artwork, error) {
6072
if common.MeilisearchClient == nil {
6173
return nil, errs.ErrSearchEngineUnavailable
6274
}
75+
76+
var filter string
77+
switch r18 {
78+
case types.R18TypeAll:
79+
filter = ""
80+
case types.R18TypeNone:
81+
filter = "r18 = false"
82+
case types.R18TypeOnly:
83+
filter = "r18 = true"
84+
}
85+
6386
index := common.MeilisearchClient.Index(config.Cfg.Search.MeiliSearch.Index)
6487
var resp meilisearch.SimilarDocumentResult
6588
if err := index.SearchSimilarDocumentsWithContext(ctx, &meilisearch.SimilarDocumentQuery{
@@ -68,6 +91,7 @@ func SearchSimilarArtworks(ctx context.Context, artworkIdStr string, offset, lim
6891
Embedder: config.Cfg.Search.MeiliSearch.Embedder,
6992
Offset: offset,
7093
Limit: limit,
94+
Filter: filter,
7195
}, &resp); err != nil {
7296
return nil, err
7397
}

‎service/change_stream.go

+12-5
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package service
22

33
import (
44
"context"
5+
"time"
56

67
"github.com/bytedance/sonic"
78
"github.com/krau/ManyACG/adapter"
@@ -99,7 +100,9 @@ func (m *artworkSyncManager) ProcessArtworkUpdateEvent(event bson.M) {
99100
common.Logger.Errorf("decode artwork from event error: %s", err)
100101
return
101102
}
102-
searchDoc, err := adapter.ConvertToSearchDoc(context.Background(), artwork)
103+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
104+
defer cancel()
105+
searchDoc, err := adapter.ConvertToSearchDoc(ctx, artwork)
103106
if err != nil {
104107
common.Logger.Errorf("convert to search doc error: %s", err)
105108
return
@@ -109,7 +112,7 @@ func (m *artworkSyncManager) ProcessArtworkUpdateEvent(event bson.M) {
109112
common.Logger.Errorf("marshal search doc error: %s", err)
110113
return
111114
}
112-
task, err := common.MeilisearchClient.Index(config.Cfg.Search.MeiliSearch.Index).UpdateDocuments(artworkJSON)
115+
task, err := common.MeilisearchClient.Index(config.Cfg.Search.MeiliSearch.Index).UpdateDocumentsWithContext(ctx, artworkJSON)
113116
if err != nil {
114117
common.Logger.Errorf("update artwork to meilisearch error: %s", err)
115118
return
@@ -119,7 +122,9 @@ func (m *artworkSyncManager) ProcessArtworkUpdateEvent(event bson.M) {
119122

120123
func (m *artworkSyncManager) ProcessArtworkDeleteEvent(event bson.M) {
121124
docID := event["documentKey"].(bson.M)["_id"].(primitive.ObjectID).Hex()
122-
task, err := common.MeilisearchClient.Index(config.Cfg.Search.MeiliSearch.Index).DeleteDocument(docID)
125+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
126+
defer cancel()
127+
task, err := common.MeilisearchClient.Index(config.Cfg.Search.MeiliSearch.Index).DeleteDocumentWithContext(ctx, docID)
123128
if err != nil {
124129
common.Logger.Errorf("delete artwork from meilisearch error: %s", err)
125130
return
@@ -134,7 +139,9 @@ func (m *artworkSyncManager) ProcessArtworkReplaceEvent(event bson.M) {
134139
common.Logger.Errorf("decode artwork from event error: %s", err)
135140
return
136141
}
137-
searchDoc, err := adapter.ConvertToSearchDoc(context.Background(), artwork)
142+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
143+
defer cancel()
144+
searchDoc, err := adapter.ConvertToSearchDoc(ctx, artwork)
138145
if err != nil {
139146
common.Logger.Errorf("convert to search doc error: %s", err)
140147
return
@@ -144,7 +151,7 @@ func (m *artworkSyncManager) ProcessArtworkReplaceEvent(event bson.M) {
144151
common.Logger.Errorf("marshal search doc error: %s", err)
145152
return
146153
}
147-
task, err := common.MeilisearchClient.Index(config.Cfg.Search.MeiliSearch.Index).UpdateDocuments(artworkJSON)
154+
task, err := common.MeilisearchClient.Index(config.Cfg.Search.MeiliSearch.Index).UpdateDocumentsWithContext(ctx, artworkJSON)
148155
if err != nil {
149156
common.Logger.Errorf("update artwork to meilisearch error: %s", err)
150157
return

0 commit comments

Comments
 (0)