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

refactor: improve knowledge base index removal and streamline vector handling logic

Browse Source
This commit is contained in:
mlogclub
2026-06-02 11:47:35 +08:00
parent c2d05863ba
commit 99d5a951be
4 changed files with 35 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/ai/rag/index.go
+2 -2
View File
@@ -157,7 +157,7 @@ func (s *index) RemoveKnowledgeBaseIndex(ctx context.Context, knowledgeBaseID in
return 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)
slog.Error("Failed to delete knowledge base vectors", "knowledge_base_id", knowledgeBaseID, "error", 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)
@@ -326,6 +326,6 @@ func joinSimilarQuestions(items []string) string {
}
func (s *index) resetKnowledgeBaseIndexStorage(ctx context.Context, knowledgeBaseID int64) error {
chunks := repositories.KnowledgeChunkRepository.Find(sqls.DB(), sqls.NewCnd().Eq("knowledge_base_id", knowledgeBaseID))
chunks := repositories.KnowledgeChunkRepository.FindByKnowledgeBaseID(sqls.DB(), knowledgeBaseID)
return s.cleanupKnowledgeBaseChunks(ctx, knowledgeBaseID, chunks)
}
internal/ai/rag/index_document_helpers.go
-12
View File
@@ -12,8 +12,6 @@ import (
"agent-desk/internal/pkg/enums"
"agent-desk/internal/ai"
"github.com/mlogclub/simple/common/strs"
)
func (s *index) buildDocumentChunkRequest(document models.KnowledgeDocument, knowledgeBase models.KnowledgeBase) *ragchunk.ChunkRequest {
@@ -45,16 +43,6 @@ func (s *index) buildDocumentChunks(ctx context.Context, document models.Knowled
return chunks, nil
}
func (s *index) collectExistingVectorIDs(chunks []models.KnowledgeChunk) []string {
ret := make([]string, 0, len(chunks))
for _, chunk := range chunks {
if strs.IsNotBlank(chunk.VectorID) {
ret = append(ret, chunk.VectorID)
}
}
return ret
}
func (s *index) prepareDocumentVectors(ctx context.Context, knowledgeBase models.KnowledgeBase, document models.KnowledgeDocument, chunks []ragchunk.ChunkResult) ([]vectordb.Vector, []models.KnowledgeChunk, int, error) {
vectors := make([]vectordb.Vector, 0, len(chunks))
chunkModels := make([]models.KnowledgeChunk, 0, len(chunks))
internal/ai/rag/index_run_helpers.go
+32 -18
View File
@@ -3,13 +3,13 @@ package rag
import (
"context"
"fmt"
"log/slog"
"agent-desk/internal/ai"
"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"
)
@@ -29,7 +29,6 @@ func (s *index) runDocumentIndex(ctx context.Context, document models.KnowledgeD
return nil, 0, fmt.Errorf("failed to get embedding model: %w", err)
}
existingVectorIDs := s.collectExistingVectorIDs(existingChunks)
vectors, chunkModels, dimension, err := s.prepareDocumentVectors(ctx, knowledgeBase, document, chunks)
if err != nil {
return nil, 0, err
@@ -37,17 +36,17 @@ func (s *index) runDocumentIndex(ctx context.Context, document models.KnowledgeD
if err := s.ensureCollection(ctx, provider, collectionName, dimension); err != nil {
return nil, 0, err
}
if len(existingVectorIDs) > 0 {
if err := provider.DeleteVectors(ctx, collectionName, existingVectorIDs); err != nil {
return nil, 0, fmt.Errorf("failed to delete old vectors: %w", err)
}
}
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 {
return nil, 0, fmt.Errorf("failed to save chunks: %w", err)
}
if staleVectorIDs := s.collectStaleVectorIDs(existingChunks, vectors); len(staleVectorIDs) > 0 {
if err := provider.DeleteVectors(ctx, collectionName, staleVectorIDs); err != nil {
slog.Error("Failed to delete stale document vectors", "document_id", document.ID, "error", err)
}
}
return vectors, len(chunks), nil
}
@@ -75,22 +74,37 @@ func (s *index) runFAQIndex(ctx context.Context, faq models.KnowledgeFAQ, knowle
return err
}
existingVectorIDs := make([]string, 0, len(existingChunks))
for _, chunk := range existingChunks {
if strs.IsNotBlank(chunk.VectorID) {
existingVectorIDs = append(existingVectorIDs, chunk.VectorID)
}
}
if len(existingVectorIDs) > 0 {
if err := provider.DeleteVectors(ctx, collectionName, existingVectorIDs); err != nil {
return fmt.Errorf("failed to delete old vectors: %w", err)
}
}
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 {
return fmt.Errorf("failed to save faq chunk: %w", err)
}
if staleVectorIDs := s.collectStaleVectorIDs(existingChunks, []vectordb.Vector{vector}); len(staleVectorIDs) > 0 {
if err := provider.DeleteVectors(ctx, collectionName, staleVectorIDs); err != nil {
slog.Error("Failed to delete stale faq vectors", "faq_id", faq.ID, "error", err)
}
}
return nil
}
func (s *index) collectStaleVectorIDs(existingChunks []models.KnowledgeChunk, currentVectors []vectordb.Vector) []string {
currentVectorIDs := make(map[string]struct{}, len(currentVectors))
for _, vector := range currentVectors {
if vector.ID == "" {
continue
}
currentVectorIDs[vector.ID] = struct{}{}
}
staleVectorIDs := make([]string, 0, len(existingChunks))
for _, chunk := range existingChunks {
if chunk.VectorID == "" {
continue
}
if _, ok := currentVectorIDs[chunk.VectorID]; ok {
continue
}
staleVectorIDs = append(staleVectorIDs, chunk.VectorID)
}
return staleVectorIDs
}
internal/services/knowledge_base_service.go
+1 -5
View File
@@ -137,10 +137,6 @@ func (s *knowledgeBaseService) DeleteKnowledgeBase(id int64) error {
return errorsx.Forbidden(fmt.Sprintf("知识库已被 %d 个 AI Agent 引用，请先解除绑定", len(referencingAgents)))
}
if err := rag.Index.RemoveKnowledgeBaseIndex(context.Background(), id); err != nil {
return err
}
if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error {
if err := repositories.KnowledgeDocumentRepository.DeleteByKnowledgeBaseID(ctx.Tx, id); err != nil {
return err
@@ -153,7 +149,7 @@ func (s *knowledgeBaseService) DeleteKnowledgeBase(id int64) error {
return err
}
return nil
return rag.Index.RemoveKnowledgeBaseIndex(context.Background(), id)
}
func (s *knowledgeBaseService) UpdateSort(ids []int64) error {