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

refactor: streamline document and FAQ index removal methods and enhance knowledge base deletion logic

Browse Source
This commit is contained in:
mlogclub
2026-06-02 11:25:30 +08:00
parent 5c4d486df0
commit abedafe911
13 changed files with 215 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/ai/rag/index.go
+26 -25
View File
@@ -119,33 +119,24 @@ func (s *index) IndexFAQByID(ctx context.Context, faqID int64) error {
}
func (s *index) RemoveDocumentIndex(ctx context.Context, documentID int64) error {
document := repositories.KnowledgeDocumentRepository.Get(sqls.DB(), documentID)
if document == nil {
return nil
}
chunks := repositories.KnowledgeChunkRepository.Find(sqls.DB(), sqls.NewCnd().Eq("document_id", documentID))
return s.removeDocumentIndexByChunks(ctx, document.KnowledgeBaseID, documentID, chunks)
return s.removeDocumentIndexByChunks(ctx, documentID, chunks)
}
func (s *index) RemoveDocumentIndexFromKnowledgeBase(ctx context.Context, knowledgeBaseID int64, documentID int64) error {
chunks := repositories.KnowledgeChunkRepository.Find(sqls.DB(), sqls.NewCnd().Eq("document_id", documentID))
return s.removeDocumentIndexByChunks(ctx, knowledgeBaseID, documentID, chunks)
func (s *index) RemoveDocumentIndexByChunkModels(ctx context.Context, documentID int64, chunks []models.KnowledgeChunk) error {
return s.removeDocumentIndexByChunks(ctx, documentID, chunks)
}
func (s *index) RemoveDocumentIndexByChunkModels(ctx context.Context, knowledgeBaseID int64, documentID int64, chunks []models.KnowledgeChunk) error {
return s.removeDocumentIndexByChunks(ctx, knowledgeBaseID, documentID, chunks)
}
func (s *index) removeDocumentIndexByChunks(ctx context.Context, knowledgeBaseID int64, documentID int64, chunks []models.KnowledgeChunk) error {
func (s *index) removeDocumentIndexByChunks(ctx context.Context, documentID int64, chunks []models.KnowledgeChunk) error {
if len(chunks) == 0 {
return nil
}
if err := s.deleteChunkVectors(ctx, collectChunkVectorIDs(chunks)); err != nil {
if err := s.deleteChunkVectors(ctx, s.collectChunkVectorIDs(chunks)); err != nil {
slog.Error("Failed to delete vectors", "error", err)
}
if err := deleteChunksByCondition("document_id", documentID); err != nil {
if err := repositories.KnowledgeChunkRepository.DeleteByDocumentID(sqls.DB(), documentID); err != nil {
return fmt.Errorf("failed to delete chunks: %w", err)
}
@@ -154,26 +145,36 @@ func (s *index) removeDocumentIndexByChunks(ctx context.Context, knowledgeBaseID
}
func (s *index) RemoveFAQIndex(ctx context.Context, faqID int64) error {
faq := repositories.KnowledgeFAQRepository.Get(sqls.DB(), faqID)
if faq == nil {
return nil
}
chunks := repositories.KnowledgeChunkRepository.FindByFaqID(sqls.DB(), faqID)
return s.removeFAQIndexByChunks(ctx, faq.KnowledgeBaseID, faqID, chunks)
return s.removeFAQIndexByChunks(ctx, faqID, chunks)
}
func (s *index) RemoveFAQIndexByChunkModels(ctx context.Context, knowledgeBaseID int64, faqID int64, chunks []models.KnowledgeChunk) error {
return s.removeFAQIndexByChunks(ctx, knowledgeBaseID, faqID, chunks)
func (s *index) RemoveFAQIndexByChunkModels(ctx context.Context, faqID int64, chunks []models.KnowledgeChunk) error {
return s.removeFAQIndexByChunks(ctx, faqID, chunks)
}
func (s *index) removeFAQIndexByChunks(ctx context.Context, knowledgeBaseID int64, faqID int64, chunks []models.KnowledgeChunk) error {
func (s *index) RemoveKnowledgeBaseIndexByChunkModels(ctx context.Context, knowledgeBaseID int64, chunks []models.KnowledgeChunk) error {
if len(chunks) == 0 {
return nil
}
if err := s.deleteChunkVectors(ctx, collectChunkVectorIDs(chunks)); err != nil {
if err := s.deleteChunkVectors(ctx, s.collectChunkVectorIDs(chunks)); err != nil {
return fmt.Errorf("failed to delete vectors for knowledge base %d: %w", knowledgeBaseID, err)
}
if err := repositories.KnowledgeChunkRepository.DeleteByKnowledgeBaseID(sqls.DB(), knowledgeBaseID); err != nil {
return fmt.Errorf("failed to delete chunks for knowledge base %d: %w", knowledgeBaseID, err)
}
slog.Info("Knowledge base index removed", "knowledge_base_id", knowledgeBaseID, "chunks_removed", len(chunks))
return nil
}
func (s *index) removeFAQIndexByChunks(ctx context.Context, faqID int64, chunks []models.KnowledgeChunk) error {
if len(chunks) == 0 {
return nil
}
if err := s.deleteChunkVectors(ctx, s.collectChunkVectorIDs(chunks)); err != nil {
slog.Error("Failed to delete faq vectors", "error", err)
}
if err := deleteChunksByCondition("faq_id", faqID); err != nil {
if err := repositories.KnowledgeChunkRepository.DeleteByFaqID(sqls.DB(), 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))
internal/ai/rag/index_cleanup_helpers.go
+4 -9
View File
@@ -7,12 +7,13 @@ import (
"agent-desk/internal/ai/rag/vectordb"
"agent-desk/internal/models"
"agent-desk/internal/repositories"
"github.com/mlogclub/simple/common/strs"
"github.com/mlogclub/simple/sqls"
)
func collectChunkVectorIDs(chunks []models.KnowledgeChunk) []string {
func (s *index) collectChunkVectorIDs(chunks []models.KnowledgeChunk) []string {
vectorIDs := make([]string, 0, len(chunks))
for _, chunk := range chunks {
if strs.IsNotBlank(chunk.VectorID) {
@@ -33,18 +34,12 @@ func (s *index) deleteChunkVectors(ctx context.Context, vectorIDs []string) erro
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)
vectorIDs := s.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 {
if err := repositories.KnowledgeChunkRepository.DeleteByKnowledgeBaseID(sqls.DB(), knowledgeBaseID); err != nil {
return fmt.Errorf("failed to clear chunks before rebuild: %w", err)
}
slog.Info("Knowledge base index storage reset",
internal/ai/rag/index_document_helpers.go
+1 -1
View File
@@ -45,7 +45,7 @@ func (s *index) buildDocumentChunks(ctx context.Context, document models.Knowled
return chunks, nil
}
func collectExistingVectorIDs(chunks []models.KnowledgeChunk) []string {
func (s *index) collectExistingVectorIDs(chunks []models.KnowledgeChunk) []string {
ret := make([]string, 0, len(chunks))
for _, chunk := range chunks {
if strs.IsNotBlank(chunk.VectorID) {
internal/ai/rag/index_run_helpers.go
+1 -1
View File
@@ -29,7 +29,7 @@ func (s *index) runDocumentIndex(ctx context.Context, document models.KnowledgeD
return nil, 0, fmt.Errorf("failed to get embedding model: %w", err)
}
existingVectorIDs := collectExistingVectorIDs(existingChunks)
existingVectorIDs := s.collectExistingVectorIDs(existingChunks)
vectors, chunkModels, dimension, err := s.prepareDocumentVectors(ctx, knowledgeBase, document, chunks)
if err != nil {
return nil, 0, err