feat(search): implement tree-based deduplication and strict modes

doITmagic · doITmagic · commit 0d4dd84a4063 · 2026-03-07T19:34:26.000+02:00
diff --git a/internal/service/tools/smart_search.go b/internal/service/tools/smart_search.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"path/filepath"
 	"sort"
 	"strings"
 	"sync"
@@ -42,7 +43,10 @@ func (t *SmartSearchTool) Description() string {
 		"high-confidence matches return full source code, exploratory results return compact summaries. " +
 		"No need to choose a search mode. Provide 'file_path' for faster workspace detection, or omit it for Auto-Discovery. " +
 		"Set 'include_full_content' to true to force full source code in all results, overriding compact mode. " +
-		"Set 'include_docs' to true to also search project documentation (README, guides, Markdown files) alongside code."
+		"Set 'include_docs' to true to also search project documentation (README, guides, Markdown files) alongside code. " +
+		"Use 'mode'=\"strict_code\" when you ONLY want to see implementation logic exactly (Go, Python, etc) and strictly ignore documentation. " +
+		"Use 'mode'=\"strict_docs\" when searching for architectural plans or summaries. " +
+		"Use 'mode'=\"all\" or omit for broad scans."
 }
 
 type SmartSearchInput struct {
@@ -51,6 +55,7 @@ type SmartSearchInput struct {
 	Limit              int    `json:"limit,omitempty"`
 	IncludeFullContent bool   `json:"include_full_content,omitempty"`
 	IncludeDocs        bool   `json:"include_docs,omitempty"`
+	Mode               string `json:"mode,omitempty"`
 }
 
 // highConfidenceThreshold: if top result score exceeds this, return full content.
@@ -172,6 +177,24 @@ func (t *SmartSearchTool) Execute(ctx context.Context, input SmartSearchInput) (
 	// Merge and deduplicate results from both strategies
 	merged := t.mergeResults(semanticRes, hybridRes, limit)
 
+	// Apply mode filtering
+	var filtered []mergedResult
+	for _, m := range merged {
+		// Strict code mode: ignore completely any markdown or documentation type
+		if input.Mode == "strict_code" && (m.symbolType == "documentation" || m.symbolType == "markdown" || m.symbolType == "code_block" || strings.HasSuffix(strings.ToLower(m.filePath), ".md") || strings.HasSuffix(strings.ToLower(m.filePath), ".html")) {
+			continue
+		}
+		// Strict docs mode: ignore anything that isn't documentation
+		if input.Mode == "strict_docs" && !(m.symbolType == "documentation" || m.symbolType == "markdown" || m.symbolType == "code_block" || strings.HasSuffix(strings.ToLower(m.filePath), ".md") || strings.HasSuffix(strings.ToLower(m.filePath), ".html")) {
+			continue
+		}
+		filtered = append(filtered, m)
+	}
+	merged = filtered
+
+	// Apply tree-based grouping for documentation chunks
+	merged = t.groupDocsByTree(merged)
+
 	if len(merged) == 0 {
 		response := ToolResponse{
 			Status:  "no_results",
@@ -457,3 +480,165 @@ func (t *SmartSearchTool) handleSearchError(err error, workspaceRoot, workspaceI
 	response.Error = fmt.Sprintf("search failed: %v", err)
 	return response.JSON()
 }
+
+// readLines reads a specific range of lines from a file.
+// Lines are 1-indexed.
+func readLines(filePath string, startLine, endLine int) (string, error) {
+	content, err := os.ReadFile(filePath)
+	if err != nil {
+		return "", err
+	}
+	
+	lines := strings.Split(string(content), "\n")
+	if startLine < 1 {
+		startLine = 1
+	}
+	if endLine > len(lines) {
+		endLine = len(lines)
+	}
+	if startLine > endLine || startLine > len(lines) {
+		return "", fmt.Errorf("invalid line range")
+	}
+	
+	return strings.Join(lines[startLine-1:endLine], "\n"), nil
+}
+
+// groupDocsByTree aggregates "documentation" and "code_block" chunks
+// from the same file and AST Signature (Markdown heading) into single unified blocks,
+// fetching the continuous text from disk to prevent Frankenstein gaps.
+func (t *SmartSearchTool) groupDocsByTree(results []mergedResult) []mergedResult {
+	if len(results) == 0 {
+		return results
+	}
+
+	var out []mergedResult
+	
+	// Groups are keyed by: filePath_|_signature -> slice of mergedResult indices in groups slice
+	type docGroup struct {
+		filePath  string
+		signature string
+		items     []*mergedResult
+		maxScore  float32
+		minLine   int
+		maxLine   int
+		source    string
+	}
+	
+	groupsMap := make(map[string]*docGroup)
+	var orderedGroups []string // keep track of the first time we see a group to maintain rough sorting
+
+	for i := range results {
+		res := &results[i]
+		
+		// Only group documentation types and files ending in .md or .html
+		ext := strings.ToLower(filepath.Ext(res.filePath))
+		isDocFile := ext == ".md" || ext == ".markdown" || ext == ".html" || ext == ".htm" || 
+			ext == ".yaml" || ext == ".yml" || ext == ".json" || ext == ".xml" || 
+			ext == ".toml" || ext == ".rst" || ext == ".css" || ext == ".scss" || ext == ".svelte" || ext == ".sql" || ext == ".sh"
+		
+		isDocType := res.symbolType == "documentation" || res.symbolType == "code_block" || res.symbolType == "markdown"
+		
+		if !isDocType || !isDocFile || res.signature == "" {
+			// Pass-through code or items without signature
+			out = append(out, *res)
+			continue
+		}
+		
+		key := fmt.Sprintf("%s_|_%s", res.filePath, res.signature)
+		if g, exists := groupsMap[key]; exists {
+			g.items = append(g.items, res)
+			if res.score > g.maxScore {
+				g.maxScore = res.score
+			}
+			if res.startLine > 0 && (g.minLine == 0 || res.startLine < g.minLine) {
+				g.minLine = res.startLine
+			}
+			if res.endLine > 0 && res.endLine > g.maxLine {
+				g.maxLine = res.endLine
+			}
+			if g.source != "both" && g.source != res.source {
+				g.source = "both"
+			}
+		} else {
+			minL := res.startLine
+			if minL == 0 {
+				minL = 1
+			}
+			maxL := res.endLine
+			if maxL == 0 {
+				maxL = 1
+			}
+			groupsMap[key] = &docGroup{
+				filePath:  res.filePath,
+				signature: res.signature,
+				items:     []*mergedResult{res},
+				maxScore:  res.score,
+				minLine:   minL,
+				maxLine:   maxL,
+				source:    res.source,
+			}
+			orderedGroups = append(orderedGroups, key)
+		}
+	}
+
+	// Reconstruct the grouped items
+	for _, key := range orderedGroups {
+		g := groupsMap[key]
+		
+		if len(g.items) == 1 {
+			// Nothing to merge, just append
+			out = append(out, *g.items[0])
+			continue
+		}
+		
+		// Multiple chunks in this group. Let's merge them!
+		// Attempt to read the full continuous block from the file
+		fullContent := ""
+		if g.minLine > 0 && g.maxLine >= g.minLine {
+			content, err := readLines(g.filePath, g.minLine, g.maxLine)
+			if err == nil {
+				fullContent = content
+			}
+		}
+		
+		// If reading from disk failed, append the contents manually with an ellipsis
+		if fullContent == "" {
+			var contents []string
+			// Sort items by line number
+			sortedItems := make([]*mergedResult, len(g.items))
+			copy(sortedItems, g.items)
+			sort.Slice(sortedItems, func(i, j int) bool {
+				return sortedItems[i].startLine < sortedItems[j].startLine
+			})
+			for _, item := range sortedItems {
+				contents = append(contents, strings.TrimSpace(item.content))
+			}
+			fullContent = strings.Join(contents, "\n\n[...]\n\n")
+		}
+
+		baseItem := g.items[0] // take the first item as a prototype
+		merged := mergedResult{
+			id:         fmt.Sprintf("merged_%s_%d_%d", baseItem.id, g.minLine, g.maxLine),
+			score:      g.maxScore,
+			filePath:   g.filePath,
+			name:       baseItem.name,
+			symbolType: "documentation_merged",
+			signature:  g.signature,
+			pkg:        baseItem.pkg,
+			docstring:  fmt.Sprintf("Merged %d chunks spanning %d lines.", len(g.items), g.maxLine-g.minLine+1),
+			content:    fullContent,
+			startLine:  g.minLine,
+			endLine:    g.maxLine,
+			source:     g.source,
+		}
+		out = append(out, merged)
+	}
+
+	// After mixing merged chunks and original unmerged items, we should re-sort by score
+	sort.Slice(out, func(i, j int) bool {
+		return out[i].score > out[j].score
+	})
+
+	return out
+}
+
diff --git a/internal/service/tools/smart_search_test.go b/internal/service/tools/smart_search_test.go
@@ -0,0 +1,106 @@
+package tools
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+)
+
+func TestGroupDocsByTree(t *testing.T) {
+	// Create a temporary file for testing disk reads
+	tempDir := t.TempDir()
+	tempFilePath := filepath.Join(tempDir, "test_doc.md")
+	
+	// Create a dummy document with 20 lines
+	lines := make([]string, 20)
+	for i := 0; i < 20; i++ {
+		lines[i] = "Line " + string(rune('A'+i)) // Line A, Line B, etc.
+	}
+	os.WriteFile(tempFilePath, []byte(strings.Join(lines, "\n")), 0644)
+
+	// Create tool instance (we only need the groupDocsByTree method)
+	tool := &SmartSearchTool{}
+
+	tests := []struct {
+		name     string
+		input    []mergedResult
+		expected int // Expected number of results after grouping
+		check    func(t *testing.T, results []mergedResult)
+	}{
+		{
+			name: "No grouping needed for code",
+			input: []mergedResult{
+				{id: "1", filePath: "main.go", symbolType: "function", score: 0.9},
+				{id: "2", filePath: "main.go", symbolType: "function", score: 0.8},
+			},
+			expected: 2,
+			check: func(t *testing.T, results []mergedResult) {
+				if len(results) != 2 {
+					t.Errorf("Expected 2 results, got %d", len(results))
+				}
+			},
+		},
+		{
+			name: "Group documentation chunks from same file and signature",
+			input: []mergedResult{
+				{
+					id: "chunk_1", filePath: tempFilePath, symbolType: "documentation", signature: "### Intro",
+					startLine: 1, endLine: 3, score: 0.8,
+				},
+				{
+					id: "chunk_2", filePath: tempFilePath, symbolType: "documentation", signature: "### Intro",
+					startLine: 4, endLine: 6, score: 0.9,
+				},
+				// This one is in a different signature
+				{
+					id: "chunk_3", filePath: tempFilePath, symbolType: "documentation", signature: "### Setup",
+					startLine: 8, endLine: 10, score: 0.5,
+				},
+				// This one is code, ignore grouping
+				{
+					id: "chunk_4", filePath: tempFilePath, symbolType: "code_block", signature: "### Intro",
+					startLine: 7, endLine: 7, score: 0.85,
+				},
+			},
+			// Expect: 1 merged block for "### Intro", 1 block for "### Setup".
+			// Note: chunk_4 has "code_block" type! Code blocks in markdown also get grouped!
+			expected: 2,
+			check: func(t *testing.T, results []mergedResult) {
+				if len(results) != 2 {
+					t.Fatalf("Expected 2 results, got %d", len(results))
+				}
+				
+				// results are sorted by score. The merged "### Intro" should have max score 0.9
+				if results[0].score != 0.9 {
+					t.Errorf("Expected max score 0.9, got %v", results[0].score)
+				}
+				
+				// Check start/end line bounds for the merged "### Intro"
+				if results[0].startLine != 1 || results[0].endLine != 7 {
+					t.Errorf("Expected lines 1-7, got %d-%d", results[0].startLine, results[0].endLine)
+				}
+				
+				// Verify the content was loaded from disk and contains Lines A to G (1 to 7)
+				if !strings.Contains(results[0].content, "Line A") || !strings.Contains(results[0].content, "Line G") {
+					t.Errorf("Merged content does not match expected lines from disk")
+				}
+				
+				// Setup should be untouched
+				if results[1].signature != "### Setup" {
+					t.Errorf("Expected second result to be '### Setup'")
+				}
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got := tool.groupDocsByTree(tc.input)
+			if len(got) != tc.expected {
+				t.Errorf("Expected %d results, got %d", tc.expected, len(got))
+			}
+			tc.check(t, got)
+		})
+	}
+}
