From 92e8eb5fe7f6b8f8c6466e1b68e2bb634391f6a8 Mon Sep 17 00:00:00 2001 From: Haruko386 Date: Thu, 2 Jul 2026 15:57:07 +0800 Subject: [PATCH] fix: add search keywords and filter for datasets-search (#16550) --- internal/dao/document.go | 173 ++++++++++++++- internal/engine/infinity/sql_test.go | 9 +- internal/handler/document.go | 300 ++++++++++++++++++++++++++- internal/handler/document_test.go | 124 ++++++++++- internal/service/document.go | 78 ++++++- 5 files changed, 661 insertions(+), 23 deletions(-) diff --git a/internal/dao/document.go b/internal/dao/document.go index af5eeb3491..d016a4d325 100644 --- a/internal/dao/document.go +++ b/internal/dao/document.go @@ -18,6 +18,9 @@ package dao import ( "ragflow/internal/entity" + "strings" + + "gorm.io/gorm" ) // DocumentDAO document data access object @@ -86,29 +89,177 @@ func (dao *DocumentDAO) List(offset, limit int) ([]*entity.Document, int64, erro return documents, total, err } +// DocumentListOptions contains filters for listing documents in a dataset. +type DocumentListOptions struct { + KbID string + Keywords string + RunStatuses []string + Types []string + Suffixes []string + Name string + DocIDs []string + DocIDFilterApplied bool + CreateTimeFrom int64 + CreateTimeTo int64 + OrderBy string + Desc bool + Offset int + Limit int +} + // ListByKBID list documents by knowledge base ID -func (dao *DocumentDAO) ListByKBID(kbID string, offset, limit int) ([]*entity.DocumentListItem, int64, error) { +func (dao *DocumentDAO) ListByKBID(kbID, keywords string, offset, limit int) ([]*entity.DocumentListItem, int64, error) { + return dao.ListByKBIDWithOptions(DocumentListOptions{ + KbID: kbID, + Keywords: keywords, + OrderBy: "create_time", + Desc: true, + Offset: offset, + Limit: limit, + }) +} + +// ListByKBIDWithOptions lists documents by knowledge base ID with filters. +func (dao *DocumentDAO) ListByKBIDWithOptions(opts DocumentListOptions) ([]*entity.DocumentListItem, int64, error) { var documents []*entity.DocumentListItem var total int64 - if err := DB.Model(&entity.Document{}).Where("kb_id = ?", kbID).Count(&total).Error; err != nil { - return nil, 0, err - } - - err := DB.Table("document"). + listQuery := DB.Table("document"). Select(`document.*, user_canvas.title as pipeline_name, user.nickname`). Joins("JOIN file2document ON file2document.document_id = document.id"). Joins("JOIN file ON file.id = file2document.file_id"). Joins("LEFT JOIN user_canvas ON document.pipeline_id = user_canvas.id"). - Joins("LEFT JOIN user ON document.created_by = user.id"). - Where("document.kb_id = ?", kbID). - Order("document.create_time DESC"). - Offset(offset). - Limit(limit). + Joins("LEFT JOIN user ON document.created_by = user.id") + + listQuery = applyDocumentListFilters(listQuery, opts, true) + countQuery := applyDocumentListFilters(DB.Model(&entity.Document{}), opts, false) + + if err := countQuery.Count(&total).Error; err != nil { + return nil, 0, err + } + + orderBy := documentListOrderColumn(opts.OrderBy) + if opts.Desc { + orderBy += " DESC" + } else { + orderBy += " ASC" + } + + err := listQuery. + Order(orderBy). + Offset(opts.Offset). + Limit(opts.Limit). Scan(&documents).Error return documents, total, err } +// GetFilterByKBID returns aggregate filter counts for documents in a dataset. +func (dao *DocumentDAO) GetFilterByKBID(opts DocumentListOptions) (map[string]interface{}, int64, error) { + var rows []struct { + ID string `gorm:"column:id"` + Run *string `gorm:"column:run"` + Suffix string `gorm:"column:suffix"` + } + + query := DB.Table("document"). + Select("document.id, document.run, document.suffix"). + Joins("JOIN file2document ON file2document.document_id = document.id"). + Joins("JOIN file ON file.id = file2document.file_id") + query = applyDocumentListFilters(query, opts, true) + + if err := query.Scan(&rows).Error; err != nil { + return nil, 0, err + } + + suffixCounter := map[string]int64{} + runStatusCounter := map[string]int64{} + for _, row := range rows { + if row.Suffix != "" { + suffixCounter[row.Suffix]++ + } + if row.Run != nil { + runStatusCounter[*row.Run]++ + } + } + + return map[string]interface{}{ + "suffix": suffixCounter, + "run_status": runStatusCounter, + "metadata": map[string]interface{}{}, + }, int64(len(rows)), nil +} + +// ListIDsByKBIDWithOptions lists matching document IDs without pagination. +func (dao *DocumentDAO) ListIDsByKBIDWithOptions(opts DocumentListOptions) ([]string, error) { + var ids []string + query := DB.Table("document"). + Select("document.id"). + Joins("JOIN file2document ON file2document.document_id = document.id"). + Joins("JOIN file ON file.id = file2document.file_id") + query = applyDocumentListFilters(query, opts, true) + if err := query.Scan(&ids).Error; err != nil { + return nil, err + } + return ids, nil +} + +func applyDocumentListFilters(query *gorm.DB, opts DocumentListOptions, qualified bool) *gorm.DB { + column := func(name string) string { + if qualified { + return "document." + name + } + return name + } + + query = query.Where(column("kb_id")+" = ?", opts.KbID) + if strings.TrimSpace(opts.Keywords) != "" { + query = query.Where("LOWER("+column("name")+") LIKE ?", "%"+strings.ToLower(strings.TrimSpace(opts.Keywords))+"%") + } + if len(opts.RunStatuses) > 0 { + query = query.Where(column("run")+" IN ?", opts.RunStatuses) + } + if len(opts.Types) > 0 { + query = query.Where(column("type")+" IN ?", opts.Types) + } + if len(opts.Suffixes) > 0 { + query = query.Where(column("suffix")+" IN ?", opts.Suffixes) + } + if opts.Name != "" { + query = query.Where(column("name")+" = ?", opts.Name) + } + if opts.DocIDFilterApplied { + if len(opts.DocIDs) == 0 { + query = query.Where("1 = 0") + } else { + query = query.Where(column("id")+" IN ?", opts.DocIDs) + } + } + if opts.CreateTimeFrom > 0 { + query = query.Where(column("create_time")+" >= ?", opts.CreateTimeFrom) + } + if opts.CreateTimeTo > 0 { + query = query.Where(column("create_time")+" <= ?", opts.CreateTimeTo) + } + return query +} + +func documentListOrderColumn(orderBy string) string { + switch orderBy { + case "update_time": + return "document.update_time" + case "name": + return "document.name" + case "size": + return "document.size" + case "type": + return "document.type" + case "run": + return "document.run" + default: + return "document.create_time" + } +} + // GetByKBID retrieves all documents in a knowledge base ordered by create time. func (dao *DocumentDAO) GetByKBID(kbID string) ([]*entity.Document, int64, error) { var documents []*entity.Document diff --git a/internal/engine/infinity/sql_test.go b/internal/engine/infinity/sql_test.go index 4c7e7b98ac..727d285a92 100644 --- a/internal/engine/infinity/sql_test.go +++ b/internal/engine/infinity/sql_test.go @@ -327,20 +327,15 @@ func TestLoadFieldMapping_MissingFileReturnsEmpty(t *testing.T) { func TestLoadFieldMapping_ParsesAliases(t *testing.T) { // Write a temporary mapping file. dir := t.TempDir() - mappingPath := filepath.Join(dir, "test_mapping.json") contents := `{ "docnm": {"type": "varchar", "comment": "docnm_kwd, title_tks, title_sm_tks"}, "content": {"type": "varchar", "comment": "content_with_weight, content_ltks"}, "plain": {"type": "varchar"} }` - if err := os.WriteFile(mappingPath, []byte(contents), 0o644); err != nil { - t.Fatalf("write mapping: %v", err) - } // Set RAG_PROJECT_BASE to the temp dir's parent so loadFieldMapping // finds the file at /conf/. - os.Setenv("RAG_PROJECT_BASE", dir) - defer os.Unsetenv("RAG_PROJECT_BASE") + t.Setenv("RAG_PROJECT_BASE", dir) // Need to create conf/ subdir. if err := os.MkdirAll(filepath.Join(dir, "conf"), 0o755); err != nil { @@ -384,6 +379,8 @@ func TestLoadFieldMapping_EmptyNameDefaultsToInfinityMappingJSON(t *testing.T) { // Empty name → defaults to "infinity_mapping.json" (line 145). // We just verify the function doesn't panic and the file-not-found // path is taken silently. + t.Setenv("RAG_PROJECT_BASE", t.TempDir()) + a2a, r2a, err := loadFieldMapping("") if err != nil { t.Fatalf("empty name: %v", err) diff --git a/internal/handler/document.go b/internal/handler/document.go index d0d2442917..0740a6b17c 100644 --- a/internal/handler/document.go +++ b/internal/handler/document.go @@ -35,6 +35,7 @@ import ( "github.com/gin-gonic/gin" + "ragflow/internal/dao" "ragflow/internal/service" ) @@ -50,7 +51,11 @@ type documentServiceIface interface { ParseDocuments(datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) StopParseDocuments(datasetID string, docIDs []string) (map[string]interface{}, error) ListDocuments(page, pageSize int) ([]*service.DocumentResponse, int64, error) - ListDocumentsByDatasetID(kbID string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) + ListDocumentsByDatasetID(kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) + ListDocumentsByDatasetIDWithOptions(opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) + ListDocumentIDsByDatasetIDWithOptions(opts dao.DocumentListOptions) ([]string, error) + GetDocumentFiltersByDatasetID(opts dao.DocumentListOptions) (map[string]interface{}, int64, error) + GetMetadataByKBs(kbIDs []string) (map[string]interface{}, error) GetDocumentsByAuthorID(authorID, page, pageSize int) ([]*service.DocumentResponse, int64, error) GetThumbnails(userID string, docIDs []string) (map[string]string, error) GetDocumentImage(imageID string) ([]byte, error) @@ -521,8 +526,48 @@ func (h *DocumentHandler) ListDocuments(c *gin.Context) { pageSize = 10 } + opts, errMsg := parseDocumentListOptions(c, datasetID) + if errMsg != "" { + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeDataError, + "message": errMsg, + "data": map[string]interface{}{"total": 0, "docs": []interface{}{}}, + }) + return + } + opts, errMsg = h.applyDocumentMetadataFilter(c, opts) + if errMsg != "" { + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeDataError, + "message": errMsg, + "data": map[string]interface{}{"total": 0, "docs": []interface{}{}}, + }) + return + } + + if c.Query("type") == "filter" { + filters, total, err := h.documentService.GetDocumentFiltersByDatasetID(opts) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeExceptionError, + "message": "failed to get document filters", + "data": map[string]interface{}{"total": 0, "filter": map[string]interface{}{}}, + }) + return + } + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeSuccess, + "message": "success", + "data": gin.H{ + "total": total, + "filter": filters, + }, + }) + return + } + // Use kbID to filter documents - documents, total, err := h.documentService.ListDocumentsByDatasetID(datasetID, page, pageSize) + documents, total, err := h.documentService.ListDocumentsByDatasetIDWithOptions(opts, page, pageSize) if err != nil { c.JSON(http.StatusOK, gin.H{ "code": 1, @@ -552,6 +597,257 @@ func (h *DocumentHandler) ListDocuments(c *gin.Context) { }) } +func parseDocumentListOptions(c *gin.Context, datasetID string) (dao.DocumentListOptions, string) { + opts := dao.DocumentListOptions{ + KbID: datasetID, + Keywords: c.Query("keywords"), + OrderBy: c.DefaultQuery("orderby", "create_time"), + Desc: strings.ToLower(strings.TrimSpace(c.DefaultQuery("desc", "true"))) != "false", + Suffixes: queryValues(c, "suffix"), + Types: queryValues(c, "types"), + } + + opts.RunStatuses = normalizeRunStatusFilter(queryValues(c, "run", "run_status")) + if len(queryValues(c, "run", "run_status")) > 0 && len(opts.RunStatuses) == 0 { + return opts, "Invalid filter run status conditions" + } + + opts.Name = c.Query("name") + docID := c.Query("id") + docIDs := queryValues(c, "ids") + if docID != "" && len(docIDs) > 0 { + return opts, fmt.Sprintf("Should not provide both 'id':%s and 'ids'%v", docID, docIDs) + } + if docID != "" { + opts.DocIDs = []string{docID} + opts.DocIDFilterApplied = true + } else if len(docIDs) > 0 { + opts.DocIDs = docIDs + opts.DocIDFilterApplied = true + } + + if v := c.Query("create_time_from"); v != "" { + createTimeFrom, err := strconv.ParseInt(v, 10, 64) + if err != nil { + return opts, "create_time_from must be an integer" + } + opts.CreateTimeFrom = createTimeFrom + } + if v := c.Query("create_time_to"); v != "" { + createTimeTo, err := strconv.ParseInt(v, 10, 64) + if err != nil { + return opts, "create_time_to must be an integer" + } + opts.CreateTimeTo = createTimeTo + } + + return opts, "" +} + +func (h *DocumentHandler) applyDocumentMetadataFilter(c *gin.Context, opts dao.DocumentListOptions) (dao.DocumentListOptions, string) { + metadata, err := parseMetadataQuery(c.Request.URL.Query()) + if err != nil { + return opts, err.Error() + } + returnEmptyMetadata := strings.ToLower(strings.TrimSpace(c.Query("return_empty_metadata"))) == "true" + if !returnEmptyMetadata && len(metadata) == 0 { + return opts, "" + } + + candidateIDs, err := h.documentService.ListDocumentIDsByDatasetIDWithOptions(opts) + if err != nil { + return opts, "failed to get documents" + } + candidateSet := stringSet(candidateIDs) + + metadataByKey, err := h.documentService.GetMetadataByKBs([]string{opts.KbID}) + if err != nil { + return opts, err.Error() + } + + docIDsWithMetadata := map[string]bool{} + matchedIDs := map[string]bool{} + firstMetadataKey := true + for key, values := range metadata { + valueMatches := map[string]bool{} + rawValues, _ := metadataByKey[key].(map[string][]string) + for _, value := range values { + for _, docID := range rawValues[value] { + valueMatches[docID] = true + docIDsWithMetadata[docID] = true + } + } + if firstMetadataKey { + matchedIDs = valueMatches + firstMetadataKey = false + } else { + matchedIDs = intersectStringSets(matchedIDs, valueMatches) + } + } + if returnEmptyMetadata { + for _, rawValue := range metadataByKey { + values, _ := rawValue.(map[string][]string) + for _, docIDs := range values { + for _, docID := range docIDs { + docIDsWithMetadata[docID] = true + } + } + } + } + + filteredIDs := make([]string, 0) + if returnEmptyMetadata { + for _, docID := range candidateIDs { + if !docIDsWithMetadata[docID] { + filteredIDs = append(filteredIDs, docID) + } + } + } else { + for docID := range matchedIDs { + if candidateSet[docID] { + filteredIDs = append(filteredIDs, docID) + } + } + } + + opts.DocIDs = filteredIDs + opts.DocIDFilterApplied = true + return opts, "" +} + +func parseMetadataQuery(values url.Values) (map[string][]string, error) { + metadata := map[string][]string{} + if raw := strings.TrimSpace(values.Get("metadata")); raw != "" { + var parsed map[string]interface{} + if err := json.Unmarshal([]byte(raw), &parsed); err != nil { + return nil, fmt.Errorf("metadata must be valid JSON") + } + for key, value := range parsed { + for _, item := range interfaceToStringSlice(value) { + metadata[key] = append(metadata[key], item) + } + } + } + + for key, vals := range values { + if !strings.HasPrefix(key, "metadata[") || !strings.HasSuffix(key, "]") { + continue + } + name := strings.TrimPrefix(key, "metadata[") + if end := strings.Index(name, "]"); end >= 0 { + name = name[:end] + } + if name == "" || name == "empty_metadata" { + continue + } + for _, value := range vals { + for _, item := range interfaceToStringSlice(value) { + metadata[name] = append(metadata[name], item) + } + } + } + return metadata, nil +} + +func interfaceToStringSlice(value interface{}) []string { + switch typed := value.(type) { + case []interface{}: + out := make([]string, 0, len(typed)) + for _, item := range typed { + if item == nil { + continue + } + if s := strings.TrimSpace(fmt.Sprintf("%v", item)); s != "" { + out = append(out, s) + } + } + return out + case []string: + out := make([]string, 0, len(typed)) + for _, item := range typed { + if s := strings.TrimSpace(item); s != "" { + out = append(out, s) + } + } + return out + case string: + if strings.TrimSpace(typed) == "" { + return nil + } + return []string{strings.TrimSpace(typed)} + default: + if value == nil { + return nil + } + return []string{fmt.Sprintf("%v", value)} + } +} + +func stringSet(values []string) map[string]bool { + out := make(map[string]bool, len(values)) + for _, value := range values { + out[value] = true + } + return out +} + +func intersectStringSets(left, right map[string]bool) map[string]bool { + out := make(map[string]bool) + for value := range left { + if right[value] { + out[value] = true + } + } + return out +} + +func queryValues(c *gin.Context, names ...string) []string { + values := make([]string, 0) + for _, name := range names { + values = append(values, c.QueryArray(name)...) + values = append(values, c.QueryArray(name+"[]")...) + } + out := make([]string, 0, len(values)) + for _, value := range values { + if trimmed := strings.TrimSpace(value); trimmed != "" { + out = append(out, trimmed) + } + } + return out +} + +func normalizeRunStatusFilter(statuses []string) []string { + if len(statuses) == 0 { + return nil + } + statusTextToNumeric := map[string]string{ + "UNSTART": string(entity.TaskStatusUnstart), + "RUNNING": string(entity.TaskStatusRunning), + "CANCEL": string(entity.TaskStatusCancel), + "DONE": string(entity.TaskStatusDone), + "FAIL": string(entity.TaskStatusFail), + } + validStatuses := map[string]bool{ + string(entity.TaskStatusUnstart): true, + string(entity.TaskStatusRunning): true, + string(entity.TaskStatusCancel): true, + string(entity.TaskStatusDone): true, + string(entity.TaskStatusFail): true, + } + out := make([]string, 0, len(statuses)) + for _, status := range statuses { + normalized := statusTextToNumeric[strings.ToUpper(status)] + if normalized == "" { + normalized = status + } + if !validStatuses[normalized] { + return nil + } + out = append(out, normalized) + } + return out +} + func (h *DocumentHandler) UploadDocuments(c *gin.Context) { user, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { diff --git a/internal/handler/document_test.go b/internal/handler/document_test.go index 22be6fdc82..ef2d06d409 100644 --- a/internal/handler/document_test.go +++ b/internal/handler/document_test.go @@ -59,6 +59,12 @@ type fakeDocumentService struct { ingestErr error ingestUserID string ingestReq *service.IngestDocumentRequest + listOpts dao.DocumentListOptions + filterOpts dao.DocumentListOptions + filterResult map[string]interface{} + filterTotal int64 + listIDs []string + metadataByKBs map[string]interface{} } func (f *fakeDocumentService) Ingest(userID string, req *service.IngestDocumentRequest) (common.ErrorCode, error) { @@ -143,9 +149,30 @@ func (f *fakeDocumentService) StopParseDocuments(datasetID string, docIDs []stri func (f *fakeDocumentService) ListDocuments(page, pageSize int) ([]*service.DocumentResponse, int64, error) { return nil, 0, nil } -func (f *fakeDocumentService) ListDocumentsByDatasetID(kbID string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { +func (f *fakeDocumentService) ListDocumentsByDatasetID(kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { return nil, 0, nil } +func (f *fakeDocumentService) ListDocumentsByDatasetIDWithOptions(opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { + f.listOpts = opts + return nil, 0, nil +} +func (f *fakeDocumentService) ListDocumentIDsByDatasetIDWithOptions(opts dao.DocumentListOptions) ([]string, error) { + f.listOpts = opts + return f.listIDs, nil +} +func (f *fakeDocumentService) GetDocumentFiltersByDatasetID(opts dao.DocumentListOptions) (map[string]interface{}, int64, error) { + f.filterOpts = opts + if f.filterResult != nil { + return f.filterResult, f.filterTotal, nil + } + return map[string]interface{}{}, 0, nil +} +func (f *fakeDocumentService) GetMetadataByKBs(kbIDs []string) (map[string]interface{}, error) { + if f.metadataByKBs != nil { + return f.metadataByKBs, nil + } + return map[string]interface{}{}, nil +} func (f *fakeDocumentService) BatchUpdateDocumentStatus(userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) { return map[string]interface{}{}, common.CodeSuccess, nil } @@ -761,6 +788,101 @@ func setupHandlerAccessDB(t *testing.T) *gorm.DB { // sptr returns a pointer to the given string (copy of service test helper). func sptr(s string) *string { return &s } +func TestListDocumentsHandler_FilterRequestUsesQueryFilters(t *testing.T) { + db := setupHandlerAccessDB(t) + orig := dao.DB + dao.DB = db + t.Cleanup(func() { dao.DB = orig }) + + gin.SetMode(gin.TestMode) + + fake := &fakeDocumentService{ + filterResult: map[string]interface{}{ + "suffix": map[string]int64{"pdf": 2}, + "run_status": map[string]int64{"3": 2}, + "metadata": map[string]interface{}{}, + }, + filterTotal: 2, + } + h := &DocumentHandler{ + documentService: fake, + datasetService: service.NewDatasetService(), + } + + c, w := setupGinContextWithUser("GET", "/api/v1/datasets/ds-1/documents?type=filter&keywords=report&suffix=pdf&run=DONE&types=doc&desc=false", "") + c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}} + + h.ListDocuments(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + if fake.filterOpts.KbID != "ds-1" { + t.Fatalf("expected dataset filter ds-1, got %q", fake.filterOpts.KbID) + } + if fake.filterOpts.Keywords != "report" { + t.Fatalf("expected keywords report, got %q", fake.filterOpts.Keywords) + } + if len(fake.filterOpts.Suffixes) != 1 || fake.filterOpts.Suffixes[0] != "pdf" { + t.Fatalf("expected suffix pdf, got %#v", fake.filterOpts.Suffixes) + } + if len(fake.filterOpts.RunStatuses) != 1 || fake.filterOpts.RunStatuses[0] != string(entity.TaskStatusDone) { + t.Fatalf("expected run DONE to map to %q, got %#v", string(entity.TaskStatusDone), fake.filterOpts.RunStatuses) + } + if len(fake.filterOpts.Types) != 1 || fake.filterOpts.Types[0] != "doc" { + t.Fatalf("expected type doc, got %#v", fake.filterOpts.Types) + } + if fake.filterOpts.Desc { + t.Fatal("expected desc=false to be parsed") + } + + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("invalid json response: %v", err) + } + data := resp["data"].(map[string]interface{}) + if data["total"] != float64(2) { + t.Fatalf("expected total 2, got %v", data["total"]) + } +} + +func TestListDocumentsHandler_MetadataFilterNarrowsDocumentIDs(t *testing.T) { + db := setupHandlerAccessDB(t) + orig := dao.DB + dao.DB = db + t.Cleanup(func() { dao.DB = orig }) + + gin.SetMode(gin.TestMode) + + fake := &fakeDocumentService{ + listIDs: []string{"doc-1", "doc-2", "doc-3"}, + metadataByKBs: map[string]interface{}{ + "author": map[string][]string{ + "Alice": []string{"doc-2", "doc-4"}, + }, + }, + } + h := &DocumentHandler{ + documentService: fake, + datasetService: service.NewDatasetService(), + } + + c, w := setupGinContextWithUser("GET", "/api/v1/datasets/ds-1/documents?metadata[author][]=Alice", "") + c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}} + + h.ListDocuments(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + if !fake.listOpts.DocIDFilterApplied { + t.Fatal("expected metadata filter to apply doc id filter") + } + if len(fake.listOpts.DocIDs) != 1 || fake.listOpts.DocIDs[0] != "doc-2" { + t.Fatalf("expected metadata filter to keep doc-2, got %#v", fake.listOpts.DocIDs) + } +} + func TestStopParseDocumentsHandler_Success(t *testing.T) { db := setupHandlerAccessDB(t) orig := dao.DB diff --git a/internal/service/document.go b/internal/service/document.go index 76281fdced..adeb55f8e1 100644 --- a/internal/service/document.go +++ b/internal/service/document.go @@ -1011,9 +1011,23 @@ func (s *DocumentService) BatchUpdateDocumentStatus(userID, datasetID, status st } // ListDocumentsByDatasetID list documents by knowledge base ID -func (s *DocumentService) ListDocumentsByDatasetID(kbID string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { - offset := (page - 1) * pageSize - documents, total, err := s.documentDAO.ListByKBID(kbID, offset, pageSize) +func (s *DocumentService) ListDocumentsByDatasetID(kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { + return s.ListDocumentsByDatasetIDWithOptions(dao.DocumentListOptions{ + KbID: kbID, + Keywords: keywords, + OrderBy: "create_time", + Desc: true, + }, page, pageSize) +} + +// ListDocumentsByDatasetIDWithOptions lists documents by knowledge base ID with filters. +func (s *DocumentService) ListDocumentsByDatasetIDWithOptions(opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) { + opts.Offset = (page - 1) * pageSize + opts.Limit = pageSize + if opts.OrderBy == "" { + opts.OrderBy = "create_time" + } + documents, total, err := s.documentDAO.ListByKBIDWithOptions(opts) if err != nil { return nil, 0, err } @@ -1026,6 +1040,64 @@ func (s *DocumentService) ListDocumentsByDatasetID(kbID string, page, pageSize i return responses, total, nil } +// GetDocumentFiltersByDatasetID returns aggregate filter values for documents in a dataset. +func (s *DocumentService) GetDocumentFiltersByDatasetID(opts dao.DocumentListOptions) (map[string]interface{}, int64, error) { + filters, total, err := s.documentDAO.GetFilterByKBID(opts) + if err != nil { + return nil, 0, err + } + docIDs, err := s.documentDAO.ListIDsByKBIDWithOptions(opts) + if err != nil { + return nil, 0, err + } + metadataFilter, err := s.getDocumentMetadataFilter(opts.KbID, docIDs) + if err != nil { + return nil, 0, err + } + filters["metadata"] = metadataFilter + return filters, total, nil +} + +func (s *DocumentService) getDocumentMetadataFilter(kbID string, docIDs []string) (map[string]interface{}, error) { + metadataByKey, err := s.GetMetadataByKBs([]string{kbID}) + if err != nil { + return nil, err + } + candidateSet := make(map[string]bool, len(docIDs)) + for _, docID := range docIDs { + candidateSet[docID] = true + } + + metadataCounter := map[string]interface{}{} + docIDsWithMetadata := map[string]bool{} + for key, rawValues := range metadataByKey { + values, ok := rawValues.(map[string][]string) + if !ok { + continue + } + valueCounter := map[string]int64{} + for value, valueDocIDs := range values { + for _, docID := range valueDocIDs { + if !candidateSet[docID] { + continue + } + valueCounter[value]++ + docIDsWithMetadata[docID] = true + } + } + if len(valueCounter) > 0 { + metadataCounter[key] = valueCounter + } + } + metadataCounter["empty_metadata"] = map[string]int64{"true": int64(len(docIDs) - len(docIDsWithMetadata))} + return metadataCounter, nil +} + +// ListDocumentIDsByDatasetIDWithOptions lists matching document IDs without pagination. +func (s *DocumentService) ListDocumentIDsByDatasetIDWithOptions(opts dao.DocumentListOptions) ([]string, error) { + return s.documentDAO.ListIDsByKBIDWithOptions(opts) +} + // GetDocumentsByAuthorID get documents by author ID func (s *DocumentService) GetDocumentsByAuthorID(authorID, page, pageSize int) ([]*DocumentResponse, int64, error) { offset := (page - 1) * pageSize