package rag import ( "context" "crypto/sha256" "encoding/hex" "encoding/json" "fmt" "log/slog" "time" "cs-agent/internal/ai" 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/repositories" "github.com/google/uuid" "github.com/mlogclub/simple/common/strs" "github.com/mlogclub/simple/sqls" ) type ChunkingConfig struct { Provider string TargetTokens int MaxTokens int OverlapTokens int EnableFallback bool } type index struct { chunkConfig ChunkingConfig registry *ragchunk.Registry } const knowledgeCollectionName = "knowledge_chunks" var Index = &index{ chunkConfig: ChunkingConfig{ Provider: string(enums.KnowledgeChunkProviderStructured), TargetTokens: 300, MaxTokens: 400, OverlapTokens: 40, EnableFallback: true, }, registry: ragchunk.NewDefaultRegistry(), } func (s *index) IndexDocumentByID(ctx context.Context, documentID int64) error { document := repositories.KnowledgeDocumentRepository.Get(sqls.DB(), documentID) if document == nil { return fmt.Errorf("document not found: %d", documentID) } return s.IndexDocument(ctx, document) } func (s *index) IndexDocument(ctx context.Context, document *models.KnowledgeDocument) error { start := time.Now() if err := s.markDocumentIndexPending(document.ID); err != nil { slog.Error("Failed to mark knowledge document index as pending", "document_id", document.ID, "error", err) } fail := func(err error) error { if updateErr := s.markDocumentIndexFailed(document.ID, err); updateErr != nil { slog.Error("Failed to mark knowledge document index as failed", "document_id", document.ID, "error", updateErr) } return err } // TODO 这里每次都查询下知识库不太友好 knowledgeBase := repositories.KnowledgeBaseRepository.Get(sqls.DB(), document.KnowledgeBaseID) if knowledgeBase == nil { return fail(fmt.Errorf("knowledge base not found: %d", document.KnowledgeBaseID)) } existingChunks := repositories.KnowledgeChunkRepository.FindByDocumentID(sqls.DB(), document.ID) chunks, err := s.registry.Chunk(ctx, &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, }, }) if err != nil { return fail(fmt.Errorf("failed to chunk document: %w", err)) } if len(chunks) == 0 { return fail(fmt.Errorf("no chunks generated from document")) } 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 := make([]string, 0, len(existingChunks)) for _, chunk := range existingChunks { if strs.IsNotBlank(chunk.VectorID) { existingVectorIDs = append(existingVectorIDs, chunk.VectorID) } } 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 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 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 := sqls.WithTransaction(func(ctx *sqls.TxContext) error { 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)) } 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.Info("Document indexed successfully", slog.Any("document_id", document.ID), slog.Any("chunks_count", len(chunks)), slog.Any("vectors_count", len(vectors)), slog.Any("time_taken", time.Since(start).String()), ) return nil } func (s *index) IndexFAQByID(ctx context.Context, faqID int64) error { faq := repositories.KnowledgeFAQRepository.Get(sqls.DB(), faqID) if faq == nil { return fmt.Errorf("faq not found: %d", faqID) } if err := s.markFAQIndexPending(faq.ID); err != nil { slog.Error("Failed to mark knowledge faq index as pending", "faq_id", faq.ID, "error", err) } fail := func(err error) error { if updateErr := s.markFAQIndexFailed(faq.ID, err); updateErr != nil { slog.Error("Failed to mark knowledge faq index as failed", "faq_id", faq.ID, "error", updateErr) } return err } knowledgeBase := repositories.KnowledgeBaseRepository.Get(sqls.DB(), faq.KnowledgeBaseID) if knowledgeBase == nil { return fail(fmt.Errorf("knowledge base not found: %d", faq.KnowledgeBaseID)) } if knowledgeBase.KnowledgeType != string(enums.KnowledgeBaseTypeFAQ) { return fail(fmt.Errorf("knowledge base %d is not faq type", knowledgeBase.ID)) } existingChunks := repositories.KnowledgeChunkRepository.FindByFaqID(sqls.DB(), faq.ID) 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)) } embeddingResult, err := ai.Embedding.GenerateEmbedding(ctx, content) if err != nil { return fail(fmt.Errorf("failed to generate embedding for faq %d: %w", faq.ID, 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() collectionInfo, err := provider.GetCollection(ctx, collectionName) if err != nil || collectionInfo == nil { 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)) 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{{ 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)) } if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error { 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)) } if err := s.markFAQIndexIndexed(faq.ID); err != nil { slog.Error("Failed to mark knowledge faq index as indexed", "faq_id", faq.ID, "error", err) } return nil } 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) } 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, 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 { if len(chunks) == 0 { return nil } collectionName := s.getCollectionName() provider := vectordb.GetProvider() if provider == nil { return fmt.Errorf("vectordb provider not initialized") } vectorIDs := make([]string, 0, len(chunks)) for _, chunk := range chunks { if chunk.VectorID != "" { vectorIDs = append(vectorIDs, chunk.VectorID) } } if len(vectorIDs) > 0 { if err := provider.DeleteVectors(ctx, collectionName, vectorIDs); err != nil { slog.Error("Failed to delete vectors", "error", err) } } if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error { return ctx.Tx.Where("document_id = ?", documentID).Delete(&models.KnowledgeChunk{}).Error }); err != nil { return fmt.Errorf("failed to delete chunks: %w", err) } slog.Info("Document index removed", "document_id", documentID, "chunks_removed", len(chunks)) return nil } 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) } 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) removeFAQIndexByChunks(ctx context.Context, knowledgeBaseID int64, faqID int64, chunks []models.KnowledgeChunk) error { if len(chunks) == 0 { return nil } collectionName := s.getCollectionName() provider := vectordb.GetProvider() if provider == nil { return fmt.Errorf("vectordb provider not initialized") } vectorIDs := make([]string, 0, len(chunks)) for _, chunk := range chunks { if chunk.VectorID != "" { vectorIDs = append(vectorIDs, chunk.VectorID) } } if len(vectorIDs) > 0 { if err := provider.DeleteVectors(ctx, collectionName, vectorIDs); err != nil { slog.Error("Failed to delete faq vectors", "error", err) } } if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error { return ctx.Tx.Where("faq_id = ?", faqID).Delete(&models.KnowledgeChunk{}).Error }); err != nil { return fmt.Errorf("failed to delete faq chunks: %w", err) } slog.Info("FAQ index removed", "faq_id", faqID, "chunks_removed", len(chunks)) return nil } func (s *index) getCollectionName() string { return knowledgeCollectionName } func buildKnowledgeChunkVectorID(knowledgeBaseID int64, documentID int64, chunkNo int) string { raw := fmt.Sprintf("kb:%d:doc:%d:chunk:%d", knowledgeBaseID, documentID, chunkNo) return uuid.NewSHA1(uuid.NameSpaceOID, []byte(raw)).String() } func buildKnowledgeFAQChunkVectorID(knowledgeBaseID int64, faqID int64, chunkNo int) string { raw := fmt.Sprintf("kb:%d:faq:%d:chunk:%d", knowledgeBaseID, faqID, chunkNo) return uuid.NewSHA1(uuid.NameSpaceOID, []byte(raw)).String() } func buildChunkContentHash(content string) string { sum := sha256.Sum256([]byte(content)) return hex.EncodeToString(sum[:]) } func firstPositiveInt(values ...int) int { for _, value := range values { if value > 0 { return value } } return 0 } func firstNonEmptyString(values ...string) string { for _, value := range values { if value != "" { return value } } return "" } func (s *index) EnsureCollection(ctx context.Context) error { dimension, err := ai.Embedding.GetDimension(ctx) if err != nil { return fmt.Errorf("failed to get embedding dimension: %w", err) } collectionName := s.getCollectionName() provider := vectordb.GetProvider() if provider == nil { return fmt.Errorf("vectordb provider not initialized") } existing, err := provider.GetCollection(ctx, collectionName) if err == nil && existing != nil { return nil } return provider.CreateCollection(ctx, collectionName, dimension) } func (s *index) RebuildKnowledgeBaseIndex(ctx context.Context, knowledgeBaseID int64) error { knowledgeBase := repositories.KnowledgeBaseRepository.Get(sqls.DB(), knowledgeBaseID) if knowledgeBase == nil { return fmt.Errorf("knowledge base not found: %d", knowledgeBaseID) } if err := s.resetKnowledgeBaseIndexStorage(ctx, knowledgeBaseID); err != nil { return err } successCount := 0 failedCount := 0 if knowledgeBase.KnowledgeType == string(enums.KnowledgeBaseTypeFAQ) { faqs := repositories.KnowledgeFAQRepository.Find(sqls.DB(), sqls.NewCnd(). Eq("knowledge_base_id", knowledgeBaseID). Where("status != ?", enums.StatusDeleted)) if len(faqs) == 0 { slog.Info("No faqs found in knowledge base, nothing to rebuild", "knowledge_base_id", knowledgeBaseID) return nil } slog.Info("Rebuilding faq knowledge base index", "knowledge_base_id", knowledgeBaseID, "faq_count", len(faqs)) for _, faq := range faqs { if err := s.IndexFAQByID(ctx, faq.ID); err != nil { slog.Error("Failed to index faq", "faq_id", faq.ID, "error", err) failedCount++ } else { successCount++ } } } else { documents := repositories.KnowledgeDocumentRepository.Find(sqls.DB(), sqls.NewCnd(). Eq("knowledge_base_id", knowledgeBaseID). Where("status != ?", enums.StatusDeleted)) if len(documents) == 0 { slog.Info("No documents found in knowledge base, nothing to rebuild", "knowledge_base_id", knowledgeBaseID) return nil } documentIDs := make([]int64, 0, len(documents)) for _, doc := range documents { documentIDs = append(documentIDs, doc.ID) } if err := s.markKnowledgeBaseDocumentsIndexPending(knowledgeBaseID, documentIDs); err != nil { slog.Error("Failed to mark knowledge base documents index as pending", "knowledge_base_id", knowledgeBaseID, "error", err) } slog.Info("Rebuilding knowledge base index", "knowledge_base_id", knowledgeBaseID, "document_count", len(documents)) for _, doc := range documents { if err := s.IndexDocumentByID(ctx, doc.ID); err != nil { slog.Error("Failed to index document", "document_id", doc.ID, "error", err) failedCount++ } else { successCount++ } } } slog.Info("Knowledge base index rebuild completed", "knowledge_base_id", knowledgeBaseID, "success_count", successCount, "failed_count", failedCount) return nil } func buildFAQChunkContent(faq *models.KnowledgeFAQ) string { if faq == nil { return "" } parts := []string{fmt.Sprintf("问题:%s", faq.Question)} var similarQuestions []string if faq.SimilarQuestions != "" { _ = json.Unmarshal([]byte(faq.SimilarQuestions), &similarQuestions) } if len(similarQuestions) > 0 { parts = append(parts, fmt.Sprintf("相似问:%s", joinSimilarQuestions(similarQuestions))) } parts = append(parts, fmt.Sprintf("回答:%s", faq.Answer)) content := "" for _, part := range parts { if part == "" { continue } if content != "" { content += "\n" } content += part } return content } func joinSimilarQuestions(items []string) string { result := "" for _, item := range items { if item == "" { continue } if result != "" { result += ";" } result += item } 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 { collectionName := s.getCollectionName() provider := vectordb.GetProvider() if provider == nil { return fmt.Errorf("vectordb provider not initialized") } chunks := repositories.KnowledgeChunkRepository.Find(sqls.DB(), sqls.NewCnd().Eq("knowledge_base_id", knowledgeBaseID)) vectorIDs := make([]string, 0, len(chunks)) for _, chunk := range chunks { if strs.IsNotBlank(chunk.VectorID) { vectorIDs = append(vectorIDs, chunk.VectorID) } } if len(vectorIDs) > 0 { if err := provider.DeleteVectors(ctx, collectionName, vectorIDs); err != nil { return fmt.Errorf("failed to delete vectors for knowledge base %d before rebuild: %w", knowledgeBaseID, err) } } if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error { return ctx.Tx.Where("knowledge_base_id = ?", knowledgeBaseID).Delete(&models.KnowledgeChunk{}).Error }); err != nil { return fmt.Errorf("failed to clear chunks before rebuild: %w", err) } slog.Info("Knowledge base index storage reset", "knowledge_base_id", knowledgeBaseID, "collection", collectionName) return nil }