From b777e5029139ba875ba5cb59a309312ed5a0be03 Mon Sep 17 00:00:00 2001 From: Haruko386 Date: Mon, 22 Jun 2026 18:14:01 +0800 Subject: [PATCH] feat[Go]: implement api /api/v1/datasets//chunks POST (#16067) ### What problem does this PR solve? As title ### Type of change - [x] New Feature (non-breaking change which adds functionality) --- internal/dao/task.go | 10 +- internal/handler/chunk.go | 58 ++- internal/handler/chunk_test.go | 3 + internal/router/router.go | 1 + internal/service/chunk/chunk.go | 607 ++++++++++++++++++++++++++- internal/service/chunk/chunk_test.go | 579 +++++++++++++++++++++++++ internal/service/chunk_types.go | 9 +- 7 files changed, 1259 insertions(+), 8 deletions(-) diff --git a/internal/dao/task.go b/internal/dao/task.go index 2f34f682bb..2870657bd6 100644 --- a/internal/dao/task.go +++ b/internal/dao/task.go @@ -33,6 +33,14 @@ func (dao *TaskDAO) Create(task *entity.Task) error { return DB.Create(task).Error } +// CreateMany creates multiple tasks in one batch. +func (dao *TaskDAO) CreateMany(tasks []*entity.Task) error { + if len(tasks) == 0 { + return nil + } + return DB.Create(&tasks).Error +} + // GetByID gets task by ID func (dao *TaskDAO) GetByID(id string) (*entity.Task, error) { var task entity.Task @@ -61,7 +69,7 @@ func (dao *TaskDAO) DeleteByTenantID(tenantID string) (int64, error) { // GetByDocID gets all tasks by document ID func (dao *TaskDAO) GetByDocID(docID string) ([]*entity.Task, error) { var tasks []*entity.Task - err := DB.Where("doc_id = ?", docID).Find(&tasks).Error + err := DB.Where("doc_id = ?", docID).Order("from_page ASC, create_time ASC").Find(&tasks).Error return tasks, err } diff --git a/internal/handler/chunk.go b/internal/handler/chunk.go index 08eab28403..ff75f155bc 100644 --- a/internal/handler/chunk.go +++ b/internal/handler/chunk.go @@ -18,8 +18,8 @@ package handler import ( "encoding/json" "net/http" - "strings" "ragflow/internal/common" + "strings" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -34,6 +34,7 @@ type chunkService interface { List(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) UpdateChunk(req *service.UpdateChunkRequest, userID string) error RemoveChunks(req *service.RemoveChunksRequest, userID string) (int64, error) + Parse(userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) } // ChunkHandler chunk handler @@ -99,7 +100,7 @@ func (h *ChunkHandler) RetrievalTest(c *gin.Context) { // an empty result for blank questions rather than an error. if strings.TrimSpace(req.Question) == "" { c.JSON(http.StatusOK, gin.H{ - "code": int(common.CodeSuccess), + "code": int(common.CodeSuccess), "data": &service.RetrievalTestResponse{ Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}, @@ -203,6 +204,59 @@ func (h *ChunkHandler) Get(c *gin.Context) { }) } +// Parse reparse the datasets' files +func (h *ChunkHandler) Parse(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + userID := strings.TrimSpace(user.ID) + if userID == "" { + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeArgumentError, + "data": nil, + "message": "user_id is required", + }) + return + } + datasetId := strings.TrimSpace(c.Param("dataset_id")) + if datasetId == "" { + c.JSON(http.StatusBadRequest, gin.H{ + "code": common.CodeBadRequest, + "message": "dataset_id is required", + }) + return + } + + var req service.ParseFileRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "code": common.CodeArgumentError, + "data": nil, + "message": err.Error(), + }) + return + } + + data, code, err := h.chunkService.Parse(userID, datasetId, &req) + if code != common.CodeSuccess { + c.JSON(http.StatusOK, gin.H{ + "code": code, + "data": data, + "message": err.Error(), + }) + return + } + + c.JSON(http.StatusOK, gin.H{ + "code": code, + "data": data, + "message": "success", + }) +} + // List retrieves chunks for a document. // @Summary List Chunks // @Description Retrieve paginated chunks for a document with optional filtering. diff --git a/internal/handler/chunk_test.go b/internal/handler/chunk_test.go index 3b6efe2a82..e82b0d63c3 100644 --- a/internal/handler/chunk_test.go +++ b/internal/handler/chunk_test.go @@ -42,6 +42,9 @@ func (m *mockChunkSvc) UpdateChunk(*service.UpdateChunkRequest, string) error { func (m *mockChunkSvc) RemoveChunks(*service.RemoveChunksRequest, string) (int64, error) { panic("not implemented") } +func (m *mockChunkSvc) Parse(string, string, *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) { + panic("not implemented") +} func setupChunkRetrievalTest(userID string) (*gin.Engine, *mockChunkSvc) { mock := &mockChunkSvc{} diff --git a/internal/router/router.go b/internal/router/router.go index f095364f0e..a3ba05992d 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -290,6 +290,7 @@ func (r *Router) Setup(engine *gin.Engine) { // Dataset document chunk datasets.GET("/:dataset_id/documents/:document_id/chunks/:chunk_id", r.chunkHandler.Get) + datasets.POST("/:dataset_id/chunks", r.chunkHandler.Parse) datasets.POST("/:dataset_id/documents/parse", r.documentHandler.StartIngestionTask) datasets.GET("/ingestion/tasks", r.documentHandler.ListIngestionTasks) datasets.PUT("/ingestion/tasks", r.documentHandler.StopIngestionTasks) diff --git a/internal/service/chunk/chunk.go b/internal/service/chunk/chunk.go index 3042f48aa2..4b608d8e09 100644 --- a/internal/service/chunk/chunk.go +++ b/internal/service/chunk/chunk.go @@ -17,26 +17,45 @@ package chunk import ( + "archive/zip" + "bytes" "context" + "encoding/csv" + "encoding/json" + "encoding/xml" "fmt" + "io" + "math/rand" + "path/filepath" "ragflow/internal/common" "ragflow/internal/entity" "ragflow/internal/entity/models" "ragflow/internal/server" + "regexp" + "sort" "strconv" "strings" + "time" + "github.com/cespare/xxhash/v2" "go.uber.org/zap" "ragflow/internal/dao" "ragflow/internal/engine" + "ragflow/internal/engine/redis" "ragflow/internal/engine/types" "ragflow/internal/service" "ragflow/internal/service/nlp" + "ragflow/internal/storage" "ragflow/internal/tokenizer" "ragflow/internal/utility" ) +const ( + maximumPageNumber = 100000 + maximumTaskPageNumber = maximumPageNumber * 1000 +) + // ChunkService chunk service type ChunkService struct { docEngine engine.DocEngine @@ -46,6 +65,14 @@ type ChunkService struct { userTenantDAO *dao.UserTenantDAO documentDAO *dao.DocumentDAO searchService *service.SearchService + + accessibleFunc func(string, string) bool + getKnowledgebaseByIDFunc func(string) (*entity.Knowledgebase, error) + getDocumentsByIDsFunc func([]string) ([]*entity.Document, error) + getDocumentStorageAddressFunc func(*entity.Document) (string, string, error) + queueParseTasksFunc func(*entity.Document, string, string, int64) error + beginParseDocumentFunc func(string) error + deleteTasksByDocIDsFunc func([]string) (int64, error) } // NewChunkService creates chunk service @@ -453,7 +480,6 @@ func hydrateChunkVectors(ctx context.Context, engine engine.DocEngine, chunks [] } } - // Get retrieves a chunk by ID func (s *ChunkService) Get(req *service.GetChunkRequest, userID string) (*service.GetChunkResponse, error) { if s.docEngine == nil { @@ -545,6 +571,584 @@ func (s *ChunkService) Get(req *service.GetChunkRequest, userID string) (*servic return &service.GetChunkResponse{Chunk: chunk}, nil } +func checkDuplicateIDs(documentIDs []string, idTypes string) ([]string, []string) { + idCount := make(map[string]int, len(documentIDs)) + duplicateMessages := make([]string, 0) + uniqueDocIDs := make([]string, 0, len(documentIDs)) + + for _, id := range documentIDs { + idCount[id]++ + } + for id, count := range idCount { + if count > 1 { + duplicateMessages = append(duplicateMessages, fmt.Sprintf("Duplicate %s ids: %s ", idTypes, id)) + } + uniqueDocIDs = append(uniqueDocIDs, id) + } + return uniqueDocIDs, duplicateMessages +} + +func (s *ChunkService) queueParseTasks(doc *entity.Document, bucket, objectName string, priority int64) error { + if s.queueParseTasksFunc != nil { + return s.queueParseTasksFunc(doc, bucket, objectName, priority) + } + tasks, err := s.buildParseTasks(doc, bucket, objectName, priority) + if err != nil { + return err + } + if len(tasks) == 0 { + return nil + } + if err := dao.NewTaskDAO().CreateMany(tasks); err != nil { + return err + } + + queueName := s.parseQueueName(doc, priority) + for _, task := range tasks { + if task.Progress >= 1 { + continue + } + message := parseTaskMessage(task) + if ok := redis.Get().QueueProduct(queueName, message); !ok { + if _, err := dao.NewTaskDAO().DeleteByDocIDs([]string{doc.ID}); err != nil { + common.Warn("Failed to clean parse tasks after Redis enqueue failure", + zap.String("docID", doc.ID), + zap.Error(err)) + } + return fmt.Errorf("Can't access Redis. Please check the Redis' status.") + } + } + return nil +} + +func (s *ChunkService) buildParseTasks(doc *entity.Document, bucket, objectName string, priority int64) ([]*entity.Task, error) { + now := time.Now() + ranges, err := s.parseTaskRanges(doc, bucket, objectName) + if err != nil { + return nil, err + } + tasks := make([]*entity.Task, 0, len(ranges)) + for _, pageRange := range ranges { + taskID := common.GenerateUUID() + progressMsg := "" + digest := s.parseTaskDigest(doc, pageRange.from, pageRange.to) + chunkIDs := "" + tasks = append(tasks, &entity.Task{ + ID: taskID, + DocID: doc.ID, + FromPage: pageRange.from, + ToPage: pageRange.to, + TaskType: "", + Priority: priority, + BeginAt: &now, + Progress: 0, + ProgressMsg: &progressMsg, + Digest: &digest, + ChunkIDs: &chunkIDs, + }) + } + return tasks, nil +} + +type parsePageRange struct { + from int64 + to int64 +} + +func (s *ChunkService) parseTaskRanges(doc *entity.Document, bucket, objectName string) ([]parsePageRange, error) { + if doc.Type == "pdf" { + return s.pdfParseTaskRanges(doc, bucket, objectName) + } + if doc.ParserID == string(entity.ParserTypeTable) { + return s.tableParseTaskRanges(doc, bucket, objectName) + } + return []parsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil +} + +func (s *ChunkService) pdfParseTaskRanges(doc *entity.Document, bucket, objectName string) ([]parsePageRange, error) { + binary, err := s.getStorageBinary(bucket, objectName) + if err != nil { + return nil, err + } + pages := estimatePDFPageCount(binary) + pageSize := int64(parserConfigInt(doc.ParserConfig, "task_page_size", 12)) + if doc.ParserID == string(entity.ParserTypePaper) { + pageSize = int64(parserConfigInt(doc.ParserConfig, "task_page_size", 22)) + } + if doc.ParserID == string(entity.ParserTypeOne) || + doc.ParserID == string(entity.ParserTypeKG) || + parserConfigString(doc.ParserConfig, "layout_recognize", "DeepDOC") != "DeepDOC" || + parserConfigBool(doc.ParserConfig, "toc_extraction", false) { + pageSize = maximumTaskPageNumber + } + if pageSize <= 0 { + pageSize = 12 + } + + pageRanges := parserConfigPageRanges(doc.ParserConfig) + ranges := make([]parsePageRange, 0) + for _, configuredRange := range pageRanges { + start := configuredRange.from - 1 + if start < 0 { + start = 0 + } + end := configuredRange.to - 1 + if pages >= 0 && end > pages { + end = pages + } + for page := start; page < end; page += pageSize { + to := page + pageSize + if to > end { + to = end + } + ranges = append(ranges, parsePageRange{from: page, to: to}) + } + } + if len(ranges) == 0 { + ranges = append(ranges, parsePageRange{from: 0, to: maximumTaskPageNumber}) + } + return ranges, nil +} + +func (s *ChunkService) tableParseTaskRanges(doc *entity.Document, bucket, objectName string) ([]parsePageRange, error) { + binary, err := s.getStorageBinary(bucket, objectName) + if err != nil { + return nil, err + } + rows := estimateTableRowCount(docName(doc), binary) + if rows <= 0 { + return []parsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil + } + ranges := make([]parsePageRange, 0, (rows+2999)/3000) + for row := int64(0); row < int64(rows); row += 3000 { + to := row + 3000 + if to > int64(rows) { + to = int64(rows) + } + ranges = append(ranges, parsePageRange{from: row, to: to}) + } + return ranges, nil +} + +func (s *ChunkService) getStorageBinary(bucket, objectName string) ([]byte, error) { + storageImpl := storage.GetStorageFactory().GetStorage() + if storageImpl == nil { + return nil, fmt.Errorf("storage not initialized") + } + return storageImpl.Get(bucket, objectName) +} + +func (s *ChunkService) beginParseDocument(docID string) error { + if s.beginParseDocumentFunc != nil { + return s.beginParseDocumentFunc(docID) + } + now := time.Now() + return dao.GetDB().Model(&entity.Document{}).Where("id = ?", docID).Updates(map[string]interface{}{ + "progress_msg": "Task is queued...", + "process_begin_at": now, + "progress": rand.Float64() * 0.01, + "run": string(entity.TaskStatusRunning), + "chunk_num": 0, + "token_num": 0, + }).Error +} + +func (s *ChunkService) getDocumentStorageAddress(doc *entity.Document) (string, string, error) { + if s.getDocumentStorageAddressFunc != nil { + return s.getDocumentStorageAddressFunc(doc) + } + return service.NewDocumentService().GetDocumentStorageAddress(doc) +} + +func (s *ChunkService) deleteTasksByDocIDs(docIDs []string) (int64, error) { + if s.deleteTasksByDocIDsFunc != nil { + return s.deleteTasksByDocIDsFunc(docIDs) + } + return dao.NewTaskDAO().DeleteByDocIDs(docIDs) +} + +func (s *ChunkService) accessible(datasetID, userID string) bool { + if s.accessibleFunc != nil { + return s.accessibleFunc(datasetID, userID) + } + return s.kbDAO.Accessible(datasetID, userID) +} + +func (s *ChunkService) getKnowledgebaseByID(datasetID string) (*entity.Knowledgebase, error) { + if s.getKnowledgebaseByIDFunc != nil { + return s.getKnowledgebaseByIDFunc(datasetID) + } + return s.kbDAO.GetByID(datasetID) +} + +func (s *ChunkService) getDocumentsByIDs(docIDs []string) ([]*entity.Document, error) { + if s.getDocumentsByIDsFunc != nil { + return s.getDocumentsByIDsFunc(docIDs) + } + return s.documentDAO.GetByIDs(docIDs) +} + +func (s *ChunkService) parseQueueName(doc *entity.Document, priority int64) string { + suffix := "common" + if doc.ParserID == string(entity.ParserTypeResume) { + suffix = "resume" + } + return fmt.Sprintf("te.%d.%s", priority, suffix) +} + +func (s *ChunkService) parseTaskDigest(doc *entity.Document, fromPage, toPage int64) string { + hasher := xxhash.New() + config := chunkingConfigForDigest(doc) + keys := make([]string, 0, len(config)) + for key := range config { + keys = append(keys, key) + } + sort.Strings(keys) + for _, key := range keys { + hasher.WriteString(stableString(config[key])) + } + hasher.WriteString(doc.ID) + hasher.WriteString(strconv.FormatInt(fromPage, 10)) + hasher.WriteString(strconv.FormatInt(toPage, 10)) + return fmt.Sprintf("%x", hasher.Sum64()) +} + +func parseTaskMessage(task *entity.Task) map[string]interface{} { + beginAt := "" + if task.BeginAt != nil { + beginAt = task.BeginAt.Format("2006-01-02 15:04:05") + } + digest := "" + if task.Digest != nil { + digest = *task.Digest + } + return map[string]interface{}{ + "id": task.ID, + "doc_id": task.DocID, + "from_page": task.FromPage, + "to_page": task.ToPage, + "progress": task.Progress, + "priority": task.Priority, + "begin_at": beginAt, + "digest": digest, + } +} + +func chunkingConfigForDigest(doc *entity.Document) map[string]interface{} { + return map[string]interface{}{ + "doc_id": doc.ID, + "kb_id": doc.KbID, + "parser_id": doc.ParserID, + "parser_config": copyParserConfigForDigest(doc.ParserConfig), + } +} + +func copyParserConfigForDigest(config map[string]interface{}) map[string]interface{} { + copied := make(map[string]interface{}, len(config)) + for key, value := range config { + if key == "raptor" || key == "graphrag" { + continue + } + copied[key] = value + } + return copied +} + +func stableString(value interface{}) string { + binary, err := json.Marshal(value) + if err != nil { + return fmt.Sprint(value) + } + return string(binary) +} + +func parserConfigInt(config map[string]interface{}, key string, fallback int) int { + value, ok := config[key] + if !ok || value == nil { + return fallback + } + switch typedValue := value.(type) { + case int: + return typedValue + case int64: + return int(typedValue) + case float64: + return int(typedValue) + case json.Number: + if intValue, err := typedValue.Int64(); err == nil { + return int(intValue) + } + case string: + if intValue, err := strconv.Atoi(strings.TrimSpace(typedValue)); err == nil { + return intValue + } + } + return fallback +} + +func parserConfigString(config map[string]interface{}, key, fallback string) string { + value, ok := config[key] + if !ok || value == nil { + return fallback + } + if stringValue, ok := value.(string); ok { + return stringValue + } + return fmt.Sprint(value) +} + +func parserConfigBool(config map[string]interface{}, key string, fallback bool) bool { + value, ok := config[key] + if !ok || value == nil { + return fallback + } + switch typedValue := value.(type) { + case bool: + return typedValue + case string: + switch strings.ToLower(strings.TrimSpace(typedValue)) { + case "true", "1", "yes", "on": + return true + case "false", "0", "no", "off": + return false + } + } + return fallback +} + +func parserConfigPageRanges(config map[string]interface{}) []parsePageRange { + defaultRanges := []parsePageRange{{from: 1, to: maximumPageNumber}} + raw, ok := config["pages"] + if !ok || raw == nil { + return defaultRanges + } + rawRanges, ok := raw.([]interface{}) + if !ok || len(rawRanges) == 0 { + return defaultRanges + } + + ranges := make([]parsePageRange, 0, len(rawRanges)) + for _, rawRange := range rawRanges { + rangeValues, ok := rawRange.([]interface{}) + if !ok || len(rangeValues) < 2 { + continue + } + from, okFrom := toInt64(rangeValues[0]) + to, okTo := toInt64(rangeValues[1]) + if okFrom && okTo && to > from { + ranges = append(ranges, parsePageRange{from: from, to: to}) + } + } + if len(ranges) == 0 { + return defaultRanges + } + return ranges +} + +func toInt64(value interface{}) (int64, bool) { + switch typedValue := value.(type) { + case int: + return int64(typedValue), true + case int64: + return typedValue, true + case float64: + return int64(typedValue), true + case json.Number: + intValue, err := typedValue.Int64() + return intValue, err == nil + case string: + intValue, err := strconv.ParseInt(strings.TrimSpace(typedValue), 10, 64) + return intValue, err == nil + default: + return 0, false + } +} + +var pdfPagePattern = regexp.MustCompile(`/Type\s*/Page\b`) + +func estimatePDFPageCount(binary []byte) int64 { + if len(binary) == 0 { + return 0 + } + return int64(len(pdfPagePattern.FindAll(binary, -1))) +} + +func estimateTableRowCount(name string, binary []byte) int { + switch strings.ToLower(filepath.Ext(name)) { + case ".xlsx": + if rows, err := countXLSXRows(binary); err == nil { + return rows + } + case ".csv", ".tsv", ".txt": + return countDelimitedRows(name, binary) + } + return 0 +} + +func countDelimitedRows(name string, binary []byte) int { + reader := csv.NewReader(bytes.NewReader(binary)) + reader.FieldsPerRecord = -1 + reader.ReuseRecord = true + if strings.EqualFold(filepath.Ext(name), ".tsv") { + reader.Comma = '\t' + } + rows := 0 + for { + _, err := reader.Read() + if err == nil { + rows++ + continue + } + if err == io.EOF { + break + } + rows += bytes.Count(binary, []byte{'\n'}) + if len(binary) > 0 && binary[len(binary)-1] != '\n' { + rows++ + } + break + } + return rows +} + +func countXLSXRows(binary []byte) (int, error) { + zipReader, err := zip.NewReader(bytes.NewReader(binary), int64(len(binary))) + if err != nil { + return 0, err + } + maxRows := 0 + for _, file := range zipReader.File { + if !strings.HasPrefix(file.Name, "xl/worksheets/") || !strings.HasSuffix(file.Name, ".xml") { + continue + } + rows, err := countWorksheetRows(file) + if err != nil { + return 0, err + } + if rows > maxRows { + maxRows = rows + } + } + return maxRows, nil +} + +func countWorksheetRows(file *zip.File) (int, error) { + reader, err := file.Open() + if err != nil { + return 0, err + } + defer reader.Close() + + decoder := xml.NewDecoder(reader) + rows := 0 + for { + token, err := decoder.Token() + if err == io.EOF { + break + } + if err != nil { + return 0, err + } + start, ok := token.(xml.StartElement) + if ok && start.Name.Local == "row" { + rows++ + } + } + return rows, nil +} + +func docName(doc *entity.Document) string { + if doc.Name == nil { + return "" + } + return *doc.Name +} + +func (s *ChunkService) Parse(userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) { + if !s.accessible(datasetID, userID) { + return nil, common.CodeOperatingError, fmt.Errorf("You don't own the dataset %s.", datasetID) + } + if req == nil || len(req.DocumentIDs) == 0 { + return nil, common.CodeDataError, fmt.Errorf("`document_ids` is required") + } + + kb, err := s.getKnowledgebaseByID(datasetID) + if err != nil || kb == nil { + return nil, common.CodeDataError, fmt.Errorf("dataset not found") + } + + docIDs, duplicateMessages := checkDuplicateIDs(req.DocumentIDs, "document") + notFound := make([]string, 0) + + docs, err := s.getDocumentsByIDs(docIDs) + if err != nil { + return nil, common.CodeServerError, err + } + docByID := make(map[string]*entity.Document, len(docs)) + for _, doc := range docs { + docByID[doc.ID] = doc + } + for _, docID := range docIDs { + doc := docByID[docID] + if doc == nil || doc.KbID != datasetID { + notFound = append(notFound, docID) + } + } + if len(notFound) > 0 { + return nil, common.CodeDataError, fmt.Errorf("Documents not found: %v", notFound) + } + for _, docID := range docIDs { + doc := docByID[docID] + if doc.Run != nil && *doc.Run == string(entity.TaskStatusRunning) { + return nil, common.CodeDataError, fmt.Errorf("Can't parse document that is currently being processed") + } + } + + successCount := 0 + + for _, docID := range docIDs { + doc := docByID[docID] + + if s.docEngine != nil { + indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) + if _, err := s.docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": docID}, indexName, datasetID); err != nil { + return nil, common.CodeServerError, err + } + } + if _, err := s.deleteTasksByDocIDs([]string{docID}); err != nil { + return nil, common.CodeServerError, err + } + + bucket, objectName, err := s.getDocumentStorageAddress(doc) + if err != nil { + return nil, common.CodeServerError, err + } + if err := s.queueParseTasks(doc, bucket, objectName, 0); err != nil { + return nil, common.CodeServerError, err + } + if err := s.beginParseDocument(doc.ID); err != nil { + if _, delErr := s.deleteTasksByDocIDs([]string{doc.ID}); delErr != nil { + common.Warn("Failed to clean parse tasks after document state update failure", + zap.String("docID", doc.ID), + zap.Error(delErr)) + } + return nil, common.CodeServerError, err + } + successCount++ + } + + if len(duplicateMessages) > 0 { + if successCount > 0 { + return map[string]interface{}{ + "success_count": successCount, + "errors": duplicateMessages, + }, common.CodeSuccess, fmt.Errorf("Partially parsed %d documents with %d errors", successCount, len(duplicateMessages)) + } + return nil, common.CodeDataError, fmt.Errorf("%s", strings.Join(duplicateMessages, ";")) + } + return nil, common.CodeSuccess, nil +} + // List retrieves chunks for a document func (s *ChunkService) List(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error) { if s.docEngine == nil { @@ -931,4 +1535,3 @@ func (s *ChunkService) RemoveChunks(req *service.RemoveChunksRequest, userID str return deletedCount, nil } - diff --git a/internal/service/chunk/chunk_test.go b/internal/service/chunk/chunk_test.go index f0d8940e3d..0871fc4ce9 100644 --- a/internal/service/chunk/chunk_test.go +++ b/internal/service/chunk/chunk_test.go @@ -2,9 +2,18 @@ package chunk import ( "context" + "errors" "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/engine/types" + "ragflow/internal/entity" + "ragflow/internal/service" "reflect" + "strings" "testing" + + "github.com/glebarez/sqlite" + "gorm.io/gorm" ) func TestIsZeroVector(t *testing.T) { @@ -60,3 +69,573 @@ func TestHydrateChunkVectors_NoDim(t *testing.T) { hydrateChunkVectors(context.Background(), nil, chunks, []string{"kb1"}, []string{"t1"}) // Empty vectors have dim=0 → early return. No crash. } + +func TestParsePrevalidatesDocumentsBeforeMutating(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + + userID := "user-1" + datasetID := "kb-1" + insertChunkTestKB(t, datasetID, userID) + insertChunkTestKB(t, "kb-2", userID) + insertChunkTestDoc(t, "doc-1", datasetID) + insertChunkTestDoc(t, "doc-2", "kb-2") + insertChunkTestTask(t, "task-1", "doc-1") + + svc := &ChunkService{ + docEngine: nil, + kbDAO: dao.NewKnowledgebaseDAO(), + documentDAO: dao.NewDocumentDAO(), + } + + _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{ + DocumentIDs: []string{"doc-1", "missing-doc", "doc-2"}, + }) + if err == nil { + t.Fatal("expected parse to fail") + } + if code != common.CodeDataError { + t.Fatalf("expected CodeDataError, got %v", code) + } + if !strings.Contains(err.Error(), "missing-doc") || !strings.Contains(err.Error(), "doc-2") { + t.Fatalf("expected missing and foreign documents in error, got %q", err.Error()) + } + + var taskCount int64 + if err := dao.DB.Model(&entity.Task{}).Where("doc_id = ?", "doc-1").Count(&taskCount).Error; err != nil { + t.Fatalf("count tasks: %v", err) + } + if taskCount != 1 { + t.Fatalf("expected existing task to remain, got %d tasks", taskCount) + } + + doc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("get doc: %v", err) + } + if doc.Run != nil { + t.Fatalf("expected doc run to remain nil, got %q", *doc.Run) + } + if doc.ChunkNum != 7 { + t.Fatalf("expected chunk_num to remain 7, got %d", doc.ChunkNum) + } +} + +func TestParseRejectsInaccessibleDataset(t *testing.T) { + svc := newParseTestService(t) + svc.accessibleFunc = func(string, string) bool { return false } + + _, code, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + if err == nil { + t.Fatal("expected parse to fail") + } + if code != common.CodeOperatingError { + t.Fatalf("expected CodeOperatingError, got %v", code) + } + if !strings.Contains(err.Error(), "don't own the dataset") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestParseRequiresDocumentIDs(t *testing.T) { + svc := newParseTestService(t) + svc.accessibleFunc = func(string, string) bool { return true } + + for _, req := range []*service.ParseFileRequest{nil, {DocumentIDs: nil}, {DocumentIDs: []string{}}} { + _, code, err := svc.Parse("user-1", "kb-1", req) + if err == nil { + t.Fatal("expected parse to fail") + } + if code != common.CodeDataError { + t.Fatalf("expected CodeDataError, got %v", code) + } + if !strings.Contains(err.Error(), "document_ids") { + t.Fatalf("unexpected error: %v", err) + } + } +} + +func TestParseReturnsDataErrorWhenDatasetMissing(t *testing.T) { + svc := newParseTestService(t) + svc.accessibleFunc = func(string, string) bool { return true } + svc.getKnowledgebaseByIDFunc = func(string) (*entity.Knowledgebase, error) { return nil, nil } + + _, code, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + if err == nil { + t.Fatal("expected parse to fail") + } + if code != common.CodeDataError { + t.Fatalf("expected CodeDataError, got %v", code) + } + if !strings.Contains(err.Error(), "dataset not found") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestParseReturnsServerErrorWhenDocumentsQueryFails(t *testing.T) { + svc := newParseTestService(t) + queryErr := errors.New("documents query failed") + svc.accessibleFunc = func(string, string) bool { return true } + svc.getKnowledgebaseByIDFunc = func(string) (*entity.Knowledgebase, error) { + return &entity.Knowledgebase{ID: "kb-1", TenantID: "user-1"}, nil + } + svc.getDocumentsByIDsFunc = func([]string) ([]*entity.Document, error) { + return nil, queryErr + } + + _, code, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + if !errors.Is(err, queryErr) { + t.Fatalf("expected query error, got %v", err) + } + if code != common.CodeServerError { + t.Fatalf("expected CodeServerError, got %v", code) + } +} + +func TestParseRejectsRunningDocument(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + + userID := "user-1" + datasetID := "kb-1" + insertChunkTestKB(t, datasetID, userID) + insertChunkTestDoc(t, "doc-1", datasetID) + running := string(entity.TaskStatusRunning) + if err := dao.DB.Model(&entity.Document{}).Where("id = ?", "doc-1").Update("run", running).Error; err != nil { + t.Fatalf("mark doc running: %v", err) + } + + svc := newParseTestService(t) + _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + if err == nil { + t.Fatal("expected parse to fail") + } + if code != common.CodeDataError { + t.Fatalf("expected CodeDataError, got %v", code) + } + if !strings.Contains(err.Error(), "currently being processed") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestParseReturnsServerErrorWhenDeleteChunksFails(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + userID, datasetID := "user-1", "kb-1" + insertChunkTestKB(t, datasetID, userID) + insertChunkTestDoc(t, "doc-1", datasetID) + + deleteErr := errors.New("delete chunks failed") + engine := &parseTestDocEngine{deleteChunksErr: deleteErr} + svc := newParseTestService(t) + svc.docEngine = engine + + _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + if !errors.Is(err, deleteErr) { + t.Fatalf("expected delete chunks error, got %v", err) + } + if code != common.CodeServerError { + t.Fatalf("expected CodeServerError, got %v", code) + } + if engine.deleteChunksCalls != 1 { + t.Fatalf("expected DeleteChunks once, got %d", engine.deleteChunksCalls) + } +} + +func TestParseReturnsServerErrorWhenDeleteTasksFails(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + userID, datasetID := "user-1", "kb-1" + insertChunkTestKB(t, datasetID, userID) + insertChunkTestDoc(t, "doc-1", datasetID) + + deleteErr := errors.New("delete tasks failed") + svc := newParseTestService(t) + svc.deleteTasksByDocIDsFunc = func([]string) (int64, error) { return 0, deleteErr } + + _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + if !errors.Is(err, deleteErr) { + t.Fatalf("expected delete tasks error, got %v", err) + } + if code != common.CodeServerError { + t.Fatalf("expected CodeServerError, got %v", code) + } +} + +func TestParseReturnsServerErrorWhenStorageAddressFails(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + userID, datasetID := "user-1", "kb-1" + insertChunkTestKB(t, datasetID, userID) + insertChunkTestDoc(t, "doc-1", datasetID) + + storageErr := errors.New("storage address failed") + svc := newParseTestService(t) + svc.getDocumentStorageAddressFunc = func(*entity.Document) (string, string, error) { + return "", "", storageErr + } + + _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + if !errors.Is(err, storageErr) { + t.Fatalf("expected storage error, got %v", err) + } + if code != common.CodeServerError { + t.Fatalf("expected CodeServerError, got %v", code) + } +} + +func TestParseReturnsServerErrorWhenQueueFails(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + userID, datasetID := "user-1", "kb-1" + insertChunkTestKB(t, datasetID, userID) + insertChunkTestDoc(t, "doc-1", datasetID) + + queueErr := errors.New("queue failed") + svc := newParseTestService(t) + svc.queueParseTasksFunc = func(*entity.Document, string, string, int64) error { + return queueErr + } + + _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + if !errors.Is(err, queueErr) { + t.Fatalf("expected queue error, got %v", err) + } + if code != common.CodeServerError { + t.Fatalf("expected CodeServerError, got %v", code) + } +} + +func TestParseCleansTasksWhenBeginParseFails(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + userID, datasetID := "user-1", "kb-1" + insertChunkTestKB(t, datasetID, userID) + insertChunkTestDoc(t, "doc-1", datasetID) + + beginErr := errors.New("begin parse failed") + deleteCalls := 0 + svc := newParseTestService(t) + svc.beginParseDocumentFunc = func(string) error { return beginErr } + svc.deleteTasksByDocIDsFunc = func([]string) (int64, error) { + deleteCalls++ + return 1, nil + } + + _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + if !errors.Is(err, beginErr) { + t.Fatalf("expected begin error, got %v", err) + } + if code != common.CodeServerError { + t.Fatalf("expected CodeServerError, got %v", code) + } + if deleteCalls != 2 { + t.Fatalf("expected initial delete and cleanup delete, got %d calls", deleteCalls) + } +} + +func TestParseReturnsPartialSuccessForDuplicateDocumentIDs(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + userID, datasetID := "user-1", "kb-1" + insertChunkTestKB(t, datasetID, userID) + insertChunkTestDoc(t, "doc-1", datasetID) + + svc := newParseTestService(t) + result, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{ + DocumentIDs: []string{"doc-1", "doc-1"}, + }) + if err == nil { + t.Fatal("expected duplicate warning error") + } + if code != common.CodeSuccess { + t.Fatalf("expected CodeSuccess, got %v", code) + } + if result["success_count"] != 1 { + t.Fatalf("expected one successful parse, got %v", result["success_count"]) + } + errorsValue, ok := result["errors"].([]string) + if !ok || len(errorsValue) != 1 || !strings.Contains(errorsValue[0], "Duplicate document ids: doc-1") { + t.Fatalf("unexpected duplicate errors: %#v", result["errors"]) + } +} + +func TestParseQueuesAndBeginsDocument(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + userID, datasetID := "user-1", "kb-1" + insertChunkTestKB(t, datasetID, userID) + insertChunkTestDoc(t, "doc-1", datasetID) + insertChunkTestTask(t, "task-1", "doc-1") + + queueCalls := 0 + svc := newParseTestService(t) + svc.queueParseTasksFunc = func(doc *entity.Document, bucket, objectName string, priority int64) error { + queueCalls++ + if doc.ID != "doc-1" || bucket != datasetID || objectName != "doc-1.txt" || priority != 0 { + t.Fatalf("unexpected queue args: doc=%s bucket=%s object=%s priority=%d", doc.ID, bucket, objectName, priority) + } + return nil + } + + result, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + if err != nil { + t.Fatalf("expected parse success, got %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("expected CodeSuccess, got %v", code) + } + if result != nil { + t.Fatalf("expected nil result, got %#v", result) + } + if queueCalls != 1 { + t.Fatalf("expected queue once, got %d", queueCalls) + } + + var taskCount int64 + if err := dao.DB.Model(&entity.Task{}).Where("doc_id = ?", "doc-1").Count(&taskCount).Error; err != nil { + t.Fatalf("count tasks: %v", err) + } + if taskCount != 0 { + t.Fatalf("expected old tasks to be deleted, got %d", taskCount) + } + + doc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("get doc: %v", err) + } + if doc.Run == nil || *doc.Run != string(entity.TaskStatusRunning) { + t.Fatalf("expected document to be running, got %v", doc.Run) + } + if doc.ChunkNum != 0 { + t.Fatalf("expected chunk_num reset to 0, got %d", doc.ChunkNum) + } +} + +func setupChunkTestDB(t *testing.T) *gorm.DB { + t.Helper() + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ + TranslateError: true, + }) + if err != nil { + t.Fatalf("failed to open sqlite: %v", err) + } + if err := db.AutoMigrate( + &entity.Knowledgebase{}, + &entity.Document{}, + &entity.Task{}, + ); err != nil { + t.Fatalf("failed to migrate: %v", err) + } + return db +} + +func pushChunkTestDB(t *testing.T, testDB *gorm.DB) { + t.Helper() + + orig := dao.DB + dao.DB = testDB + t.Cleanup(func() { + dao.DB = orig + }) +} + +func newParseTestService(t *testing.T) *ChunkService { + t.Helper() + + return &ChunkService{ + docEngine: nil, + kbDAO: dao.NewKnowledgebaseDAO(), + documentDAO: dao.NewDocumentDAO(), + getDocumentStorageAddressFunc: func(doc *entity.Document) (string, string, error) { + if doc.Location == nil { + return doc.KbID, "", nil + } + return doc.KbID, *doc.Location, nil + }, + queueParseTasksFunc: func(*entity.Document, string, string, int64) error { + return nil + }, + } +} + +func insertChunkTestKB(t *testing.T, id, tenantID string) { + t.Helper() + + status := string(entity.StatusValid) + kb := &entity.Knowledgebase{ + ID: id, + TenantID: tenantID, + Name: id, + EmbdID: "embedding-model", + Permission: string(entity.TenantPermissionMe), + CreatedBy: tenantID, + ParserConfig: entity.JSONMap{}, + Status: &status, + } + if err := dao.DB.Create(kb).Error; err != nil { + t.Fatalf("insert kb: %v", err) + } +} + +func insertChunkTestDoc(t *testing.T, id, kbID string) { + t.Helper() + + location := id + ".txt" + name := id + ".txt" + status := string(entity.StatusValid) + doc := &entity.Document{ + ID: id, + KbID: kbID, + ParserID: string(entity.ParserTypeNaive), + ParserConfig: entity.JSONMap{}, + SourceType: string(entity.FileSourceLocal), + Type: "txt", + CreatedBy: "user-1", + Name: &name, + Location: &location, + ChunkNum: 7, + Suffix: ".txt", + Status: &status, + } + if err := dao.DB.Create(doc).Error; err != nil { + t.Fatalf("insert doc: %v", err) + } +} + +func insertChunkTestTask(t *testing.T, id, docID string) { + t.Helper() + + digest := "digest" + task := &entity.Task{ + ID: id, + DocID: docID, + Digest: &digest, + } + if err := dao.DB.Create(task).Error; err != nil { + t.Fatalf("insert task: %v", err) + } +} + +type parseTestDocEngine struct { + deleteChunksErr error + deleteChunksCalls int +} + +func (e *parseTestDocEngine) CreateChunkStore(context.Context, string, string, int, string) error { + return nil +} + +func (e *parseTestDocEngine) InsertChunks(context.Context, []map[string]interface{}, string, string) ([]string, error) { + return nil, nil +} + +func (e *parseTestDocEngine) UpdateChunks(context.Context, map[string]interface{}, map[string]interface{}, string, string) error { + return nil +} + +func (e *parseTestDocEngine) DeleteChunks(context.Context, map[string]interface{}, string, string) (int64, error) { + e.deleteChunksCalls++ + return 0, e.deleteChunksErr +} + +func (e *parseTestDocEngine) Search(context.Context, *types.SearchRequest) (*types.SearchResult, error) { + return nil, nil +} + +func (e *parseTestDocEngine) GetChunk(context.Context, string, string, []string) (interface{}, error) { + return nil, nil +} + +func (e *parseTestDocEngine) DropChunkStore(context.Context, string, string) error { + return nil +} + +func (e *parseTestDocEngine) ChunkStoreExists(context.Context, string, string) (bool, error) { + return false, nil +} + +func (e *parseTestDocEngine) CreateMetadataStore(context.Context, string) error { + return nil +} + +func (e *parseTestDocEngine) InsertMetadata(context.Context, []map[string]interface{}, string) ([]string, error) { + return nil, nil +} + +func (e *parseTestDocEngine) UpdateMetadata(context.Context, string, string, map[string]interface{}, string) error { + return nil +} + +func (e *parseTestDocEngine) DeleteMetadata(context.Context, map[string]interface{}, string) (int64, error) { + return 0, nil +} + +func (e *parseTestDocEngine) DeleteMetadataKeys(context.Context, string, string, []string, string) error { + return nil +} + +func (e *parseTestDocEngine) DropMetadataStore(context.Context, string) error { + return nil +} + +func (e *parseTestDocEngine) MetadataStoreExists(context.Context, string) (bool, error) { + return false, nil +} + +func (e *parseTestDocEngine) SearchMetadata(context.Context, *types.SearchMetadataRequest) (*types.SearchMetadataResult, error) { + return nil, nil +} + +func (e *parseTestDocEngine) IndexDocument(context.Context, string, string, interface{}) error { + return nil +} + +func (e *parseTestDocEngine) DeleteDocument(context.Context, string, string) error { + return nil +} + +func (e *parseTestDocEngine) BulkIndex(context.Context, string, []interface{}) (interface{}, error) { + return nil, nil +} + +func (e *parseTestDocEngine) GetFields([]map[string]interface{}, []string) map[string]map[string]interface{} { + return nil +} + +func (e *parseTestDocEngine) GetAggregation([]map[string]interface{}, string) []map[string]interface{} { + return nil +} + +func (e *parseTestDocEngine) GetHighlight([]map[string]interface{}, []string, string) map[string]string { + return nil +} + +func (e *parseTestDocEngine) GetChunkIDs([]map[string]interface{}) []string { + return nil +} + +func (e *parseTestDocEngine) KNNScores(context.Context, []map[string]interface{}, []float64, int) (map[string]interface{}, error) { + return nil, nil +} + +func (e *parseTestDocEngine) GetScores(map[string]interface{}) map[string]float64 { + return nil +} + +func (e *parseTestDocEngine) Ping(context.Context) error { + return nil +} + +func (e *parseTestDocEngine) Close() error { + return nil +} + +func (e *parseTestDocEngine) GetType() string { + return "test" +} + +func (e *parseTestDocEngine) FilterDocIdsByMetaPushdown(context.Context, []string, []map[string]interface{}, string) []string { + return nil +} diff --git a/internal/service/chunk_types.go b/internal/service/chunk_types.go index 2fe429cd65..244770eb07 100644 --- a/internal/service/chunk_types.go +++ b/internal/service/chunk_types.go @@ -24,7 +24,6 @@ import ( "ragflow/internal/server" "strings" - "ragflow/internal/dao" "ragflow/internal/engine" "ragflow/internal/engine/types" @@ -43,10 +42,9 @@ type ChunkService struct { searchService *SearchService } - // RetrievalTestRequest retrieval test request type RetrievalTestRequest struct { - Datasets common.StringSlice `json:"dataset_ids" binding:"required"` // string or []string + Datasets common.StringSlice `json:"dataset_ids" binding:"required"` // string or []string Question string `json:"question"` Page *int `json:"page,omitempty"` Size *int `json:"size,omitempty"` @@ -81,6 +79,11 @@ type GetChunkResponse struct { Chunk map[string]interface{} `json:"chunk"` } +// ParseFileRequest is the request body for reparsing documents in a dataset. +type ParseFileRequest struct { + DocumentIDs []string `json:"document_ids"` +} + // Get retrieves a chunk by ID func (s *ChunkService) Get(req *GetChunkRequest, userID string) (*GetChunkResponse, error) { if s.docEngine == nil {