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
docs
+1 -1
Submodule docs updated: 9526a7c086...f657db0533
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
internal/repositories/ai_agent_repository.go
+16
View File
@@ -1,7 +1,10 @@
package repositories
import (
"strconv"
"agent-desk/internal/models"
"agent-desk/internal/pkg/enums"
"agent-desk/internal/pkg/httpx/params"
@@ -109,3 +112,16 @@ func (r *aIAgentRepository) FindByIds(db *gorm.DB, ids []int64) []models.AIAgent
db.Where("id IN ?", ids).Find(&list)
return list
}
func (r *aIAgentRepository) FindByKnowledgeBaseID(db *gorm.DB, knowledgeBaseID int64) (list []models.AIAgent) {
id := strconv.FormatInt(knowledgeBaseID, 10)
db.Where(
"(knowledge_ids = ? OR knowledge_ids LIKE ? OR knowledge_ids LIKE ? OR knowledge_ids LIKE ?) AND status <> ?",
id,
id+",%",
"%,"+id,
"%,"+id+",%",
enums.StatusDeleted,
).Find(&list)
return
}
internal/repositories/knowledge_chunk_repository.go
+8 -4
View File
@@ -96,12 +96,16 @@ func (r *knowledgeChunkRepository) Delete(db *gorm.DB, id int64) {
db.Delete(&models.KnowledgeChunk{}, "id = ?", id)
}
func (r *knowledgeChunkRepository) DeleteByDocumentID(db *gorm.DB, documentID int64) {
db.Delete(&models.KnowledgeChunk{}, "document_id = ?", documentID)
func (r *knowledgeChunkRepository) DeleteByDocumentID(db *gorm.DB, documentID int64) error {
return db.Delete(&models.KnowledgeChunk{}, "document_id = ?", documentID).Error
}
func (r *knowledgeChunkRepository) DeleteByFaqID(db *gorm.DB, faqID int64) {
db.Delete(&models.KnowledgeChunk{}, "faq_id = ?", faqID)
func (r *knowledgeChunkRepository) DeleteByFaqID(db *gorm.DB, faqID int64) error {
return db.Delete(&models.KnowledgeChunk{}, "faq_id = ?", faqID).Error
}
func (r *knowledgeChunkRepository) DeleteByKnowledgeBaseID(db *gorm.DB, knowledgeBaseID int64) error {
return db.Delete(&models.KnowledgeChunk{}, "knowledge_base_id = ?", knowledgeBaseID).Error
}
func (r *knowledgeChunkRepository) FindByDocumentID(db *gorm.DB, documentID int64) (list []models.KnowledgeChunk) {
internal/repositories/knowledge_document_repository.go
+6 -1
View File
@@ -2,6 +2,7 @@ package repositories
import (
"agent-desk/internal/models"
"agent-desk/internal/pkg/enums"
"agent-desk/internal/pkg/httpx/params"
@@ -103,6 +104,10 @@ func (r *knowledgeDocumentRepository) Delete(db *gorm.DB, id int64) {
db.Delete(&models.KnowledgeDocument{}, "id = ?", id)
}
func (r *knowledgeDocumentRepository) DeleteByKnowledgeBaseID(db *gorm.DB, knowledgeBaseID int64) error {
return db.Delete(&models.KnowledgeDocument{}, "knowledge_base_id = ?", knowledgeBaseID).Error
}
func (r *knowledgeDocumentRepository) FindByIDs(db *gorm.DB, ids []int64) (list []models.KnowledgeDocument) {
if len(ids) == 0 {
return nil
@@ -113,6 +118,6 @@ func (r *knowledgeDocumentRepository) FindByIDs(db *gorm.DB, ids []int64) (list
func (r *knowledgeDocumentRepository) CountByKnowledgeBaseID(db *gorm.DB, knowledgeBaseID int64) int64 {
var count int64
db.Model(&models.KnowledgeDocument{}).Where("knowledge_base_id = ?", knowledgeBaseID).Count(&count)
db.Model(&models.KnowledgeDocument{}).Where("knowledge_base_id = ? AND status <> ?", knowledgeBaseID, enums.StatusDeleted).Count(&count)
return count
}
internal/repositories/knowledge_faq_repository.go
+5 -1
View File
@@ -58,6 +58,10 @@ func (r *knowledgeFAQRepository) Delete(db *gorm.DB, id int64) {
db.Delete(&models.KnowledgeFAQ{}, "id = ?", id)
}
func (r *knowledgeFAQRepository) DeleteByKnowledgeBaseID(db *gorm.DB, knowledgeBaseID int64) error {
return db.Delete(&models.KnowledgeFAQ{}, "knowledge_base_id = ?", knowledgeBaseID).Error
}
func (r *knowledgeFAQRepository) FindByIDs(db *gorm.DB, ids []int64) (list []models.KnowledgeFAQ) {
if len(ids) == 0 {
return nil
@@ -81,6 +85,6 @@ func (r *knowledgeFAQRepository) FindByKnowledgeBaseIDAndQuestions(db *gorm.DB,
func (r *knowledgeFAQRepository) CountByKnowledgeBaseID(db *gorm.DB, knowledgeBaseID int64) int64 {
var count int64
db.Model(&models.KnowledgeFAQ{}).Where("knowledge_base_id = ?", knowledgeBaseID).Count(&count)
db.Model(&models.KnowledgeFAQ{}).Where("knowledge_base_id = ? AND status <> ?", knowledgeBaseID, enums.StatusDeleted).Count(&count)
return count
}
internal/services/knowledge_base_service.go
+27 -8
View File
@@ -1,8 +1,11 @@
package services
import (
"context"
"fmt"
"time"
"agent-desk/internal/ai/rag"
"agent-desk/internal/models"
"agent-desk/internal/pkg/dto"
"agent-desk/internal/pkg/dto/request"
@@ -125,16 +128,32 @@ func (s *knowledgeBaseService) DeleteKnowledgeBase(id int64) error {
if current == nil {
return errorsx.InvalidParam("知识库不存在")
}
docCount := repositories.KnowledgeDocumentRepository.CountByKnowledgeBaseID(sqls.DB(), id)
if docCount > 0 {
return errorsx.InvalidParam("知识库下存在文档，无法删除")
referencingAgents := repositories.AIAgentRepository.FindByKnowledgeBaseID(sqls.DB(), id)
if len(referencingAgents) > 0 {
if len(referencingAgents) == 1 {
return errorsx.Forbidden(fmt.Sprintf("知识库已被 AI Agent「%s」引用，请先解除绑定", referencingAgents[0].Name))
}
return errorsx.Forbidden(fmt.Sprintf("知识库已被 %d 个 AI Agent 引用，请先解除绑定", len(referencingAgents)))
}
faqCount := repositories.KnowledgeFAQRepository.CountByKnowledgeBaseID(sqls.DB(), id)
if faqCount > 0 {
return errorsx.InvalidParam("知识库下存在FAQ，无法删除")
chunks := repositories.KnowledgeChunkRepository.Find(sqls.DB(), sqls.NewCnd().Eq("knowledge_base_id", id))
if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error {
if err := repositories.KnowledgeChunkRepository.DeleteByKnowledgeBaseID(ctx.Tx, id); err != nil {
return err
}
if err := repositories.KnowledgeDocumentRepository.DeleteByKnowledgeBaseID(ctx.Tx, id); err != nil {
return err
}
if err := repositories.KnowledgeFAQRepository.DeleteByKnowledgeBaseID(ctx.Tx, id); err != nil {
return err
}
return ctx.Tx.Delete(&models.KnowledgeBase{}, "id = ?", id).Error
}); err != nil {
return err
}
repositories.KnowledgeBaseRepository.Delete(sqls.DB(), id)
return nil
return rag.Index.RemoveKnowledgeBaseIndexByChunkModels(context.Background(), id, chunks)
}
func (s *knowledgeBaseService) UpdateSort(ids []int64) error {
internal/services/knowledge_base_service_test.go
+117
View File
@@ -1,9 +1,18 @@
package services
import (
"fmt"
"strings"
"testing"
"agent-desk/internal/models"
"agent-desk/internal/pkg/dto/request"
"agent-desk/internal/pkg/enums"
"agent-desk/internal/repositories"
"github.com/glebarez/sqlite"
"github.com/mlogclub/simple/sqls"
"gorm.io/gorm"
)
func TestBuildKnowledgeBaseModelUsesLowerDefaultScoreThreshold(t *testing.T) {
@@ -15,3 +24,111 @@ func TestBuildKnowledgeBaseModelUsesLowerDefaultScoreThreshold(t *testing.T) {
t.Fatalf("expected default score threshold 0.2, got %v", item.DefaultScoreThreshold)
}
}
func TestDeleteKnowledgeBaseRejectsAIAgentReference(t *testing.T) {
setupKnowledgeBaseServiceTestDB(t)
kb := createKnowledgeBaseServiceTestBase(t, "Referenced KB")
otherKB := createKnowledgeBaseServiceTestBase(t, "Other KB")
if err := repositories.AIAgentRepository.Create(sqls.DB(), &models.AIAgent{
Name: "Support Agent",
Status: enums.StatusOk,
KnowledgeIDs: "12",
}); err != nil {
t.Fatalf("create unrelated ai agent: %v", err)
}
if err := repositories.AIAgentRepository.Create(sqls.DB(), &models.AIAgent{
Name: "Knowledge Agent",
Status: enums.StatusOk,
KnowledgeIDs: fmt.Sprintf("12,%d,%d", kb.ID, otherKB.ID),
}); err != nil {
t.Fatalf("create ai agent: %v", err)
}
err := KnowledgeBaseService.DeleteKnowledgeBase(kb.ID)
if err == nil {
t.Fatal("DeleteKnowledgeBase() error is nil, want referenced knowledge base error")
}
if got := err.Error(); !strings.Contains(got, "Knowledge Agent") {
t.Fatalf("DeleteKnowledgeBase() error = %q, want agent name", got)
}
if repositories.KnowledgeBaseRepository.Get(sqls.DB(), kb.ID) == nil {
t.Fatal("knowledge base was deleted despite ai agent reference")
}
}
func TestDeleteKnowledgeBaseCascadesContentWhenNotReferenced(t *testing.T) {
setupKnowledgeBaseServiceTestDB(t)
kb := createKnowledgeBaseServiceTestBase(t, "Delete KB")
document := &models.KnowledgeDocument{
KnowledgeBaseID: kb.ID,
Title: "Doc",
ContentType: enums.KnowledgeDocumentContentTypeMarkdown,
Content: "content",
Status: enums.StatusOk,
IndexStatus: enums.KnowledgeDocumentIndexStatusIndexed,
}
if err := repositories.KnowledgeDocumentRepository.Create(sqls.DB(), document); err != nil {
t.Fatalf("create document: %v", err)
}
faq := &models.KnowledgeFAQ{
KnowledgeBaseID: kb.ID,
Question: "Question",
Answer: "Answer",
Status: enums.StatusOk,
IndexStatus: enums.KnowledgeDocumentIndexStatusIndexed,
}
if err := repositories.KnowledgeFAQRepository.Create(sqls.DB(), faq); err != nil {
t.Fatalf("create faq: %v", err)
}
if err := repositories.KnowledgeChunkRepository.BatchCreate(sqls.DB(), []models.KnowledgeChunk{
{KnowledgeBaseID: kb.ID, DocumentID: document.ID, ChunkNo: 1, Status: enums.StatusOk},
{KnowledgeBaseID: kb.ID, FaqID: faq.ID, ChunkNo: 1, Status: enums.StatusOk},
}); err != nil {
t.Fatalf("create chunks: %v", err)
}
if err := KnowledgeBaseService.DeleteKnowledgeBase(kb.ID); err != nil {
t.Fatalf("DeleteKnowledgeBase() error = %v", err)
}
assertKnowledgeBaseServiceTestCount(t, &models.KnowledgeBase{}, "id = ?", kb.ID, 0)
assertKnowledgeBaseServiceTestCount(t, &models.KnowledgeDocument{}, "knowledge_base_id = ?", kb.ID, 0)
assertKnowledgeBaseServiceTestCount(t, &models.KnowledgeFAQ{}, "knowledge_base_id = ?", kb.ID, 0)
assertKnowledgeBaseServiceTestCount(t, &models.KnowledgeChunk{}, "knowledge_base_id = ?", kb.ID, 0)
}
func setupKnowledgeBaseServiceTestDB(t *testing.T) {
t.Helper()
db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
if err != nil {
t.Fatalf("open sqlite db: %v", err)
}
if err := db.AutoMigrate(&models.KnowledgeBase{}, &models.KnowledgeDocument{}, &models.KnowledgeFAQ{}, &models.KnowledgeChunk{}, &models.AIAgent{}); err != nil {
t.Fatalf("auto migrate: %v", err)
}
sqls.SetDB(db)
}
func createKnowledgeBaseServiceTestBase(t *testing.T, name string) *models.KnowledgeBase {
t.Helper()
item := &models.KnowledgeBase{
Name: name,
KnowledgeType: string(enums.KnowledgeBaseTypeDocument),
Status: enums.StatusOk,
}
if err := repositories.KnowledgeBaseRepository.Create(sqls.DB(), item); err != nil {
t.Fatalf("create knowledge base: %v", err)
}
return item
}
func assertKnowledgeBaseServiceTestCount(t *testing.T, model any, query string, arg any, want int64) {
t.Helper()
var count int64
if err := sqls.DB().Model(model).Where(query, arg).Count(&count).Error; err != nil {
t.Fatalf("count %T: %v", model, err)
}
if count != want {
t.Fatalf("count %T = %d, want %d", model, count, want)
}
}
internal/services/knowledge_document_service.go
+2 -6
View File
@@ -155,7 +155,7 @@ func (s *knowledgeDocumentService) UpdateKnowledgeDocument(req request.UpdateKno
}
if oldKnowledgeBaseID != item.KnowledgeBaseID {
if err := rag.Index.RemoveDocumentIndexFromKnowledgeBase(context.Background(), oldKnowledgeBaseID, req.ID); err != nil {
if err := rag.Index.RemoveDocumentIndex(context.Background(), req.ID); err != nil {
slog.Error("failed to remove old document index after knowledge base change", "document_id", req.ID, "knowledge_base_id", oldKnowledgeBaseID, "error", err)
}
}
@@ -167,10 +167,6 @@ func (s *knowledgeDocumentService) UpdateKnowledgeDocument(req request.UpdateKno
}
func (s *knowledgeDocumentService) DeleteKnowledgeDocument(id int64) error {
current := s.Get(id)
if current == nil {
return errorsx.InvalidParam("文档不存在")
}
chunks := repositories.KnowledgeChunkRepository.FindByDocumentID(sqls.DB(), id)
if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error {
_ = repositories.KnowledgeDocumentRepository.Updates(ctx.Tx, id, map[string]any{
@@ -182,7 +178,7 @@ func (s *knowledgeDocumentService) DeleteKnowledgeDocument(id int64) error {
}); err != nil {
return err
}
return rag.Index.RemoveDocumentIndexByChunkModels(context.Background(), current.KnowledgeBaseID, id, chunks)
return rag.Index.RemoveDocumentIndexByChunkModels(context.Background(), id, chunks)
}
func (s *knowledgeDocumentService) buildKnowledgeDocumentModel(req request.CreateKnowledgeDocumentRequest) (*models.KnowledgeDocument, error) {
internal/services/knowledge_faq_service.go
+1 -1
View File
@@ -113,7 +113,7 @@ func (s *knowledgeFAQService) DeleteKnowledgeFAQ(id int64) error {
}); err != nil {
return err
}
return rag.Index.RemoveFAQIndexByChunkModels(context.Background(), current.KnowledgeBaseID, id, chunks)
return rag.Index.RemoveFAQIndexByChunkModels(context.Background(), id, chunks)
}
func (s *knowledgeFAQService) buildKnowledgeFAQModel(req request.CreateKnowledgeFAQRequest) (*models.KnowledgeFAQ, error) {