feat: refactor indexing logic by introducing runDocumentIndex and runFAQIndex helper methods
This commit is contained in:
@@ -17,7 +17,6 @@ import (
|
|||||||
"cs-agent/internal/repositories"
|
"cs-agent/internal/repositories"
|
||||||
|
|
||||||
"github.com/google/uuid"
|
"github.com/google/uuid"
|
||||||
"github.com/mlogclub/simple/common/strs"
|
|
||||||
"github.com/mlogclub/simple/sqls"
|
"github.com/mlogclub/simple/sqls"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -73,55 +72,18 @@ func (s *index) IndexDocument(ctx context.Context, document *models.KnowledgeDoc
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return fail(err)
|
return fail(err)
|
||||||
}
|
}
|
||||||
|
vectors, chunkCount, err := s.runDocumentIndex(ctx, document, knowledgeBase)
|
||||||
existingChunks := repositories.KnowledgeChunkRepository.FindByDocumentID(sqls.DB(), document.ID)
|
|
||||||
|
|
||||||
chunks, err := s.buildDocumentChunks(ctx, document, knowledgeBase)
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fail(err)
|
return fail(err)
|
||||||
}
|
}
|
||||||
|
|
||||||
collectionName := s.getCollectionName()
|
|
||||||
provider := vectordb.GetProvider()
|
|
||||||
if provider == nil {
|
|
||||||
return fail(fmt.Errorf("vectordb provider not initialized"))
|
|
||||||
}
|
|
||||||
|
|
||||||
if _, err := ai.Embedding.GetModel(ctx); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to get embedding model: %w", err))
|
|
||||||
}
|
|
||||||
|
|
||||||
existingVectorIDs := collectExistingVectorIDs(existingChunks)
|
|
||||||
vectors, chunkModels, dimension, err := s.prepareDocumentVectors(ctx, knowledgeBase, document, chunks)
|
|
||||||
if err != nil {
|
|
||||||
return fail(err)
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := s.ensureCollection(ctx, provider, collectionName, dimension); err != nil {
|
|
||||||
return fail(err)
|
|
||||||
}
|
|
||||||
|
|
||||||
if len(existingVectorIDs) > 0 {
|
|
||||||
if err := provider.DeleteVectors(ctx, collectionName, existingVectorIDs); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to delete old vectors: %w", err))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := provider.UpsertVectors(ctx, collectionName, vectors); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to upsert vectors: %w", err))
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := s.replaceDocumentChunks(document.ID, chunkModels); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to save chunks: %w", err))
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := s.markDocumentIndexIndexed(document.ID); err != nil {
|
if err := s.markDocumentIndexIndexed(document.ID); err != nil {
|
||||||
slog.Error("Failed to mark knowledge document index as indexed", "document_id", document.ID, "error", err)
|
slog.Error("Failed to mark knowledge document index as indexed", "document_id", document.ID, "error", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
slog.Info("Document indexed successfully",
|
slog.Info("Document indexed successfully",
|
||||||
slog.Any("document_id", document.ID),
|
slog.Any("document_id", document.ID),
|
||||||
slog.Any("chunks_count", len(chunks)),
|
slog.Any("chunks_count", chunkCount),
|
||||||
slog.Any("vectors_count", len(vectors)),
|
slog.Any("vectors_count", len(vectors)),
|
||||||
slog.Any("time_taken", time.Since(start).String()),
|
slog.Any("time_taken", time.Since(start).String()),
|
||||||
)
|
)
|
||||||
@@ -147,48 +109,9 @@ func (s *index) IndexFAQByID(ctx context.Context, faqID int64) error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return fail(err)
|
return fail(err)
|
||||||
}
|
}
|
||||||
existingChunks := repositories.KnowledgeChunkRepository.FindByFaqID(sqls.DB(), faq.ID)
|
if err := s.runFAQIndex(ctx, faq, knowledgeBase); err != nil {
|
||||||
content := buildFAQChunkContent(faq)
|
|
||||||
if content == "" {
|
|
||||||
return fail(fmt.Errorf("faq content is empty"))
|
|
||||||
}
|
|
||||||
|
|
||||||
provider := vectordb.GetProvider()
|
|
||||||
if provider == nil {
|
|
||||||
return fail(fmt.Errorf("vectordb provider not initialized"))
|
|
||||||
}
|
|
||||||
if _, err := ai.Embedding.GetModel(ctx); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to get embedding model: %w", err))
|
|
||||||
}
|
|
||||||
vector, chunkModel, dimension, err := s.prepareFAQVector(ctx, knowledgeBase, faq, content)
|
|
||||||
if err != nil {
|
|
||||||
return fail(err)
|
return fail(err)
|
||||||
}
|
}
|
||||||
|
|
||||||
collectionName := s.getCollectionName()
|
|
||||||
if err := s.ensureCollection(ctx, provider, collectionName, dimension); err != nil {
|
|
||||||
return fail(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 fail(fmt.Errorf("failed to delete old vectors: %w", err))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := provider.UpsertVectors(ctx, collectionName, []vectordb.Vector{vector}); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to upsert vectors: %w", err))
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := s.replaceFAQChunk(faq.ID, &chunkModel); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to save faq chunk: %w", err))
|
|
||||||
}
|
|
||||||
if err := s.markFAQIndexIndexed(faq.ID); err != nil {
|
if err := s.markFAQIndexIndexed(faq.ID); err != nil {
|
||||||
slog.Error("Failed to mark knowledge faq index as indexed", "faq_id", faq.ID, "error", err)
|
slog.Error("Failed to mark knowledge faq index as indexed", "faq_id", faq.ID, "error", err)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,96 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"cs-agent/internal/ai"
|
||||||
|
"cs-agent/internal/ai/rag/vectordb"
|
||||||
|
"cs-agent/internal/models"
|
||||||
|
"cs-agent/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/common/strs"
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *index) runDocumentIndex(ctx context.Context, document *models.KnowledgeDocument, knowledgeBase *models.KnowledgeBase) ([]vectordb.Vector, int, error) {
|
||||||
|
existingChunks := repositories.KnowledgeChunkRepository.FindByDocumentID(sqls.DB(), document.ID)
|
||||||
|
chunks, err := s.buildDocumentChunks(ctx, document, knowledgeBase)
|
||||||
|
if err != nil {
|
||||||
|
return nil, 0, err
|
||||||
|
}
|
||||||
|
|
||||||
|
collectionName := s.getCollectionName()
|
||||||
|
provider := vectordb.GetProvider()
|
||||||
|
if provider == nil {
|
||||||
|
return nil, 0, fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
if _, err := ai.Embedding.GetModel(ctx); err != nil {
|
||||||
|
return nil, 0, fmt.Errorf("failed to get embedding model: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
existingVectorIDs := collectExistingVectorIDs(existingChunks)
|
||||||
|
vectors, chunkModels, dimension, err := s.prepareDocumentVectors(ctx, knowledgeBase, document, chunks)
|
||||||
|
if err != nil {
|
||||||
|
return nil, 0, err
|
||||||
|
}
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
return vectors, len(chunks), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) runFAQIndex(ctx context.Context, faq *models.KnowledgeFAQ, knowledgeBase *models.KnowledgeBase) error {
|
||||||
|
existingChunks := repositories.KnowledgeChunkRepository.FindByFaqID(sqls.DB(), faq.ID)
|
||||||
|
content := buildFAQChunkContent(faq)
|
||||||
|
if content == "" {
|
||||||
|
return fmt.Errorf("faq content is empty")
|
||||||
|
}
|
||||||
|
|
||||||
|
provider := vectordb.GetProvider()
|
||||||
|
if provider == nil {
|
||||||
|
return fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
if _, err := ai.Embedding.GetModel(ctx); err != nil {
|
||||||
|
return fmt.Errorf("failed to get embedding model: %w", err)
|
||||||
|
}
|
||||||
|
vector, chunkModel, dimension, err := s.prepareFAQVector(ctx, knowledgeBase, faq, content)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
collectionName := s.getCollectionName()
|
||||||
|
if err := s.ensureCollection(ctx, provider, collectionName, dimension); err != nil {
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user