d546406a15
- Updated the formatKnowledgeBaseList function to return knowledge base information as XML, improving data structure for integration. - Modified the formatOutput function in GrepChunksTool to format search results as XML, enhancing readability and consistency. - Adjusted the formatOutput method in KnowledgeSearchTool to output search results in XML format, providing a structured response. - Refactored the buildOutput function in ListKnowledgeChunksTool to generate XML output, improving the clarity of chunk information. These changes standardize output formats across multiple tools, facilitating better data handling and integration.
749 lines
23 KiB
Go
749 lines
23 KiB
Go
package tools
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"math"
|
||
"sort"
|
||
"strings"
|
||
|
||
"github.com/Tencent/WeKnora/internal/logger"
|
||
"github.com/Tencent/WeKnora/internal/searchutil"
|
||
"github.com/Tencent/WeKnora/internal/types"
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
var grepChunksTool = BaseTool{
|
||
name: ToolGrepChunks,
|
||
description: `Unix-style text pattern matching tool for knowledge base chunks.
|
||
|
||
Searches for text patterns in chunk content using strict literal text matching (fixed-string search). This tool performs exact keyword lookup, not semantic search.
|
||
|
||
## Core Function
|
||
Performs exact, literal text pattern matching. Accepts multiple patterns and returns chunks matching any of them (OR logic).
|
||
|
||
## CRITICAL – Keyword Extraction Rules
|
||
This tool MUST receive **short, high-value keywords** only.
|
||
**Do NOT use long phrases, sentences, or multi-word expressions.**
|
||
|
||
Provide only the **minimal core entities** extracted from user query, such as:
|
||
- Proper nouns
|
||
- Key concepts
|
||
- Domain terms
|
||
- Distinct entities that define the query
|
||
|
||
### Requirements
|
||
- Keywords should be **1–3 words maximum**
|
||
- Focus exclusively on **core entities**, not descriptions
|
||
- Break complex input into individual, essential keywords
|
||
- Avoid phrases, explanations, or anything that reduces match probability
|
||
- Preserve precision details embedded in the query (e.g., version numbers, build IDs) when they materially define the entity being matched.
|
||
|
||
Long phrases dramatically reduce recall because chunks rarely contain identical wording.
|
||
Only short, atomic keywords ensure accurate matching and avoid unrelated retrieval.
|
||
|
||
|
||
## Usage
|
||
grep_chunks scans enabled chunks across the specified knowledge bases and returns those containing any provided keyword. Matching is case-insensitive, with chunk indices and local context included.
|
||
|
||
## When to Use
|
||
- Extracting core entities from user input
|
||
- Exact keyword presence checks
|
||
- Fast preliminary filtering before semantic search
|
||
- Situations requiring deterministic text search
|
||
`,
|
||
schema: json.RawMessage(`{
|
||
"type": "object",
|
||
"properties": {
|
||
"patterns": {
|
||
"type": "array",
|
||
"description": "REQUIRED: Text patterns to search for. Can be a single pattern or multiple patterns. Treated as literal text (fixed string matching). Results match any of the patterns (OR logic).",
|
||
"items": {
|
||
"type": "string"
|
||
},
|
||
"minItems": 1
|
||
},
|
||
"knowledge_base_ids": {
|
||
"type": "array",
|
||
"description": "Filter by knowledge base IDs. If empty, searches all allowed KBs.",
|
||
"items": {
|
||
"type": "string"
|
||
}
|
||
},
|
||
"max_results": {
|
||
"type": "integer",
|
||
"description": "Maximum number of matching chunks to return (default: 50, max: 200)",
|
||
"default": 50,
|
||
"minimum": 1,
|
||
"maximum": 200
|
||
}
|
||
},
|
||
"required": ["patterns"]
|
||
}`),
|
||
}
|
||
|
||
// GrepChunksInput defines the input parameters for grep chunks tool
|
||
type GrepChunksInput struct {
|
||
Patterns []string `json:"patterns" `
|
||
KnowledgeBaseIDs []string `json:"knowledge_base_ids,omitempty"`
|
||
MaxResults int `json:"max_results,omitempty"`
|
||
}
|
||
|
||
// GrepChunksTool performs text pattern matching in knowledge base chunks
|
||
// Similar to grep command in Unix-like systems, but operates on knowledge base content
|
||
type GrepChunksTool struct {
|
||
BaseTool
|
||
db *gorm.DB
|
||
searchTargets types.SearchTargets // Pre-computed unified search targets with KB-tenant mapping
|
||
}
|
||
|
||
// NewGrepChunksTool creates a new grep chunks tool
|
||
func NewGrepChunksTool(db *gorm.DB, searchTargets types.SearchTargets) *GrepChunksTool {
|
||
return &GrepChunksTool{
|
||
BaseTool: grepChunksTool,
|
||
db: db,
|
||
searchTargets: searchTargets,
|
||
}
|
||
}
|
||
|
||
// Execute executes the grep chunks tool
|
||
func (t *GrepChunksTool) Execute(ctx context.Context, args json.RawMessage) (*types.ToolResult, error) {
|
||
logger.Infof(ctx, "[Tool][GrepChunks] Execute started")
|
||
|
||
// Parse args from json.RawMessage
|
||
var input GrepChunksInput
|
||
if err := json.Unmarshal(args, &input); err != nil {
|
||
logger.Errorf(ctx, "[Tool][GrepChunks] Failed to parse args: %v", err)
|
||
return &types.ToolResult{
|
||
Success: false,
|
||
Error: fmt.Sprintf("Failed to parse args: %v", err),
|
||
}, err
|
||
}
|
||
|
||
// Parse pattern parameter (required) - support multiple patterns
|
||
patterns := input.Patterns
|
||
|
||
// Validate patterns
|
||
if len(patterns) == 0 {
|
||
logger.Errorf(ctx, "[Tool][GrepChunks] Missing or invalid patterns parameter")
|
||
return &types.ToolResult{
|
||
Success: false,
|
||
Error: "pattern parameter is required and must contain at least one non-empty pattern",
|
||
}, fmt.Errorf("missing pattern parameter")
|
||
}
|
||
|
||
// Use default values for all options
|
||
countOnly := false // default: show results
|
||
|
||
maxResults := 50
|
||
if input.MaxResults > 0 {
|
||
maxResults = input.MaxResults
|
||
if maxResults < 1 {
|
||
maxResults = 1
|
||
} else if maxResults > 200 {
|
||
maxResults = 200
|
||
}
|
||
}
|
||
|
||
// Get allowed KBs from searchTargets
|
||
allowedKBIDs := t.searchTargets.GetAllKnowledgeBaseIDs()
|
||
kbTenantMap := t.searchTargets.GetKBTenantMap()
|
||
|
||
// Collect all specific knowledge IDs from searchTargets
|
||
var allowedKnowledgeIDs []string
|
||
for _, target := range t.searchTargets {
|
||
if target.Type == types.SearchTargetTypeKnowledge && len(target.KnowledgeIDs) > 0 {
|
||
allowedKnowledgeIDs = append(allowedKnowledgeIDs, target.KnowledgeIDs...)
|
||
}
|
||
}
|
||
|
||
// Parse knowledge_base_ids filter from input
|
||
kbIDs := input.KnowledgeBaseIDs
|
||
if len(kbIDs) == 0 {
|
||
kbIDs = allowedKBIDs
|
||
} else {
|
||
// Validate input KBs against allowed KBs
|
||
validKBs := make([]string, 0)
|
||
for _, kbID := range kbIDs {
|
||
if t.searchTargets.ContainsKB(kbID) {
|
||
validKBs = append(validKBs, kbID)
|
||
}
|
||
}
|
||
kbIDs = validKBs
|
||
}
|
||
|
||
logger.Infof(ctx, "[Tool][GrepChunks] Patterns: %v, MaxResults: %d, KBs: %v, KnowledgeIDs: %v, KBTenantMap: %v",
|
||
patterns, maxResults, kbIDs, allowedKnowledgeIDs, kbTenantMap)
|
||
|
||
// Build and execute query with tenant info
|
||
results, totalCount, err := t.searchChunks(ctx, patterns, kbIDs, allowedKnowledgeIDs, kbTenantMap)
|
||
if err != nil {
|
||
logger.Errorf(ctx, "[Tool][GrepChunks] Search failed: %v", err)
|
||
return &types.ToolResult{
|
||
Success: false,
|
||
Error: fmt.Sprintf("Search failed: %v", err),
|
||
}, err
|
||
}
|
||
|
||
logger.Infof(ctx, "[Tool][GrepChunks] Found %d matching chunks", len(results))
|
||
|
||
// Apply deduplication to remove duplicate or near-duplicate chunks
|
||
deduplicatedResults := t.deduplicateChunks(ctx, results)
|
||
logger.Infof(ctx, "[Tool][GrepChunks] After deduplication: %d chunks (from %d)",
|
||
len(deduplicatedResults), len(results))
|
||
|
||
// Calculate match scores for sorting (based on match count and position)
|
||
scoredResults := t.scoreChunks(ctx, deduplicatedResults, patterns)
|
||
|
||
// Apply MMR to reduce redundancy if we have many results
|
||
finalResults := scoredResults
|
||
if len(scoredResults) > 10 {
|
||
// Use MMR when we have more than 10 results
|
||
mmrK := len(scoredResults)
|
||
if maxResults > 0 && mmrK > maxResults {
|
||
mmrK = maxResults
|
||
}
|
||
logger.Debugf(
|
||
ctx,
|
||
"[Tool][GrepChunks] Applying MMR: k=%d, lambda=0.7, input=%d results",
|
||
mmrK,
|
||
len(scoredResults),
|
||
)
|
||
mmrResults := t.applyMMR(ctx, scoredResults, patterns, mmrK, 0.7)
|
||
if len(mmrResults) > 0 {
|
||
finalResults = mmrResults
|
||
logger.Infof(ctx, "[Tool][GrepChunks] MMR completed: %d results selected", len(finalResults))
|
||
}
|
||
}
|
||
|
||
// Sort by match score (descending), then by chunk index
|
||
sort.Slice(finalResults, func(i, j int) bool {
|
||
if finalResults[i].MatchedPatterns != finalResults[j].MatchedPatterns {
|
||
return finalResults[i].MatchedPatterns > finalResults[j].MatchedPatterns
|
||
}
|
||
if finalResults[i].MatchScore != finalResults[j].MatchScore {
|
||
return finalResults[i].MatchScore > finalResults[j].MatchScore
|
||
}
|
||
return finalResults[i].ChunkIndex < finalResults[j].ChunkIndex
|
||
})
|
||
|
||
aggregatedResults := t.aggregateByKnowledge(finalResults, patterns)
|
||
|
||
totalKnowledge := len(aggregatedResults)
|
||
|
||
if len(aggregatedResults) > 20 {
|
||
aggregatedResults = aggregatedResults[:20]
|
||
}
|
||
|
||
logger.Infof(ctx, "[Tool][GrepChunks] Aggregated results: %d", len(aggregatedResults))
|
||
|
||
// Format output
|
||
output := t.formatOutput(ctx, aggregatedResults, totalCount, patterns, countOnly)
|
||
|
||
return &types.ToolResult{
|
||
Success: true,
|
||
Output: output,
|
||
Data: map[string]interface{}{
|
||
"patterns": patterns,
|
||
"knowledge_results": aggregatedResults,
|
||
"result_count": len(aggregatedResults),
|
||
"total_matches": totalKnowledge,
|
||
"knowledge_base_ids": kbIDs,
|
||
"max_results": maxResults,
|
||
"display_type": "grep_results",
|
||
},
|
||
}, nil
|
||
}
|
||
|
||
type chunkWithTitle struct {
|
||
types.Chunk
|
||
KnowledgeTitle string `json:"knowledge_title" gorm:"column:knowledge_title"`
|
||
MatchScore float64 `json:"match_score" gorm:"column:match_score"` // Score based on match count and position
|
||
MatchedPatterns int `json:"matched_patterns"` // Number of unique patterns matched
|
||
TotalChunkCount int `json:"total_chunk_count" gorm:"column:total_chunk_count"`
|
||
}
|
||
|
||
// searchChunks performs the database search with pattern matching
|
||
// kbTenantMap provides KB-to-tenant mapping for cross-tenant queries
|
||
func (t *GrepChunksTool) searchChunks(
|
||
ctx context.Context,
|
||
patterns []string,
|
||
kbIDs []string,
|
||
knowledgeIDs []string,
|
||
kbTenantMap map[string]uint64,
|
||
) ([]chunkWithTitle, int64, error) {
|
||
if len(kbIDs) == 0 && len(knowledgeIDs) == 0 {
|
||
logger.Warnf(ctx, "[Tool][GrepChunks] No kbIDs or knowledgeIDs specified, returning empty results")
|
||
return nil, 0, nil
|
||
}
|
||
|
||
// PostgreSQL uses ILIKE for case-insensitive matching;
|
||
// MySQL and SQLite LIKE is already case-insensitive under default collation.
|
||
likeOp := "LIKE"
|
||
if t.db.Dialector.Name() == "postgres" {
|
||
likeOp = "ILIKE"
|
||
}
|
||
|
||
query := t.db.Debug().WithContext(ctx).Table("chunks").
|
||
Select("chunks.id, chunks.content, chunks.chunk_index, chunks.knowledge_id, "+
|
||
"chunks.knowledge_base_id, chunks.chunk_type, chunks.created_at, "+
|
||
"knowledges.title as knowledge_title").
|
||
Joins("JOIN knowledges ON chunks.knowledge_id = knowledges.id").
|
||
Where("chunks.is_enabled = ?", true).
|
||
Where("chunks.deleted_at IS NULL").
|
||
Where("knowledges.deleted_at IS NULL")
|
||
|
||
if len(knowledgeIDs) > 0 {
|
||
query = query.Where("chunks.knowledge_id IN ?", knowledgeIDs)
|
||
logger.Infof(ctx, "[Tool][GrepChunks] Filtering by %d specific knowledge IDs", len(knowledgeIDs))
|
||
} else if len(kbIDs) > 0 {
|
||
var conditions []string
|
||
var args []interface{}
|
||
for _, kbID := range kbIDs {
|
||
tenantID := kbTenantMap[kbID]
|
||
if tenantID > 0 {
|
||
conditions = append(conditions, "(chunks.knowledge_base_id = ? AND chunks.tenant_id = ?)")
|
||
args = append(args, kbID, tenantID)
|
||
}
|
||
}
|
||
if len(conditions) > 0 {
|
||
query = query.Where("("+strings.Join(conditions, " OR ")+")", args...)
|
||
} else {
|
||
logger.Warnf(ctx, "[Tool][GrepChunks] No valid KB-tenant pairs found")
|
||
return nil, 0, nil
|
||
}
|
||
}
|
||
|
||
if len(patterns) == 1 {
|
||
query = query.Where("chunks.content "+likeOp+" ?", "%"+patterns[0]+"%")
|
||
} else {
|
||
var conditions []string
|
||
var args []interface{}
|
||
for _, pattern := range patterns {
|
||
conditions = append(conditions, "chunks.content "+likeOp+" ?")
|
||
args = append(args, "%"+pattern+"%")
|
||
}
|
||
query = query.Where("("+strings.Join(conditions, " OR ")+")", args...)
|
||
}
|
||
|
||
const maxFetchLimit = 500
|
||
|
||
var results []chunkWithTitle
|
||
if err := query.Order("chunks.created_at DESC").Limit(maxFetchLimit).Find(&results).Error; err != nil {
|
||
logger.Errorf(ctx, "[Tool][GrepChunks] Failed to fetch results: %v", err)
|
||
return nil, 0, err
|
||
}
|
||
|
||
if len(results) > 0 {
|
||
knowledgeIDSet := make(map[string]struct{})
|
||
for _, r := range results {
|
||
if r.KnowledgeID != "" {
|
||
knowledgeIDSet[r.KnowledgeID] = struct{}{}
|
||
}
|
||
}
|
||
uniqueKnowledgeIDs := make([]string, 0, len(knowledgeIDSet))
|
||
for kid := range knowledgeIDSet {
|
||
uniqueKnowledgeIDs = append(uniqueKnowledgeIDs, kid)
|
||
}
|
||
|
||
type countRow struct {
|
||
KnowledgeID string `gorm:"column:knowledge_id"`
|
||
Count int `gorm:"column:cnt"`
|
||
}
|
||
var counts []countRow
|
||
if err := t.db.WithContext(ctx).Table("chunks").
|
||
Select("knowledge_id, COUNT(*) AS cnt").
|
||
Where("knowledge_id IN ?", uniqueKnowledgeIDs).
|
||
Where("is_enabled = ?", true).
|
||
Where("deleted_at IS NULL").
|
||
Group("knowledge_id").
|
||
Find(&counts).Error; err != nil {
|
||
logger.Warnf(ctx, "[Tool][GrepChunks] Failed to fetch chunk counts, skipping: %v", err)
|
||
} else {
|
||
countMap := make(map[string]int, len(counts))
|
||
for _, c := range counts {
|
||
countMap[c.KnowledgeID] = c.Count
|
||
}
|
||
for i := range results {
|
||
results[i].TotalChunkCount = countMap[results[i].KnowledgeID]
|
||
}
|
||
}
|
||
}
|
||
|
||
return results, int64(len(results)), nil
|
||
}
|
||
|
||
// formatOutput formats the search results as XML
|
||
func (t *GrepChunksTool) formatOutput(
|
||
ctx context.Context,
|
||
results []knowledgeAggregation,
|
||
totalCount int64,
|
||
patterns []string,
|
||
countOnly bool,
|
||
) string {
|
||
var b strings.Builder
|
||
|
||
if countOnly {
|
||
b.WriteString(fmt.Sprintf("<grep_results count=\"%d\" />\n", totalCount))
|
||
return b.String()
|
||
}
|
||
|
||
b.WriteString(fmt.Sprintf("<grep_results match_count=\"%d\">\n", len(results)))
|
||
for _, p := range patterns {
|
||
b.WriteString(fmt.Sprintf("<pattern>%s</pattern>\n", p))
|
||
}
|
||
|
||
if len(results) == 0 {
|
||
b.WriteString("</grep_results>")
|
||
return b.String()
|
||
}
|
||
|
||
for _, result := range results {
|
||
b.WriteString(fmt.Sprintf("<match knowledge_id=\"%s\" title=\"%s\" chunk_hits=\"%d\" chunk_total=\"%d\">\n",
|
||
result.KnowledgeID, result.KnowledgeTitle, result.ChunkHitCount, result.TotalChunkCount))
|
||
for _, pattern := range patterns {
|
||
count := result.PatternCounts[pattern]
|
||
if count > 0 {
|
||
b.WriteString(fmt.Sprintf("<pattern_hit pattern=\"%s\" count=\"%d\" />\n", pattern, count))
|
||
}
|
||
}
|
||
b.WriteString("</match>\n")
|
||
}
|
||
|
||
b.WriteString("</grep_results>")
|
||
return b.String()
|
||
}
|
||
|
||
type knowledgeAggregation struct {
|
||
KnowledgeID string `json:"knowledge_id"`
|
||
KnowledgeBaseID string `json:"knowledge_base_id"`
|
||
KnowledgeTitle string `json:"knowledge_title"`
|
||
ChunkHitCount int `json:"chunk_hit_count"`
|
||
TotalChunkCount int `json:"total_chunk_count"`
|
||
PatternCounts map[string]int `json:"pattern_counts"`
|
||
TotalPatternHits int `json:"total_pattern_hits"`
|
||
DistinctPatterns int `json:"distinct_patterns"`
|
||
}
|
||
|
||
func (t *GrepChunksTool) aggregateByKnowledge(results []chunkWithTitle, patterns []string) []knowledgeAggregation {
|
||
if len(results) == 0 {
|
||
return nil
|
||
}
|
||
|
||
patternKeys := make([]string, 0, len(patterns))
|
||
for _, p := range patterns {
|
||
if strings.TrimSpace(p) == "" {
|
||
continue
|
||
}
|
||
patternKeys = append(patternKeys, p)
|
||
}
|
||
|
||
aggregated := make(map[string]*knowledgeAggregation)
|
||
for _, chunk := range results {
|
||
knowledgeID := chunk.KnowledgeID
|
||
if knowledgeID == "" {
|
||
knowledgeID = fmt.Sprintf("chunk-%s", chunk.ID)
|
||
}
|
||
|
||
if _, ok := aggregated[knowledgeID]; !ok {
|
||
title := chunk.KnowledgeTitle
|
||
if strings.TrimSpace(title) == "" {
|
||
title = "Untitled"
|
||
}
|
||
aggregated[knowledgeID] = &knowledgeAggregation{
|
||
KnowledgeID: knowledgeID,
|
||
KnowledgeBaseID: chunk.KnowledgeBaseID,
|
||
KnowledgeTitle: title,
|
||
TotalChunkCount: chunk.TotalChunkCount,
|
||
PatternCounts: make(map[string]int, len(patternKeys)),
|
||
}
|
||
for _, pKey := range patternKeys {
|
||
aggregated[knowledgeID].PatternCounts[pKey] = 0
|
||
}
|
||
}
|
||
|
||
entry := aggregated[knowledgeID]
|
||
entry.ChunkHitCount++
|
||
|
||
patternOccurrences := t.countPatternOccurrences(chunk.Content, patternKeys)
|
||
for _, p := range patternKeys {
|
||
count := patternOccurrences[p]
|
||
if count == 0 {
|
||
continue
|
||
}
|
||
entry.PatternCounts[p] += count
|
||
entry.TotalPatternHits += count
|
||
}
|
||
}
|
||
|
||
resultSlice := make([]knowledgeAggregation, 0, len(aggregated))
|
||
for _, entry := range aggregated {
|
||
distinct := 0
|
||
for _, count := range entry.PatternCounts {
|
||
if count > 0 {
|
||
distinct++
|
||
}
|
||
}
|
||
entry.DistinctPatterns = distinct
|
||
resultSlice = append(resultSlice, *entry)
|
||
}
|
||
|
||
sort.Slice(resultSlice, func(i, j int) bool {
|
||
if resultSlice[i].DistinctPatterns != resultSlice[j].DistinctPatterns {
|
||
return resultSlice[i].DistinctPatterns > resultSlice[j].DistinctPatterns
|
||
}
|
||
if resultSlice[i].TotalPatternHits != resultSlice[j].TotalPatternHits {
|
||
return resultSlice[i].TotalPatternHits > resultSlice[j].TotalPatternHits
|
||
}
|
||
if resultSlice[i].ChunkHitCount != resultSlice[j].ChunkHitCount {
|
||
return resultSlice[i].ChunkHitCount > resultSlice[j].ChunkHitCount
|
||
}
|
||
return resultSlice[i].KnowledgeTitle < resultSlice[j].KnowledgeTitle
|
||
})
|
||
return resultSlice
|
||
}
|
||
|
||
func (t *GrepChunksTool) countPatternOccurrences(content string, patterns []string) map[string]int {
|
||
counts := make(map[string]int, len(patterns))
|
||
if content == "" || len(patterns) == 0 {
|
||
return counts
|
||
}
|
||
|
||
contentLower := strings.ToLower(content)
|
||
for _, pattern := range patterns {
|
||
p := strings.ToLower(pattern)
|
||
if strings.TrimSpace(p) == "" {
|
||
continue
|
||
}
|
||
counts[pattern] = countOccurrences(contentLower, p)
|
||
}
|
||
return counts
|
||
}
|
||
|
||
func countOccurrences(text string, pattern string) int {
|
||
if pattern == "" {
|
||
return 0
|
||
}
|
||
count := 0
|
||
index := 0
|
||
for index < len(text) {
|
||
pos := strings.Index(text[index:], pattern)
|
||
if pos == -1 {
|
||
break
|
||
}
|
||
count++
|
||
index += pos + len(pattern)
|
||
}
|
||
return count
|
||
}
|
||
|
||
// deduplicateChunks removes duplicate or near-duplicate chunks using content signature
|
||
func (t *GrepChunksTool) deduplicateChunks(ctx context.Context, results []chunkWithTitle) []chunkWithTitle {
|
||
seen := make(map[string]bool)
|
||
contentSig := make(map[string]bool)
|
||
uniqueResults := make([]chunkWithTitle, 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, keep the first one
|
||
seenByID := make(map[string]bool)
|
||
deduplicated := make([]chunkWithTitle, 0)
|
||
for _, r := range uniqueResults {
|
||
if !seenByID[r.ID] {
|
||
seenByID[r.ID] = true
|
||
deduplicated = append(deduplicated, r)
|
||
}
|
||
}
|
||
|
||
return deduplicated
|
||
}
|
||
|
||
// buildContentSignature creates a normalized signature for content to detect near-duplicates
|
||
func (t *GrepChunksTool) buildContentSignature(content string) string {
|
||
return searchutil.BuildContentSignature(content)
|
||
}
|
||
|
||
// scoreChunks calculates match scores for chunks based on pattern matches
|
||
func (t *GrepChunksTool) scoreChunks(
|
||
ctx context.Context,
|
||
results []chunkWithTitle,
|
||
patterns []string,
|
||
) []chunkWithTitle {
|
||
scored := make([]chunkWithTitle, len(results))
|
||
for i := range results {
|
||
scored[i] = results[i]
|
||
score, patternCount := t.calculateMatchScore(results[i].Content, patterns)
|
||
scored[i].MatchScore = score
|
||
scored[i].MatchedPatterns = patternCount
|
||
}
|
||
return scored
|
||
}
|
||
|
||
// calculateMatchScore calculates a score based on how many patterns match and their positions
|
||
func (t *GrepChunksTool) calculateMatchScore(content string, patterns []string) (float64, int) {
|
||
if content == "" || len(patterns) == 0 {
|
||
return 0.0, 0
|
||
}
|
||
|
||
contentLower := strings.ToLower(content)
|
||
matchCount := 0
|
||
earliestPos := len(content)
|
||
|
||
// Count how many patterns match and find earliest position
|
||
for _, pattern := range patterns {
|
||
patternLower := strings.ToLower(pattern)
|
||
if strings.Contains(contentLower, patternLower) {
|
||
matchCount++
|
||
// Find position of first match
|
||
pos := strings.Index(contentLower, patternLower)
|
||
if pos >= 0 && pos < earliestPos {
|
||
earliestPos = pos
|
||
}
|
||
}
|
||
}
|
||
|
||
// Score: higher for more matches, slightly higher for earlier positions
|
||
// Base score: match ratio (0.0 to 1.0)
|
||
baseScore := float64(matchCount) / float64(len(patterns))
|
||
|
||
// Position bonus: earlier matches get slight boost (max 0.1)
|
||
positionBonus := 0.0
|
||
if earliestPos < len(content) {
|
||
// Normalize position to [0, 1] and apply small bonus
|
||
positionRatio := 1.0 - float64(earliestPos)/float64(len(content))
|
||
positionBonus = positionRatio * 0.1
|
||
}
|
||
|
||
return math.Min(baseScore+positionBonus, 1.0), matchCount
|
||
}
|
||
|
||
// applyMMR applies Maximal Marginal Relevance algorithm to reduce redundancy
|
||
func (t *GrepChunksTool) applyMMR(
|
||
ctx context.Context,
|
||
results []chunkWithTitle,
|
||
patterns []string,
|
||
k int,
|
||
lambda float64,
|
||
) []chunkWithTitle {
|
||
if k <= 0 || len(results) == 0 {
|
||
return nil
|
||
}
|
||
|
||
logger.Debugf(ctx, "[Tool][GrepChunks] Applying MMR: lambda=%.2f, k=%d, candidates=%d",
|
||
lambda, k, len(results))
|
||
|
||
selected := make([]chunkWithTitle, 0, k)
|
||
selectedTokenSets := make([]map[string]struct{}, 0, k) // cache of token sets
|
||
|
||
candidates := make([]chunkWithTitle, 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(r.Content)
|
||
}
|
||
|
||
// MMR selection loop
|
||
for len(selected) < k && len(candidates) > 0 {
|
||
bestIdx := 0
|
||
bestScore := -1.0
|
||
|
||
for i, r := range candidates {
|
||
relevance := r.MatchScore
|
||
redundancy := 0.0
|
||
|
||
// Calculate maximum redundancy with already selected results
|
||
for _, selectedTS := range selectedTokenSets {
|
||
redundancy = math.Max(redundancy, t.jaccard(tokenSets[i], selectedTS))
|
||
}
|
||
|
||
// MMR score: balance relevance and diversity
|
||
mmr := lambda*relevance - (1.0-lambda)*redundancy
|
||
if mmr > bestScore {
|
||
bestScore = mmr
|
||
bestIdx = i
|
||
}
|
||
}
|
||
|
||
// Add best candidate to selected and remove from candidates
|
||
selected = append(selected, candidates[bestIdx])
|
||
selectedTokenSets = append(selectedTokenSets, tokenSets[bestIdx])
|
||
|
||
// Remove corresponding token set. Use swap deletion
|
||
last := len(candidates) - 1
|
||
candidates[bestIdx] = candidates[last]
|
||
tokenSets[bestIdx] = tokenSets[last]
|
||
candidates = candidates[:last]
|
||
tokenSets = tokenSets[:last]
|
||
}
|
||
|
||
// Compute average redundancy among selected results
|
||
avgRed := 0.0
|
||
if len(selected) > 1 {
|
||
pairs := 0
|
||
for i := 0; i < len(selected); i++ {
|
||
for j := i + 1; j < len(selected); j++ {
|
||
avgRed += t.jaccard(selectedTokenSets[i], selectedTokenSets[j]) // read token from cache
|
||
pairs++
|
||
}
|
||
}
|
||
if pairs > 0 {
|
||
avgRed /= float64(pairs)
|
||
}
|
||
}
|
||
|
||
logger.Debugf(ctx, "[Tool][GrepChunks] 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 *GrepChunksTool) tokenizeSimple(text string) map[string]struct{} {
|
||
return searchutil.TokenizeSimple(text)
|
||
}
|
||
|
||
// jaccard calculates Jaccard similarity between two token sets
|
||
func (t *GrepChunksTool) jaccard(a, b map[string]struct{}) float64 {
|
||
return searchutil.Jaccard(a, b)
|
||
}
|