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)
}