Add files via upload

internal/knowledge/indexer.go

@@ -213,9 +213,13 @@ func (idx *Indexer) HasIndex() (bool, error) {
 	return count > 0, nil
 }
 
-// RebuildIndex 重建所有索引
-func (idx *Indexer) RebuildIndex(ctx context.Context) error {
+func (idx *Indexer) beginIndexRun() error {
 	idx.rebuildMu.Lock()
+	defer idx.rebuildMu.Unlock()
+
+	if idx.isRebuilding {
+		return fmt.Errorf("索引任务已在进行中")
+	}
 	idx.isRebuilding = true
 	idx.rebuildTotalItems = 0
 	idx.rebuildCurrent = 0
@@ -223,41 +227,124 @@ func (idx *Indexer) RebuildIndex(ctx context.Context) error {
 	idx.rebuildStartTime = time.Now()
 	idx.rebuildLastItemID = ""
 	idx.rebuildLastChunks = 0
-	idx.rebuildMu.Unlock()
+	return nil
+}
 
+// TryBeginIndexRun 同步占用索引任务槽位；调用方必须在后台任务结束时调用 FinishIndexRun。
+func (idx *Indexer) TryBeginIndexRun() error {
+	return idx.beginIndexRun()
+}
+
+func (idx *Indexer) FinishIndexRun() {
+	idx.rebuildMu.Lock()
+	idx.isRebuilding = false
+	idx.rebuildMu.Unlock()
+}
+
+func (idx *Indexer) resetLastError() {
 	idx.mu.Lock()
 	idx.lastError = ""
 	idx.lastErrorTime = time.Time{}
 	idx.errorCount = 0
 	idx.mu.Unlock()
+}
 
-	rows, err := idx.db.Query("SELECT id FROM knowledge_base_items")
+func (idx *Indexer) setIndexRunTotal(total int) {
+	idx.rebuildMu.Lock()
+	idx.rebuildTotalItems = total
+	idx.rebuildMu.Unlock()
+}
+
+// IndexMissing 为尚无向量的知识项构建索引（默认推荐路径，适合冷启动与中断续跑）。
+func (idx *Indexer) IndexMissing(ctx context.Context) error {
+	if err := idx.beginIndexRun(); err != nil {
+		return err
+	}
+	defer idx.FinishIndexRun()
+	return idx.runIndexMissing(ctx)
+}
+
+// RebuildIndex 全量重建所有知识项索引（显式 opt-in，成本更高）。
+func (idx *Indexer) RebuildIndex(ctx context.Context) error {
+	if err := idx.beginIndexRun(); err != nil {
+		return err
+	}
+	defer idx.FinishIndexRun()
+	return idx.runRebuildIndex(ctx)
+}
+
+// RunRebuildIndex 在已占用索引任务槽位后执行全量重建（供 HTTP handler 后台任务使用）。
+func (idx *Indexer) RunRebuildIndex(ctx context.Context) error {
+	return idx.runRebuildIndex(ctx)
+}
+
+// RunIndexMissing 在已占用索引任务槽位后执行缺失索引补齐（供 HTTP handler 后台任务使用）。
+func (idx *Indexer) RunIndexMissing(ctx context.Context) error {
+	return idx.runIndexMissing(ctx)
+}
+
+func (idx *Indexer) runRebuildIndex(ctx context.Context) error {
+	idx.resetLastError()
+
+	rows, err := idx.db.QueryContext(ctx, "SELECT id FROM knowledge_base_items ORDER BY updated_at ASC, id ASC")
 	if err != nil {
-		idx.rebuildMu.Lock()
-		idx.isRebuilding = false
-		idx.rebuildMu.Unlock()
 		return fmt.Errorf("查询知识项失败：%w", err)
 	}
 	defer rows.Close()
 
+	itemIDs, err := scanKnowledgeItemIDs(rows)
+	if err != nil {
+		return err
+	}
+
+	idx.setIndexRunTotal(len(itemIDs))
+	idx.logger.Info("开始重建索引", zap.Int("totalItems", len(itemIDs)))
+
+	return idx.indexItemIDs(ctx, itemIDs, "索引重建完成")
+}
+
+func (idx *Indexer) runIndexMissing(ctx context.Context) error {
+	idx.resetLastError()
+
+	rows, err := idx.db.QueryContext(ctx, `
+		SELECT i.id
+		FROM knowledge_base_items i
+		LEFT JOIN knowledge_embeddings e ON e.item_id = i.id
+		WHERE e.item_id IS NULL
+		ORDER BY i.updated_at ASC, i.id ASC
+	`)
+	if err != nil {
+		return fmt.Errorf("查询未索引知识项失败：%w", err)
+	}
+	defer rows.Close()
+
+	itemIDs, err := scanKnowledgeItemIDs(rows)
+	if err != nil {
+		return fmt.Errorf("扫描未索引知识项 ID 失败：%w", err)
+	}
+
+	idx.setIndexRunTotal(len(itemIDs))
+	idx.logger.Info("开始补齐缺失索引", zap.Int("totalItems", len(itemIDs)))
+
+	return idx.indexItemIDs(ctx, itemIDs, "索引构建完成")
+}
+
+func scanKnowledgeItemIDs(rows *sql.Rows) ([]string, error) {
 	var itemIDs []string
 	for rows.Next() {
 		var id string
 		if err := rows.Scan(&id); err != nil {
-			idx.rebuildMu.Lock()
-			idx.isRebuilding = false
-			idx.rebuildMu.Unlock()
-			return fmt.Errorf("扫描知识项 ID 失败：%w", err)
+			return nil, fmt.Errorf("扫描知识项 ID 失败：%w", err)
 		}
 		itemIDs = append(itemIDs, id)
 	}
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("扫描知识项 ID 失败：%w", err)
+	}
+	return itemIDs, nil
+}
 
-	idx.rebuildMu.Lock()
-	idx.rebuildTotalItems = len(itemIDs)
-	idx.rebuildMu.Unlock()
-
-	idx.logger.Info("开始重建索引", zap.Int("totalItems", len(itemIDs)))
-
+func (idx *Indexer) indexItemIDs(ctx context.Context, itemIDs []string, doneMessage string) error {
 	failedCount := 0
 	consecutiveFailures := 0
 	maxConsecutiveFailures := 5
@@ -329,11 +416,7 @@ func (idx *Indexer) RebuildIndex(ctx context.Context) error {
 		}
 	}
 
-	idx.rebuildMu.Lock()
-	idx.isRebuilding = false
-	idx.rebuildMu.Unlock()
-
-	idx.logger.Info("索引重建完成", zap.Int("totalItems", len(itemIDs)), zap.Int("failedCount", failedCount))
+	idx.logger.Info(doneMessage, zap.Int("totalItems", len(itemIDs)), zap.Int("failedCount", failedCount))
 	return nil
 }
 

internal/knowledge/indexer_rebuild_state_test.go

@@ -0,0 +1,20 @@
+package knowledge
+
+import "testing"
+
+func TestIndexerRejectsConcurrentIndexRuns(t *testing.T) {
+	idx := &Indexer{}
+
+	if err := idx.beginIndexRun(); err != nil {
+		t.Fatalf("first index run should start: %v", err)
+	}
+	if err := idx.beginIndexRun(); err == nil {
+		t.Fatal("second index run should be rejected while one is active")
+	}
+
+	idx.FinishIndexRun()
+	if err := idx.beginIndexRun(); err != nil {
+		t.Fatalf("index run should start again after finish: %v", err)
+	}
+	idx.FinishIndexRun()
+}