feat(agent): enhance chunk search functionality to include total chunk counts

- Updated the searchChunks method to calculate and include the total chunk count for each knowledge ID in the results.
- Introduced logic to fetch chunk counts based on unique knowledge IDs, improving the efficiency of data retrieval.
- Enhanced error handling for chunk count fetching, ensuring robust logging in case of failures.

internal/agent/tools/grep_chunks.go

@@ -288,8 +288,7 @@ func (t *GrepChunksTool) searchChunks(
 	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, "+
-			"COUNT(*) OVER (PARTITION BY chunks.knowledge_id) AS total_chunk_count").
+			"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").
@@ -336,6 +335,42 @@ func (t *GrepChunksTool) searchChunks(
 		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
 }
 

migrations/versioned/000026_chunks_query_indexes.down.sql

@@ -0,0 +1,4 @@
+DO $$ BEGIN RAISE NOTICE '[Migration 000026] Dropping chunks query indexes'; END $$;
+
+DROP INDEX IF EXISTS idx_chunks_kb_tenant;
+DROP INDEX IF EXISTS idx_chunks_knowledge_enabled;

migrations/versioned/000026_chunks_query_indexes.up.sql

@@ -0,0 +1,18 @@
+-- ============================================================================
+-- Migration 000026: Add missing indexes for chunks query performance
+-- ============================================================================
+-- Covers grep_chunks tool queries that filter by (knowledge_base_id, tenant_id)
+-- and aggregate by (knowledge_id) with is_enabled/deleted_at filters.
+-- ============================================================================
+
+DO $$ BEGIN RAISE NOTICE '[Migration 000026] Adding chunks query indexes'; END $$;
+
+-- Index for main query: WHERE knowledge_base_id = ? AND tenant_id = ?
+-- Also benefits JOIN and OR conditions across multiple KB-tenant pairs
+CREATE INDEX IF NOT EXISTS idx_chunks_kb_tenant
+    ON chunks(knowledge_base_id, tenant_id);
+
+-- Index for count query: WHERE knowledge_id IN (...) AND is_enabled = true AND deleted_at IS NULL
+-- GROUP BY knowledge_id
+CREATE INDEX IF NOT EXISTS idx_chunks_knowledge_enabled
+    ON chunks(knowledge_id, is_enabled, deleted_at);