WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

2026-09-22 16:19:06 +08:00
commit f7fc992f57
3929 changed files with 1006607 additions and 0 deletions
@@ -0,0 +1,652 @@
+package chatpipeline
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/Tencent/WeKnora/internal/config"
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/Tencent/WeKnora/internal/searchutil"
+	"github.com/Tencent/WeKnora/internal/tracing/langfuse"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+)
+
+// PluginSearch implements search functionality for chat pipeline
+type PluginSearch struct {
+	knowledgeBaseService  interfaces.KnowledgeBaseService
+	knowledgeService      interfaces.KnowledgeService
+	chunkService          interfaces.ChunkService
+	config                *config.Config
+	webSearchService      interfaces.WebSearchService
+	tenantService         interfaces.TenantService
+	sessionService        interfaces.SessionService
+	webSearchStateService interfaces.WebSearchStateService
+	webSearchProviderRepo interfaces.WebSearchProviderRepository
+}
+
+func NewPluginSearch(eventManager *EventManager,
+	knowledgeBaseService interfaces.KnowledgeBaseService,
+	knowledgeService interfaces.KnowledgeService,
+	chunkService interfaces.ChunkService,
+	config *config.Config,
+	webSearchService interfaces.WebSearchService,
+	tenantService interfaces.TenantService,
+	sessionService interfaces.SessionService,
+	webSearchStateService interfaces.WebSearchStateService,
+	webSearchProviderRepo interfaces.WebSearchProviderRepository,
+) *PluginSearch {
+	res := &PluginSearch{
+		knowledgeBaseService:  knowledgeBaseService,
+		knowledgeService:      knowledgeService,
+		chunkService:          chunkService,
+		config:                config,
+		webSearchService:      webSearchService,
+		tenantService:         tenantService,
+		sessionService:        sessionService,
+		webSearchStateService: webSearchStateService,
+		webSearchProviderRepo: webSearchProviderRepo,
+	}
+	eventManager.Register(res)
+	return res
+}
+
+// ActivationEvents returns the event types this plugin handles
+func (p *PluginSearch) ActivationEvents() []types.EventType {
+	return []types.EventType{types.CHUNK_SEARCH}
+}
+
+// OnEvent handles search events in the chat pipeline
+func (p *PluginSearch) OnEvent(ctx context.Context,
+	eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
+) *PluginError {
+	// Check if we have search targets or web search enabled
+	hasKBTargets := types.HasKnowledgeRetrievalScope(
+		chatManage.SearchTargets,
+		chatManage.KnowledgeBaseIDs,
+		chatManage.KnowledgeIDs,
+	)
+	if !hasKBTargets && !chatManage.WebSearchEnabled {
+		pipelineError(ctx, "Search", "kb_not_found", map[string]interface{}{
+			"session_id": chatManage.SessionID,
+		})
+		return nil
+	}
+
+	pipelineInfo(ctx, "Search", "input", map[string]interface{}{
+		"session_id":     chatManage.SessionID,
+		"rewrite_query":  chatManage.RewriteQuery,
+		"search_targets": len(chatManage.SearchTargets),
+		"tenant_id":      chatManage.TenantID,
+		"web_enabled":    chatManage.WebSearchEnabled,
+	})
+
+	// Run KB search and web search concurrently
+	pipelineInfo(ctx, "Search", "plan", map[string]interface{}{
+		"search_targets":    len(chatManage.SearchTargets),
+		"embedding_top_k":   chatManage.EmbeddingTopK,
+		"vector_threshold":  chatManage.VectorThreshold,
+		"keyword_threshold": chatManage.KeywordThreshold,
+	})
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	allResults := make([]*types.SearchResult, 0)
+	var kbSearchErr error
+
+	wg.Add(2)
+	// Goroutine 1: Knowledge base search using SearchTargets
+	go func() {
+		defer wg.Done()
+		kbResults, err := p.searchByTargets(ctx, chatManage)
+		kbSearchErr = err
+		if len(kbResults) > 0 {
+			mu.Lock()
+			allResults = append(allResults, kbResults...)
+			mu.Unlock()
+		}
+	}()
+
+	// Goroutine 2: Web search (if enabled)
+	go func() {
+		defer wg.Done()
+		webResults := p.searchWebIfEnabled(ctx, chatManage)
+		if len(webResults) > 0 {
+			mu.Lock()
+			allResults = append(allResults, webResults...)
+			mu.Unlock()
+		}
+	}()
+
+	wg.Wait()
+	if kbSearchErr != nil && len(allResults) == 0 {
+		pipelineError(ctx, "Search", "kb_search_failed", map[string]interface{}{
+			"error": kbSearchErr.Error(),
+		})
+		return ErrSearch.WithError(kbSearchErr)
+	}
+	if kbSearchErr != nil {
+		pipelineWarn(ctx, "Search", "kb_search_partial_failure", map[string]interface{}{
+			"error":        kbSearchErr.Error(),
+			"result_count": len(allResults),
+		})
+	}
+
+	chatManage.SearchResult = allResults
+
+	logSearchScoreSample(ctx, "result_score_before_normalize", chatManage.SearchResult)
+
+	// If recall is low, attempt query expansion with keyword-focused search
+	if chatManage.EnableQueryExpansion && len(chatManage.SearchResult) < max(1, chatManage.EmbeddingTopK) {
+		expResults := p.runQueryExpansion(ctx, chatManage)
+		if len(expResults) > 0 {
+			chatManage.SearchResult = append(chatManage.SearchResult, expResults...)
+		}
+	}
+
+	logSearchScoreSample(ctx, "final_score", chatManage.SearchResult)
+
+	// Return if we have results
+	if len(chatManage.SearchResult) != 0 {
+		pipelineInfo(ctx, "Search", "output", map[string]interface{}{
+			"session_id":   chatManage.SessionID,
+			"result_count": len(chatManage.SearchResult),
+		})
+		return next()
+	}
+	pipelineWarn(ctx, "Search", "output", map[string]interface{}{
+		"session_id":   chatManage.SessionID,
+		"result_count": 0,
+	})
+	return ErrSearchNothing
+}
+
+// getSearchResultFromHistory retrieves relevant knowledge references from chat history
+func getSearchResultFromHistory(chatManage *types.ChatManage) []*types.SearchResult {
+	if len(chatManage.History) == 0 {
+		return nil
+	}
+	// Search history in reverse chronological order
+	for i := len(chatManage.History) - 1; i >= 0; i-- {
+		if len(chatManage.History[i].KnowledgeReferences) > 0 {
+			// Mark all references as history matches
+			for _, reference := range chatManage.History[i].KnowledgeReferences {
+				reference.MatchType = types.MatchTypeHistory
+			}
+			return chatManage.History[i].KnowledgeReferences
+		}
+	}
+	return nil
+}
+
+func removeDuplicateResults(results []*types.SearchResult) []*types.SearchResult {
+	seen := make(map[string]bool)
+	contentSig := make(map[string]string) // sig -> first chunk ID
+	var uniqueResults []*types.SearchResult
+	for _, r := range results {
+		// Only deduplicate by exact chunk ID — do NOT treat shared ParentChunkID
+		// as duplicates, because different child chunks of the same parent carry
+		// different content segments that may all be relevant.
+		if seen[r.ID] {
+			logger.Debugf(context.Background(), "Dedup: chunk %s removed due to duplicate ID", r.ID)
+			continue
+		}
+		sig := buildContentSignature(r.Content)
+		if sig != "" {
+			if firstChunk, exists := contentSig[sig]; exists {
+				logger.Debugf(context.Background(), "Dedup: chunk %s removed due to content signature (dup of %s, sig prefix: %.50s...)", r.ID, firstChunk, sig)
+				continue
+			}
+			contentSig[sig] = r.ID
+		}
+		seen[r.ID] = true
+		uniqueResults = append(uniqueResults, r)
+	}
+	return uniqueResults
+}
+
+func buildContentSignature(content string) string {
+	return searchutil.BuildContentSignature(content)
+}
+
+// removePartialOverlaps drops chunks whose content is largely contained within
+// a higher-scored chunk, even across different knowledge sources. This catches
+// cross-KB duplicates and near-duplicates that exact-signature dedup misses.
+//
+// Two thresholds are used:
+//   - Substring containment: if the normalized short text is a literal substring
+//     of the normalized long text, the shorter chunk is removed.
+//   - Token overlap coefficient >= 0.85: if 85%+ of the smaller chunk's tokens
+//     appear in the larger chunk, the smaller one is redundant.
+//
+// The input slice MUST already be deduplicated by ID/signature. Within each
+// pair the chunk with the lower score is the candidate for removal; ties are
+// broken by content length (longer wins).
+func removePartialOverlaps(ctx context.Context, results []*types.SearchResult) []*types.SearchResult {
+	const overlapThreshold = 0.85
+
+	if len(results) <= 1 {
+		return results
+	}
+
+	type normEntry struct {
+		norm   string
+		result *types.SearchResult
+	}
+
+	entries := make([]normEntry, 0, len(results))
+	for _, r := range results {
+		entries = append(entries, normEntry{
+			norm:   searchutil.NormalizeContent(r.Content),
+			result: r,
+		})
+	}
+
+	removed := make(map[int]bool)
+
+	for i := 0; i < len(entries); i++ {
+		if removed[i] {
+			continue
+		}
+		for j := i + 1; j < len(entries); j++ {
+			if removed[j] {
+				continue
+			}
+
+			a, b := entries[i], entries[j]
+
+			shortIdx, longIdx := i, j
+			if len(a.norm) > len(b.norm) {
+				shortIdx, longIdx = j, i
+			}
+
+			contained := searchutil.IsContentContained(
+				entries[shortIdx].norm, entries[longIdx].norm,
+			)
+
+			if !contained {
+				ratio := searchutil.ContentOverlapRatio(
+					entries[shortIdx].result.Content,
+					entries[longIdx].result.Content,
+				)
+				if ratio < overlapThreshold {
+					continue
+				}
+			}
+
+			victim := shortIdx
+			if entries[shortIdx].result.Score > entries[longIdx].result.Score {
+				victim = longIdx
+			}
+			removed[victim] = true
+
+			keptIdx := i
+			if victim == i {
+				keptIdx = j
+			}
+			pipelineInfo(ctx, "Merge", "partial_overlap_drop", map[string]interface{}{
+				"kept_id":    entries[keptIdx].result.ID,
+				"dropped_id": entries[victim].result.ID,
+				"contained":  contained,
+			})
+		}
+	}
+
+	out := make([]*types.SearchResult, 0, len(results)-len(removed))
+	for i, e := range entries {
+		if !removed[i] {
+			out = append(out, e.result)
+		}
+	}
+	return out
+}
+
+func logSearchScoreSample(ctx context.Context, action string, results []*types.SearchResult) {
+	const maxLogRows = 8
+	limit := min(maxLogRows, len(results))
+	for i := 0; i < limit; i++ {
+		r := results[i]
+		pipelineInfo(ctx, "Search", action, map[string]interface{}{
+			"index":      i,
+			"chunk_id":   r.ID,
+			"score":      fmt.Sprintf("%.4f", r.Score),
+			"match_type": r.MatchType,
+		})
+	}
+	if len(results) > limit {
+		pipelineInfo(ctx, "Search", action+"_summary", map[string]interface{}{
+			"total":     len(results),
+			"logged":    limit,
+			"truncated": len(results) - limit,
+		})
+	}
+}
+
+// targetReportsEmbedFailure reports whether an embedding failure for a KB should
+// be recorded as a retrieval error. Wiki/graph-only KBs have no vector or keyword
+// index to degrade into; HybridSearch returns empty without error. When KB metadata
+// is unavailable, callers still attempt keyword-degraded search.
+func targetReportsEmbedFailure(kb *types.KnowledgeBase) bool {
+	if kb == nil {
+		return false
+	}
+	if kb.Type == types.KnowledgeBaseTypeFAQ {
+		return true
+	}
+	if kb.IsKeywordEnabled() {
+		return false
+	}
+	return kb.IsVectorEnabled()
+}
+
+// searchByTargets performs KB searches using pre-computed SearchTargets.
+// Targets sharing the same underlying embedding model (identified by model
+// name + endpoint, not just model ID) are grouped so the query embedding is
+// computed once per model AND all full-KB targets in a group are combined into
+// a single retrieval call, reducing both embedding API calls and DB round-trips.
+func (p *PluginSearch) searchByTargets(
+	ctx context.Context,
+	chatManage *types.ChatManage,
+) ([]*types.SearchResult, error) {
+	if len(chatManage.SearchTargets) == 0 {
+		return nil, nil
+	}
+
+	queryText := strings.TrimSpace(chatManage.RewriteQuery)
+
+	// Batch-fetch KB records to determine embedding model grouping.
+	// On failure, all targets fall into an empty-key group and HybridSearch
+	// computes the embedding per-KB (graceful degradation).
+	kbIDs := make([]string, 0, len(chatManage.SearchTargets))
+	for _, t := range chatManage.SearchTargets {
+		kbIDs = append(kbIDs, t.KnowledgeBaseID)
+	}
+	var kbList []*types.KnowledgeBase
+	kbMap := make(map[string]*types.KnowledgeBase)
+	if kbs, err := p.knowledgeBaseService.GetKnowledgeBasesByIDsOnly(ctx, kbIDs); err == nil {
+		kbList = kbs
+		for _, kb := range kbs {
+			if kb != nil {
+				kbMap[kb.ID] = kb
+			}
+		}
+	} else {
+		pipelineWarn(ctx, "Search", "batch_kb_fetch_error", map[string]interface{}{
+			"error": err.Error(),
+		})
+	}
+
+	// Resolve actual model identities (name + endpoint) so that cross-tenant
+	// KBs backed by the same physical model share one embedding computation.
+	modelKeyMap := p.knowledgeBaseService.ResolveEmbeddingModelKeys(ctx, kbList)
+
+	groups := make(map[string][]*types.SearchTarget)
+	for _, t := range chatManage.SearchTargets {
+		key := modelKeyMap[t.KnowledgeBaseID] // empty string if unresolved
+		groups[key] = append(groups[key], t)
+	}
+
+	pipelineInfo(ctx, "Search", "embedding_groups", map[string]interface{}{
+		"total_targets": len(chatManage.SearchTargets),
+		"unique_models": len(groups),
+	})
+
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	var results []*types.SearchResult
+	var firstErr error
+	var errOnce sync.Once
+	recordError := func(err error) {
+		if err != nil {
+			errOnce.Do(func() { firstErr = err })
+		}
+	}
+
+	for modelKey, targets := range groups {
+		wg.Add(1)
+		go func(modelKey string, targets []*types.SearchTarget) {
+			defer wg.Done()
+
+			// Compute embedding once for this model group. When that fails, retain
+			// only targets that have a real keyword index; vector-only targets must
+			// propagate the embedding failure instead of looking like empty recall.
+			var queryEmbedding []float32
+			disableVector := false
+			searchableTargets := targets
+			if modelKey != "" {
+				emb, err := p.knowledgeBaseService.GetQueryEmbedding(ctx, targets[0].KnowledgeBaseID, queryText)
+				if err != nil {
+					searchableTargets = make([]*types.SearchTarget, 0, len(targets))
+					for _, target := range targets {
+						kb := kbMap[target.KnowledgeBaseID]
+						if !targetReportsEmbedFailure(kb) {
+							searchableTargets = append(searchableTargets, target)
+							continue
+						}
+						recordError(fmt.Errorf("knowledge base %s has no keyword fallback: %w", target.KnowledgeBaseID, err))
+					}
+					pipelineWarn(ctx, "Search", "group_embed_degrade_keyword", map[string]interface{}{
+						"model_key":        modelKey,
+						"kb_id":            targets[0].KnowledgeBaseID,
+						"error":            err.Error(),
+						"fallback_targets": len(searchableTargets),
+						"failed_targets":   len(targets) - len(searchableTargets),
+					})
+					disableVector = true
+				} else {
+					queryEmbedding = emb
+				}
+			}
+
+			// Separate full-KB targets (can be combined into one retrieval)
+			// from specific-knowledge targets (need per-target direct loading).
+			var fullKBIDs []string
+			var knowledgeTargets []*types.SearchTarget
+			for _, t := range searchableTargets {
+				if t.Type == types.SearchTargetTypeKnowledgeBase && len(t.TagIDs) == 0 {
+					fullKBIDs = append(fullKBIDs, t.KnowledgeBaseID)
+				} else {
+					knowledgeTargets = append(knowledgeTargets, t)
+				}
+			}
+
+			pipelineInfo(ctx, "Search", "group_plan", map[string]interface{}{
+				"model_key":          modelKey,
+				"combined_kb_count":  len(fullKBIDs),
+				"individual_targets": len(knowledgeTargets),
+				"vector_len":         len(queryEmbedding),
+			})
+
+			var innerWg sync.WaitGroup
+
+			// Combined search: one HybridSearch call spanning all full-KB targets
+			if len(fullKBIDs) > 0 {
+				innerWg.Add(1)
+				go func() {
+					defer innerWg.Done()
+
+					params := types.SearchParams{
+						QueryText:             queryText,
+						QueryEmbedding:        queryEmbedding,
+						KnowledgeBaseIDs:      fullKBIDs,
+						VectorThreshold:       chatManage.VectorThreshold,
+						KeywordThreshold:      chatManage.KeywordThreshold,
+						MatchCount:            chatManage.EmbeddingTopK,
+						SkipContextEnrichment: true,
+						DisableVectorMatch:    disableVector,
+					}
+					res, err := p.knowledgeBaseService.HybridSearch(ctx, fullKBIDs[0], params)
+					if err != nil {
+						pipelineWarn(ctx, "Search", "combined_kb_search_error", map[string]interface{}{
+							"kb_ids": fullKBIDs,
+							"error":  err.Error(),
+						})
+						recordError(err)
+						return
+					}
+					pipelineInfo(ctx, "Search", "combined_kb_result", map[string]interface{}{
+						"kb_ids":    fullKBIDs,
+						"hit_count": len(res),
+					})
+					mu.Lock()
+					results = append(results, res...)
+					mu.Unlock()
+				}()
+			}
+
+			// Individual search: per-target handling for specific-knowledge targets
+			for _, target := range knowledgeTargets {
+				innerWg.Add(1)
+				go func(t *types.SearchTarget) {
+					defer innerWg.Done()
+					recordError(p.searchSingleTarget(
+						ctx, chatManage, t, queryText, queryEmbedding, disableVector, &mu, &results,
+					))
+				}(target)
+			}
+
+			innerWg.Wait()
+		}(modelKey, targets)
+	}
+
+	wg.Wait()
+
+	pipelineInfo(ctx, "Search", "kb_result_summary", map[string]interface{}{
+		"total_hits": len(results),
+	})
+	return results, firstErr
+}
+
+// searchSingleTarget performs hybrid retrieval inside one constrained target.
+func (p *PluginSearch) searchSingleTarget(
+	ctx context.Context,
+	chatManage *types.ChatManage,
+	t *types.SearchTarget,
+	queryText string,
+	queryEmbedding []float32,
+	disableVector bool,
+	mu *sync.Mutex,
+	results *[]*types.SearchResult,
+) error {
+	if t.Type == types.SearchTargetTypeKnowledge && len(t.KnowledgeIDs) == 0 {
+		return nil
+	}
+
+	vectorThreshold, keywordThreshold := t.RecallThresholds(
+		chatManage.VectorThreshold,
+		chatManage.KeywordThreshold,
+	)
+	if t.DisableRecallThresholds {
+		pipelineInfo(ctx, "Search", "explicit_scope_threshold_override", map[string]interface{}{
+			"kb_id":              t.KnowledgeBaseID,
+			"knowledge_id_count": len(t.KnowledgeIDs),
+			"tag_id_count":       len(t.TagIDs),
+		})
+	}
+	params := types.SearchParams{
+		QueryText:             queryText,
+		QueryEmbedding:        queryEmbedding,
+		VectorThreshold:       vectorThreshold,
+		KeywordThreshold:      keywordThreshold,
+		MatchCount:            chatManage.EmbeddingTopK,
+		TagIDs:                t.TagIDs,
+		ScopeTagIDs:           t.ScopeTagIDs,
+		SkipContextEnrichment: true,
+		DisableVectorMatch:    disableVector,
+	}
+	if t.Type == types.SearchTargetTypeKnowledge {
+		params.KnowledgeIDs = t.KnowledgeIDs
+	}
+	res, err := p.knowledgeBaseService.HybridSearch(ctx, t.KnowledgeBaseID, params)
+	if err != nil {
+		pipelineWarn(ctx, "Search", "kb_search_error", map[string]interface{}{
+			"kb_id":       t.KnowledgeBaseID,
+			"target_type": t.Type,
+			"query":       params.QueryText,
+			"error":       err.Error(),
+		})
+		return err
+	}
+	pipelineInfo(ctx, "Search", "kb_result", map[string]interface{}{
+		"kb_id":       t.KnowledgeBaseID,
+		"target_type": t.Type,
+		"hit_count":   len(res),
+	})
+	mu.Lock()
+	*results = append(*results, res...)
+	mu.Unlock()
+	return nil
+}
+
+// searchWebIfEnabled executes web search when enabled and returns converted results
+func (p *PluginSearch) searchWebIfEnabled(ctx context.Context, chatManage *types.ChatManage) []*types.SearchResult {
+	if !chatManage.WebSearchEnabled || p.webSearchService == nil || p.tenantService == nil {
+		return nil
+	}
+	tenant, _ := types.TenantInfoFromContext(ctx)
+	providerID := chatManage.WebSearchProviderID
+
+	if providerID == "" {
+		pipelineWarn(ctx, "Search", "web_config_missing", map[string]interface{}{
+			"tenant_id": chatManage.TenantID,
+		})
+		return nil
+	}
+
+	webConfig := types.EffectiveWebSearchConfig(nil)
+	if tenant != nil {
+		webConfig = types.EffectiveWebSearchConfig(tenant.WebSearchConfig)
+	}
+
+	// Apply agent-level web search overrides
+	if chatManage.WebSearchMaxResults > 0 {
+		webConfig.MaxResults = chatManage.WebSearchMaxResults
+	}
+
+	pipelineInfo(ctx, "Search", "web_request", map[string]interface{}{
+		"tenant_id":   chatManage.TenantID,
+		"provider_id": providerID,
+	})
+	webCtx, webSpan := langfuse.GetManager().StartSpan(ctx, langfuse.SpanOptions{
+		Name: "web_search",
+		Input: map[string]interface{}{
+			"provider_id": providerID,
+			"query":       chatManage.RewriteQuery,
+			"max_results": webConfig.MaxResults,
+		},
+	})
+	webResults, err := p.webSearchService.Search(webCtx, providerID, webConfig, chatManage.RewriteQuery)
+	webSpan.Finish(map[string]interface{}{
+		"hit_count": len(webResults),
+	}, nil, err)
+	if err != nil {
+		pipelineWarn(ctx, "Search", "web_search_error", map[string]interface{}{
+			"tenant_id": chatManage.TenantID,
+			"error":     err.Error(),
+		})
+		return nil
+	}
+	// Build questions using RewriteQuery only
+	// questions := []string{strings.TrimSpace(chatManage.RewriteQuery)}
+	// Load session-scoped temp KB state from Redis using WebSearchStateRepository
+	// tempKBID, seen, ids := p.webSearchStateService.GetWebSearchTempKBState(ctx, chatManage.SessionID)
+	// compressed, kbID, newSeen, newIDs, err := p.webSearchService.CompressWithRAG(
+	// 	ctx, chatManage.SessionID, tempKBID, questions, webResults, webConfig,
+	// 	p.knowledgeBaseService, p.knowledgeService, seen, ids,
+	// )
+	// if err != nil {
+	// 	pipelineWarn(ctx, "Search", "web_compress_error", map[string]interface{}{
+	// 		"error": err.Error(),
+	// 	})
+	// } else {
+	// 	webResults = compressed
+	// 	// Persist temp KB state back into Redis using WebSearchStateRepository
+	// 	p.webSearchStateService.SaveWebSearchTempKBState(ctx, chatManage.SessionID, kbID, newSeen, newIDs)
+	// }
+	res := searchutil.ConvertWebSearchResults(webResults)
+	pipelineInfo(ctx, "Search", "web_hits", map[string]interface{}{
+		"hit_count": len(res),
+	})
+	return res
+}