package tools import ( "context" "encoding/json" "fmt" "math" "sort" "strings" "sync" "github.com/Tencent/WeKnora/internal/config" "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/models/rerank" "github.com/Tencent/WeKnora/internal/searchutil" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" ) var knowledgeSearchTool = BaseTool{ name: ToolKnowledgeSearch, description: `Semantic/vector search tool for retrieving knowledge by meaning, intent, and conceptual relevance. This tool uses embeddings to understand the user's query and find semantically similar content across knowledge base chunks. ## Purpose Designed for high-level understanding tasks, such as: - conceptual explanations - topic overviews - reasoning-based information needs - contextual or intent-driven retrieval - queries that cannot be answered with literal keyword matching The tool searches by MEANING rather than exact text. It identifies chunks that are conceptually relevant even when the wording differs. ## What the Tool Does NOT Do - Does NOT perform exact keyword matching - Does NOT search for specific named entities - Should NOT be used for literal lookup tasks - Should NOT receive long raw text or user messages as queries - Should NOT be used to locate specific strings or error codes For literal/keyword/entity search, another tool should be used. ## Required Input Behavior "queries" must contain **1–5 short, well-formed semantic questions or conceptual statements** that clearly express the meaning the model is trying to retrieve. Each query should represent a **concept, idea, topic, explanation, or intent**, such as: - abstract topics - definitions - mechanisms - best practices - comparisons - how/why questions Avoid: - keyword lists - raw text from user messages - full paragraphs - unprocessed input ## Examples of valid query shapes (not content): - "What is the main idea of..." - "How does X work in general?" - "Explain the purpose of..." - "What are the key principles behind..." - "Overview of ..." ## Parameters - queries (required): 1–5 semantic questions or conceptual statements. These should reflect the meaning or topic you want embeddings to capture. - knowledge_base_ids (optional): limit the search scope. ## Output Returns chunks ranked by semantic similarity, reranked when applicable. Each chunk has a short cN source ID and belongs to a dN document ID. Results represent conceptual relevance, not literal keyword overlap. Use dN for document-level follow-up tool calls.`, schema: json.RawMessage(`{ "type": "object", "properties": { "queries": { "type": "array", "description": "REQUIRED: 1-5 semantic questions/topics (e.g., ['What is RAG?', 'RAG benefits'])", "items": { "type": "string" }, "minItems": 1, "maxItems": 5 }, "knowledge_base_ids": { "type": "array", "description": "Optional: bound knowledge-base IDs (the short bN values shown in runtime context)", "items": { "type": "string" }, "minItems": 0, "maxItems": 10 } }, "required": ["queries"] }`), } // KnowledgeSearchInput defines the input parameters for knowledge search tool type KnowledgeSearchInput struct { Queries []string `json:"queries"` KnowledgeBaseIDs []string `json:"knowledge_base_ids,omitempty"` } // searchResultWithMeta wraps search result with metadata about which query matched it type searchResultWithMeta struct { *types.SearchResult SourceQuery string QueryType string // "vector" or "keyword" KnowledgeBaseID string // ID of the knowledge base this result came from KnowledgeBaseType string // Type of the knowledge base (document, faq, etc.) } // KnowledgeSearchTool searches knowledge bases with flexible query modes. // seenChunks lets repeated calls in the same session surface previously- // returned chunks in a compact form (mirroring wiki_search's de-duping UX) // so the LLM doesn't burn tokens re-reading identical content. type KnowledgeSearchTool struct { BaseTool knowledgeBaseService interfaces.KnowledgeBaseService knowledgeService interfaces.KnowledgeService chunkService interfaces.ChunkService searchTargets types.SearchTargets // Pre-computed unified search targets rerankModel rerank.Reranker config *config.Config // Global config for fallback values seenMu sync.Mutex seenChunks map[string]bool } // NewKnowledgeSearchTool creates a new knowledge search tool func NewKnowledgeSearchTool( knowledgeBaseService interfaces.KnowledgeBaseService, knowledgeService interfaces.KnowledgeService, chunkService interfaces.ChunkService, searchTargets types.SearchTargets, rerankModel rerank.Reranker, cfg *config.Config, ) *KnowledgeSearchTool { return &KnowledgeSearchTool{ BaseTool: knowledgeSearchTool, knowledgeBaseService: knowledgeBaseService, knowledgeService: knowledgeService, chunkService: chunkService, searchTargets: searchTargets, rerankModel: rerankModel, config: cfg, seenChunks: make(map[string]bool), } } // Execute executes the knowledge search tool func (t *KnowledgeSearchTool) Execute(ctx context.Context, args json.RawMessage) (*types.ToolResult, error) { logger.Infof(ctx, "[Tool][KnowledgeSearch] Execute started") // Parse args from json.RawMessage var input KnowledgeSearchInput if err := json.Unmarshal(args, &input); err != nil { logger.Errorf(ctx, "[Tool][KnowledgeSearch] Failed to parse args: %v", err) return &types.ToolResult{ Success: false, Error: fmt.Sprintf("Failed to parse args: %v", err), }, err } // Log input arguments argsJSON, _ := json.MarshalIndent(input, "", " ") logger.Debugf(ctx, "[Tool][KnowledgeSearch] Input args:\n%s", string(argsJSON)) // Determine which KBs to search - user can optionally filter to specific KBs var userSpecifiedKBs []string if len(input.KnowledgeBaseIDs) > 0 { userSpecifiedKBs = input.KnowledgeBaseIDs if err := validateKnowledgeBaseIDsInSearchTargets(t.searchTargets, userSpecifiedKBs); err != nil { return &types.ToolResult{Success: false, Error: err.Error()}, err } logger.Infof(ctx, "[Tool][KnowledgeSearch] User specified %d knowledge bases: %v", len(userSpecifiedKBs), userSpecifiedKBs) } // Use pre-computed search targets, optionally filtered by user-specified KBs searchTargets := t.searchTargets if len(userSpecifiedKBs) > 0 { // Filter search targets to only include user-specified KBs userKBSet := make(map[string]bool) for _, kbID := range userSpecifiedKBs { userKBSet[kbID] = true } var filteredTargets types.SearchTargets for _, target := range t.searchTargets { if target == nil { continue } if userKBSet[target.KnowledgeBaseID] { filteredTargets = append(filteredTargets, target) } } searchTargets = filteredTargets } // Validate search targets if len(searchTargets) == 0 { logger.Errorf(ctx, "[Tool][KnowledgeSearch] No search targets available") return &types.ToolResult{ Success: false, Error: "no knowledge bases specified and no search targets configured", }, fmt.Errorf("no search targets available") } kbIDs := searchTargets.GetAllKnowledgeBaseIDs() logger.Infof(ctx, "[Tool][KnowledgeSearch] Using %d search targets across %d KBs", len(searchTargets), len(kbIDs)) // Parse query parameter queries := input.Queries // Validate: query must be provided if len(queries) == 0 { logger.Errorf(ctx, "[Tool][KnowledgeSearch] No queries provided") return &types.ToolResult{ Success: false, Error: "queries parameter is required", }, fmt.Errorf("no queries provided") } logger.Infof(ctx, "[Tool][KnowledgeSearch] Queries: %v", queries) // Search parameters: fall back to global config, then to hardcoded defaults. // We used to read tenant.ConversationConfig here as the first source of // truth, but that field was removed when the chat pipeline moved to // CustomAgent — tenant-level KV settings now live on the agent itself. var topK int var vectorThreshold, keywordThreshold, minScore float64 // Fallback to global config if not set if topK == 0 && t.config != nil { topK = t.config.Conversation.EmbeddingTopK } if vectorThreshold == 0 && t.config != nil { vectorThreshold = t.config.Conversation.VectorThreshold } if keywordThreshold == 0 && t.config != nil { keywordThreshold = t.config.Conversation.KeywordThreshold } // Final fallback to hardcoded defaults if config is not available if topK == 0 { topK = 5 } if vectorThreshold == 0 { vectorThreshold = 0.6 } if keywordThreshold == 0 { keywordThreshold = 0.5 } if minScore == 0 { minScore = 0.3 } logger.Infof( ctx, "[Tool][KnowledgeSearch] Search params: top_k=%d, vector_threshold=%.2f, keyword_threshold=%.2f, min_score=%.2f", topK, vectorThreshold, keywordThreshold, minScore, ) // Execute concurrent search using pre-computed search targets logger.Infof(ctx, "[Tool][KnowledgeSearch] Starting concurrent search with %d search targets", len(searchTargets)) kbTypeMap := t.getKnowledgeBaseTypes(ctx, kbIDs) allResults := t.concurrentSearchByTargets(ctx, queries, searchTargets, topK, vectorThreshold, keywordThreshold, kbTypeMap) logger.Infof(ctx, "[Tool][KnowledgeSearch] Concurrent search completed: %d raw results", len(allResults)) // Note: HybridSearch now uses RRF (Reciprocal Rank Fusion) which produces normalized scores // RRF scores are in range [0, ~0.033] (max when rank=1 on both sides: 2/(60+1)) // Threshold filtering is already done inside HybridSearch before RRF, so we skip it here // Deduplicate before reranking to reduce processing overhead deduplicatedBeforeRerank := t.deduplicateResults(allResults) // Apply ReRank if model is configured. // Use first query for reranking (or combine all queries if needed) rerankQuery := "" if len(queries) > 0 { rerankQuery = queries[0] if len(queries) > 1 { // Combine multiple queries for reranking rerankQuery = strings.Join(queries, " ") } } // Variable to hold results through reranking and MMR stages var filteredResults []*searchResultWithMeta if t.rerankModel != nil && len(deduplicatedBeforeRerank) > 0 && rerankQuery != "" { logger.Infof(ctx, "[Tool][KnowledgeSearch] Applying rerank, input: %d results, threshold: %.2f, queries: %v", len(deduplicatedBeforeRerank), t.rerankThreshold(), queries) rerankedResults, err := t.rerankResults(ctx, rerankQuery, deduplicatedBeforeRerank) if err != nil { logger.Warnf(ctx, "[Tool][KnowledgeSearch] Rerank failed, using original results: %v", err) filteredResults = deduplicatedBeforeRerank } else { filteredResults = rerankedResults logger.Infof(ctx, "[Tool][KnowledgeSearch] Rerank completed successfully: %d results", len(filteredResults)) } } else { // No reranking model available, use deduplicated results filteredResults = deduplicatedBeforeRerank } // Apply MMR (Maximal Marginal Relevance) to reduce redundancy and improve diversity // Note: composite scoring is already applied inside rerankResults if len(filteredResults) > 0 { // Calculate k for MMR: use min(len(results), max(1, topK)) mmrK := len(filteredResults) if topK > 0 && mmrK > topK { mmrK = topK } if mmrK < 1 { mmrK = 1 } // Apply MMR with lambda=0.7 (balance between relevance and diversity) logger.Debugf( ctx, "[Tool][KnowledgeSearch] Applying MMR: k=%d, lambda=0.7, input=%d results", mmrK, len(filteredResults), ) mmrResults := t.applyMMR(ctx, filteredResults, mmrK, 0.7) if len(mmrResults) > 0 { filteredResults = mmrResults logger.Infof(ctx, "[Tool][KnowledgeSearch] MMR completed: %d results selected", len(filteredResults)) } else { logger.Warnf(ctx, "[Tool][KnowledgeSearch] MMR returned no results, using original results") } } // Note: minScore filter is skipped because HybridSearch now uses RRF scores // RRF scores are in range [0, ~0.033], not [0, 1], so old thresholds don't apply // Threshold filtering is already done inside HybridSearch before RRF fusion // Final deduplication after rerank (in case rerank changed scores/order but duplicates remain) logger.Debugf(ctx, "[Tool][KnowledgeSearch] Final deduplication after rerank...") deduplicatedResults := t.deduplicateResults(filteredResults) logger.Infof(ctx, "[Tool][KnowledgeSearch] After final deduplication: %d results (from %d)", len(deduplicatedResults), len(filteredResults)) // Sort results by score (descending) sort.Slice(deduplicatedResults, func(i, j int) bool { if deduplicatedResults[i].Score != deduplicatedResults[j].Score { return deduplicatedResults[i].Score > deduplicatedResults[j].Score } // If scores are equal, sort by knowledge ID for consistency return deduplicatedResults[i].KnowledgeID < deduplicatedResults[j].KnowledgeID }) // Log all ranked results (including lower ranks for rerank debugging) if len(deduplicatedResults) > 0 { total := len(deduplicatedResults) for i, r := range deduplicatedResults { logger.Infof(ctx, "[Tool][KnowledgeSearch][Rank %d/%d] score=%.3f, type=%s, kb=%s, chunk_id=%s", i+1, total, r.Score, r.QueryType, r.KnowledgeID, r.ID) } } // Enrich image info for search results (lazy-loaded from child image chunks) if t.chunkService != nil && len(deduplicatedResults) > 0 { byTenant := make(map[uint64][]*types.SearchResult) for _, r := range deduplicatedResults { tid := t.searchTargets.GetTenantIDForKB(r.KnowledgeBaseID) if tid == 0 { continue } byTenant[tid] = append(byTenant[tid], r.SearchResult) } for tid, batch := range byTenant { searchutil.EnrichSearchResultsImageInfo(ctx, t.chunkService.GetRepository(), tid, batch) } } // Build output logger.Infof(ctx, "[Tool][KnowledgeSearch] Formatting output with %d final results", len(deduplicatedResults)) result, err := t.formatOutput(ctx, deduplicatedResults, kbIDs, queries) if err != nil { logger.Errorf(ctx, "[Tool][KnowledgeSearch] Failed to format output: %v", err) return result, err } logger.Infof(ctx, "[Tool][KnowledgeSearch] Output: %s", result.Output) return result, nil } // getKnowledgeBaseTypes fetches knowledge base types for the given IDs func (t *KnowledgeSearchTool) getKnowledgeBaseTypes(ctx context.Context, kbIDs []string) map[string]string { kbTypeMap := make(map[string]string, len(kbIDs)) for _, kbID := range kbIDs { if kbID == "" { continue } if _, exists := kbTypeMap[kbID]; exists { continue } kb, err := t.knowledgeBaseService.GetKnowledgeBaseByID(ctx, kbID) if err != nil { logger.Warnf(ctx, "[Tool][KnowledgeSearch] Failed to fetch knowledge base %s info: %v", kbID, err) continue } kbTypeMap[kbID] = kb.Type } return kbTypeMap } // concurrentSearchByTargets executes hybrid search using pre-computed search targets. // Targets sharing the same underlying embedding model (identified by model name + endpoint) // are grouped so the query embedding is computed once per (model, query) pair, and all // full-KB targets in a group are combined into a single retrieval call. func (t *KnowledgeSearchTool) concurrentSearchByTargets( ctx context.Context, queries []string, searchTargets types.SearchTargets, topK int, vectorThreshold, keywordThreshold float64, kbTypeMap map[string]string, ) []*searchResultWithMeta { // Batch-fetch KB records for embedding model grouping kbIDs := searchTargets.GetAllKnowledgeBaseIDs() var kbList []*types.KnowledgeBase if kbs, err := t.knowledgeBaseService.GetKnowledgeBasesByIDsOnly(ctx, kbIDs); err == nil { kbList = kbs } // Filter out non-searchable KBs (wiki-only / graph-only). knowledge_search // can only serve KBs with vector or keyword indexing; feeding a wiki-only // KB into HybridSearch causes spurious "model ID cannot be empty" errors // because such KBs have no EmbeddingModelID configured. Such scopes // should be queried via wiki_search / graph tools instead. // // KBs that we couldn't fetch from the repo (not in kbList) are kept so // the downstream HybridSearch path can still surface the real error. searchableKBs := make(map[string]bool, len(kbList)) knownKBs := make(map[string]bool, len(kbList)) for _, kb := range kbList { if kb == nil { continue } knownKBs[kb.ID] = true if kb.IsVectorEnabled() || kb.IsKeywordEnabled() { searchableKBs[kb.ID] = true } } filteredTargets := make(types.SearchTargets, 0, len(searchTargets)) for _, st := range searchTargets { if st == nil || st.KnowledgeBaseID == "" { continue } if searchableKBs[st.KnowledgeBaseID] { filteredTargets = append(filteredTargets, st) continue } if knownKBs[st.KnowledgeBaseID] { logger.Infof(ctx, "[Tool][KnowledgeSearch] Skipping non-searchable KB %s (no vector/keyword index, likely wiki/graph-only)", st.KnowledgeBaseID) continue } // KB record unavailable; keep so downstream can surface real errors. filteredTargets = append(filteredTargets, st) } if len(filteredTargets) == 0 { logger.Infof(ctx, "[Tool][KnowledgeSearch] No searchable KBs in scope (all wiki/graph-only); skipping retrieval") return nil } searchTargets = filteredTargets // Resolve actual model identities (name + endpoint) for cross-tenant grouping modelKeyMap := t.knowledgeBaseService.ResolveEmbeddingModelKeys(ctx, kbList) groups := make(map[string][]*types.SearchTarget) for _, st := range searchTargets { if st == nil || st.KnowledgeBaseID == "" { continue } key := modelKeyMap[st.KnowledgeBaseID] groups[key] = append(groups[key], st) } var wg sync.WaitGroup var mu sync.Mutex allResults := make([]*searchResultWithMeta, 0) for _, query := range queries { q := query for modelKey, targets := range groups { wg.Add(1) go func(q string, modelKey string, targets []*types.SearchTarget) { defer wg.Done() // Compute embedding once for this (model, query) pair var queryEmbedding []float32 if modelKey != "" { emb, err := t.knowledgeBaseService.GetQueryEmbedding(ctx, targets[0].KnowledgeBaseID, q) if err != nil { logger.Warnf(ctx, "[Tool][KnowledgeSearch] Failed to pre-compute embedding for model %s: %v", modelKey, err) } else { queryEmbedding = emb } } // Separate full-KB targets (combinable) from specific-knowledge targets var fullKBIDs []string var knowledgeTargets []*types.SearchTarget for _, st := range targets { if st.Type == types.SearchTargetTypeKnowledgeBase && len(st.TagIDs) == 0 { fullKBIDs = append(fullKBIDs, st.KnowledgeBaseID) } else { knowledgeTargets = append(knowledgeTargets, st) } } var innerWg sync.WaitGroup // Combined retrieval for all full-KB targets in this group if len(fullKBIDs) > 0 { innerWg.Add(1) go func() { defer innerWg.Done() searchParams := types.SearchParams{ QueryText: q, QueryEmbedding: queryEmbedding, KnowledgeBaseIDs: fullKBIDs, MatchCount: topK, VectorThreshold: vectorThreshold, KeywordThreshold: keywordThreshold, } kbResults, err := t.knowledgeBaseService.HybridSearch(ctx, fullKBIDs[0], searchParams) if err != nil { logger.Warnf(ctx, "[Tool][KnowledgeSearch] Combined search failed for KBs %v: %v", fullKBIDs, err) return } mu.Lock() for _, r := range kbResults { allResults = append(allResults, &searchResultWithMeta{ SearchResult: r, SourceQuery: q, QueryType: "hybrid", KnowledgeBaseID: r.KnowledgeBaseID, KnowledgeBaseType: kbTypeMap[r.KnowledgeBaseID], }) } mu.Unlock() }() } // Individual retrieval for specific-knowledge targets for _, target := range knowledgeTargets { st := target innerWg.Add(1) go func() { defer innerWg.Done() stVectorThreshold, stKeywordThreshold := st.RecallThresholds( vectorThreshold, keywordThreshold, ) searchParams := types.SearchParams{ QueryText: q, QueryEmbedding: queryEmbedding, MatchCount: topK, VectorThreshold: stVectorThreshold, KeywordThreshold: stKeywordThreshold, KnowledgeIDs: st.KnowledgeIDs, TagIDs: st.TagIDs, ScopeTagIDs: st.ScopeTagIDs, } kbResults, err := t.knowledgeBaseService.HybridSearch(ctx, st.KnowledgeBaseID, searchParams) if err != nil { logger.Warnf(ctx, "[Tool][KnowledgeSearch] Failed to search KB %s: %v", st.KnowledgeBaseID, err) return } mu.Lock() for _, r := range kbResults { allResults = append(allResults, &searchResultWithMeta{ SearchResult: r, SourceQuery: q, QueryType: "hybrid", KnowledgeBaseID: r.KnowledgeBaseID, KnowledgeBaseType: kbTypeMap[r.KnowledgeBaseID], }) } mu.Unlock() }() } innerWg.Wait() }(q, modelKey, targets) } } wg.Wait() return allResults } // rerankResults applies reranking to all search results (including FAQ entries) // using the configured rerank model, then filters by threshold and applies // composite scoring so MMR/sorting uses a single score scale. // // A failed rerank call degrades to the raw retrieval order, mirroring the chat // pipeline's api_error_fallback. An empty result after threshold filtering is // kept empty: filterRerankRankResults already preserves the top candidate down // to agentRerankFallbackMinScore, so reaching zero means even the best match is // below that floor. There is deliberately no chat-model re-scoring path here — // it mixed the reranker's [0,1] scale with raw RRF scores and could resurrect // candidates the reranker had already rejected. func (t *KnowledgeSearchTool) rerankResults( ctx context.Context, query string, results []*searchResultWithMeta, ) ([]*searchResultWithMeta, error) { if len(results) == 0 || t.rerankModel == nil { return results, nil } rankResults, err := t.rerankScores(ctx, query, results) if err != nil { logger.Warnf(ctx, "[Tool][KnowledgeSearch] Rerank model failed, using raw retrieval results: %v", err) return results, nil } threshold := t.rerankThreshold() reranked := t.applyModelRerankScores( results, rankResults, threshold, t.searchTargets.HasRecallThresholdOverride(), ) logger.Infof(ctx, "[Tool][KnowledgeSearch] Reranked %d/%d results above threshold %.2f", len(reranked), len(results), threshold) return reranked, nil } func (t *KnowledgeSearchTool) getFAQMetadata( ctx context.Context, chunkID string, cache map[string]*types.FAQChunkMetadata, ) (*types.FAQChunkMetadata, error) { if chunkID == "" || t.chunkService == nil { return nil, nil } if meta, ok := cache[chunkID]; ok { return meta, nil } chunk, err := t.chunkService.GetChunkByID(ctx, chunkID) if err != nil { cache[chunkID] = nil return nil, err } if chunk == nil { cache[chunkID] = nil return nil, nil } meta, err := chunk.FAQMetadata() if err != nil { cache[chunkID] = nil return nil, err } cache[chunkID] = meta return meta, nil } // rerankScores scores the candidates with the configured rerank model and // returns the raw relevance scores, leaving threshold filtering and composite // scoring to the caller. func (t *KnowledgeSearchTool) rerankScores( ctx context.Context, query string, results []*searchResultWithMeta, ) ([]rerank.RankResult, error) { passages := make([]string, len(results)) for i, result := range results { passages[i] = t.getEnrichedPassage(ctx, result.SearchResult) } rerankResp, err := t.rerankModel.Rerank(ctx, query, passages) if err != nil { return nil, fmt.Errorf("rerank call failed: %w", err) } return rerankResp, nil } func (t *KnowledgeSearchTool) rerankThreshold() float64 { if t.config != nil && t.config.Conversation != nil && t.config.Conversation.RerankThreshold > 0 { return t.config.Conversation.RerankThreshold } return 0.3 } const agentRerankFallbackMinScore = 0.15 func filterRerankRankResults( rankResults []rerank.RankResult, threshold float64, preserveTop bool, ) []rerank.RankResult { if len(rankResults) == 0 { return nil } filtered := make([]rerank.RankResult, 0, len(rankResults)) for _, r := range rankResults { if r.RelevanceScore >= threshold { filtered = append(filtered, r) } } if len(filtered) == 0 { top := rankResults[0] for _, r := range rankResults[1:] { if r.RelevanceScore > top.RelevanceScore { top = r } } if preserveTop || top.RelevanceScore >= agentRerankFallbackMinScore { return []rerank.RankResult{top} } } return filtered } func (t *KnowledgeSearchTool) applyModelRerankScores( originals []*searchResultWithMeta, rankResults []rerank.RankResult, threshold float64, preserveTop bool, ) []*searchResultWithMeta { filtered := filterRerankRankResults(rankResults, threshold, preserveTop) out := make([]*searchResultWithMeta, 0, len(filtered)) for _, rr := range filtered { if rr.Index < 0 || rr.Index >= len(originals) { continue } newResult := *originals[rr.Index] baseScore := newResult.Score modelScore := rr.RelevanceScore newResult.Score = t.compositeScore(&newResult, modelScore, baseScore) out = append(out, &newResult) } sort.Slice(out, func(i, j int) bool { return out[i].Score > out[j].Score }) return out } // deduplicateResults removes duplicate chunks, keeping the highest score // Uses multiple keys (ID, parent chunk ID, knowledge+index) and content signature for deduplication func (t *KnowledgeSearchTool) deduplicateResults(results []*searchResultWithMeta) []*searchResultWithMeta { seen := make(map[string]bool) contentSig := make(map[string]bool) uniqueResults := make([]*searchResultWithMeta, 0) for _, r := range results { // Build multiple keys for deduplication keys := []string{r.ID} if r.ParentChunkID != "" { keys = append(keys, "parent:"+r.ParentChunkID) } if r.KnowledgeID != "" { keys = append(keys, fmt.Sprintf("kb:%s#%d", r.KnowledgeID, r.ChunkIndex)) } // Check if any key is already seen dup := false for _, k := range keys { if seen[k] { dup = true break } } if dup { continue } // Check content signature for near-duplicate content sig := t.buildContentSignature(r.Content) if sig != "" { if contentSig[sig] { continue } contentSig[sig] = true } // Mark all keys as seen for _, k := range keys { seen[k] = true } uniqueResults = append(uniqueResults, r) } // If we have duplicates by ID but different scores, keep the highest score // This handles cases where the same chunk appears multiple times with different scores seenByID := make(map[string]*searchResultWithMeta) for _, r := range uniqueResults { if existing, ok := seenByID[r.ID]; ok { // Keep the result with higher score if r.Score > existing.Score { seenByID[r.ID] = r } } else { seenByID[r.ID] = r } } // Convert back to slice deduplicated := make([]*searchResultWithMeta, 0, len(seenByID)) for _, r := range seenByID { deduplicated = append(deduplicated, r) } return deduplicated } // buildContentSignature creates a normalized signature for content to detect near-duplicates func (t *KnowledgeSearchTool) buildContentSignature(content string) string { return searchutil.BuildContentSignature(content) } // writeKnowledgeMetadataHeader emits document-scoped metadata once per // knowledge item. Chunk entries keep only chunk-specific content so repeated // results from the same document do not waste model context. func writeKnowledgeMetadataHeader(ob *strings.Builder, results []*searchResultWithMeta) { seen := make(map[string]struct{}, len(results)) hasMetadata := false var documents strings.Builder for _, result := range results { if result == nil || result.SearchResult == nil || result.KnowledgeID == "" || result.KnowledgeCustomMetadata == "" { continue } if _, ok := seen[result.KnowledgeID]; ok { continue } seen[result.KnowledgeID] = struct{}{} hasMetadata = true documents.WriteString(fmt.Sprintf( "\n", xmlEscape(result.KnowledgeID), xmlEscape(result.KnowledgeBaseID), xmlEscape(result.KnowledgeTitle), )) documents.WriteString(fmt.Sprintf("%s\n", xmlEscape(result.KnowledgeCustomMetadata))) documents.WriteString("\n") } if !hasMetadata { return } ob.WriteString("\n") ob.WriteString(documents.String()) ob.WriteString("\n") } // formatOutput formats the search results for display func (t *KnowledgeSearchTool) formatOutput( ctx context.Context, results []*searchResultWithMeta, kbsToSearch []string, queries []string, ) (*types.ToolResult, error) { if len(results) == 0 { data := map[string]interface{}{ "knowledge_base_ids": kbsToSearch, "results": []interface{}{}, "count": 0, } if len(queries) > 0 { data["queries"] = queries } output := fmt.Sprintf("No relevant content found in %d knowledge base(s).\n\n", len(kbsToSearch)) output += "=== ⚠️ CRITICAL - Next Steps ===\n" output += "- ❌ DO NOT use training data or general knowledge to answer\n" output += "- ✅ If web_search is enabled: You MUST use web_search to find information\n" output += "- ✅ If web_search is disabled: State 'I couldn't find relevant information in the knowledge base'\n" output += "- NEVER fabricate or infer answers - ONLY use retrieved content\n" return &types.ToolResult{ Success: true, Output: output, Data: data, }, nil } // Count results by KB kbCounts := make(map[string]int) for _, r := range results { kbCounts[r.KnowledgeBaseID]++ } // Format individual results as XML. Tag names are kept in sync with // wiki_search (``, per-entry element, ``) so that // agents and downstream consumers see a single consistent shape across // all retrieval tools. var ob strings.Builder ob.WriteString(fmt.Sprintf("\n", len(results))) for _, q := range queries { ob.WriteString(fmt.Sprintf("%s\n", xmlEscape(q))) } writeKnowledgeMetadataHeader(&ob, results) formattedResults := make([]map[string]interface{}, 0, len(results)) enabled := true faqMetadataCache := make(map[string]*types.FAQChunkMetadata) knowledgeChunkMap := make(map[string]map[int]bool) knowledgeTotalMap := make(map[string]int64) knowledgeTitleMap := make(map[string]string) for i, result := range results { var faqMeta *types.FAQChunkMetadata if result.KnowledgeBaseType == types.KnowledgeBaseTypeFAQ { meta, err := t.getFAQMetadata(ctx, result.ID, faqMetadataCache) if err != nil { logger.Warnf(ctx, "[Tool][KnowledgeSearch] Failed to load FAQ metadata for chunk %s: %v", result.ID, err) } else { faqMeta = meta } } if knowledgeChunkMap[result.KnowledgeID] == nil { knowledgeChunkMap[result.KnowledgeID] = make(map[int]bool) } knowledgeChunkMap[result.KnowledgeID][result.ChunkIndex] = true knowledgeTitleMap[result.KnowledgeID] = result.KnowledgeTitle // Cache total chunk count per knowledge if _, exists := knowledgeTotalMap[result.KnowledgeID]; !exists { effectiveTenantID := t.searchTargets.GetTenantIDForKB(result.KnowledgeBaseID) if effectiveTenantID == 0 { logger.Warnf(ctx, "[Tool][KnowledgeSearch] KB %s not found in searchTargets, skipping chunk count", result.KnowledgeBaseID) knowledgeTotalMap[result.KnowledgeID] = 0 } else { // Use the same chunk-type filter as list_knowledge_chunks so the // total reported here matches what list_knowledge_chunks can page // over. Mismatched filters previously let LLMs compute offsets // against an inflated/deflated total and page past the end. _, total, err := t.chunkService.GetRepository().ListPagedChunksByKnowledgeID(ctx, effectiveTenantID, result.KnowledgeID, &types.Pagination{Page: 1, PageSize: 1}, []types.ChunkType{types.ChunkTypeText, types.ChunkTypeFAQ}, nil, "", "", "", "", &enabled, ) if err != nil { logger.Warnf(ctx, "[Tool][KnowledgeSearch] Failed to get total chunks for knowledge %s: %v", result.KnowledgeID, err) knowledgeTotalMap[result.KnowledgeID] = 0 } else { knowledgeTotalMap[result.KnowledgeID] = total } } } t.seenMu.Lock() seen := t.seenChunks[result.ID] t.seenChunks[result.ID] = true t.seenMu.Unlock() isFAQ := faqMeta != nil if seen { // Compact rendering for chunks we already returned in a previous // knowledge_search call during this session. The model has the // content in context already, so re-emitting it only burns tokens. if isFAQ { ob.WriteString(fmt.Sprintf( "\n", i+1, xmlEscape(result.ID), result.ChunkIndex, xmlEscape(result.KnowledgeBaseID), xmlEscape(result.KnowledgeTitle), result.Score, xmlEscape(result.SourceQuery), )) } else { ob.WriteString(fmt.Sprintf( "\n", i+1, xmlEscape(result.ID), result.ChunkIndex, xmlEscape(result.KnowledgeID), xmlEscape(result.KnowledgeBaseID), xmlEscape(result.KnowledgeTitle), result.Score, xmlEscape(result.SourceQuery), )) } ob.WriteString("(content omitted, already returned in a previous knowledge_search call this session)\n") if isFAQ { ob.WriteString("\n") } else { ob.WriteString("\n") } } else { if isFAQ { ob.WriteString(fmt.Sprintf( "\n", i+1, xmlEscape(result.ID), result.ChunkIndex, xmlEscape(result.KnowledgeBaseID), xmlEscape(result.KnowledgeTitle), result.Score, xmlEscape(result.SourceQuery), )) } else { ob.WriteString(fmt.Sprintf( "\n", i+1, xmlEscape(result.ID), result.ChunkIndex, xmlEscape(result.KnowledgeID), xmlEscape(result.KnowledgeBaseID), xmlEscape(result.KnowledgeTitle), result.Score, xmlEscape(result.SourceQuery), )) } snippet := "" if faqMeta != nil { snippet = faqMatchSnippetFromQueries(faqMeta, queries) } if snippet == "" { snippet = extractSnippetForQueries(result.Content, queries) } if snippet != "" { ob.WriteString(fmt.Sprintf("%s\n", xmlEscape(snippet))) } ob.WriteString(fmt.Sprintf("%s\n", result.Content)) if result.ImageInfo != "" { var imageInfos []types.ImageInfo if err := json.Unmarshal([]byte(result.ImageInfo), &imageInfos); err == nil && len(imageInfos) > 0 { for _, img := range imageInfos { if imageMarkdown := searchutil.BuildImageInfoMarkdownWithURL(img.URL, &img); imageMarkdown != "" { ob.WriteString(imageMarkdown) ob.WriteString("\n") } } } } if isFAQ { writeFAQFieldsXML(&ob, faqMeta) ob.WriteString("\n") } else { ob.WriteString("\n") } } formattedResults = append(formattedResults, map[string]interface{}{ "result_index": i + 1, "content": result.Content, "knowledge_id": result.KnowledgeID, "knowledge_base_id": result.KnowledgeBaseID, "knowledge_title": result.KnowledgeTitle, "knowledge_metadata": result.KnowledgeCustomMetadata, "match_type": result.MatchType, "source_query": result.SourceQuery, "query_type": result.QueryType, "knowledge_base_type": result.KnowledgeBaseType, }) last := formattedResults[len(formattedResults)-1] if result.ImageInfo != "" { var imageInfos []types.ImageInfo if err := json.Unmarshal([]byte(result.ImageInfo), &imageInfos); err == nil && len(imageInfos) > 0 { imageList := make([]map[string]string, 0, len(imageInfos)) for _, img := range imageInfos { imgData := make(map[string]string) if img.URL != "" { imgData["url"] = img.URL } if img.Caption != "" { imgData["caption"] = img.Caption } if img.OCRText != "" { imgData["ocr_text"] = img.OCRText } if len(imgData) > 0 { imageList = append(imageList, imgData) } } if len(imageList) > 0 { last["images"] = imageList } } } if faqMeta != nil { last["faq_id"] = result.ID last["index"] = result.ChunkIndex if faqMeta.StandardQuestion != "" { last["faq_standard_question"] = faqMeta.StandardQuestion } appendSimilarQuestionsToChunkData(last, faqMeta.SimilarQuestions) if len(faqMeta.Answers) > 0 { last["faq_answers"] = faqMeta.Answers } } else { last["chunk_id"] = result.ID last["chunk_index"] = result.ChunkIndex } } // Retrieval statistics ob.WriteString("\n") for knowledgeID, retrievedChunks := range knowledgeChunkMap { totalChunks := knowledgeTotalMap[knowledgeID] retrievedCount := len(retrievedChunks) title := knowledgeTitleMap[knowledgeID] if totalChunks > 0 { remaining := totalChunks - int64(retrievedCount) percentage := float64(retrievedCount) / float64(totalChunks) * 100 ob.WriteString(fmt.Sprintf("\n", xmlEscape(knowledgeID), xmlEscape(title), totalChunks, retrievedCount, remaining, percentage)) } } ob.WriteString("\n") ob.WriteString("") output := ob.String() data := map[string]interface{}{ "knowledge_base_ids": kbsToSearch, "results": formattedResults, "count": len(formattedResults), "kb_counts": kbCounts, "display_type": "search_results", } if len(queries) > 0 { data["queries"] = queries } return &types.ToolResult{ Success: true, Output: output, Data: data, }, nil } // chunkRange represents a continuous range of chunk indices type chunkRange struct { start int end int } // getEnrichedPassage 合并Content和ImageInfo的文本内容 func (t *KnowledgeSearchTool) getEnrichedPassage(ctx context.Context, result *types.SearchResult) string { if result.ImageInfo == "" { return result.Content } // 解析ImageInfo var imageInfos []types.ImageInfo err := json.Unmarshal([]byte(result.ImageInfo), &imageInfos) if err != nil { logger.Warnf(ctx, "[Tool][KnowledgeSearch] Failed to parse image info: %v", err) return result.Content } if len(imageInfos) == 0 { return result.Content } // 提取所有图片的描述和OCR文本 var imageTexts []string for _, img := range imageInfos { if img.Caption != "" { imageTexts = append(imageTexts, fmt.Sprintf("Image Caption: %s", img.Caption)) } if img.OCRText != "" { imageTexts = append(imageTexts, fmt.Sprintf("Image Text: %s", img.OCRText)) } } if len(imageTexts) == 0 { return result.Content } // 组合内容和图片信息 combinedText := result.Content if combinedText != "" { combinedText += "\n\n" } combinedText += strings.Join(imageTexts, "\n") logger.Debugf(ctx, "[Tool][KnowledgeSearch] Enriched passage: content_len=%d, image_texts=%d", len(result.Content), len(imageTexts)) return combinedText } // compositeScore calculates a composite score considering multiple factors func (t *KnowledgeSearchTool) compositeScore( result *searchResultWithMeta, modelScore, baseScore float64, ) float64 { // Source weight: web_search results get slightly lower weight sourceWeight := 1.0 if strings.ToLower(result.KnowledgeSource) == "web_search" { sourceWeight = 0.95 } // Position prior: slightly favor chunks earlier in the document positionPrior := 1.0 if result.StartAt >= 0 && result.EndAt > result.StartAt { // Calculate position ratio and apply small boost for earlier positions positionRatio := 1.0 - float64(result.StartAt)/float64(result.EndAt+1) positionPrior += t.clampFloat(positionRatio, -0.05, 0.05) } // Composite formula: weighted combination of model score, base score, and source weight composite := 0.6*modelScore + 0.3*baseScore + 0.1*sourceWeight composite *= positionPrior // Clamp to [0, 1] if composite < 0 { composite = 0 } if composite > 1 { composite = 1 } return composite } // clampFloat clamps a float value to the specified range func (t *KnowledgeSearchTool) clampFloat(v, minV, maxV float64) float64 { return searchutil.ClampFloat(v, minV, maxV) } // applyMMR applies Maximal Marginal Relevance algorithm to reduce redundancy func (t *KnowledgeSearchTool) applyMMR( ctx context.Context, results []*searchResultWithMeta, k int, lambda float64, ) []*searchResultWithMeta { if k <= 0 || len(results) == 0 { return nil } logger.Infof(ctx, "[Tool][KnowledgeSearch] Applying MMR: lambda=%.2f, k=%d, candidates=%d", lambda, k, len(results)) selected := make([]*searchResultWithMeta, 0, k) candidates := make([]*searchResultWithMeta, len(results)) copy(candidates, results) // Pre-compute token sets for all candidates tokenSets := make([]map[string]struct{}, len(candidates)) for i, r := range candidates { tokenSets[i] = t.tokenizeSimple(t.getEnrichedPassage(ctx, r.SearchResult)) } // MMR selection loop, incremental form: maxRedundancy[i] caches candidate i's // maximum jaccard against everything selected so far, so each round only needs // one comparison per remaining candidate instead of one per (candidate, selected) // pair. Selection output is identical to the naive form, including tie-breaking, // because the candidate iteration order is unchanged. selectedTokenSets := make([]map[string]struct{}, 0, k) maxRedundancy := make([]float64, len(candidates)) for len(selected) < k && len(candidates) > 0 { bestIdx := 0 bestScore := -1.0 for i, r := range candidates { // MMR score: balance relevance and diversity mmr := lambda*r.Score - (1.0-lambda)*maxRedundancy[i] if mmr > bestScore { bestScore = mmr bestIdx = i } } // Add best candidate to selected and remove from candidates selected = append(selected, candidates[bestIdx]) chosenTokens := tokenSets[bestIdx] selectedTokenSets = append(selectedTokenSets, chosenTokens) candidates = append(candidates[:bestIdx], candidates[bestIdx+1:]...) // Remove corresponding token set and cached redundancy tokenSets = append(tokenSets[:bestIdx], tokenSets[bestIdx+1:]...) maxRedundancy = append(maxRedundancy[:bestIdx], maxRedundancy[bestIdx+1:]...) // Fold the freshly selected result into every remaining candidate's cache for i := range candidates { maxRedundancy[i] = math.Max(maxRedundancy[i], t.jaccard(tokenSets[i], chosenTokens)) } } // Compute average redundancy among selected results, reusing the cached token // sets instead of re-tokenizing every pair avgRed := 0.0 if len(selectedTokenSets) > 1 { pairs := 0 for i := 0; i < len(selectedTokenSets); i++ { for j := i + 1; j < len(selectedTokenSets); j++ { avgRed += t.jaccard(selectedTokenSets[i], selectedTokenSets[j]) pairs++ } } if pairs > 0 { avgRed /= float64(pairs) } } logger.Infof(ctx, "[Tool][KnowledgeSearch] MMR completed: selected=%d, avg_redundancy=%.4f", len(selected), avgRed) return selected } // tokenizeSimple tokenizes text into a set of words (simple whitespace-based) func (t *KnowledgeSearchTool) tokenizeSimple(text string) map[string]struct{} { return searchutil.TokenizeSimple(text) } // extractSnippetForQueries tries to produce a short contextual snippet around // the first occurrence of any token extracted from the provided queries. // When no token matches (common for fully paraphrased semantic queries) it // falls back to the leading 160 runes of content so callers always get // something to scan. The snippet is single-lined and bounded in length to // keep the rendered XML compact. func extractSnippetForQueries(content string, queries []string) string { content = strings.TrimSpace(content) if content == "" { return "" } tokens := searchQueryTokens(queries) lowered := strings.ToLower(content) earliest := -1 earliestEnd := -1 for _, tok := range tokens { idx := strings.Index(lowered, tok) if idx < 0 { continue } end := idx + len(tok) if earliest < 0 || idx < earliest { earliest = idx earliestEnd = end } } if earliest < 0 { runes := []rune(content) if len(runes) > snippetContextRunes*2 { return strings.TrimSpace(string(runes[:snippetContextRunes*2])) + " ..." } return content } matchStr := content[earliest:earliestEnd] before := content[:earliest] after := content[earliestEnd:] beforeRunes := []rune(before) if len(beforeRunes) > snippetContextRunes { beforeRunes = beforeRunes[len(beforeRunes)-snippetContextRunes:] } afterRunes := []rune(after) if len(afterRunes) > snippetContextRunes { afterRunes = afterRunes[:snippetContextRunes] } snippet := string(beforeRunes) + matchStr + string(afterRunes) snippet = strings.ReplaceAll(snippet, "\n", " ") for strings.Contains(snippet, " ") { snippet = strings.ReplaceAll(snippet, " ", " ") } return "... " + strings.TrimSpace(snippet) + " ..." } // jaccard calculates Jaccard similarity between two token sets func (t *KnowledgeSearchTool) jaccard(a, b map[string]struct{}) float64 { return searchutil.Jaccard(a, b) }