fix: add search keywords and filter for datasets-search (#16550)

This commit is contained in:
Haruko386
2026-07-02 15:57:07 +08:00
committed by GitHub
parent 4130091b69
commit 92e8eb5fe7
5 changed files with 661 additions and 23 deletions

View File

@@ -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 {

View File

@@ -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