diff --git a/internal/handler/knowledge.go b/internal/handler/knowledge.go index eee106ac..f08882c8 100644 --- a/internal/handler/knowledge.go +++ b/internal/handler/knowledge.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "net/http" + "strings" "time" "cyberstrike-ai/internal/audit" @@ -316,20 +317,55 @@ func (h *KnowledgeHandler) DeleteItem(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"message": "删除成功"}) } -// RebuildIndex 重建索引 -func (h *KnowledgeHandler) RebuildIndex(c *gin.Context) { - // 异步重建索引 +// StartIndex 构建知识库向量索引。默认仅补齐尚无向量的知识项；mode=full 时全量重建。 +func (h *KnowledgeHandler) StartIndex(c *gin.Context) { + if err := h.indexer.TryBeginIndexRun(); err != nil { + c.JSON(http.StatusConflict, gin.H{"error": "已有索引任务正在进行，请等待完成"}) + return + } + + mode := strings.TrimSpace(c.Query("mode")) + if mode == "" { + mode = "missing" + } + if mode != "full" && mode != "missing" { + h.indexer.FinishIndexRun() + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的 mode 参数，可选值：missing、full"}) + return + } + + fullRebuild := mode == "full" + message := "索引构建已开始，将在后台进行" + auditAction := "index_build" + auditDetail := "构建知识库索引" + if fullRebuild { + message = "全量索引重建已开始，将在后台进行" + auditAction = "index_rebuild_full" + auditDetail = "全量重建知识库索引" + } + go func() { + defer h.indexer.FinishIndexRun() ctx := context.Background() - if err := h.indexer.RebuildIndex(ctx); err != nil { - h.logger.Error("重建索引失败", zap.Error(err)) + var err error + if fullRebuild { + err = h.indexer.RunRebuildIndex(ctx) + } else { + err = h.indexer.RunIndexMissing(ctx) + } + if err != nil { + if fullRebuild { + h.logger.Error("全量重建索引失败", zap.Error(err)) + } else { + h.logger.Error("构建知识库索引失败", zap.Error(err)) + } } }() if h.audit != nil { - h.audit.RecordOK(c, "knowledge", "index_rebuild", "重建知识库索引", "knowledge", "", nil) + h.audit.RecordOK(c, "knowledge", auditAction, auditDetail, "knowledge", "", nil) } - c.JSON(http.StatusOK, gin.H{"message": "索引重建已开始，将在后台进行"}) + c.JSON(http.StatusOK, gin.H{"message": message, "mode": mode}) } // ScanKnowledgeBase 扫描知识库 diff --git a/internal/handler/openapi.go b/internal/handler/openapi.go index ba8f4eca..4860e89c 100644 --- a/internal/handler/openapi.go +++ b/internal/handler/openapi.go @@ -4416,16 +4416,35 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) { "/api/knowledge/index": map[string]interface{}{ "post": map[string]interface{}{ "tags": []string{"知识库"}, - "summary": "重建索引", - "description": "重新构建知识库索引", - "operationId": "rebuildIndex", + "summary": "构建索引", + "description": "构建知识库向量索引。默认仅处理尚无向量的知识项；mode=full 时全量重建。", + "operationId": "startKnowledgeIndex", + "parameters": []map[string]interface{}{ + { + "name": "mode", + "in": "query", + "required": false, + "description": "索引模式：missing（默认，补齐缺失向量）或 full（全量重建）", + "schema": map[string]interface{}{ + "type": "string", + "enum": []string{"missing", "full"}, + "default": "missing", + }, + }, + }, "responses": map[string]interface{}{ "200": map[string]interface{}{ - "description": "重建索引任务已启动", + "description": "索引任务已启动", + }, + "400": map[string]interface{}{ + "description": "无效的 mode 参数", }, "401": map[string]interface{}{ "description": "未授权", }, + "409": map[string]interface{}{ + "description": "已有索引任务正在进行", + }, }, }, }, diff --git a/internal/handler/openapi_i18n.go b/internal/handler/openapi_i18n.go index efdbb757..11ce5733 100644 --- a/internal/handler/openapi_i18n.go +++ b/internal/handler/openapi_i18n.go @@ -43,7 +43,7 @@ var apiDocI18nSummaryToKey = map[string]string{ "设置对话置顶": "pinConversation", "设置分组置顶": "pinGroup", "设置分组中对话的置顶": "pinGroupConversation", "获取分类": "getCategories", "列出知识项": "listKnowledgeItems", "创建知识项": "createKnowledgeItem", "获取知识项": "getKnowledgeItem", "更新知识项": "updateKnowledgeItem", "删除知识项": "deleteKnowledgeItem", - "获取索引状态": "getIndexStatus", "重建索引": "rebuildIndex", "扫描知识库": "scanKnowledgeBase", + "获取索引状态": "getIndexStatus", "构建索引": "startKnowledgeIndex", "扫描知识库": "scanKnowledgeBase", "搜索知识库": "searchKnowledgeBase", "基础搜索": "basicSearch", "按风险类型搜索": "searchByRiskType", "获取检索日志": "getRetrievalLogs", "删除检索日志": "deleteRetrievalLog", "MCP端点": "mcpEndpoint", "列出所有工具": "listAllTools", "调用工具": "invokeTool", "初始化连接": "initConnection", diff --git a/internal/knowledge/indexer.go b/internal/knowledge/indexer.go index aeb6b9ff..1c1e7cbe 100644 --- a/internal/knowledge/indexer.go +++ b/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 } diff --git a/internal/knowledge/indexer_rebuild_state_test.go b/internal/knowledge/indexer_rebuild_state_test.go new file mode 100644 index 00000000..1e04480e --- /dev/null +++ b/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() +}