diff --git a/internal/handler/dify_retrieval_handler.go b/internal/handler/dify_retrieval_handler.go index 4f7a0c7843..af4bdb2bf5 100644 --- a/internal/handler/dify_retrieval_handler.go +++ b/internal/handler/dify_retrieval_handler.go @@ -55,6 +55,7 @@ type ModelServiceIface interface { // MetadataServiceIface abstracts MetadataService for the Dify handler. type MetadataServiceIface interface { GetFlattedMetaByKBs(kbIDs []string) (common.MetaData, error) + SearchMetadataByKBs(kbIDs []string, size int) (*service.SearchMetadataResponse, error) LabelQuestion(question string, kbs []*entity.Knowledgebase) map[string]float64 } @@ -333,6 +334,21 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) { } } + metaMap := make(map[string]map[string]interface{}) + metaResult, err := h.metadataSvc.SearchMetadataByKBs([]string{kb.ID}, 10000) + if err == nil { + for _, metadata := range metaResult.MetadataRecords { + docID, ok := service.ExtractDocumentID(metadata) + if !ok { + continue + } + metaFields, err := service.ExtractMetaFields(metadata) + if err == nil { + metaMap[docID] = metaFields + } + } + } + // Build response records := make([]difyRecord, 0, len(chunks)) for _, ch := range chunks { @@ -346,8 +362,8 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) { delete(ch, "vector") meta := make(map[string]interface{}) - if doc.MetaFields != nil { - for k, v := range *doc.MetaFields { + if metaFields := metaMap[docID]; metaFields != nil { + for k, v := range metaFields { meta[k] = v } } diff --git a/internal/handler/dify_retrieval_handler_test.go b/internal/handler/dify_retrieval_handler_test.go index f5d196d39a..7dc46cec78 100644 --- a/internal/handler/dify_retrieval_handler_test.go +++ b/internal/handler/dify_retrieval_handler_test.go @@ -31,6 +31,7 @@ import ( "ragflow/internal/engine/types" "ragflow/internal/entity" modelModule "ragflow/internal/entity/models" + "ragflow/internal/service" "ragflow/internal/service/nlp" "github.com/gin-gonic/gin" @@ -82,8 +83,9 @@ func (m *mockModelService) GetChatModel(tenantID, llmID string) (*modelModule.Ch type mockMetadataService struct { MetadataServiceIface - getFlattedMetaFn func(kbIDs []string) (common.MetaData, error) - labelQuestionFn func(question string, kbs []*entity.Knowledgebase) map[string]float64 + getFlattedMetaFn func(kbIDs []string) (common.MetaData, error) + searchMetadataByKBs func(kbIDs []string, size int) (*service.SearchMetadataResponse, error) + labelQuestionFn func(question string, kbs []*entity.Knowledgebase) map[string]float64 } func (m *mockMetadataService) GetFlattedMetaByKBs(kbIDs []string) (common.MetaData, error) { @@ -93,6 +95,17 @@ func (m *mockMetadataService) GetFlattedMetaByKBs(kbIDs []string) (common.MetaDa return common.MetaData{}, nil } +func (m *mockMetadataService) SearchMetadataByKBs(kbIDs []string, size int) (*service.SearchMetadataResponse, error) { + if m.searchMetadataByKBs != nil { + return m.searchMetadataByKBs(kbIDs, size) + } + return &service.SearchMetadataResponse{ + MetadataRecords: []map[string]interface{}{ + {"id": "doc1", "meta_fields": map[string]interface{}{"author": "Zhang San"}}, + }, + }, nil +} + func (m *mockMetadataService) LabelQuestion(question string, kbs []*entity.Knowledgebase) map[string]float64 { if m.labelQuestionFn != nil { return m.labelQuestionFn(question, kbs) @@ -126,7 +139,7 @@ func (m *mockDocDAO) GetByIDs(ids []string) ([]*entity.Document, error) { return m.getByIDsFn(ids) } return []*entity.Document{ - {ID: "doc1", Name: strPtr("Test Doc"), MetaFields: &entity.JSONMap{"author": "Zhang San"}}, + {ID: "doc1", Name: strPtr("Test Doc")}, }, nil }