From 79e834e8e9cc9a0099b28177126ff2bc81e09505 Mon Sep 17 00:00:00 2001 From: Blizzard Date: Tue, 30 Jun 2026 13:37:45 +0800 Subject: [PATCH] =?UTF-8?q?feat(rag):=20=E5=A4=A7=E6=96=87=E4=BB=B6=20RAG?= =?UTF-8?q?=20=E5=81=9A=E6=B7=B1=20=E2=80=94=E2=80=94=20=E5=B9=B6=E5=8F=91?= =?UTF-8?q?=E5=85=A5=E5=BA=93/=E5=9B=BE=E8=B0=B1=E7=AA=97=E5=8F=A3?= =?UTF-8?q?=E5=8C=96=20+=20Bleve=E8=90=BD=E7=9B=98=20+=20file=5Fid?= =?UTF-8?q?=E6=B2=BB=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 几十万字文件从"广而浅"到准生产级： - 向量化串行→并发分批保序(embedAll)；几十万字上千块快数倍 - 图谱整篇喂LLM(爆上下文只抽开头)→窗口化并发抽(extractGraphWindowed)， 全覆盖；窗口/封顶/并发 env 可配；图谱可单配便宜模型(GRAPH_CHAT_*，未配回退主chat) - Bleve 内存索引(重启即丢、三路退两路)→落盘 scorch(env BLEVE_PATH，失败退内存兜底) - 下游键改稳定 file_id：Neo4j 关系打 file_id(实体仍 kb+name 共享)； 新增 kb_delete 工具 + Engine.DeleteDoc 级联删 Milvus/Bleve/Neo4j - 单测：窗口化/去重/env可配/落盘持久/图谱模型回退 Co-Authored-By: Claude Opus 4.8 --- sundynix-mcp-go/internal/mcp/gateway.go | 14 ++ sundynix-mcp-go/internal/rag/bleve.go | 46 ++++- sundynix-mcp-go/internal/rag/cjk_test.go | 22 +++ sundynix-mcp-go/internal/rag/graph.go | 28 ++- sundynix-mcp-go/internal/rag/ingest_large.go | 163 ++++++++++++++++++ .../internal/rag/ingest_large_test.go | 108 ++++++++++++ sundynix-mcp-go/internal/rag/milvus.go | 11 ++ sundynix-mcp-go/internal/rag/rag.go | 116 ++++++++----- 8 files changed, 462 insertions(+), 46 deletions(-) create mode 100644 sundynix-mcp-go/internal/rag/ingest_large.go create mode 100644 sundynix-mcp-go/internal/rag/ingest_large_test.go diff --git a/sundynix-mcp-go/internal/mcp/gateway.go b/sundynix-mcp-go/internal/mcp/gateway.go index 6d64931..d622f73 100644 --- a/sundynix-mcp-go/internal/mcp/gateway.go +++ b/sundynix-mcp-go/internal/mcp/gateway.go @@ -162,6 +162,7 @@ func (g *Gateway) buildRegistry() map[string]toolDef { // —— 仅内部/流水线/管理用，不暴露给自主 agent —— "kb_ingest": {cn: "知识入库", desc: "文本切块 → 向量化 → 写入 Milvus / Bleve", handler: g.kbIngest}, + "kb_delete": {cn: "知识删除", desc: "按 file_id 级联删某文档的向量/全文/图谱", handler: g.kbDelete}, "kb_search": {cn: "检索台查询", desc: "结构化返回命中内容与相似度分数", handler: g.kbSearch}, "kb_graph": {cn: "知识图谱", desc: "取某库的实体关系三元组（Neo4j）", handler: g.kbGraph}, "report_render": {cn: "报告渲染", desc: "把结构化报告渲染为 Word（.docx）", handler: g.reportRender}, @@ -507,3 +508,16 @@ func (g *Gateway) kbIngest(ctx context.Context, call *contract.ToolCall) *contra } return &contract.ToolResult{OK: true, Content: fmt.Sprintf("已入库 %d 块到知识库 %q", n, kb)} } + +// kbDelete 按 file_id 级联删某文档在三库的痕迹（向量/全文/图谱）。 +func (g *Gateway) kbDelete(ctx context.Context, call *contract.ToolCall) *contract.ToolResult { + kb, _ := call.Args["kb"].(string) + fileID, _ := call.Args["file_id"].(string) + if fileID == "" { + return &contract.ToolResult{OK: false, Error: "kb_delete: file_id 必填"} + } + if err := g.rag.DeleteDoc(ctx, kb, fileID); err != nil { + return &contract.ToolResult{OK: false, Error: "kb_delete: " + err.Error()} + } + return &contract.ToolResult{OK: true, Content: fmt.Sprintf("已删除文档 %s 的向量/全文/图谱", fileID)} +} diff --git a/sundynix-mcp-go/internal/rag/bleve.go b/sundynix-mcp-go/internal/rag/bleve.go index c41a778..11a18f6 100644 --- a/sundynix-mcp-go/internal/rag/bleve.go +++ b/sundynix-mcp-go/internal/rag/bleve.go @@ -4,6 +4,8 @@ import ( "fmt" "hash/fnv" "log" + "os" + "path/filepath" "github.com/blevesearch/bleve/v2" "github.com/blevesearch/bleve/v2/analysis/analyzer/keyword" @@ -12,8 +14,9 @@ import ( "github.com/blevesearch/bleve/v2/search/query" ) -// bleveStore 是全文（BM25）检索路。内存索引：随 ingest 写入，进程重启重建。 -// 真实生产应落盘(bleve.New(path,...))；此处内存优先求简。 +// bleveStore 是全文（BM25）检索路。落盘索引（scorch）：随 ingest 写入并持久，进程重启不丢， +// 与 Milvus(向量)/Neo4j(图谱) 同为持久存储——保证"三路混合"重启后仍是三路。 +// 落盘不可用时退回内存索引（功能在、重启会丢），好过全文路全哑。 type bleveStore struct { idx bleve.Index } @@ -34,17 +37,54 @@ func bleveMapping() mapping.IndexMapping { return im } +// blevePath 返回全文索引落盘目录（env BLEVE_PATH，默认 ./.data/bleve）。 +// 生产部署应指向持久卷；开发期落在工作目录下的 .data。 +func blevePath() string { + if p := os.Getenv("BLEVE_PATH"); p != "" { + return p + } + return ".data/bleve" +} + +// openBleve 打开（或新建）落盘全文索引；落盘失败退回内存索引兜底。 func openBleve() *bleveStore { + path := blevePath() + idx, err := bleve.Open(path) + if err == bleve.ErrorIndexPathDoesNotExist { + if mkerr := os.MkdirAll(filepath.Dir(path), 0o755); mkerr != nil { + log.Printf("[rag] bleve 目录创建失败，全文路退内存（重启会丢）: %v", mkerr) + return memBleve() + } + idx, err = bleve.New(path, bleveMapping()) + } + if err != nil { + log.Printf("[rag] bleve 落盘打开失败，全文路退内存（重启会丢）: %v", err) + return memBleve() + } + log.Printf("[rag] bleve 全文索引落盘就绪: %s", path) + return &bleveStore{idx: idx} +} + +// memBleve 内存兜底：落盘不可用时退回内存索引（功能在、重启会丢），好过全文路全哑。 +func memBleve() *bleveStore { idx, err := bleve.NewMemOnly(bleveMapping()) if err != nil { - log.Printf("[rag] bleve 初始化失败，全文路降级: %v", err) + log.Printf("[rag] bleve 内存兜底也失败，全文路降级: %v", err) return &bleveStore{} } + log.Printf("[rag] bleve 退回内存索引（非持久）") return &bleveStore{idx: idx} } func (b *bleveStore) ready() bool { return b != nil && b.idx != nil } +// close 关闭索引（落盘版释放锁并刷盘）；优雅停机时调。 +func (b *bleveStore) close() { + if b.ready() { + _ = b.idx.Close() + } +} + // index 把 (kb, doc, texts) 写入全文索引（id 含 kb+doc+文本哈希，幂等）。 func (b *bleveStore) index(kb, doc string, texts []string) error { if !b.ready() { diff --git a/sundynix-mcp-go/internal/rag/cjk_test.go b/sundynix-mcp-go/internal/rag/cjk_test.go index 650709f..783e295 100644 --- a/sundynix-mcp-go/internal/rag/cjk_test.go +++ b/sundynix-mcp-go/internal/rag/cjk_test.go @@ -5,7 +5,9 @@ import "testing" // TestBleve_ChineseSearch 钉死中文全文检索：CJK 分词器让"星云一号的总设计师"能命中含 // "星云一号""总设计师"的中文块。修复前默认标准分词器把整段中文当一个 token → 永远 0 命中。 func TestBleve_ChineseSearch(t *testing.T) { + t.Setenv("BLEVE_PATH", t.TempDir()+"/bleve") // 隔离落盘索引，测试互不污染 b := openBleve() + defer b.close() if !b.ready() { t.Skip("bleve 不可用") } @@ -39,6 +41,26 @@ func TestQueryNgrams(t *testing.T) { } } +// TestBleve_PersistsAcrossReopen 钉死落盘：写入→关闭→同路径重开，数据仍在。 +// 修复"NewMemOnly 进程重启即丢、三路静默退化两路"的回归——这是 P0 的核心验收。 +func TestBleve_PersistsAcrossReopen(t *testing.T) { + t.Setenv("BLEVE_PATH", t.TempDir()+"/bleve") + b1 := openBleve() + if !b1.ready() { + t.Skip("bleve 不可用") + } + if err := b1.index("kp", "d1", []string{"持久化测试：龙渊号探测器由总师陈思远负责。"}); err != nil { + t.Fatal(err) + } + b1.close() // 模拟进程退出（落盘刷盘） + + b2 := openBleve() // 模拟重启：同路径重开 + defer b2.close() + if hits := b2.search("kp", "龙渊号探测器", 5); len(hits) == 0 { + t.Fatal("重开后应仍能检索到（落盘持久），got 0 —— 退回内存了？") + } +} + func contains(s, sub string) bool { return len(sub) == 0 || (len(s) >= len(sub) && indexOf(s, sub) >= 0) } diff --git a/sundynix-mcp-go/internal/rag/graph.go b/sundynix-mcp-go/internal/rag/graph.go index 8306b6d..b6b7b66 100644 --- a/sundynix-mcp-go/internal/rag/graph.go +++ b/sundynix-mcp-go/internal/rag/graph.go @@ -56,7 +56,9 @@ func (g *graphStore) close(ctx context.Context) { } // store 把三元组 MERGE 进 Neo4j（实体 + 关系，按 kb 隔离）。 -func (g *graphStore) store(ctx context.Context, kb string, triples []Triple) (int, error) { +// 关系带 file_id（来源文档稳定 ID）→ 同一三元组来自不同文档各成一条边，供按文档级联删； +// 实体仍按 (kb,name) 共享去重（一个实体可被多篇文档提及）。 +func (g *graphStore) store(ctx context.Context, kb, fileID string, triples []Triple) (int, error) { if !g.ready() { return 0, nil } @@ -68,8 +70,8 @@ func (g *graphStore) store(ctx context.Context, kb string, triples []Triple) (in _, err := neo4j.ExecuteQuery(ctx, g.driver, `MERGE (a:Entity {kb:$kb, name:$s}) MERGE (b:Entity {kb:$kb, name:$o}) - MERGE (a)-[r:REL {type:$p}]->(b)`, - map[string]any{"kb": kb, "s": t.S, "o": t.O, "p": t.P}, + MERGE (a)-[r:REL {type:$p, file_id:$fid}]->(b)`, + map[string]any{"kb": kb, "s": t.S, "o": t.O, "p": t.P, "fid": fileID}, neo4j.EagerResultTransformer, neo4j.ExecuteQueryWithDatabase("neo4j")) if err != nil { return n, err @@ -79,6 +81,26 @@ func (g *graphStore) store(ctx context.Context, kb string, triples []Triple) (in return n, nil } +// deleteByFile 删除某文档（file_id）贡献的全部关系，再清掉因此变孤儿的实体（无任何关系）。 +// 共享实体（仍被其它文档的关系连着）保留。kb 作用域隔离。 +func (g *graphStore) deleteByFile(ctx context.Context, kb, fileID string) error { + if !g.ready() || fileID == "" { + return nil + } + if _, err := neo4j.ExecuteQuery(ctx, g.driver, + `MATCH (:Entity {kb:$kb})-[r:REL {file_id:$fid}]->(:Entity {kb:$kb}) DELETE r`, + map[string]any{"kb": kb, "fid": fileID}, + neo4j.EagerResultTransformer, neo4j.ExecuteQueryWithDatabase("neo4j")); err != nil { + return err + } + // 孤儿实体清理：本 kb 内不再连任何关系的实体删除。 + _, err := neo4j.ExecuteQuery(ctx, g.driver, + `MATCH (e:Entity {kb:$kb}) WHERE NOT (e)--() DELETE e`, + map[string]any{"kb": kb}, + neo4j.EagerResultTransformer, neo4j.ExecuteQueryWithDatabase("neo4j")) + return err +} + // search 图谱召回：找查询里提到的实体，返回其相连三元组（文本化）。 func (g *graphStore) search(ctx context.Context, kb, query string, limit int) []Hit { if !g.ready() || query == "" { diff --git a/sundynix-mcp-go/internal/rag/ingest_large.go b/sundynix-mcp-go/internal/rag/ingest_large.go new file mode 100644 index 0000000..ec0fb07 --- /dev/null +++ b/sundynix-mcp-go/internal/rag/ingest_large.go @@ -0,0 +1,163 @@ +package rag + +import ( + "context" + "fmt" + "log" + "strings" + "sync" + + "github.com/sundynix/sundynix-shared/contract" +) + +// embedAll 并发分批把 chunks 向量化，保序返回（vecs[i] 对应 chunks[i]）。 +// 串行版几十万字上千块 = 上百次顺序 HTTP，慢且阻塞；这里限并发 embedConcurrency 跑， +// 各批写入预分配切片的对应区间保持顺序，任一批失败即整体失败（上层会清理重入）。 +func (e *Engine) embedAll(ctx context.Context, chunks []string, emit func(contract.IngestEvent)) ([][]float32, error) { + vecs := make([][]float32, len(chunks)) + sem := make(chan struct{}, embedConcurrency) + var wg sync.WaitGroup + var mu sync.Mutex + var firstErr error + done := 0 + + for start := 0; start < len(chunks); start += embedBatch { + end := min(start+embedBatch, len(chunks)) + wg.Add(1) + sem <- struct{}{} + go func(start, end int) { + defer wg.Done() + defer func() { <-sem }() + mu.Lock() + stop := firstErr != nil + mu.Unlock() + if stop { + return // 已有批失败，不再发起新请求 + } + bv, err := e.embed().Embed(ctx, chunks[start:end]) + mu.Lock() + defer mu.Unlock() + if err != nil { + if firstErr == nil { + firstErr = err + } + return + } + copy(vecs[start:end], bv) + done += end - start + emit(contract.IngestEvent{Stage: "向量化", Done: done, Total: len(chunks)}) + }(start, end) + } + wg.Wait() + if firstErr != nil { + return nil, firstErr + } + return vecs, nil +} + +// packGraphWindows 把已切好的语义块合并成抽取窗口（窗内拼接，每窗 ~target 字）。 +// 复用语义块边界，避免再切一遍；超出 maxWindows 截断并回报 truncated=true（调用方告警）。 +func packGraphWindows(chunks []string, target, maxWindows int) (windows []string, truncated bool) { + var cur strings.Builder + curLen := 0 + flush := func() { + if curLen > 0 { + windows = append(windows, cur.String()) + cur.Reset() + curLen = 0 + } + } + for _, c := range chunks { + cl := runeLen(c) + if curLen > 0 && curLen+cl > target { + flush() + if len(windows) >= maxWindows { + return windows, true + } + } + if curLen > 0 { + cur.WriteByte('\n') + curLen++ + } + cur.WriteString(c) + curLen += cl + } + flush() + return windows, false +} + +// extractGraphWindowed 把全文按窗口并发抽三元组并 MERGE 进 Neo4j。 +// 取代"整篇喂 LLM"——几十万字也能全覆盖；实体按 kb+name 在 Neo4j 天然去重， +// 这里再做一次内存去重减少重复 MERGE。单窗失败只丢该窗，不影响其余。 +func (e *Engine) extractGraphWindowed(ctx context.Context, kb, fileID string, chunks []string, emit func(contract.IngestEvent)) { + windows, truncated := packGraphWindows(chunks, graphWindowRunes(), graphMaxWindows()) + if len(windows) == 0 { + return + } + if truncated { + log.Printf("[rag] 文档过大，图谱仅抽取前 %d 窗（其余略过）kb=%s", len(windows), kb) + emit(contract.IngestEvent{Stage: "抽实体", Msg: fmt.Sprintf("文档较大，图谱抽取前 %d 段", len(windows))}) + } + emit(contract.IngestEvent{Stage: "抽实体", Done: 0, Total: len(windows), Msg: fmt.Sprintf("LLM 分 %d 段抽取知识三元组…", len(windows))}) + + chat := e.graphChatClient() + sem := make(chan struct{}, graphConcurrency()) + var wg sync.WaitGroup + var mu sync.Mutex + var all []Triple + doneWin := 0 + + for _, w := range windows { + wg.Add(1) + sem <- struct{}{} + go func(w string) { + defer wg.Done() + defer func() { <-sem }() + ts, err := extractTriples(ctx, chat, w) + mu.Lock() + defer mu.Unlock() + doneWin++ + if err != nil { + log.Printf("[rag] 窗口三元组抽取失败（忽略该窗）: %v", err) + } else { + all = append(all, ts...) + } + emit(contract.IngestEvent{Stage: "抽实体", Done: doneWin, Total: len(windows), Msg: fmt.Sprintf("图谱抽取 %d/%d 段", doneWin, len(windows))}) + }(w) + } + wg.Wait() + + all = dedupeTriples(all) + if len(all) == 0 { + return + } + // 实时回流给 UI（边出现边渲染图谱）。 + tv := make([]contract.TripleView, len(all)) + for i, t := range all { + tv[i] = contract.TripleView{S: t.S, P: t.P, O: t.O} + } + emit(contract.IngestEvent{Stage: "抽实体", Total: len(all), Triples: tv, Msg: fmt.Sprintf("抽出 %d 条知识三元组", len(all))}) + emit(contract.IngestEvent{Stage: "写Neo4j", Total: len(all), Msg: fmt.Sprintf("%d 条三元组写入图谱", len(all))}) + if n, gerr := e.graph.store(ctx, kb, fileID, all); gerr != nil { + log.Printf("[rag] 写 Neo4j 失败（图谱降级）: %v", gerr) + } else { + log.Printf("[rag] 图谱: 写入 %d 条三元组到 kb=%s（%d 窗）", n, kb, len(windows)) + } +} + +// dedupeTriples 去掉空项与重复的 (s,p,o)，减少跨窗重复 MERGE。 +func dedupeTriples(ts []Triple) []Triple { + seen := make(map[string]bool, len(ts)) + out := ts[:0] + for _, t := range ts { + if t.S == "" || t.P == "" || t.O == "" { + continue + } + k := t.S + "\x00" + t.P + "\x00" + t.O + if !seen[k] { + seen[k] = true + out = append(out, t) + } + } + return out +} diff --git a/sundynix-mcp-go/internal/rag/ingest_large_test.go b/sundynix-mcp-go/internal/rag/ingest_large_test.go new file mode 100644 index 0000000..5e9bd97 --- /dev/null +++ b/sundynix-mcp-go/internal/rag/ingest_large_test.go @@ -0,0 +1,108 @@ +package rag + +import ( + "strings" + "testing" +) + +// TestPackGraphWindows_Groups 把多块合并成 ~target 的窗口，窗口数远少于块数。 +func TestPackGraphWindows_Groups(t *testing.T) { + var chunks []string + for i := 0; i < 50; i++ { + chunks = append(chunks, strings.Repeat("字", 200)) // 每块 200 字 + } + // target=4000 → 每窗约 20 块 → ~3 窗（不再是 50 次 LLM 调用）。 + windows, truncated := packGraphWindows(chunks, 4000, 60) + if truncated { + t.Fatal("50×200=1万字不该触发封顶") + } + if len(windows) >= len(chunks) { + t.Fatalf("窗口应远少于块数：windows=%d chunks=%d", len(windows), len(chunks)) + } + for i, w := range windows { + if rl := runeLen(w); rl > 4000+200 { // 容一块的溢出 + t.Fatalf("窗口 %d 超目标过多：%d 字", i, rl) + } + } +} + +// TestPackGraphWindows_Truncate 超大文档封顶窗口数并报 truncated。 +func TestPackGraphWindows_Truncate(t *testing.T) { + var chunks []string + for i := 0; i < 1000; i++ { + chunks = append(chunks, strings.Repeat("字", 500)) + } + windows, truncated := packGraphWindows(chunks, 4000, 60) + if !truncated { + t.Fatal("50万字应触发封顶 truncated=true") + } + if len(windows) > 60 { + t.Fatalf("窗口数应封顶 ≤60，got %d", len(windows)) + } +} + +// TestDedupeTriples 去重 + 去空项。 +func TestDedupeTriples(t *testing.T) { + in := []Triple{ + {S: "李明华", P: "是", O: "总设计师"}, + {S: "李明华", P: "是", O: "总设计师"}, // 重复 + {S: "星云一号", P: "发射于", O: "2023"}, + {S: "", P: "x", O: "y"}, // 空主体 → 丢 + {S: "a", P: "", O: "b"}, // 空关系 → 丢 + {S: "李明华", P: "领导", O: "团队"}, // 同主体不同关系 → 保留 + } + out := dedupeTriples(in) + if len(out) != 3 { + t.Fatalf("应剩 3 条（去 1 重复 + 2 空），got %d: %+v", len(out), out) + } +} + +// TestEnvInt env 正整数读取：缺省/0/非法都退 def。 +func TestEnvInt(t *testing.T) { + if envInt("NO_SUCH_ENV_XZ", 7) != 7 { + t.Fatal("缺省应返回 def") + } + t.Setenv("X_TEST_INT", "12") + if envInt("X_TEST_INT", 7) != 12 { + t.Fatal("应读 env=12") + } + t.Setenv("X_TEST_INT", "0") + if envInt("X_TEST_INT", 7) != 7 { + t.Fatal("非正数应退 def") + } + t.Setenv("X_TEST_INT", "abc") + if envInt("X_TEST_INT", 7) != 7 { + t.Fatal("非法应退 def") + } +} + +// TestGraphKnobs_EnvConfigurable 图谱封顶/窗口可经 env 调（瓶颈①）。 +func TestGraphKnobs_EnvConfigurable(t *testing.T) { + t.Setenv("GRAPH_MAX_WINDOWS", "2") + if graphMaxWindows() != 2 { + t.Fatalf("env 应覆盖封顶为 2，got %d", graphMaxWindows()) + } + if graphWindowRunes() != 4000 { + t.Fatalf("未设时默认窗口应 4000，got %d", graphWindowRunes()) + } +} + +// TestGraphChatClient_FallsBackToMain 未配专用图谱模型 → 回退主 chat；配了 → 用专用（瓶颈⑤）。 +func TestGraphChatClient_FallsBackToMain(t *testing.T) { + e := &Engine{chat: newChatClient("http://main", "k", "m")} + if e.graphChatClient() != e.chat { + t.Fatal("未配专用模型应回退主 chat") + } + e.graphChat = newChatClient("http://cheap", "k", "cheap") + if e.graphChatClient() != e.graphChat { + t.Fatal("配了专用模型应用专用") + } +} + +// TestPackGraphWindows_Empty 空输入不 panic。 +func TestPackGraphWindows_Empty(t *testing.T) { + w, trunc := packGraphWindows(nil, 4000, 60) + if len(w) != 0 || trunc { + t.Fatalf("空输入应得空窗口、不截断，got %d %v", len(w), trunc) + } +} diff --git a/sundynix-mcp-go/internal/rag/milvus.go b/sundynix-mcp-go/internal/rag/milvus.go index e9267bb..129cde2 100644 --- a/sundynix-mcp-go/internal/rag/milvus.go +++ b/sundynix-mcp-go/internal/rag/milvus.go @@ -115,6 +115,17 @@ func (m *milvusStore) deleteDoc(ctx context.Context, kb, doc string, dim int) { } } +// deleteByFile 按 file_id 删某文档的全部向量块（级联删用，无需 dim）。集合不存在/未加载时静默忽略。 +func (m *milvusStore) deleteByFile(ctx context.Context, kb, fileID string) { + if fileID == "" { + return + } + expr := fmt.Sprintf("kb == %q && doc == %q", kb, fileID) + if err := m.cli.Delete(ctx, collection, "", expr); err != nil { + log.Printf("[rag] 按 file_id 删除向量失败（忽略）: %v", err) + } +} + // insert 写入若干 (kb, doc, text, vector)。 // 若集合在运行期被丢失（如 Milvus 重启）→ 清缓存、重建集合后重试一次，避免必须重启进程才能恢复。 func (m *milvusStore) insert(ctx context.Context, kb, doc string, texts []string, vecs [][]float32) error { diff --git a/sundynix-mcp-go/internal/rag/rag.go b/sundynix-mcp-go/internal/rag/rag.go index f535fde..e3ea870 100644 --- a/sundynix-mcp-go/internal/rag/rag.go +++ b/sundynix-mcp-go/internal/rag/rag.go @@ -6,13 +6,33 @@ import ( "context" "errors" "log" + "os" + "strconv" "sync" "github.com/sundynix/sundynix-shared/contract" ) -// embedBatch 是每批向量化的块数（让大文件的入库进度可观测）。 -const embedBatch = 10 +const ( + embedBatch = 10 // 每批向量化的块数（兼顾进度可观测 + provider 单批上限） + embedConcurrency = 4 // 并发批数：几十万字上千块串行太慢，限并发跑快数倍且不打爆 provider 速率 +) + +// 图谱抽取窗口化参数（env 可配，便于按文档体量/成本调）：取代"整篇喂 LLM"（几十万字必爆上下文）。 +// 把已切的语义块合并成 ~graphWindowRunes 的窗口，逐窗并发抽三元组；实体按 kb+name 在 Neo4j 去重。 +func graphWindowRunes() int { return envInt("GRAPH_WINDOW_RUNES", 4000) } // 每窗字数（远小于模型上下文） +func graphMaxWindows() int { return envInt("GRAPH_MAX_WINDOWS", 60) } // 封顶窗口数（控成本；大文档超出略过+告警） +func graphConcurrency() int { return envInt("GRAPH_CONCURRENCY", 3) } // 并发抽取窗口数 + +// envInt 读正整数 env，缺省/非法时返回 def。 +func envInt(key string, def int) int { + if v := os.Getenv(key); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + return n + } + } + return def +} // Config 是 RAG 引擎的初始化配置。 type Config struct { @@ -25,13 +45,14 @@ type Config struct { // Engine 聚合 embedding + Milvus(向量) + Bleve(全文) + Neo4j(图谱) → RRF 融合 + 可选 rerank。 // embedding 与 chat(图谱抽取用)可热更新（控制面下发）。 type Engine struct { - mu sync.RWMutex - emb *embedClient - chat *chatClient - mv *milvusStore - bleve *bleveStore - rerank *rerankClient - graph *graphStore + mu sync.RWMutex + emb *embedClient + chat *chatClient // 控制面下发的主对话模型 + graphChat *chatClient // 可选：图谱抽取专用便宜模型（env GRAPH_CHAT_*）；未配则回退主 chat + mv *milvusStore + bleve *bleveStore + rerank *rerankClient + graph *graphStore } // SetEmbedding 热更新 embedding 配置（控制面变更时调用）。空配置=关闭向量检索。 @@ -68,6 +89,17 @@ func (e *Engine) chatClient() *chatClient { return e.chat } +// graphChatClient 返回图谱抽取用的模型：优先专用便宜模型(env GRAPH_CHAT_*)，未配则回退主 chat。 +// 让"每篇最多 60 次图谱抽取"走便宜模型，主对话仍用好模型——成本与质量解耦。 +func (e *Engine) graphChatClient() *chatClient { + e.mu.RLock() + defer e.mu.RUnlock() + if e.graphChat.ready() { + return e.graphChat + } + return e.chat +} + // Open 建立 RAG 引擎。各路连不上 → 降级（不阻断工具服务）。 func Open(ctx context.Context, cfg Config) *Engine { e := &Engine{ @@ -75,6 +107,13 @@ func Open(ctx context.Context, cfg Config) *Engine { rerank: newRerankClient(cfg.RerankBase, cfg.RerankKey, cfg.RerankModel), graph: openGraph(ctx, cfg.Neo4jURI, cfg.Neo4jUser, cfg.Neo4jPass), } + // 图谱抽取专用便宜模型（可选，env GRAPH_CHAT_*）：未配则图谱抽取回退控制面主 chat。 + if gb := os.Getenv("GRAPH_CHAT_BASE"); gb != "" { + e.graphChat = newChatClient(gb, os.Getenv("GRAPH_CHAT_KEY"), os.Getenv("GRAPH_CHAT_MODEL")) + if e.graphChat.ready() { + log.Printf("[rag] 图谱抽取专用便宜模型: %s model=%s", gb, os.Getenv("GRAPH_CHAT_MODEL")) + } + } if e.rerank.ready() { log.Printf("[rag] rerank: %s model=%s", cfg.RerankBase, cfg.RerankModel) } @@ -130,17 +169,11 @@ func (e *Engine) Ingest(ctx context.Context, kb, doc, text string, onProgress fu } emit(contract.IngestEvent{Stage: "切块", Total: len(chunks), Chunks: previews(chunks), Msg: "拆为 " + itoa(len(chunks)) + " 块"}) - // 分批向量化（逐批回报进度）。 - vecs := make([][]float32, 0, len(chunks)) - for i := 0; i < len(chunks); i += embedBatch { - end := min(i+embedBatch, len(chunks)) - bv, err := e.embed().Embed(ctx, chunks[i:end]) - if err != nil { - emit(contract.IngestEvent{Stage: "失败", Error: "向量化: " + err.Error()}) - return 0, err - } - vecs = append(vecs, bv...) - emit(contract.IngestEvent{Stage: "向量化", Done: end, Total: len(chunks)}) + // 并发分批向量化（保序，逐批回报进度）—— 几十万字上千块时比串行快数倍。 + vecs, err := e.embedAll(ctx, chunks, emit) + if err != nil { + emit(contract.IngestEvent{Stage: "失败", Error: "向量化: " + err.Error()}) + return 0, err } emit(contract.IngestEvent{Stage: "写Milvus", Msg: "向量库写入中"}) @@ -155,26 +188,10 @@ func (e *Engine) Ingest(ctx context.Context, kb, doc, text string, onProgress fu e.bleve.deleteDoc(kb, doc) _ = e.bleve.index(kb, doc, chunks) // 同步写全文索引（失败不阻断向量入库） - // 图谱路：LLM 抽实体/关系 → Neo4j（可降级，不阻断向量入库）。 - if e.graph.ready() && e.chatClient().ready() { - emit(contract.IngestEvent{Stage: "抽实体", Msg: "LLM 正在抽取知识三元组…"}) - triples, terr := extractTriples(ctx, e.chatClient(), text) - if terr != nil { - log.Printf("[rag] 三元组抽取失败（图谱降级）: %v", terr) - } else if len(triples) > 0 { - // 把抽出的三元组实时回流给 UI（边出现边渲染图谱）。 - tv := make([]contract.TripleView, len(triples)) - for i, t := range triples { - tv[i] = contract.TripleView{S: t.S, P: t.P, O: t.O} - } - emit(contract.IngestEvent{Stage: "抽实体", Total: len(triples), Triples: tv, Msg: "抽出 " + itoa(len(triples)) + " 条知识三元组"}) - emit(contract.IngestEvent{Stage: "写Neo4j", Total: len(triples), Msg: itoa(len(triples)) + " 条三元组写入图谱"}) - if n, gerr := e.graph.store(ctx, kb, triples); gerr != nil { - log.Printf("[rag] 写 Neo4j 失败（图谱降级）: %v", gerr) - } else { - log.Printf("[rag] 图谱: 写入 %d 条三元组到 kb=%s", n, kb) - } - } + // 图谱路：窗口化并发抽实体/关系 → Neo4j（可降级，不阻断向量入库）。 + // 几十万字按 ~graphWindowRunes 分窗逐窗抽，全覆盖；实体在 Neo4j 按 kb+name 去重。 + if e.graph.ready() && e.graphChatClient().ready() { + e.extractGraphWindowed(ctx, kb, doc, chunks, emit) // doc=file_id：图谱关系按它标源，供级联删 } return len(chunks), nil @@ -244,10 +261,29 @@ func (e *Engine) Search(ctx context.Context, kb, query string, topK int) ([]Hit, return cand, nil } +// DeleteDoc 按 file_id 级联删某文档在三库的痕迹：Milvus 向量块 + Bleve 全文块 + Neo4j 关系。 +// Milvus/Bleve 尽力删（void），Neo4j 返回错误供上层感知；任一失败不阻断其余。 +func (e *Engine) DeleteDoc(ctx context.Context, kb, fileID string) error { + if fileID == "" { + return errors.New("file_id 必填") + } + e.mv.deleteByFile(ctx, kb, fileID) + e.bleve.deleteDoc(kb, fileID) + if err := e.graph.deleteByFile(ctx, kb, fileID); err != nil { + log.Printf("[rag] 图谱按 file_id 删除失败: %v", err) + return err + } + log.Printf("[rag] 已删除文档痕迹 kb=%s file_id=%s（向量/全文/图谱）", kb, fileID) + return nil +} + func (e *Engine) Close() { if e.mv != nil { e.mv.close() } + if e.bleve != nil { + e.bleve.close() // 落盘版释放锁并刷盘 + } e.graph.close(context.Background()) }