From 42aba36c1b6d38035b2f4d75c6856c9f60a5edd9 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Thu, 2 Jul 2026 19:54:42 +0800 Subject: [PATCH] fix(go): chunk stats after chunk deletion (#16553) ## Summary - Decrement document and knowledgebase chunk counts after chunks are deleted - Keep token counts unchanged because deleted chunk token totals are not available - Add tests for stats update, zero-delete behavior, error handling, and transaction rollback --- internal/service/chunk/chunk.go | 42 +++++ internal/service/chunk/chunk_test.go | 228 ++++++++++++++++++++++++++- internal/service/chunk_types.go | 46 ++++++ internal/service/chunk_types_test.go | 85 ++++++++++ 4 files changed, 397 insertions(+), 4 deletions(-) diff --git a/internal/service/chunk/chunk.go b/internal/service/chunk/chunk.go index b16543913..83b2f74f2 100644 --- a/internal/service/chunk/chunk.go +++ b/internal/service/chunk/chunk.go @@ -108,6 +108,7 @@ type ChunkService struct { deleteTasksByDocIDsFunc func([]string) (int64, error) getEmbeddingModelFunc func(string, string) (*models.EmbeddingModel, error) incrementChunkStatsFunc func(string, string, int64, int64, float64) error + decrementChunkStatsFunc func(string, string, int64, int64, float64) error storeChunkImageFunc func(string, string, []byte) error tokenizeFunc func(string) (string, error) fineGrainedTokenizeFunc func(string) (string, error) @@ -1766,6 +1767,12 @@ func (s *ChunkService) RemoveChunks(req *service.RemoveChunksRequest, userID str return 0, fmt.Errorf("failed to delete chunks: %w", err) } + if deletedCount > 0 { + if err := s.decrementChunkStats(req.DocID, doc.KbID, 0, deletedCount, 0); err != nil { + return deletedCount, fmt.Errorf("failed to update chunk stats: %w", err) + } + } + return deletedCount, nil } @@ -2071,6 +2078,41 @@ func (s *ChunkService) incrementChunkStats(docID, kbID string, tokenNum, chunkNu }) } +func (s *ChunkService) decrementChunkStats(docID, kbID string, tokenNum, chunkNum int64, duration float64) error { + if s.decrementChunkStatsFunc != nil { + return s.decrementChunkStatsFunc(docID, kbID, tokenNum, chunkNum, duration) + } + return dao.DB.Transaction(func(tx *gorm.DB) error { + result := tx.Model(&entity.Document{}). + Where("id = ? AND kb_id = ?", docID, kbID). + Updates(map[string]interface{}{ + "token_num": gorm.Expr("CASE WHEN token_num - ? >= 0 THEN token_num - ? ELSE 0 END", tokenNum, tokenNum), + "chunk_num": gorm.Expr("CASE WHEN chunk_num - ? >= 0 THEN chunk_num - ? ELSE 0 END", chunkNum, chunkNum), + "process_duration": gorm.Expr("CASE WHEN process_duration + ? >= 0 THEN process_duration + ? ELSE 0 END", duration, duration), + }) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return fmt.Errorf("document not found") + } + + result = tx.Model(&entity.Knowledgebase{}). + Where("id = ?", kbID). + Updates(map[string]interface{}{ + "token_num": gorm.Expr("CASE WHEN token_num - ? >= 0 THEN token_num - ? ELSE 0 END", tokenNum, tokenNum), + "chunk_num": gorm.Expr("CASE WHEN chunk_num - ? >= 0 THEN chunk_num - ? ELSE 0 END", chunkNum, chunkNum), + }) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return fmt.Errorf("knowledgebase not found") + } + return nil + }) +} + func (s *ChunkService) storeChunkImage(bucket, chunkID string, imageBinary []byte) error { if s.storeChunkImageFunc != nil { return s.storeChunkImageFunc(bucket, chunkID, imageBinary) diff --git a/internal/service/chunk/chunk_test.go b/internal/service/chunk/chunk_test.go index cf77c990f..3397fcd67 100644 --- a/internal/service/chunk/chunk_test.go +++ b/internal/service/chunk/chunk_test.go @@ -767,6 +767,186 @@ func TestStoreChunkImageMergesExistingImage(t *testing.T) { } } +func TestRemoveChunksDecrementsStatsAfterDelete(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + insertChunkTestUserTenant(t, "user-1", "tenant-1") + insertChunkTestKB(t, "kb-1", "tenant-1") + insertChunkTestDoc(t, "doc-1", "kb-1") + setChunkTestStats(t, "doc-1", "kb-1", 70, 7, 100, 10) + + engine := &parseTestDocEngine{deleteChunksCount: 3} + svc := &ChunkService{ + docEngine: engine, + kbDAO: dao.NewKnowledgebaseDAO(), + userTenantDAO: dao.NewUserTenantDAO(), + } + + deletedCount, err := svc.RemoveChunks(&service.RemoveChunksRequest{ + DocID: "doc-1", + ChunkIDs: []string{"chunk-1", "chunk-2", "chunk-3"}, + }, "user-1") + if err != nil { + t.Fatalf("RemoveChunks() error = %v", err) + } + if deletedCount != 3 { + t.Fatalf("deleted count = %d, want 3", deletedCount) + } + if engine.deleteChunksCalls != 1 { + t.Fatalf("DeleteChunks calls = %d, want 1", engine.deleteChunksCalls) + } + if engine.deleteIndexName != "ragflow_tenant-1" || engine.deleteDatasetID != "kb-1" { + t.Fatalf("unexpected delete target index=%q dataset=%q", engine.deleteIndexName, engine.deleteDatasetID) + } + if engine.deleteChunksCondition["doc_id"] != "doc-1" { + t.Fatalf("delete doc_id condition = %#v, want doc-1", engine.deleteChunksCondition["doc_id"]) + } + if !reflect.DeepEqual(engine.deleteChunksCondition["id"], []interface{}{"chunk-1", "chunk-2", "chunk-3"}) { + t.Fatalf("delete id condition = %#v", engine.deleteChunksCondition["id"]) + } + + doc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("get doc: %v", err) + } + if doc.TokenNum != 70 { + t.Fatalf("document token_num = %d, want 70", doc.TokenNum) + } + if doc.ChunkNum != 4 { + t.Fatalf("document chunk_num = %d, want 4", doc.ChunkNum) + } + + kb, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + if err != nil { + t.Fatalf("get kb: %v", err) + } + if kb.TokenNum != 100 { + t.Fatalf("knowledgebase token_num = %d, want 100", kb.TokenNum) + } + if kb.ChunkNum != 7 { + t.Fatalf("knowledgebase chunk_num = %d, want 7", kb.ChunkNum) + } +} + +func TestRemoveChunksSkipsStatsWhenNothingDeleted(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + insertChunkTestUserTenant(t, "user-1", "tenant-1") + insertChunkTestKB(t, "kb-1", "tenant-1") + insertChunkTestDoc(t, "doc-1", "kb-1") + setChunkTestStats(t, "doc-1", "kb-1", 70, 7, 100, 10) + + var decrementCalls int + svc := &ChunkService{ + docEngine: &parseTestDocEngine{deleteChunksCount: 0}, + kbDAO: dao.NewKnowledgebaseDAO(), + userTenantDAO: dao.NewUserTenantDAO(), + decrementChunkStatsFunc: func(string, string, int64, int64, float64) error { + decrementCalls++ + return nil + }, + } + + deletedCount, err := svc.RemoveChunks(&service.RemoveChunksRequest{ + DocID: "doc-1", + DeleteAll: true, + }, "user-1") + if err != nil { + t.Fatalf("RemoveChunks() error = %v", err) + } + if deletedCount != 0 { + t.Fatalf("deleted count = %d, want 0", deletedCount) + } + if decrementCalls != 0 { + t.Fatalf("decrement calls = %d, want 0", decrementCalls) + } +} + +func TestRemoveChunksReturnsStatsError(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + insertChunkTestUserTenant(t, "user-1", "tenant-1") + insertChunkTestKB(t, "kb-1", "tenant-1") + insertChunkTestDoc(t, "doc-1", "kb-1") + + svc := &ChunkService{ + docEngine: &parseTestDocEngine{deleteChunksCount: 2}, + kbDAO: dao.NewKnowledgebaseDAO(), + userTenantDAO: dao.NewUserTenantDAO(), + decrementChunkStatsFunc: func(docID, kbID string, tokenNum, chunkNum int64, duration float64) error { + if docID != "doc-1" || kbID != "kb-1" || tokenNum != 0 || chunkNum != 2 || duration != 0 { + t.Fatalf("unexpected decrement args doc=%s kb=%s token=%d chunk=%d duration=%v", docID, kbID, tokenNum, chunkNum, duration) + } + return errors.New("stats update failed") + }, + } + + deletedCount, err := svc.RemoveChunks(&service.RemoveChunksRequest{ + DocID: "doc-1", + ChunkIDs: []string{"chunk-1", "chunk-2"}, + }, "user-1") + if err == nil || !strings.Contains(err.Error(), "failed to update chunk stats") { + t.Fatalf("expected stats update error, got count=%d err=%v", deletedCount, err) + } + if deletedCount != 2 { + t.Fatalf("deleted count on error = %d, want 2", deletedCount) + } +} + +func TestDecrementChunkStatsClampsCounters(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + insertChunkTestKB(t, "kb-1", "tenant-1") + insertChunkTestDoc(t, "doc-1", "kb-1") + setChunkTestStats(t, "doc-1", "kb-1", 2, 1, 3, 2) + if err := dao.DB.Model(&entity.Document{}). + Where("id = ?", "doc-1"). + Update("process_duration", 0.5).Error; err != nil { + t.Fatalf("set process duration: %v", err) + } + + svc := &ChunkService{} + if err := svc.decrementChunkStats("doc-1", "kb-1", 5, 7, -1); err != nil { + t.Fatalf("decrementChunkStats() error = %v", err) + } + + doc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("get doc: %v", err) + } + if doc.TokenNum != 0 || doc.ChunkNum != 0 || doc.ProcessDuration != 0 { + t.Fatalf("document stats token=%d chunk=%d duration=%v, want zeros", doc.TokenNum, doc.ChunkNum, doc.ProcessDuration) + } + + kb, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + if err != nil { + t.Fatalf("get kb: %v", err) + } + if kb.TokenNum != 0 || kb.ChunkNum != 0 { + t.Fatalf("knowledgebase stats token=%d chunk=%d, want zeros", kb.TokenNum, kb.ChunkNum) + } +} + +func TestDecrementChunkStatsRollsBackWhenKnowledgebaseMissing(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + insertChunkTestDoc(t, "doc-1", "missing-kb") + + svc := &ChunkService{} + err := svc.decrementChunkStats("doc-1", "missing-kb", 0, 3, 0) + if err == nil || !strings.Contains(err.Error(), "knowledgebase not found") { + t.Fatalf("expected missing knowledgebase error, got %v", err) + } + + doc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("get doc: %v", err) + } + if doc.ChunkNum != 7 { + t.Fatalf("document chunk_num after rollback = %d, want 7", doc.ChunkNum) + } +} + func setupChunkTestDB(t *testing.T) *gorm.DB { t.Helper() @@ -777,6 +957,7 @@ func setupChunkTestDB(t *testing.T) *gorm.DB { t.Fatalf("failed to open sqlite: %v", err) } if err := db.AutoMigrate( + &entity.UserTenant{}, &entity.Knowledgebase{}, &entity.Document{}, &entity.Task{}, @@ -815,6 +996,23 @@ func newParseTestService(t *testing.T) *ChunkService { } } +func insertChunkTestUserTenant(t *testing.T, userID, tenantID string) { + t.Helper() + + status := "1" + userTenant := &entity.UserTenant{ + ID: userID + "-" + tenantID, + UserID: userID, + TenantID: tenantID, + Role: "owner", + InvitedBy: userID, + Status: &status, + } + if err := dao.DB.Create(userTenant).Error; err != nil { + t.Fatalf("insert user tenant: %v", err) + } +} + func insertChunkTestKB(t *testing.T, id, tenantID string) { t.Helper() @@ -859,6 +1057,21 @@ func insertChunkTestDoc(t *testing.T, id, kbID string) { } } +func setChunkTestStats(t *testing.T, docID, kbID string, docTokenNum, docChunkNum, kbTokenNum, kbChunkNum int64) { + t.Helper() + + if err := dao.DB.Model(&entity.Document{}). + Where("id = ?", docID). + Updates(map[string]interface{}{"token_num": docTokenNum, "chunk_num": docChunkNum}).Error; err != nil { + t.Fatalf("set doc stats: %v", err) + } + if err := dao.DB.Model(&entity.Knowledgebase{}). + Where("id = ?", kbID). + Updates(map[string]interface{}{"token_num": kbTokenNum, "chunk_num": kbChunkNum}).Error; err != nil { + t.Fatalf("set kb stats: %v", err) + } +} + func insertChunkTestTask(t *testing.T, id, docID string) { t.Helper() @@ -874,8 +1087,12 @@ func insertChunkTestTask(t *testing.T, id, docID string) { } type parseTestDocEngine struct { - deleteChunksErr error - deleteChunksCalls int + deleteChunksErr error + deleteChunksCount int64 + deleteChunksCalls int + deleteChunksCondition map[string]interface{} + deleteIndexName string + deleteDatasetID string } func (e *parseTestDocEngine) CreateChunkStore(context.Context, string, string, int, string) error { @@ -890,9 +1107,12 @@ func (e *parseTestDocEngine) UpdateChunks(context.Context, map[string]interface{ return nil } -func (e *parseTestDocEngine) DeleteChunks(context.Context, map[string]interface{}, string, string) (int64, error) { +func (e *parseTestDocEngine) DeleteChunks(_ context.Context, condition map[string]interface{}, indexName string, datasetID string) (int64, error) { e.deleteChunksCalls++ - return 0, e.deleteChunksErr + e.deleteChunksCondition = copyMap(condition) + e.deleteIndexName = indexName + e.deleteDatasetID = datasetID + return e.deleteChunksCount, e.deleteChunksErr } func (e *parseTestDocEngine) Search(context.Context, *types.SearchRequest) (*types.SearchResult, error) { diff --git a/internal/service/chunk_types.go b/internal/service/chunk_types.go index 2e147895f..7e2dbc2cf 100644 --- a/internal/service/chunk_types.go +++ b/internal/service/chunk_types.go @@ -22,6 +22,7 @@ import ( "fmt" "ragflow/internal/common" "ragflow/internal/engine/redis" + "ragflow/internal/entity" "ragflow/internal/server" "strings" @@ -30,6 +31,8 @@ import ( "ragflow/internal/engine/types" "ragflow/internal/tokenizer" "ragflow/internal/utility" + + "gorm.io/gorm" ) var ( @@ -51,6 +54,8 @@ type ChunkService struct { userTenantDAO *dao.UserTenantDAO documentDAO *dao.DocumentDAO searchService *SearchService + + decrementChunkStatsFunc func(string, string, int64, int64, float64) error } // RetrievalTestRequest retrieval test request @@ -727,9 +732,50 @@ func (s *ChunkService) RemoveChunks(req *RemoveChunksRequest, userID string) (in return 0, fmt.Errorf("failed to delete chunks: %w", err) } + if deletedCount > 0 { + if err := s.decrementChunkStats(req.DocID, doc.KbID, 0, deletedCount, 0); err != nil { + return deletedCount, fmt.Errorf("failed to update chunk stats: %w", err) + } + } + return deletedCount, nil } +func (s *ChunkService) decrementChunkStats(docID, kbID string, tokenNum, chunkNum int64, duration float64) error { + if s.decrementChunkStatsFunc != nil { + return s.decrementChunkStatsFunc(docID, kbID, tokenNum, chunkNum, duration) + } + return dao.DB.Transaction(func(tx *gorm.DB) error { + result := tx.Model(&entity.Document{}). + Where("id = ? AND kb_id = ?", docID, kbID). + Updates(map[string]interface{}{ + "token_num": gorm.Expr("CASE WHEN token_num - ? >= 0 THEN token_num - ? ELSE 0 END", tokenNum, tokenNum), + "chunk_num": gorm.Expr("CASE WHEN chunk_num - ? >= 0 THEN chunk_num - ? ELSE 0 END", chunkNum, chunkNum), + "process_duration": gorm.Expr("CASE WHEN process_duration + ? >= 0 THEN process_duration + ? ELSE 0 END", duration, duration), + }) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return fmt.Errorf("document not found") + } + + result = tx.Model(&entity.Knowledgebase{}). + Where("id = ?", kbID). + Updates(map[string]interface{}{ + "token_num": gorm.Expr("CASE WHEN token_num - ? >= 0 THEN token_num - ? ELSE 0 END", tokenNum, tokenNum), + "chunk_num": gorm.Expr("CASE WHEN chunk_num - ? >= 0 THEN chunk_num - ? ELSE 0 END", chunkNum, chunkNum), + }) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return fmt.Errorf("knowledgebase not found") + } + return nil + }) +} + // SourcedChunk is a typed, normalized view over a retrieval result chunk. // It decouples the ask pipeline (KbPrompt, ChunksFormat) from the raw // map[string]interface{} that flows through the retrieval engine. diff --git a/internal/service/chunk_types_test.go b/internal/service/chunk_types_test.go index ad8cefd42..f5e8d1f04 100644 --- a/internal/service/chunk_types_test.go +++ b/internal/service/chunk_types_test.go @@ -18,7 +18,12 @@ package service import ( "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" "testing" + + "github.com/glebarez/sqlite" + "gorm.io/gorm" ) func TestNewSourcedChunks_Empty(t *testing.T) { @@ -221,6 +226,86 @@ func TestSourcedChunk_ZeroValue(t *testing.T) { } } +func TestChunkTypesDecrementChunkStats(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.AutoMigrate(&entity.Knowledgebase{}, &entity.Document{}); err != nil { + t.Fatalf("migrate sqlite: %v", err) + } + previousDB := dao.DB + dao.DB = db + t.Cleanup(func() { dao.DB = previousDB }) + + status := string(entity.StatusValid) + if err := dao.DB.Create(&entity.Knowledgebase{ + ID: "kb-1", + TenantID: "tenant-1", + Name: "kb-1", + EmbdID: "embed", + Permission: string(entity.TenantPermissionMe), + CreatedBy: "user-1", + ParserConfig: entity.JSONMap{}, + TokenNum: 20, + ChunkNum: 5, + Status: &status, + }).Error; err != nil { + t.Fatalf("create kb: %v", err) + } + if err := dao.DB.Create(&entity.Document{ + ID: "doc-1", + KbID: "kb-1", + ParserID: string(entity.ParserTypeNaive), + ParserConfig: entity.JSONMap{}, + SourceType: string(entity.FileSourceLocal), + Type: "txt", + CreatedBy: "user-1", + TokenNum: 10, + ChunkNum: 3, + Suffix: ".txt", + Status: &status, + }).Error; err != nil { + t.Fatalf("create doc: %v", err) + } + + svc := &ChunkService{} + if err := svc.decrementChunkStats("doc-1", "kb-1", 0, 2, 0); err != nil { + t.Fatalf("decrementChunkStats() error = %v", err) + } + + var doc entity.Document + if err := dao.DB.First(&doc, "id = ?", "doc-1").Error; err != nil { + t.Fatalf("get doc: %v", err) + } + if doc.TokenNum != 10 || doc.ChunkNum != 1 { + t.Fatalf("document stats token=%d chunk=%d, want token=10 chunk=1", doc.TokenNum, doc.ChunkNum) + } + var kb entity.Knowledgebase + if err := dao.DB.First(&kb, "id = ?", "kb-1").Error; err != nil { + t.Fatalf("get kb: %v", err) + } + if kb.TokenNum != 20 || kb.ChunkNum != 3 { + t.Fatalf("knowledgebase stats token=%d chunk=%d, want token=20 chunk=3", kb.TokenNum, kb.ChunkNum) + } + + if err := svc.decrementChunkStats("doc-1", "kb-1", 30, 30, -1); err != nil { + t.Fatalf("decrementChunkStats() clamp error = %v", err) + } + if err := dao.DB.First(&doc, "id = ?", "doc-1").Error; err != nil { + t.Fatalf("get doc after clamp: %v", err) + } + if doc.TokenNum != 0 || doc.ChunkNum != 0 || doc.ProcessDuration != 0 { + t.Fatalf("document clamped stats token=%d chunk=%d duration=%v, want zeros", doc.TokenNum, doc.ChunkNum, doc.ProcessDuration) + } + if err := dao.DB.First(&kb, "id = ?", "kb-1").Error; err != nil { + t.Fatalf("get kb after clamp: %v", err) + } + if kb.TokenNum != 0 || kb.ChunkNum != 0 { + t.Fatalf("knowledgebase clamped stats token=%d chunk=%d, want zeros", kb.TokenNum, kb.ChunkNum) + } +} + func TestNewSourcedChunks_RoundTrip(t *testing.T) { original := []SourcedChunk{{ ID: "id1",