From 229282f558c725dfac2d49723409939b39bd43d5 Mon Sep 17 00:00:00 2001 From: mlogclub Date: Mon, 13 Apr 2026 18:17:29 +0800 Subject: [PATCH] feat: refactor document and FAQ index removal logic, introducing cleanup helper functions for improved structure and readability --- internal/ai/rag/index.go | 73 +++--------------------- internal/ai/rag/index_cleanup_helpers.go | 54 ++++++++++++++++++ 2 files changed, 61 insertions(+), 66 deletions(-) create mode 100644 internal/ai/rag/index_cleanup_helpers.go diff --git a/internal/ai/rag/index.go b/internal/ai/rag/index.go index 72b1422..46f6685 100644 --- a/internal/ai/rag/index.go +++ b/internal/ai/rag/index.go @@ -221,28 +221,11 @@ func (s *index) removeDocumentIndexByChunks(ctx context.Context, knowledgeBaseID return nil } - collectionName := s.getCollectionName() - provider := vectordb.GetProvider() - if provider == nil { - return fmt.Errorf("vectordb provider not initialized") + if err := s.deleteChunkVectors(ctx, collectChunkVectorIDs(chunks)); err != nil { + slog.Error("Failed to delete vectors", "error", err) } - vectorIDs := make([]string, 0, len(chunks)) - for _, chunk := range chunks { - if chunk.VectorID != "" { - vectorIDs = append(vectorIDs, chunk.VectorID) - } - } - - if len(vectorIDs) > 0 { - if err := provider.DeleteVectors(ctx, collectionName, vectorIDs); err != nil { - slog.Error("Failed to delete vectors", "error", err) - } - } - - if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error { - return ctx.Tx.Where("document_id = ?", documentID).Delete(&models.KnowledgeChunk{}).Error - }); err != nil { + if err := deleteChunksByCondition("document_id", documentID); err != nil { return fmt.Errorf("failed to delete chunks: %w", err) } @@ -267,25 +250,10 @@ func (s *index) removeFAQIndexByChunks(ctx context.Context, knowledgeBaseID int6 if len(chunks) == 0 { return nil } - collectionName := s.getCollectionName() - provider := vectordb.GetProvider() - if provider == nil { - return fmt.Errorf("vectordb provider not initialized") + if err := s.deleteChunkVectors(ctx, collectChunkVectorIDs(chunks)); err != nil { + slog.Error("Failed to delete faq vectors", "error", err) } - vectorIDs := make([]string, 0, len(chunks)) - for _, chunk := range chunks { - if chunk.VectorID != "" { - vectorIDs = append(vectorIDs, chunk.VectorID) - } - } - if len(vectorIDs) > 0 { - if err := provider.DeleteVectors(ctx, collectionName, vectorIDs); err != nil { - slog.Error("Failed to delete faq vectors", "error", err) - } - } - if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error { - return ctx.Tx.Where("faq_id = ?", faqID).Delete(&models.KnowledgeChunk{}).Error - }); err != nil { + if err := deleteChunksByCondition("faq_id", faqID); err != nil { return fmt.Errorf("failed to delete faq chunks: %w", err) } slog.Info("FAQ index removed", "faq_id", faqID, "chunks_removed", len(chunks)) @@ -455,33 +423,6 @@ func joinSimilarQuestions(items []string) string { } func (s *index) resetKnowledgeBaseIndexStorage(ctx context.Context, knowledgeBaseID int64) error { - collectionName := s.getCollectionName() - provider := vectordb.GetProvider() - if provider == nil { - return fmt.Errorf("vectordb provider not initialized") - } - chunks := repositories.KnowledgeChunkRepository.Find(sqls.DB(), sqls.NewCnd().Eq("knowledge_base_id", knowledgeBaseID)) - vectorIDs := make([]string, 0, len(chunks)) - for _, chunk := range chunks { - if strs.IsNotBlank(chunk.VectorID) { - vectorIDs = append(vectorIDs, chunk.VectorID) - } - } - if len(vectorIDs) > 0 { - if err := provider.DeleteVectors(ctx, collectionName, vectorIDs); err != nil { - return fmt.Errorf("failed to delete vectors for knowledge base %d before rebuild: %w", knowledgeBaseID, err) - } - } - - if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error { - return ctx.Tx.Where("knowledge_base_id = ?", knowledgeBaseID).Delete(&models.KnowledgeChunk{}).Error - }); err != nil { - return fmt.Errorf("failed to clear chunks before rebuild: %w", err) - } - - slog.Info("Knowledge base index storage reset", - "knowledge_base_id", knowledgeBaseID, - "collection", collectionName) - return nil + return s.cleanupKnowledgeBaseChunks(ctx, knowledgeBaseID, chunks) } diff --git a/internal/ai/rag/index_cleanup_helpers.go b/internal/ai/rag/index_cleanup_helpers.go new file mode 100644 index 0000000..2ed34c4 --- /dev/null +++ b/internal/ai/rag/index_cleanup_helpers.go @@ -0,0 +1,54 @@ +package rag + +import ( + "context" + "fmt" + "log/slog" + + "cs-agent/internal/ai/rag/vectordb" + "cs-agent/internal/models" + + "github.com/mlogclub/simple/common/strs" + "github.com/mlogclub/simple/sqls" +) + +func collectChunkVectorIDs(chunks []models.KnowledgeChunk) []string { + vectorIDs := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + if strs.IsNotBlank(chunk.VectorID) { + vectorIDs = append(vectorIDs, chunk.VectorID) + } + } + return vectorIDs +} + +func (s *index) deleteChunkVectors(ctx context.Context, vectorIDs []string) error { + if len(vectorIDs) == 0 { + return nil + } + provider := vectordb.GetProvider() + if provider == nil { + return fmt.Errorf("vectordb provider not initialized") + } + return provider.DeleteVectors(ctx, s.getCollectionName(), vectorIDs) +} + +func deleteChunksByCondition(column string, value int64) error { + return sqls.WithTransaction(func(ctx *sqls.TxContext) error { + return ctx.Tx.Where(column+" = ?", value).Delete(&models.KnowledgeChunk{}).Error + }) +} + +func (s *index) cleanupKnowledgeBaseChunks(ctx context.Context, knowledgeBaseID int64, chunks []models.KnowledgeChunk) error { + vectorIDs := collectChunkVectorIDs(chunks) + if err := s.deleteChunkVectors(ctx, vectorIDs); err != nil { + return fmt.Errorf("failed to delete vectors for knowledge base %d before rebuild: %w", knowledgeBaseID, err) + } + if err := deleteChunksByCondition("knowledge_base_id", knowledgeBaseID); err != nil { + return fmt.Errorf("failed to clear chunks before rebuild: %w", err) + } + slog.Info("Knowledge base index storage reset", + "knowledge_base_id", knowledgeBaseID, + "collection", s.getCollectionName()) + return nil +}