diff --git a/internal/agent/tools/grep_chunks.go b/internal/agent/tools/grep_chunks.go index 77d7f40c..8fd2fda2 100644 --- a/internal/agent/tools/grep_chunks.go +++ b/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 } diff --git a/migrations/versioned/000026_chunks_query_indexes.down.sql b/migrations/versioned/000026_chunks_query_indexes.down.sql new file mode 100644 index 00000000..a76d5553 --- /dev/null +++ b/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; diff --git a/migrations/versioned/000026_chunks_query_indexes.up.sql b/migrations/versioned/000026_chunks_query_indexes.up.sql new file mode 100644 index 00000000..6ae01534 --- /dev/null +++ b/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);