Files
Enterprise-Knowledge-Base/internal/application/service/knowledgebase.go
T
wizardchen da8ba98622 feat: Refactor knowledge base service and router for improved wiki handling
- Updated the knowledge base service to ensure WikiConfig is created when wiki indexing is enabled, enhancing backward compatibility.
- Removed legacy syncing of WikiConfig.Enabled to IndexingStrategy.WikiEnabled, streamlining the logic for wiki enablement.
- Introduced a custom retry delay function for asynq tasks to handle wiki ingest lock conflicts more effectively, improving task management.
- Cleaned up the container file by removing unnecessary whitespace, contributing to better code readability.
2026-04-22 21:18:32 +08:00

637 lines
21 KiB
Go

package service
import (
"context"
"encoding/json"
"errors"
"time"
"github.com/Tencent/WeKnora/internal/application/service/retriever"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/google/uuid"
"github.com/hibiken/asynq"
)
// ErrInvalidTenantID represents an error for invalid tenant ID
var ErrInvalidTenantID = errors.New("invalid tenant ID")
// knowledgeBaseService implements the knowledge base service interface
type knowledgeBaseService struct {
repo interfaces.KnowledgeBaseRepository
kgRepo interfaces.KnowledgeRepository
chunkRepo interfaces.ChunkRepository
shareRepo interfaces.KBShareRepository
kbShareService interfaces.KBShareService
modelService interfaces.ModelService
retrieveEngine interfaces.RetrieveEngineRegistry
tenantRepo interfaces.TenantRepository
fileSvc interfaces.FileService
graphEngine interfaces.RetrieveGraphRepository
asynqClient interfaces.TaskEnqueuer
}
// NewKnowledgeBaseService creates a new knowledge base service
func NewKnowledgeBaseService(repo interfaces.KnowledgeBaseRepository,
kgRepo interfaces.KnowledgeRepository,
chunkRepo interfaces.ChunkRepository,
shareRepo interfaces.KBShareRepository,
kbShareService interfaces.KBShareService,
modelService interfaces.ModelService,
retrieveEngine interfaces.RetrieveEngineRegistry,
tenantRepo interfaces.TenantRepository,
fileSvc interfaces.FileService,
graphEngine interfaces.RetrieveGraphRepository,
asynqClient interfaces.TaskEnqueuer,
) interfaces.KnowledgeBaseService {
return &knowledgeBaseService{
repo: repo,
kgRepo: kgRepo,
chunkRepo: chunkRepo,
shareRepo: shareRepo,
kbShareService: kbShareService,
modelService: modelService,
retrieveEngine: retrieveEngine,
tenantRepo: tenantRepo,
fileSvc: fileSvc,
graphEngine: graphEngine,
asynqClient: asynqClient,
}
}
// GetRepository gets the knowledge base repository
// Parameters:
// - ctx: Context with authentication and request information
//
// Returns:
// - interfaces.KnowledgeBaseRepository: Knowledge base repository
func (s *knowledgeBaseService) GetRepository() interfaces.KnowledgeBaseRepository {
return s.repo
}
// CreateKnowledgeBase creates a new knowledge base
func (s *knowledgeBaseService) CreateKnowledgeBase(ctx context.Context,
kb *types.KnowledgeBase,
) (*types.KnowledgeBase, error) {
// Generate UUID and set creation timestamps
if kb.ID == "" {
kb.ID = uuid.New().String()
}
kb.CreatedAt = time.Now()
kb.TenantID = types.MustTenantIDFromContext(ctx)
kb.UpdatedAt = time.Now()
kb.EnsureDefaults()
logger.Infof(ctx, "Creating knowledge base, ID: %s, tenant ID: %d, name: %s", kb.ID, kb.TenantID, kb.Name)
if err := s.repo.CreateKnowledgeBase(ctx, kb); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"knowledge_base_id": kb.ID,
"tenant_id": kb.TenantID,
})
return nil, err
}
logger.Infof(ctx, "Knowledge base created successfully, ID: %s, name: %s", kb.ID, kb.Name)
return kb, nil
}
// GetKnowledgeBaseByID retrieves a knowledge base by its ID
func (s *knowledgeBaseService) GetKnowledgeBaseByID(ctx context.Context, id string) (*types.KnowledgeBase, error) {
if id == "" {
logger.Error(ctx, "Knowledge base ID is empty")
return nil, errors.New("knowledge base ID cannot be empty")
}
kb, err := s.repo.GetKnowledgeBaseByID(ctx, id)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"knowledge_base_id": id,
})
return nil, err
}
kb.EnsureDefaults()
return kb, nil
}
// GetKnowledgeBaseByIDOnly retrieves knowledge base by ID without tenant filter
// Used for cross-tenant shared KB access where permission is checked elsewhere
func (s *knowledgeBaseService) GetKnowledgeBaseByIDOnly(ctx context.Context, id string) (*types.KnowledgeBase, error) {
if id == "" {
logger.Error(ctx, "Knowledge base ID is empty")
return nil, errors.New("knowledge base ID cannot be empty")
}
kb, err := s.repo.GetKnowledgeBaseByID(ctx, id)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"knowledge_base_id": id,
})
return nil, err
}
kb.EnsureDefaults()
return kb, nil
}
// GetKnowledgeBasesByIDsOnly retrieves knowledge bases by IDs without tenant filter (batch).
func (s *knowledgeBaseService) GetKnowledgeBasesByIDsOnly(ctx context.Context, ids []string) ([]*types.KnowledgeBase, error) {
if len(ids) == 0 {
return nil, nil
}
kbs, err := s.repo.GetKnowledgeBaseByIDs(ctx, ids)
if err != nil {
return nil, err
}
for _, kb := range kbs {
if kb != nil {
kb.EnsureDefaults()
}
}
return kbs, nil
}
// ListKnowledgeBases returns all knowledge bases for a tenant
func (s *knowledgeBaseService) ListKnowledgeBases(ctx context.Context) ([]*types.KnowledgeBase, error) {
tenantID := types.MustTenantIDFromContext(ctx)
kbs, err := s.repo.ListKnowledgeBasesByTenantID(ctx, tenantID)
if err != nil {
for _, kb := range kbs {
kb.EnsureDefaults()
}
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"tenant_id": tenantID,
})
return nil, err
}
// Query knowledge count and chunk count for each knowledge base
for _, kb := range kbs {
kb.EnsureDefaults()
// Get knowledge count
switch kb.Type {
case types.KnowledgeBaseTypeDocument:
knowledgeCount, err := s.kgRepo.CountKnowledgeByKnowledgeBaseID(ctx, tenantID, kb.ID)
if err != nil {
logger.Warnf(ctx, "Failed to get knowledge count for knowledge base %s: %v", kb.ID, err)
} else {
kb.KnowledgeCount = knowledgeCount
}
case types.KnowledgeBaseTypeFAQ:
// Get chunk count
chunkCount, err := s.chunkRepo.CountChunksByKnowledgeBaseID(ctx, tenantID, kb.ID)
if err != nil {
logger.Warnf(ctx, "Failed to get chunk count for knowledge base %s: %v", kb.ID, err)
} else {
kb.ChunkCount = chunkCount
}
}
// Check if there is a processing import task
processingCount, err := s.kgRepo.CountKnowledgeByStatus(
ctx,
tenantID,
kb.ID,
[]string{"pending", "processing"},
)
if err != nil {
logger.Warnf(ctx, "Failed to check processing status for knowledge base %s: %v", kb.ID, err)
} else {
kb.IsProcessing = processingCount > 0
kb.ProcessingCount = processingCount
}
}
return kbs, nil
}
// ListKnowledgeBasesByTenantID returns all knowledge bases for the given tenant (e.g. for shared agent context).
func (s *knowledgeBaseService) ListKnowledgeBasesByTenantID(ctx context.Context, tenantID uint64) ([]*types.KnowledgeBase, error) {
kbs, err := s.repo.ListKnowledgeBasesByTenantID(ctx, tenantID)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"tenant_id": tenantID,
})
return nil, err
}
for _, kb := range kbs {
kb.EnsureDefaults()
switch kb.Type {
case types.KnowledgeBaseTypeDocument:
if cnt, err := s.kgRepo.CountKnowledgeByKnowledgeBaseID(ctx, tenantID, kb.ID); err == nil {
kb.KnowledgeCount = cnt
}
case types.KnowledgeBaseTypeFAQ:
if cnt, err := s.chunkRepo.CountChunksByKnowledgeBaseID(ctx, tenantID, kb.ID); err == nil {
kb.ChunkCount = cnt
}
}
if processingCount, err := s.kgRepo.CountKnowledgeByStatus(ctx, tenantID, kb.ID, []string{"pending", "processing"}); err == nil {
kb.IsProcessing = processingCount > 0
kb.ProcessingCount = processingCount
}
}
return kbs, nil
}
// FillKnowledgeBaseCounts fills KnowledgeCount, ChunkCount, IsProcessing, ProcessingCount for the given KB using kb.TenantID.
func (s *knowledgeBaseService) FillKnowledgeBaseCounts(ctx context.Context, kb *types.KnowledgeBase) error {
if kb == nil {
return nil
}
tenantID := kb.TenantID
kb.EnsureDefaults()
switch kb.Type {
case types.KnowledgeBaseTypeDocument:
if cnt, err := s.kgRepo.CountKnowledgeByKnowledgeBaseID(ctx, tenantID, kb.ID); err == nil {
kb.KnowledgeCount = cnt
}
case types.KnowledgeBaseTypeFAQ:
if cnt, err := s.chunkRepo.CountChunksByKnowledgeBaseID(ctx, tenantID, kb.ID); err == nil {
kb.ChunkCount = cnt
}
}
if processingCount, err := s.kgRepo.CountKnowledgeByStatus(ctx, tenantID, kb.ID, []string{"pending", "processing"}); err == nil {
kb.IsProcessing = processingCount > 0
kb.ProcessingCount = processingCount
}
return nil
}
// UpdateKnowledgeBase updates a knowledge base's properties
func (s *knowledgeBaseService) UpdateKnowledgeBase(ctx context.Context,
id string,
name string,
description string,
config *types.KnowledgeBaseConfig,
) (*types.KnowledgeBase, error) {
if id == "" {
logger.Error(ctx, "Knowledge base ID is empty")
return nil, errors.New("knowledge base ID cannot be empty")
}
logger.Infof(ctx, "Updating knowledge base, ID: %s, name: %s", id, name)
// Get existing knowledge base
kb, err := s.repo.GetKnowledgeBaseByID(ctx, id)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"knowledge_base_id": id,
})
return nil, err
}
// Update the knowledge base properties
kb.Name = name
kb.Description = description
if config != nil {
kb.ChunkingConfig = config.ChunkingConfig
kb.ImageProcessingConfig = config.ImageProcessingConfig
if config.FAQConfig != nil {
kb.FAQConfig = config.FAQConfig
}
if config.WikiConfig != nil {
kb.WikiConfig = config.WikiConfig
}
// Update indexing strategy — syncs to ExtractConfig for backward compat
if config.IndexingStrategy != nil {
if !config.IndexingStrategy.HasAnyIndexing() {
return nil, errors.New("at least one indexing strategy must be enabled")
}
kb.IndexingStrategy = *config.IndexingStrategy
// Ensure WikiConfig exists when wiki indexing is enabled so that
// wiki-specific tunables (synthesis model, granularity, …) have a home.
if kb.WikiConfig == nil && config.IndexingStrategy.WikiEnabled {
kb.WikiConfig = &types.WikiConfig{}
}
// Sync GraphEnabled → ExtractConfig
if kb.ExtractConfig != nil {
kb.ExtractConfig.Enabled = config.IndexingStrategy.GraphEnabled
} else if config.IndexingStrategy.GraphEnabled {
kb.ExtractConfig = &types.ExtractConfig{Enabled: true}
}
}
}
kb.UpdatedAt = time.Now()
kb.EnsureDefaults()
logger.Info(ctx, "Saving knowledge base update")
if err := s.repo.UpdateKnowledgeBase(ctx, kb); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"knowledge_base_id": id,
})
return nil, err
}
logger.Infof(ctx, "Knowledge base updated successfully, ID: %s, name: %s", kb.ID, kb.Name)
return kb, nil
}
// TogglePinKnowledgeBase toggles the pin status of a knowledge base
func (s *knowledgeBaseService) TogglePinKnowledgeBase(ctx context.Context, id string) (*types.KnowledgeBase, error) {
if id == "" {
return nil, errors.New("knowledge base ID cannot be empty")
}
tenantID := types.MustTenantIDFromContext(ctx)
kb, err := s.repo.TogglePinKnowledgeBase(ctx, id, tenantID)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"knowledge_base_id": id,
})
return nil, err
}
logger.Infof(ctx, "Knowledge base pin toggled, ID: %s, is_pinned: %v", id, kb.IsPinned)
return kb, nil
}
// DeleteKnowledgeBase deletes a knowledge base by its ID
// This method marks the knowledge base as deleted and enqueues an async task
// to handle the heavy cleanup operations (embeddings, chunks, files, graph data)
func (s *knowledgeBaseService) DeleteKnowledgeBase(ctx context.Context, id string) error {
if id == "" {
logger.Error(ctx, "Knowledge base ID is empty")
return errors.New("knowledge base ID cannot be empty")
}
logger.Infof(ctx, "Deleting knowledge base, ID: %s", id)
// Get tenant ID from context
tenantID := types.MustTenantIDFromContext(ctx)
tenantInfo, _ := types.TenantInfoFromContext(ctx)
// Step 1: Delete the knowledge base record first (mark as deleted)
logger.Infof(ctx, "Deleting knowledge base from database")
err := s.repo.DeleteKnowledgeBase(ctx, id)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"knowledge_base_id": id,
})
return err
}
// Step 1b: Remove all organization shares for this KB so org settings no longer show them
if delErr := s.shareRepo.DeleteByKnowledgeBaseID(ctx, id); delErr != nil {
logger.Warnf(ctx, "Failed to delete KB shares for knowledge base %s: %v", id, delErr)
}
// Step 2: Enqueue async task for heavy cleanup operations
payload := types.KBDeletePayload{
TenantID: tenantID,
KnowledgeBaseID: id,
EffectiveEngines: tenantInfo.GetEffectiveEngines(),
}
payloadBytes, err := json.Marshal(payload)
if err != nil {
logger.Warnf(ctx, "Failed to marshal KB delete payload: %v", err)
// Don't fail the request, the KB record is already deleted
return nil
}
task := asynq.NewTask(types.TypeKBDelete, payloadBytes, asynq.Queue("low"), asynq.MaxRetry(3))
info, err := s.asynqClient.Enqueue(task)
if err != nil {
logger.Warnf(ctx, "Failed to enqueue KB delete task: %v", err)
// Don't fail the request, the KB record is already deleted
return nil
}
logger.Infof(ctx, "KB delete task enqueued: %s, knowledge base ID: %s", info.ID, id)
logger.Infof(ctx, "Knowledge base deleted successfully, ID: %s", id)
return nil
}
// ProcessKBDelete handles async knowledge base deletion task
// This method performs heavy cleanup operations: deleting embeddings, chunks, files, and graph data
func (s *knowledgeBaseService) ProcessKBDelete(ctx context.Context, t *asynq.Task) error {
var payload types.KBDeletePayload
if err := json.Unmarshal(t.Payload(), &payload); err != nil {
logger.Errorf(ctx, "Failed to unmarshal KB delete payload: %v", err)
return err
}
tenantID := payload.TenantID
kbID := payload.KnowledgeBaseID
// Set tenant context for downstream services
ctx = context.WithValue(ctx, types.TenantIDContextKey, tenantID)
logger.Infof(ctx, "Processing KB delete task for knowledge base: %s", kbID)
// Step 1: Get all knowledge entries in this knowledge base
logger.Infof(ctx, "Fetching all knowledge entries in knowledge base, ID: %s", kbID)
knowledgeList, err := s.kgRepo.ListKnowledgeByKnowledgeBaseID(ctx, tenantID, kbID)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"knowledge_base_id": kbID,
})
return err
}
logger.Infof(ctx, "Found %d knowledge entries to delete", len(knowledgeList))
// Step 2: Delete all knowledge entries and their resources
if len(knowledgeList) > 0 {
knowledgeIDs := make([]string, 0, len(knowledgeList))
for _, knowledge := range knowledgeList {
knowledgeIDs = append(knowledgeIDs, knowledge.ID)
}
logger.Infof(ctx, "Deleting all knowledge entries and their resources")
// Delete embeddings from vector store
logger.Infof(ctx, "Deleting embeddings from vector store")
retrieveEngine, err := retriever.NewCompositeRetrieveEngine(
s.retrieveEngine,
payload.EffectiveEngines,
)
if err != nil {
logger.Warnf(ctx, "Failed to create retrieve engine: %v", err)
} else {
// Group knowledge by embedding model and type
type groupKey struct {
EmbeddingModelID string
Type string
}
embeddingGroups := make(map[groupKey][]string)
for _, knowledge := range knowledgeList {
key := groupKey{EmbeddingModelID: knowledge.EmbeddingModelID, Type: knowledge.Type}
embeddingGroups[key] = append(embeddingGroups[key], knowledge.ID)
}
for key, knowledgeGroup := range embeddingGroups {
embeddingModel, err := s.modelService.GetEmbeddingModel(ctx, key.EmbeddingModelID)
if err != nil {
logger.Warnf(ctx, "Failed to get embedding model %s: %v", key.EmbeddingModelID, err)
continue
}
if err := retrieveEngine.DeleteByKnowledgeIDList(ctx, knowledgeGroup, embeddingModel.GetDimensions(), key.Type); err != nil {
logger.Warnf(ctx, "Failed to delete embeddings for model %s: %v", key.EmbeddingModelID, err)
}
}
}
// Collect image URLs before chunks are deleted
chunkImageInfos, imgErr := s.chunkRepo.ListImageInfoByKnowledgeIDs(ctx, tenantID, knowledgeIDs)
if imgErr != nil {
logger.Warnf(ctx, "Failed to collect image URLs for KB delete: %v", imgErr)
}
var imageInfoStrs []string
for _, ci := range chunkImageInfos {
imageInfoStrs = append(imageInfoStrs, ci.ImageInfo)
}
imageURLs := collectImageURLs(ctx, imageInfoStrs)
// Delete all chunks
logger.Infof(ctx, "Deleting all chunks in knowledge base")
for _, knowledgeID := range knowledgeIDs {
if err := s.chunkRepo.DeleteChunksByKnowledgeID(ctx, tenantID, knowledgeID); err != nil {
logger.Warnf(ctx, "Failed to delete chunks for knowledge %s: %v", knowledgeID, err)
}
}
// Delete physical files, extracted images, and adjust storage
logger.Infof(ctx, "Deleting physical files and extracted images")
storageAdjust := int64(0)
for _, knowledge := range knowledgeList {
if knowledge.FilePath != "" {
if err := s.fileSvc.DeleteFile(ctx, knowledge.FilePath); err != nil {
logger.Warnf(ctx, "Failed to delete file %s: %v", knowledge.FilePath, err)
}
}
storageAdjust -= knowledge.StorageSize
}
deleteExtractedImages(ctx, s.fileSvc, imageURLs)
if storageAdjust != 0 {
if err := s.tenantRepo.AdjustStorageUsed(ctx, tenantID, storageAdjust); err != nil {
logger.Warnf(ctx, "Failed to adjust tenant storage: %v", err)
}
}
// Delete knowledge graph data
logger.Infof(ctx, "Deleting knowledge graph data")
namespaces := make([]types.NameSpace, 0, len(knowledgeList))
for _, knowledge := range knowledgeList {
namespaces = append(namespaces, types.NameSpace{
KnowledgeBase: knowledge.KnowledgeBaseID,
Knowledge: knowledge.ID,
})
}
if s.graphEngine != nil && len(namespaces) > 0 {
if err := s.graphEngine.DelGraph(ctx, namespaces); err != nil {
logger.Warnf(ctx, "Failed to delete knowledge graph: %v", err)
}
}
// Delete all knowledge entries from database
logger.Infof(ctx, "Deleting knowledge entries from database")
if err := s.kgRepo.DeleteKnowledgeList(ctx, tenantID, knowledgeIDs); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"knowledge_base_id": kbID,
})
return err
}
}
logger.Infof(ctx, "KB delete task completed successfully, knowledge base ID: %s", kbID)
return nil
}
// SetEmbeddingModel sets the embedding model for a knowledge base
func (s *knowledgeBaseService) SetEmbeddingModel(ctx context.Context, id string, modelID string) error {
if id == "" {
logger.Error(ctx, "Knowledge base ID is empty")
return errors.New("knowledge base ID cannot be empty")
}
if modelID == "" {
logger.Error(ctx, "Model ID is empty")
return errors.New("model ID cannot be empty")
}
logger.Infof(ctx, "Setting embedding model for knowledge base, knowledge base ID: %s, model ID: %s", id, modelID)
// Get the knowledge base
kb, err := s.repo.GetKnowledgeBaseByID(ctx, id)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"knowledge_base_id": id,
})
return err
}
// Update the knowledge base's embedding model
kb.EmbeddingModelID = modelID
kb.UpdatedAt = time.Now()
logger.Info(ctx, "Saving knowledge base embedding model update")
err = s.repo.UpdateKnowledgeBase(ctx, kb)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"knowledge_base_id": id,
"embedding_model_id": modelID,
})
return err
}
logger.Infof(
ctx,
"Knowledge base embedding model set successfully, knowledge base ID: %s, model ID: %s",
id,
modelID,
)
return nil
}
// CopyKnowledgeBase copies a knowledge base to a new knowledge base (shallow copy).
// Source and target must belong to the tenant in context; cross-tenant access is rejected.
func (s *knowledgeBaseService) CopyKnowledgeBase(ctx context.Context,
srcKB string, dstKB string,
) (*types.KnowledgeBase, *types.KnowledgeBase, error) {
tenantID := types.MustTenantIDFromContext(ctx)
// Load source KB with tenant scope to prevent cross-tenant cloning
sourceKB, err := s.repo.GetKnowledgeBaseByIDAndTenant(ctx, srcKB, tenantID)
if err != nil {
logger.Errorf(ctx, "Get source knowledge base failed: %v", err)
return nil, nil, err
}
sourceKB.EnsureDefaults()
var targetKB *types.KnowledgeBase
if dstKB != "" {
// Load target KB with tenant scope so we only clone into the caller's tenant
targetKB, err = s.repo.GetKnowledgeBaseByIDAndTenant(ctx, dstKB, tenantID)
if err != nil {
return nil, nil, err
}
} else {
var faqConfig *types.FAQConfig
if sourceKB.FAQConfig != nil {
cfg := *sourceKB.FAQConfig
faqConfig = &cfg
}
targetKB = &types.KnowledgeBase{
ID: uuid.New().String(),
Name: sourceKB.Name,
Type: sourceKB.Type,
Description: sourceKB.Description,
TenantID: tenantID,
ChunkingConfig: sourceKB.ChunkingConfig,
ImageProcessingConfig: sourceKB.ImageProcessingConfig,
EmbeddingModelID: sourceKB.EmbeddingModelID,
SummaryModelID: sourceKB.SummaryModelID,
VLMConfig: sourceKB.VLMConfig,
StorageProviderConfig: sourceKB.StorageProviderConfig,
StorageConfig: sourceKB.StorageConfig,
FAQConfig: faqConfig,
}
targetKB.EnsureDefaults()
if err := s.repo.CreateKnowledgeBase(ctx, targetKB); err != nil {
return nil, nil, err
}
}
return sourceKB, targetKB, nil
}