feat: refactor document and FAQ indexing logic, introducing helper functions for improved structure and readability
This commit is contained in:
+15
-230
@@ -76,26 +76,9 @@ func (s *index) IndexDocument(ctx context.Context, document *models.KnowledgeDoc
|
|||||||
|
|
||||||
existingChunks := repositories.KnowledgeChunkRepository.FindByDocumentID(sqls.DB(), document.ID)
|
existingChunks := repositories.KnowledgeChunkRepository.FindByDocumentID(sqls.DB(), document.ID)
|
||||||
|
|
||||||
chunks, err := s.registry.Chunk(ctx, &ragchunk.ChunkRequest{
|
chunks, err := s.buildDocumentChunks(ctx, document, knowledgeBase)
|
||||||
KnowledgeBaseID: document.KnowledgeBaseID,
|
|
||||||
DocumentID: document.ID,
|
|
||||||
DocumentTitle: document.Title,
|
|
||||||
ContentType: document.ContentType,
|
|
||||||
Content: document.Content,
|
|
||||||
PlainText: ExtractPlainText(document.Content, document.ContentType),
|
|
||||||
Options: ragchunk.ChunkOptions{
|
|
||||||
Provider: firstNonEmptyString(knowledgeBase.ChunkProvider, s.chunkConfig.Provider),
|
|
||||||
TargetTokens: firstPositiveInt(knowledgeBase.ChunkTargetTokens, s.chunkConfig.TargetTokens),
|
|
||||||
MaxTokens: firstPositiveInt(knowledgeBase.ChunkMaxTokens, s.chunkConfig.MaxTokens),
|
|
||||||
OverlapTokens: firstPositiveInt(knowledgeBase.ChunkOverlapTokens, s.chunkConfig.OverlapTokens),
|
|
||||||
EnableFallback: s.chunkConfig.EnableFallback,
|
|
||||||
},
|
|
||||||
})
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fail(fmt.Errorf("failed to chunk document: %w", err))
|
return fail(err)
|
||||||
}
|
|
||||||
if len(chunks) == 0 {
|
|
||||||
return fail(fmt.Errorf("no chunks generated from document"))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
collectionName := s.getCollectionName()
|
collectionName := s.getCollectionName()
|
||||||
@@ -108,83 +91,14 @@ func (s *index) IndexDocument(ctx context.Context, document *models.KnowledgeDoc
|
|||||||
return fail(fmt.Errorf("failed to get embedding model: %w", err))
|
return fail(fmt.Errorf("failed to get embedding model: %w", err))
|
||||||
}
|
}
|
||||||
|
|
||||||
existingVectorIDs := make([]string, 0, len(existingChunks))
|
existingVectorIDs := collectExistingVectorIDs(existingChunks)
|
||||||
for _, chunk := range existingChunks {
|
vectors, chunkModels, dimension, err := s.prepareDocumentVectors(ctx, knowledgeBase, document, chunks)
|
||||||
if strs.IsNotBlank(chunk.VectorID) {
|
if err != nil {
|
||||||
existingVectorIDs = append(existingVectorIDs, chunk.VectorID)
|
return fail(err)
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
vectors := make([]vectordb.Vector, 0, len(chunks))
|
if err := s.ensureCollection(ctx, provider, collectionName, dimension); err != nil {
|
||||||
chunkModels := make([]models.KnowledgeChunk, 0, len(chunks))
|
return fail(err)
|
||||||
dimension := 0
|
|
||||||
|
|
||||||
for i, chunk := range chunks {
|
|
||||||
embeddingResult, err := ai.Embedding.GenerateEmbedding(ctx, chunk.Content)
|
|
||||||
if err != nil {
|
|
||||||
slog.Error("Failed to generate embedding for chunk", "document_id", document.ID, "chunk_index", i, "error", err)
|
|
||||||
return fail(fmt.Errorf("failed to generate embedding for chunk %d: %w", i, err))
|
|
||||||
}
|
|
||||||
if dimension == 0 {
|
|
||||||
dimension = embeddingResult.Dimension
|
|
||||||
}
|
|
||||||
|
|
||||||
chunkID := buildKnowledgeChunkVectorID(knowledgeBase.ID, document.ID, chunk.ChunkNo)
|
|
||||||
providerName := ""
|
|
||||||
if chunk.Metadata != nil {
|
|
||||||
if value, ok := chunk.Metadata["provider"].(string); ok {
|
|
||||||
providerName = value
|
|
||||||
}
|
|
||||||
}
|
|
||||||
chunkModel := models.KnowledgeChunk{
|
|
||||||
KnowledgeBaseID: knowledgeBase.ID,
|
|
||||||
DocumentID: document.ID,
|
|
||||||
ChunkNo: chunk.ChunkNo,
|
|
||||||
Title: chunk.Title,
|
|
||||||
Content: chunk.Content,
|
|
||||||
ContentHash: buildChunkContentHash(chunk.Content),
|
|
||||||
CharCount: chunk.CharCount,
|
|
||||||
TokenCount: chunk.TokenCount,
|
|
||||||
ChunkType: string(chunk.ChunkType),
|
|
||||||
SectionPath: chunk.SectionPath,
|
|
||||||
Provider: providerName,
|
|
||||||
VectorID: chunkID,
|
|
||||||
Status: enums.StatusOk,
|
|
||||||
CreatedAt: time.Now(),
|
|
||||||
UpdatedAt: time.Now(),
|
|
||||||
}
|
|
||||||
chunkModels = append(chunkModels, chunkModel)
|
|
||||||
|
|
||||||
vectors = append(vectors, vectordb.Vector{
|
|
||||||
ID: chunkID,
|
|
||||||
Vector: embeddingResult.Vector,
|
|
||||||
Payload: vectordb.ChunkPayload{
|
|
||||||
KnowledgeBaseID: knowledgeBase.ID,
|
|
||||||
DocumentID: document.ID,
|
|
||||||
DocumentTitle: document.Title,
|
|
||||||
ChunkNo: chunk.ChunkNo,
|
|
||||||
ChunkType: string(chunk.ChunkType),
|
|
||||||
SectionPath: chunk.SectionPath,
|
|
||||||
Content: chunk.Content,
|
|
||||||
Title: chunk.Title,
|
|
||||||
Provider: providerName,
|
|
||||||
},
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
if len(vectors) == 0 {
|
|
||||||
return fail(fmt.Errorf("no vectors generated"))
|
|
||||||
}
|
|
||||||
|
|
||||||
collectionInfo, err := provider.GetCollection(ctx, collectionName)
|
|
||||||
if err != nil || collectionInfo == nil {
|
|
||||||
if dimension <= 0 {
|
|
||||||
return fail(fmt.Errorf("invalid embedding dimension: %d", dimension))
|
|
||||||
}
|
|
||||||
if err := provider.CreateCollection(ctx, collectionName, dimension); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to create collection: %w", err))
|
|
||||||
}
|
|
||||||
slog.Info("Created collection for knowledge base", "collection", collectionName, "dimension", dimension)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if len(existingVectorIDs) > 0 {
|
if len(existingVectorIDs) > 0 {
|
||||||
@@ -197,17 +111,7 @@ func (s *index) IndexDocument(ctx context.Context, document *models.KnowledgeDoc
|
|||||||
return fail(fmt.Errorf("failed to upsert vectors: %w", err))
|
return fail(fmt.Errorf("failed to upsert vectors: %w", err))
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error {
|
if err := s.replaceDocumentChunks(document.ID, chunkModels); err != nil {
|
||||||
if err := ctx.Tx.Where("document_id = ?", document.ID).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
|
|
||||||
}); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to save chunks: %w", err))
|
return fail(fmt.Errorf("failed to save chunks: %w", err))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -259,35 +163,14 @@ func (s *index) IndexFAQByID(ctx context.Context, faqID int64) error {
|
|||||||
if _, err := ai.Embedding.GetModel(ctx); err != nil {
|
if _, err := ai.Embedding.GetModel(ctx); err != nil {
|
||||||
return fail(fmt.Errorf("failed to get embedding model: %w", err))
|
return fail(fmt.Errorf("failed to get embedding model: %w", err))
|
||||||
}
|
}
|
||||||
embeddingResult, err := ai.Embedding.GenerateEmbedding(ctx, content)
|
vector, chunkModel, dimension, err := s.prepareFAQVector(ctx, knowledgeBase, faq, content)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fail(fmt.Errorf("failed to generate embedding for faq %d: %w", faq.ID, err))
|
return fail(err)
|
||||||
}
|
|
||||||
|
|
||||||
chunkID := buildKnowledgeFAQChunkVectorID(knowledgeBase.ID, faq.ID, 0)
|
|
||||||
chunkModel := models.KnowledgeChunk{
|
|
||||||
KnowledgeBaseID: knowledgeBase.ID,
|
|
||||||
FaqID: faq.ID,
|
|
||||||
ChunkNo: 0,
|
|
||||||
Title: faq.Question,
|
|
||||||
Content: content,
|
|
||||||
ContentHash: buildChunkContentHash(content),
|
|
||||||
CharCount: len([]rune(content)),
|
|
||||||
TokenCount: len([]rune(content)) / 2,
|
|
||||||
ChunkType: string(enums.KnowledgeChunkTypeFAQ),
|
|
||||||
Provider: string(enums.KnowledgeChunkProviderFAQ),
|
|
||||||
VectorID: chunkID,
|
|
||||||
Status: enums.StatusOk,
|
|
||||||
CreatedAt: time.Now(),
|
|
||||||
UpdatedAt: time.Now(),
|
|
||||||
}
|
}
|
||||||
|
|
||||||
collectionName := s.getCollectionName()
|
collectionName := s.getCollectionName()
|
||||||
collectionInfo, err := provider.GetCollection(ctx, collectionName)
|
if err := s.ensureCollection(ctx, provider, collectionName, dimension); err != nil {
|
||||||
if err != nil || collectionInfo == nil {
|
return fail(err)
|
||||||
if err := provider.CreateCollection(ctx, collectionName, embeddingResult.Dimension); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to create collection: %w", err))
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
existingVectorIDs := make([]string, 0, len(existingChunks))
|
existingVectorIDs := make([]string, 0, len(existingChunks))
|
||||||
@@ -302,29 +185,11 @@ func (s *index) IndexFAQByID(ctx context.Context, faqID int64) error {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := provider.UpsertVectors(ctx, collectionName, []vectordb.Vector{{
|
if err := provider.UpsertVectors(ctx, collectionName, []vectordb.Vector{vector}); err != nil {
|
||||||
ID: chunkID,
|
|
||||||
Vector: embeddingResult.Vector,
|
|
||||||
Payload: vectordb.ChunkPayload{
|
|
||||||
KnowledgeBaseID: knowledgeBase.ID,
|
|
||||||
FaqID: faq.ID,
|
|
||||||
FaqQuestion: faq.Question,
|
|
||||||
ChunkNo: 0,
|
|
||||||
ChunkType: string(enums.KnowledgeChunkTypeFAQ),
|
|
||||||
Content: content,
|
|
||||||
Title: faq.Question,
|
|
||||||
Provider: string(enums.KnowledgeChunkProviderFAQ),
|
|
||||||
},
|
|
||||||
}}); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to upsert vectors: %w", err))
|
return fail(fmt.Errorf("failed to upsert vectors: %w", err))
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error {
|
if err := s.replaceFAQChunk(faq.ID, &chunkModel); err != nil {
|
||||||
if err := ctx.Tx.Where("faq_id = ?", faq.ID).Delete(&models.KnowledgeChunk{}).Error; err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
return ctx.Tx.Create(&chunkModel).Error
|
|
||||||
}); err != nil {
|
|
||||||
return fail(fmt.Errorf("failed to save faq chunk: %w", err))
|
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 {
|
||||||
@@ -589,86 +454,6 @@ func joinSimilarQuestions(items []string) string {
|
|||||||
return result
|
return result
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *index) markDocumentIndexPending(documentID int64) error {
|
|
||||||
return repositories.KnowledgeDocumentRepository.Updates(sqls.DB(), documentID, map[string]any{
|
|
||||||
"index_status": enums.KnowledgeDocumentIndexStatusPending,
|
|
||||||
"indexed_at": nil,
|
|
||||||
"index_error": "",
|
|
||||||
"updated_at": time.Now(),
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *index) markDocumentIndexIndexed(documentID int64) error {
|
|
||||||
now := time.Now()
|
|
||||||
return repositories.KnowledgeDocumentRepository.Updates(sqls.DB(), documentID, map[string]any{
|
|
||||||
"index_status": enums.KnowledgeDocumentIndexStatusIndexed,
|
|
||||||
"indexed_at": &now,
|
|
||||||
"index_error": "",
|
|
||||||
"updated_at": now,
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *index) markDocumentIndexFailed(documentID int64, err error) error {
|
|
||||||
return repositories.KnowledgeDocumentRepository.Updates(sqls.DB(), documentID, map[string]any{
|
|
||||||
"index_status": enums.KnowledgeDocumentIndexStatusFailed,
|
|
||||||
"index_error": truncateIndexError(err),
|
|
||||||
"updated_at": time.Now(),
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *index) markKnowledgeBaseDocumentsIndexPending(knowledgeBaseID int64, documentIDs []int64) error {
|
|
||||||
if len(documentIDs) == 0 {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
return sqls.DB().Model(&models.KnowledgeDocument{}).
|
|
||||||
Where("knowledge_base_id = ?", knowledgeBaseID).
|
|
||||||
Where("id IN ?", documentIDs).
|
|
||||||
Updates(map[string]any{
|
|
||||||
"index_status": enums.KnowledgeDocumentIndexStatusPending,
|
|
||||||
"indexed_at": nil,
|
|
||||||
"index_error": "",
|
|
||||||
"updated_at": time.Now(),
|
|
||||||
}).Error
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *index) markFAQIndexPending(faqID int64) error {
|
|
||||||
return repositories.KnowledgeFAQRepository.Updates(sqls.DB(), faqID, map[string]any{
|
|
||||||
"index_status": enums.KnowledgeDocumentIndexStatusPending,
|
|
||||||
"indexed_at": nil,
|
|
||||||
"index_error": "",
|
|
||||||
"updated_at": time.Now(),
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *index) markFAQIndexIndexed(faqID int64) error {
|
|
||||||
now := time.Now()
|
|
||||||
return repositories.KnowledgeFAQRepository.Updates(sqls.DB(), faqID, map[string]any{
|
|
||||||
"index_status": enums.KnowledgeDocumentIndexStatusIndexed,
|
|
||||||
"indexed_at": &now,
|
|
||||||
"index_error": "",
|
|
||||||
"updated_at": now,
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *index) markFAQIndexFailed(faqID int64, err error) error {
|
|
||||||
return repositories.KnowledgeFAQRepository.Updates(sqls.DB(), faqID, map[string]any{
|
|
||||||
"index_status": enums.KnowledgeDocumentIndexStatusFailed,
|
|
||||||
"index_error": truncateIndexError(err),
|
|
||||||
"updated_at": time.Now(),
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
func truncateIndexError(err error) string {
|
|
||||||
if err == nil {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
message := err.Error()
|
|
||||||
if len(message) <= 1000 {
|
|
||||||
return message
|
|
||||||
}
|
|
||||||
return message[:1000]
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *index) resetKnowledgeBaseIndexStorage(ctx context.Context, knowledgeBaseID int64) error {
|
func (s *index) resetKnowledgeBaseIndexStorage(ctx context.Context, knowledgeBaseID int64) error {
|
||||||
collectionName := s.getCollectionName()
|
collectionName := s.getCollectionName()
|
||||||
provider := vectordb.GetProvider()
|
provider := vectordb.GetProvider()
|
||||||
|
|||||||
@@ -0,0 +1,119 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
ragchunk "cs-agent/internal/ai/rag/chunk"
|
||||||
|
"cs-agent/internal/ai/rag/vectordb"
|
||||||
|
"cs-agent/internal/models"
|
||||||
|
"cs-agent/internal/pkg/enums"
|
||||||
|
|
||||||
|
"cs-agent/internal/ai"
|
||||||
|
"github.com/mlogclub/simple/common/strs"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *index) buildDocumentChunkRequest(document *models.KnowledgeDocument, knowledgeBase *models.KnowledgeBase) *ragchunk.ChunkRequest {
|
||||||
|
return &ragchunk.ChunkRequest{
|
||||||
|
KnowledgeBaseID: document.KnowledgeBaseID,
|
||||||
|
DocumentID: document.ID,
|
||||||
|
DocumentTitle: document.Title,
|
||||||
|
ContentType: document.ContentType,
|
||||||
|
Content: document.Content,
|
||||||
|
PlainText: ExtractPlainText(document.Content, document.ContentType),
|
||||||
|
Options: ragchunk.ChunkOptions{
|
||||||
|
Provider: firstNonEmptyString(knowledgeBase.ChunkProvider, s.chunkConfig.Provider),
|
||||||
|
TargetTokens: firstPositiveInt(knowledgeBase.ChunkTargetTokens, s.chunkConfig.TargetTokens),
|
||||||
|
MaxTokens: firstPositiveInt(knowledgeBase.ChunkMaxTokens, s.chunkConfig.MaxTokens),
|
||||||
|
OverlapTokens: firstPositiveInt(knowledgeBase.ChunkOverlapTokens, s.chunkConfig.OverlapTokens),
|
||||||
|
EnableFallback: s.chunkConfig.EnableFallback,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) buildDocumentChunks(ctx context.Context, document *models.KnowledgeDocument, knowledgeBase *models.KnowledgeBase) ([]ragchunk.ChunkResult, error) {
|
||||||
|
chunks, err := s.registry.Chunk(ctx, s.buildDocumentChunkRequest(document, knowledgeBase))
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to chunk document: %w", err)
|
||||||
|
}
|
||||||
|
if len(chunks) == 0 {
|
||||||
|
return nil, fmt.Errorf("no chunks generated from document")
|
||||||
|
}
|
||||||
|
return chunks, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func 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))
|
||||||
|
dimension := 0
|
||||||
|
|
||||||
|
for i, chunk := range chunks {
|
||||||
|
embeddingResult, err := ai.Embedding.GenerateEmbedding(ctx, chunk.Content)
|
||||||
|
if err != nil {
|
||||||
|
slog.Error("Failed to generate embedding for chunk", "document_id", document.ID, "chunk_index", i, "error", err)
|
||||||
|
return nil, nil, 0, fmt.Errorf("failed to generate embedding for chunk %d: %w", i, err)
|
||||||
|
}
|
||||||
|
if dimension == 0 {
|
||||||
|
dimension = embeddingResult.Dimension
|
||||||
|
}
|
||||||
|
|
||||||
|
chunkID := buildKnowledgeChunkVectorID(knowledgeBase.ID, document.ID, chunk.ChunkNo)
|
||||||
|
providerName := ""
|
||||||
|
if chunk.Metadata != nil {
|
||||||
|
if value, ok := chunk.Metadata["provider"].(string); ok {
|
||||||
|
providerName = value
|
||||||
|
}
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
chunkModels = append(chunkModels, models.KnowledgeChunk{
|
||||||
|
KnowledgeBaseID: knowledgeBase.ID,
|
||||||
|
DocumentID: document.ID,
|
||||||
|
ChunkNo: chunk.ChunkNo,
|
||||||
|
Title: chunk.Title,
|
||||||
|
Content: chunk.Content,
|
||||||
|
ContentHash: buildChunkContentHash(chunk.Content),
|
||||||
|
CharCount: chunk.CharCount,
|
||||||
|
TokenCount: chunk.TokenCount,
|
||||||
|
ChunkType: string(chunk.ChunkType),
|
||||||
|
SectionPath: chunk.SectionPath,
|
||||||
|
Provider: providerName,
|
||||||
|
VectorID: chunkID,
|
||||||
|
Status: enums.StatusOk,
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
})
|
||||||
|
|
||||||
|
vectors = append(vectors, vectordb.Vector{
|
||||||
|
ID: chunkID,
|
||||||
|
Vector: embeddingResult.Vector,
|
||||||
|
Payload: vectordb.ChunkPayload{
|
||||||
|
KnowledgeBaseID: knowledgeBase.ID,
|
||||||
|
DocumentID: document.ID,
|
||||||
|
DocumentTitle: document.Title,
|
||||||
|
ChunkNo: chunk.ChunkNo,
|
||||||
|
ChunkType: string(chunk.ChunkType),
|
||||||
|
SectionPath: chunk.SectionPath,
|
||||||
|
Content: chunk.Content,
|
||||||
|
Title: chunk.Title,
|
||||||
|
Provider: providerName,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(vectors) == 0 {
|
||||||
|
return nil, nil, 0, fmt.Errorf("no vectors generated")
|
||||||
|
}
|
||||||
|
return vectors, chunkModels, dimension, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,57 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"cs-agent/internal/ai"
|
||||||
|
"cs-agent/internal/ai/rag/vectordb"
|
||||||
|
"cs-agent/internal/models"
|
||||||
|
"cs-agent/internal/pkg/enums"
|
||||||
|
)
|
||||||
|
|
||||||
|
func buildFAQChunkModel(knowledgeBase *models.KnowledgeBase, faq *models.KnowledgeFAQ, content string) (models.KnowledgeChunk, string) {
|
||||||
|
chunkID := buildKnowledgeFAQChunkVectorID(knowledgeBase.ID, faq.ID, 0)
|
||||||
|
now := time.Now()
|
||||||
|
return models.KnowledgeChunk{
|
||||||
|
KnowledgeBaseID: knowledgeBase.ID,
|
||||||
|
FaqID: faq.ID,
|
||||||
|
ChunkNo: 0,
|
||||||
|
Title: faq.Question,
|
||||||
|
Content: content,
|
||||||
|
ContentHash: buildChunkContentHash(content),
|
||||||
|
CharCount: len([]rune(content)),
|
||||||
|
TokenCount: len([]rune(content)) / 2,
|
||||||
|
ChunkType: string(enums.KnowledgeChunkTypeFAQ),
|
||||||
|
Provider: string(enums.KnowledgeChunkProviderFAQ),
|
||||||
|
VectorID: chunkID,
|
||||||
|
Status: enums.StatusOk,
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
}, chunkID
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) prepareFAQVector(ctx context.Context, knowledgeBase *models.KnowledgeBase, faq *models.KnowledgeFAQ, content string) (vectordb.Vector, models.KnowledgeChunk, int, error) {
|
||||||
|
embeddingResult, err := ai.Embedding.GenerateEmbedding(ctx, content)
|
||||||
|
if err != nil {
|
||||||
|
return vectordb.Vector{}, models.KnowledgeChunk{}, 0, fmt.Errorf("failed to generate embedding for faq %d: %w", faq.ID, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
chunkModel, chunkID := buildFAQChunkModel(knowledgeBase, faq, content)
|
||||||
|
vector := vectordb.Vector{
|
||||||
|
ID: chunkID,
|
||||||
|
Vector: embeddingResult.Vector,
|
||||||
|
Payload: vectordb.ChunkPayload{
|
||||||
|
KnowledgeBaseID: knowledgeBase.ID,
|
||||||
|
FaqID: faq.ID,
|
||||||
|
FaqQuestion: faq.Question,
|
||||||
|
ChunkNo: 0,
|
||||||
|
ChunkType: string(enums.KnowledgeChunkTypeFAQ),
|
||||||
|
Content: content,
|
||||||
|
Title: faq.Question,
|
||||||
|
Provider: string(enums.KnowledgeChunkProviderFAQ),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
return vector, chunkModel, embeddingResult.Dimension, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,91 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"cs-agent/internal/models"
|
||||||
|
"cs-agent/internal/pkg/enums"
|
||||||
|
"cs-agent/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *index) markDocumentIndexPending(documentID int64) error {
|
||||||
|
return repositories.KnowledgeDocumentRepository.Updates(sqls.DB(), documentID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusPending,
|
||||||
|
"indexed_at": nil,
|
||||||
|
"index_error": "",
|
||||||
|
"updated_at": time.Now(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markDocumentIndexIndexed(documentID int64) error {
|
||||||
|
now := time.Now()
|
||||||
|
return repositories.KnowledgeDocumentRepository.Updates(sqls.DB(), documentID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusIndexed,
|
||||||
|
"indexed_at": &now,
|
||||||
|
"index_error": "",
|
||||||
|
"updated_at": now,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markDocumentIndexFailed(documentID int64, err error) error {
|
||||||
|
return repositories.KnowledgeDocumentRepository.Updates(sqls.DB(), documentID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusFailed,
|
||||||
|
"index_error": truncateIndexError(err),
|
||||||
|
"updated_at": time.Now(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markKnowledgeBaseDocumentsIndexPending(knowledgeBaseID int64, documentIDs []int64) error {
|
||||||
|
if len(documentIDs) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return sqls.DB().Model(&models.KnowledgeDocument{}).
|
||||||
|
Where("knowledge_base_id = ?", knowledgeBaseID).
|
||||||
|
Where("id IN ?", documentIDs).
|
||||||
|
Updates(map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusPending,
|
||||||
|
"indexed_at": nil,
|
||||||
|
"index_error": "",
|
||||||
|
"updated_at": time.Now(),
|
||||||
|
}).Error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markFAQIndexPending(faqID int64) error {
|
||||||
|
return repositories.KnowledgeFAQRepository.Updates(sqls.DB(), faqID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusPending,
|
||||||
|
"indexed_at": nil,
|
||||||
|
"index_error": "",
|
||||||
|
"updated_at": time.Now(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markFAQIndexIndexed(faqID int64) error {
|
||||||
|
now := time.Now()
|
||||||
|
return repositories.KnowledgeFAQRepository.Updates(sqls.DB(), faqID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusIndexed,
|
||||||
|
"indexed_at": &now,
|
||||||
|
"index_error": "",
|
||||||
|
"updated_at": now,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markFAQIndexFailed(faqID int64, err error) error {
|
||||||
|
return repositories.KnowledgeFAQRepository.Updates(sqls.DB(), faqID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusFailed,
|
||||||
|
"index_error": truncateIndexError(err),
|
||||||
|
"updated_at": time.Now(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func truncateIndexError(err error) string {
|
||||||
|
if err == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
message := err.Error()
|
||||||
|
if len(message) <= 1000 {
|
||||||
|
return message
|
||||||
|
}
|
||||||
|
return message[:1000]
|
||||||
|
}
|
||||||
@@ -0,0 +1,53 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
|
||||||
|
"cs-agent/internal/ai/rag/vectordb"
|
||||||
|
"cs-agent/internal/models"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *index) ensureCollection(ctx context.Context, provider vectordb.Provider, collectionName string, dimension int) error {
|
||||||
|
collectionInfo, err := provider.GetCollection(ctx, collectionName)
|
||||||
|
if err == nil && collectionInfo != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if dimension <= 0 {
|
||||||
|
return fmt.Errorf("invalid embedding dimension: %d", dimension)
|
||||||
|
}
|
||||||
|
if err := provider.CreateCollection(ctx, collectionName, dimension); err != nil {
|
||||||
|
return fmt.Errorf("failed to create collection: %w", err)
|
||||||
|
}
|
||||||
|
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
|
||||||
|
})
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user