From c09ee0e596ff4b44dc694b1836ad0de37b9edc45 Mon Sep 17 00:00:00 2001
From: mlogclub <mlog1@qq.com>
Date: Tue, 2 Jun 2026 12:01:01 +0800
Subject: [PATCH] refactor: replace document and FAQ chunk handling with
 repository methods for improved clarity and maintainability

---
 internal/ai/rag/index_run_helpers.go          |  4 +--
 internal/ai/rag/index_storage_helpers.go      | 29 -------------------
 .../knowledge_chunk_repository.go             | 24 +++++++++++++++
 3 files changed, 26 insertions(+), 31 deletions(-)

diff --git a/internal/ai/rag/index_run_helpers.go b/internal/ai/rag/index_run_helpers.go
index 91fba07..327db8c 100644
--- a/internal/ai/rag/index_run_helpers.go
+++ b/internal/ai/rag/index_run_helpers.go
@@ -39,7 +39,7 @@ func (s *index) runDocumentIndex(ctx context.Context, document models.KnowledgeD
 	if err := provider.UpsertVectors(ctx, collectionName, vectors); err != nil {
 		return nil, 0, fmt.Errorf("failed to upsert vectors: %w", err)
 	}
-	if err := s.replaceDocumentChunks(document.ID, chunkModels); err != nil {
+	if err := repositories.KnowledgeChunkRepository.ReplaceByDocumentID(sqls.DB(), document.ID, chunkModels); err != nil {
 		return nil, 0, fmt.Errorf("failed to save chunks: %w", err)
 	}
 	if staleVectorIDs := s.collectStaleVectorIDs(existingChunks, vectors); len(staleVectorIDs) > 0 {
@@ -77,7 +77,7 @@ func (s *index) runFAQIndex(ctx context.Context, faq models.KnowledgeFAQ, knowle
 	if err := provider.UpsertVectors(ctx, collectionName, []vectordb.Vector{vector}); err != nil {
 		return fmt.Errorf("failed to upsert vectors: %w", err)
 	}
-	if err := s.replaceFAQChunk(faq.ID, &chunkModel); err != nil {
+	if err := repositories.KnowledgeChunkRepository.ReplaceByFaqID(sqls.DB(), faq.ID, &chunkModel); err != nil {
 		return fmt.Errorf("failed to save faq chunk: %w", err)
 	}
 	if staleVectorIDs := s.collectStaleVectorIDs(existingChunks, []vectordb.Vector{vector}); len(staleVectorIDs) > 0 {
diff --git a/internal/ai/rag/index_storage_helpers.go b/internal/ai/rag/index_storage_helpers.go
index 8c13228..9a4e7b2 100644
--- a/internal/ai/rag/index_storage_helpers.go
+++ b/internal/ai/rag/index_storage_helpers.go
@@ -6,9 +6,6 @@ import (
 	"log/slog"
 
 	"agent-desk/internal/ai/rag/vectordb"
-	"agent-desk/internal/models"
-
-	"github.com/mlogclub/simple/sqls"
 )
 
 func (s *index) ensureCollection(ctx context.Context, provider vectordb.Provider, collectionName string, dimension int) error {
@@ -25,29 +22,3 @@ func (s *index) ensureCollection(ctx context.Context, provider vectordb.Provider
 	slog.Info("Created collection for knowledge base", "collection", collectionName, "dimension", dimension)
 	return nil
 }
-
-func (s *index) replaceDocumentChunks(documentID int64, chunkModels []models.KnowledgeChunk) error {
-	return sqls.WithTransaction(func(ctx *sqls.TxContext) error {
-		if err := ctx.Tx.Where("document_id = ?", documentID).Delete(&models.KnowledgeChunk{}).Error; err != nil {
-			return err
-		}
-		for _, chunk := range chunkModels {
-			if err := ctx.Tx.Create(&chunk).Error; err != nil {
-				return err
-			}
-		}
-		return nil
-	})
-}
-
-func (s *index) replaceFAQChunk(faqID int64, chunkModel *models.KnowledgeChunk) error {
-	if chunkModel == nil {
-		return nil
-	}
-	return sqls.WithTransaction(func(ctx *sqls.TxContext) error {
-		if err := ctx.Tx.Where("faq_id = ?", faqID).Delete(&models.KnowledgeChunk{}).Error; err != nil {
-			return err
-		}
-		return ctx.Tx.Create(chunkModel).Error
-	})
-}
diff --git a/internal/repositories/knowledge_chunk_repository.go b/internal/repositories/knowledge_chunk_repository.go
index 0304e71..a19b7b7 100644
--- a/internal/repositories/knowledge_chunk_repository.go
+++ b/internal/repositories/knowledge_chunk_repository.go
@@ -77,6 +77,30 @@ func (r *knowledgeChunkRepository) BatchCreate(db *gorm.DB, list []models.Knowle
 	return
 }
 
+func (r *knowledgeChunkRepository) ReplaceByDocumentID(db *gorm.DB, documentID int64, list []models.KnowledgeChunk) error {
+	return sqls.WithTransaction(func(ctx *sqls.TxContext) error {
+		if err := ctx.Tx.Where("document_id = ?", documentID).Delete(&models.KnowledgeChunk{}).Error; err != nil {
+			return err
+		}
+		if len(list) == 0 {
+			return nil
+		}
+		return ctx.Tx.Create(&list).Error
+	})
+}
+
+func (r *knowledgeChunkRepository) ReplaceByFaqID(db *gorm.DB, faqID int64, item *models.KnowledgeChunk) error {
+	return sqls.WithTransaction(func(ctx *sqls.TxContext) error {
+		if err := ctx.Tx.Where("faq_id = ?", faqID).Delete(&models.KnowledgeChunk{}).Error; err != nil {
+			return err
+		}
+		if item == nil {
+			return nil
+		}
+		return ctx.Tx.Create(item).Error
+	})
+}
+
 func (r *knowledgeChunkRepository) Update(db *gorm.DB, t *models.KnowledgeChunk) (err error) {
 	err = db.Save(t).Error
 	return