mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-25 12:22:54 +02:00
Add files via upload
This commit is contained in:
@@ -4,6 +4,7 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"cyberstrike-ai/internal/audit"
|
"cyberstrike-ai/internal/audit"
|
||||||
@@ -316,20 +317,55 @@ func (h *KnowledgeHandler) DeleteItem(c *gin.Context) {
|
|||||||
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
|
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
|
||||||
}
|
}
|
||||||
|
|
||||||
// RebuildIndex 重建索引
|
// StartIndex 构建知识库向量索引。默认仅补齐尚无向量的知识项;mode=full 时全量重建。
|
||||||
func (h *KnowledgeHandler) RebuildIndex(c *gin.Context) {
|
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() {
|
go func() {
|
||||||
|
defer h.indexer.FinishIndexRun()
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
if err := h.indexer.RebuildIndex(ctx); err != nil {
|
var err error
|
||||||
h.logger.Error("重建索引失败", zap.Error(err))
|
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 {
|
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 扫描知识库
|
// ScanKnowledgeBase 扫描知识库
|
||||||
|
|||||||
@@ -4416,16 +4416,35 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"/api/knowledge/index": map[string]interface{}{
|
"/api/knowledge/index": map[string]interface{}{
|
||||||
"post": map[string]interface{}{
|
"post": map[string]interface{}{
|
||||||
"tags": []string{"知识库"},
|
"tags": []string{"知识库"},
|
||||||
"summary": "重建索引",
|
"summary": "构建索引",
|
||||||
"description": "重新构建知识库索引",
|
"description": "构建知识库向量索引。默认仅处理尚无向量的知识项;mode=full 时全量重建。",
|
||||||
"operationId": "rebuildIndex",
|
"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{}{
|
"responses": map[string]interface{}{
|
||||||
"200": map[string]interface{}{
|
"200": map[string]interface{}{
|
||||||
"description": "重建索引任务已启动",
|
"description": "索引任务已启动",
|
||||||
|
},
|
||||||
|
"400": map[string]interface{}{
|
||||||
|
"description": "无效的 mode 参数",
|
||||||
},
|
},
|
||||||
"401": map[string]interface{}{
|
"401": map[string]interface{}{
|
||||||
"description": "未授权",
|
"description": "未授权",
|
||||||
},
|
},
|
||||||
|
"409": map[string]interface{}{
|
||||||
|
"description": "已有索引任务正在进行",
|
||||||
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -43,7 +43,7 @@ var apiDocI18nSummaryToKey = map[string]string{
|
|||||||
"设置对话置顶": "pinConversation", "设置分组置顶": "pinGroup", "设置分组中对话的置顶": "pinGroupConversation",
|
"设置对话置顶": "pinConversation", "设置分组置顶": "pinGroup", "设置分组中对话的置顶": "pinGroupConversation",
|
||||||
"获取分类": "getCategories", "列出知识项": "listKnowledgeItems", "创建知识项": "createKnowledgeItem",
|
"获取分类": "getCategories", "列出知识项": "listKnowledgeItems", "创建知识项": "createKnowledgeItem",
|
||||||
"获取知识项": "getKnowledgeItem", "更新知识项": "updateKnowledgeItem", "删除知识项": "deleteKnowledgeItem",
|
"获取知识项": "getKnowledgeItem", "更新知识项": "updateKnowledgeItem", "删除知识项": "deleteKnowledgeItem",
|
||||||
"获取索引状态": "getIndexStatus", "重建索引": "rebuildIndex", "扫描知识库": "scanKnowledgeBase",
|
"获取索引状态": "getIndexStatus", "构建索引": "startKnowledgeIndex", "扫描知识库": "scanKnowledgeBase",
|
||||||
"搜索知识库": "searchKnowledgeBase", "基础搜索": "basicSearch", "按风险类型搜索": "searchByRiskType",
|
"搜索知识库": "searchKnowledgeBase", "基础搜索": "basicSearch", "按风险类型搜索": "searchByRiskType",
|
||||||
"获取检索日志": "getRetrievalLogs", "删除检索日志": "deleteRetrievalLog",
|
"获取检索日志": "getRetrievalLogs", "删除检索日志": "deleteRetrievalLog",
|
||||||
"MCP端点": "mcpEndpoint", "列出所有工具": "listAllTools", "调用工具": "invokeTool", "初始化连接": "initConnection",
|
"MCP端点": "mcpEndpoint", "列出所有工具": "listAllTools", "调用工具": "invokeTool", "初始化连接": "initConnection",
|
||||||
|
|||||||
+105
-22
@@ -213,9 +213,13 @@ func (idx *Indexer) HasIndex() (bool, error) {
|
|||||||
return count > 0, nil
|
return count > 0, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// RebuildIndex 重建所有索引
|
func (idx *Indexer) beginIndexRun() error {
|
||||||
func (idx *Indexer) RebuildIndex(ctx context.Context) error {
|
|
||||||
idx.rebuildMu.Lock()
|
idx.rebuildMu.Lock()
|
||||||
|
defer idx.rebuildMu.Unlock()
|
||||||
|
|
||||||
|
if idx.isRebuilding {
|
||||||
|
return fmt.Errorf("索引任务已在进行中")
|
||||||
|
}
|
||||||
idx.isRebuilding = true
|
idx.isRebuilding = true
|
||||||
idx.rebuildTotalItems = 0
|
idx.rebuildTotalItems = 0
|
||||||
idx.rebuildCurrent = 0
|
idx.rebuildCurrent = 0
|
||||||
@@ -223,41 +227,124 @@ func (idx *Indexer) RebuildIndex(ctx context.Context) error {
|
|||||||
idx.rebuildStartTime = time.Now()
|
idx.rebuildStartTime = time.Now()
|
||||||
idx.rebuildLastItemID = ""
|
idx.rebuildLastItemID = ""
|
||||||
idx.rebuildLastChunks = 0
|
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.mu.Lock()
|
||||||
idx.lastError = ""
|
idx.lastError = ""
|
||||||
idx.lastErrorTime = time.Time{}
|
idx.lastErrorTime = time.Time{}
|
||||||
idx.errorCount = 0
|
idx.errorCount = 0
|
||||||
idx.mu.Unlock()
|
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 {
|
if err != nil {
|
||||||
idx.rebuildMu.Lock()
|
|
||||||
idx.isRebuilding = false
|
|
||||||
idx.rebuildMu.Unlock()
|
|
||||||
return fmt.Errorf("查询知识项失败:%w", err)
|
return fmt.Errorf("查询知识项失败:%w", err)
|
||||||
}
|
}
|
||||||
defer rows.Close()
|
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
|
var itemIDs []string
|
||||||
for rows.Next() {
|
for rows.Next() {
|
||||||
var id string
|
var id string
|
||||||
if err := rows.Scan(&id); err != nil {
|
if err := rows.Scan(&id); err != nil {
|
||||||
idx.rebuildMu.Lock()
|
return nil, fmt.Errorf("扫描知识项 ID 失败:%w", err)
|
||||||
idx.isRebuilding = false
|
|
||||||
idx.rebuildMu.Unlock()
|
|
||||||
return fmt.Errorf("扫描知识项 ID 失败:%w", err)
|
|
||||||
}
|
}
|
||||||
itemIDs = append(itemIDs, id)
|
itemIDs = append(itemIDs, id)
|
||||||
}
|
}
|
||||||
|
if err := rows.Err(); err != nil {
|
||||||
|
return nil, fmt.Errorf("扫描知识项 ID 失败:%w", err)
|
||||||
|
}
|
||||||
|
return itemIDs, nil
|
||||||
|
}
|
||||||
|
|
||||||
idx.rebuildMu.Lock()
|
func (idx *Indexer) indexItemIDs(ctx context.Context, itemIDs []string, doneMessage string) error {
|
||||||
idx.rebuildTotalItems = len(itemIDs)
|
|
||||||
idx.rebuildMu.Unlock()
|
|
||||||
|
|
||||||
idx.logger.Info("开始重建索引", zap.Int("totalItems", len(itemIDs)))
|
|
||||||
|
|
||||||
failedCount := 0
|
failedCount := 0
|
||||||
consecutiveFailures := 0
|
consecutiveFailures := 0
|
||||||
maxConsecutiveFailures := 5
|
maxConsecutiveFailures := 5
|
||||||
@@ -329,11 +416,7 @@ func (idx *Indexer) RebuildIndex(ctx context.Context) error {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
idx.rebuildMu.Lock()
|
idx.logger.Info(doneMessage, zap.Int("totalItems", len(itemIDs)), zap.Int("failedCount", failedCount))
|
||||||
idx.isRebuilding = false
|
|
||||||
idx.rebuildMu.Unlock()
|
|
||||||
|
|
||||||
idx.logger.Info("索引重建完成", zap.Int("totalItems", len(itemIDs)), zap.Int("failedCount", failedCount))
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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()
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user