wlw/ai-agent
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: refactor document and FAQ index removal logic, introducing cleanup helper functions for improved structure and readability

Browse Source
This commit is contained in:
mlogclub
2026-04-13 18:17:29 +08:00
parent 227ccedee6
commit 229282f558
2 changed files with 61 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/ai/rag/index.go
+7 -66
View File
@@ -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)
}
internal/ai/rag/index_cleanup_helpers.go
+54
View File
@@ -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
}