xinmi/internal/agent/tools/knowledge_search.go

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(
			"<document knowledge_id=\"%s\" knowledge_base_id=\"%s\" title=\"%s\">\n",
			xmlEscape(result.KnowledgeID),
			xmlEscape(result.KnowledgeBaseID),
			xmlEscape(result.KnowledgeTitle),
		))
		documents.WriteString(fmt.Sprintf("<metadata>%s</metadata>\n", xmlEscape(result.KnowledgeCustomMetadata)))
		documents.WriteString("</document>\n")
	}
	if !hasMetadata {
		return
	}
	ob.WriteString("<documents>\n")
	ob.WriteString(documents.String())
	ob.WriteString("</documents>\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 (`<search_results>`, per-entry element, `<query>`) so that
	// agents and downstream consumers see a single consistent shape across
	// all retrieval tools.
	var ob strings.Builder
	ob.WriteString(fmt.Sprintf("<search_results count=\"%d\">\n", len(results)))
	for _, q := range queries {
		ob.WriteString(fmt.Sprintf("<query>%s</query>\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(
					"<faq rank=\"%d\" faq_id=\"%s\" index=\"%d\" knowledge_base_id=\"%s\" knowledge_title=\"%s\" score=\"%.3f\" source_query=\"%s\" already_seen=\"true\">\n",
					i+1,
					xmlEscape(result.ID),
					result.ChunkIndex,
					xmlEscape(result.KnowledgeBaseID),
					xmlEscape(result.KnowledgeTitle),
					result.Score,
					xmlEscape(result.SourceQuery),
				))
			} else {
				ob.WriteString(fmt.Sprintf(
					"<chunk rank=\"%d\" chunk_id=\"%s\" chunk_index=\"%d\" knowledge_id=\"%s\" knowledge_base_id=\"%s\" knowledge_title=\"%s\" score=\"%.3f\" source_query=\"%s\" already_seen=\"true\">\n",
					i+1,
					xmlEscape(result.ID),
					result.ChunkIndex,
					xmlEscape(result.KnowledgeID),
					xmlEscape(result.KnowledgeBaseID),
					xmlEscape(result.KnowledgeTitle),
					result.Score,
					xmlEscape(result.SourceQuery),
				))
			}
			ob.WriteString("<note>(content omitted, already returned in a previous knowledge_search call this session)</note>\n")
			if isFAQ {
				ob.WriteString("</faq>\n")
			} else {
				ob.WriteString("</chunk>\n")
			}
		} else {
			if isFAQ {
				ob.WriteString(fmt.Sprintf(
					"<faq rank=\"%d\" faq_id=\"%s\" index=\"%d\" knowledge_base_id=\"%s\" knowledge_title=\"%s\" score=\"%.3f\" source_query=\"%s\">\n",
					i+1,
					xmlEscape(result.ID),
					result.ChunkIndex,
					xmlEscape(result.KnowledgeBaseID),
					xmlEscape(result.KnowledgeTitle),
					result.Score,
					xmlEscape(result.SourceQuery),
				))
			} else {
				ob.WriteString(fmt.Sprintf(
					"<chunk rank=\"%d\" chunk_id=\"%s\" chunk_index=\"%d\" knowledge_id=\"%s\" knowledge_base_id=\"%s\" knowledge_title=\"%s\" score=\"%.3f\" source_query=\"%s\">\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("<match_snippet>%s</match_snippet>\n", xmlEscape(snippet)))
			}
			ob.WriteString(fmt.Sprintf("<content>%s</content>\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("</faq>\n")
			} else {
				ob.WriteString("</chunk>\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("<retrieval_statistics>\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("<document_stat knowledge_id=\"%s\" title=\"%s\" total_chunks=\"%d\" retrieved=\"%d\" remaining=\"%d\" coverage=\"%.1f%%\" />\n",
				xmlEscape(knowledgeID), xmlEscape(title), totalChunks, retrievedCount, remaining, percentage))
		}
	}
	ob.WriteString("</retrieval_statistics>\n")
	ob.WriteString("</search_results>")

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