feat(rag): 大文件 RAG 做深 —— 并发入库/图谱窗口化 + Bleve落盘 + file_id治理

几十万字文件从"广而浅"到准生产级:
- 向量化串行→并发分批保序(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 <noreply@anthropic.com>
This commit is contained in:
Blizzard
2026-06-30 13:37:45 +08:00
parent 6f16062dfc
commit 79e834e8e9
8 changed files with 462 additions and 46 deletions
+14
View File
@@ -162,6 +162,7 @@ func (g *Gateway) buildRegistry() map[string]toolDef {
// —— 仅内部/流水线/管理用,不暴露给自主 agent —— // —— 仅内部/流水线/管理用,不暴露给自主 agent ——
"kb_ingest": {cn: "知识入库", desc: "文本切块 → 向量化 → 写入 Milvus / Bleve", handler: g.kbIngest}, "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_search": {cn: "检索台查询", desc: "结构化返回命中内容与相似度分数", handler: g.kbSearch},
"kb_graph": {cn: "知识图谱", desc: "取某库的实体关系三元组(Neo4j)", handler: g.kbGraph}, "kb_graph": {cn: "知识图谱", desc: "取某库的实体关系三元组(Neo4j)", handler: g.kbGraph},
"report_render": {cn: "报告渲染", desc: "把结构化报告渲染为 Word.docx", handler: g.reportRender}, "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)} 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)}
}
+43 -3
View File
@@ -4,6 +4,8 @@ import (
"fmt" "fmt"
"hash/fnv" "hash/fnv"
"log" "log"
"os"
"path/filepath"
"github.com/blevesearch/bleve/v2" "github.com/blevesearch/bleve/v2"
"github.com/blevesearch/bleve/v2/analysis/analyzer/keyword" "github.com/blevesearch/bleve/v2/analysis/analyzer/keyword"
@@ -12,8 +14,9 @@ import (
"github.com/blevesearch/bleve/v2/search/query" "github.com/blevesearch/bleve/v2/search/query"
) )
// bleveStore 是全文(BM25)检索路。内存索引:随 ingest 写入,进程重启重建。 // bleveStore 是全文(BM25)检索路。落盘索引(scorch:随 ingest 写入并持久,进程重启不丢,
// 真实生产应落盘(bleve.New(path,...));此处内存优先求简 // 与 Milvus(向量)/Neo4j(图谱) 同为持久存储——保证"三路混合"重启后仍是三路
// 落盘不可用时退回内存索引(功能在、重启会丢),好过全文路全哑。
type bleveStore struct { type bleveStore struct {
idx bleve.Index idx bleve.Index
} }
@@ -34,17 +37,54 @@ func bleveMapping() mapping.IndexMapping {
return im 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 { 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()) idx, err := bleve.NewMemOnly(bleveMapping())
if err != nil { if err != nil {
log.Printf("[rag] bleve 初始化失败,全文路降级: %v", err) log.Printf("[rag] bleve 内存兜底也失败,全文路降级: %v", err)
return &bleveStore{} return &bleveStore{}
} }
log.Printf("[rag] bleve 退回内存索引(非持久)")
return &bleveStore{idx: idx} return &bleveStore{idx: idx}
} }
func (b *bleveStore) ready() bool { return b != nil && b.idx != nil } 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+文本哈希,幂等)。 // index 把 (kb, doc, texts) 写入全文索引(id 含 kb+doc+文本哈希,幂等)。
func (b *bleveStore) index(kb, doc string, texts []string) error { func (b *bleveStore) index(kb, doc string, texts []string) error {
if !b.ready() { if !b.ready() {
+22
View File
@@ -5,7 +5,9 @@ import "testing"
// TestBleve_ChineseSearch 钉死中文全文检索:CJK 分词器让"星云一号的总设计师"能命中含 // TestBleve_ChineseSearch 钉死中文全文检索:CJK 分词器让"星云一号的总设计师"能命中含
// "星云一号""总设计师"的中文块。修复前默认标准分词器把整段中文当一个 token → 永远 0 命中。 // "星云一号""总设计师"的中文块。修复前默认标准分词器把整段中文当一个 token → 永远 0 命中。
func TestBleve_ChineseSearch(t *testing.T) { func TestBleve_ChineseSearch(t *testing.T) {
t.Setenv("BLEVE_PATH", t.TempDir()+"/bleve") // 隔离落盘索引,测试互不污染
b := openBleve() b := openBleve()
defer b.close()
if !b.ready() { if !b.ready() {
t.Skip("bleve 不可用") 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 { func contains(s, sub string) bool {
return len(sub) == 0 || (len(s) >= len(sub) && indexOf(s, sub) >= 0) return len(sub) == 0 || (len(s) >= len(sub) && indexOf(s, sub) >= 0)
} }
+25 -3
View File
@@ -56,7 +56,9 @@ func (g *graphStore) close(ctx context.Context) {
} }
// store 把三元组 MERGE 进 Neo4j(实体 + 关系,按 kb 隔离)。 // 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() { if !g.ready() {
return 0, nil 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, _, err := neo4j.ExecuteQuery(ctx, g.driver,
`MERGE (a:Entity {kb:$kb, name:$s}) `MERGE (a:Entity {kb:$kb, name:$s})
MERGE (b:Entity {kb:$kb, name:$o}) MERGE (b:Entity {kb:$kb, name:$o})
MERGE (a)-[r:REL {type:$p}]->(b)`, MERGE (a)-[r:REL {type:$p, file_id:$fid}]->(b)`,
map[string]any{"kb": kb, "s": t.S, "o": t.O, "p": t.P}, map[string]any{"kb": kb, "s": t.S, "o": t.O, "p": t.P, "fid": fileID},
neo4j.EagerResultTransformer, neo4j.ExecuteQueryWithDatabase("neo4j")) neo4j.EagerResultTransformer, neo4j.ExecuteQueryWithDatabase("neo4j"))
if err != nil { if err != nil {
return n, err return n, err
@@ -79,6 +81,26 @@ func (g *graphStore) store(ctx context.Context, kb string, triples []Triple) (in
return n, nil 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 图谱召回:找查询里提到的实体,返回其相连三元组(文本化)。 // search 图谱召回:找查询里提到的实体,返回其相连三元组(文本化)。
func (g *graphStore) search(ctx context.Context, kb, query string, limit int) []Hit { func (g *graphStore) search(ctx context.Context, kb, query string, limit int) []Hit {
if !g.ready() || query == "" { if !g.ready() || query == "" {
@@ -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
}
@@ -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("窗口数应封顶 ≤60got %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 应覆盖封顶为 2got %d", graphMaxWindows())
}
if graphWindowRunes() != 4000 {
t.Fatalf("未设时默认窗口应 4000got %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)
}
}
+11
View File
@@ -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)。 // insert 写入若干 (kb, doc, text, vector)。
// 若集合在运行期被丢失(如 Milvus 重启)→ 清缓存、重建集合后重试一次,避免必须重启进程才能恢复。 // 若集合在运行期被丢失(如 Milvus 重启)→ 清缓存、重建集合后重试一次,避免必须重启进程才能恢复。
func (m *milvusStore) insert(ctx context.Context, kb, doc string, texts []string, vecs [][]float32) error { func (m *milvusStore) insert(ctx context.Context, kb, doc string, texts []string, vecs [][]float32) error {
+76 -40
View File
@@ -6,13 +6,33 @@ import (
"context" "context"
"errors" "errors"
"log" "log"
"os"
"strconv"
"sync" "sync"
"github.com/sundynix/sundynix-shared/contract" "github.com/sundynix/sundynix-shared/contract"
) )
// embedBatch 是每批向量化的块数(让大文件的入库进度可观测)。 const (
const embedBatch = 10 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 引擎的初始化配置。 // Config 是 RAG 引擎的初始化配置。
type Config struct { type Config struct {
@@ -25,13 +45,14 @@ type Config struct {
// Engine 聚合 embedding + Milvus(向量) + Bleve(全文) + Neo4j(图谱) → RRF 融合 + 可选 rerank。 // Engine 聚合 embedding + Milvus(向量) + Bleve(全文) + Neo4j(图谱) → RRF 融合 + 可选 rerank。
// embedding 与 chat(图谱抽取用)可热更新(控制面下发)。 // embedding 与 chat(图谱抽取用)可热更新(控制面下发)。
type Engine struct { type Engine struct {
mu sync.RWMutex mu sync.RWMutex
emb *embedClient emb *embedClient
chat *chatClient chat *chatClient // 控制面下发的主对话模型
mv *milvusStore graphChat *chatClient // 可选:图谱抽取专用便宜模型(env GRAPH_CHAT_*);未配则回退主 chat
bleve *bleveStore mv *milvusStore
rerank *rerankClient bleve *bleveStore
graph *graphStore rerank *rerankClient
graph *graphStore
} }
// SetEmbedding 热更新 embedding 配置(控制面变更时调用)。空配置=关闭向量检索。 // SetEmbedding 热更新 embedding 配置(控制面变更时调用)。空配置=关闭向量检索。
@@ -68,6 +89,17 @@ func (e *Engine) chatClient() *chatClient {
return e.chat 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 引擎。各路连不上 → 降级(不阻断工具服务)。 // Open 建立 RAG 引擎。各路连不上 → 降级(不阻断工具服务)。
func Open(ctx context.Context, cfg Config) *Engine { func Open(ctx context.Context, cfg Config) *Engine {
e := &Engine{ e := &Engine{
@@ -75,6 +107,13 @@ func Open(ctx context.Context, cfg Config) *Engine {
rerank: newRerankClient(cfg.RerankBase, cfg.RerankKey, cfg.RerankModel), rerank: newRerankClient(cfg.RerankBase, cfg.RerankKey, cfg.RerankModel),
graph: openGraph(ctx, cfg.Neo4jURI, cfg.Neo4jUser, cfg.Neo4jPass), 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() { if e.rerank.ready() {
log.Printf("[rag] rerank: %s model=%s", cfg.RerankBase, cfg.RerankModel) 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)) + " 块"}) emit(contract.IngestEvent{Stage: "切块", Total: len(chunks), Chunks: previews(chunks), Msg: "拆为 " + itoa(len(chunks)) + " 块"})
// 分批向量化(逐批回报进度)。 // 并发分批向量化(保序,逐批回报进度)—— 几十万字上千块时比串行快数倍
vecs := make([][]float32, 0, len(chunks)) vecs, err := e.embedAll(ctx, chunks, emit)
for i := 0; i < len(chunks); i += embedBatch { if err != nil {
end := min(i+embedBatch, len(chunks)) emit(contract.IngestEvent{Stage: "失败", Error: "向量化: " + err.Error()})
bv, err := e.embed().Embed(ctx, chunks[i:end]) return 0, err
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)})
} }
emit(contract.IngestEvent{Stage: "写Milvus", Msg: "向量库写入中"}) 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.deleteDoc(kb, doc)
_ = e.bleve.index(kb, doc, chunks) // 同步写全文索引(失败不阻断向量入库) _ = e.bleve.index(kb, doc, chunks) // 同步写全文索引(失败不阻断向量入库)
// 图谱路:LLM 抽实体/关系 → Neo4j(可降级,不阻断向量入库)。 // 图谱路:窗口化并发抽实体/关系 → Neo4j(可降级,不阻断向量入库)。
if e.graph.ready() && e.chatClient().ready() { // 几十万字按 ~graphWindowRunes 分窗逐窗抽,全覆盖;实体在 Neo4j 按 kb+name 去重。
emit(contract.IngestEvent{Stage: "抽实体", Msg: "LLM 正在抽取知识三元组…"}) if e.graph.ready() && e.graphChatClient().ready() {
triples, terr := extractTriples(ctx, e.chatClient(), text) e.extractGraphWindowed(ctx, kb, doc, chunks, emit) // doc=file_id:图谱关系按它标源,供级联删
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)
}
}
} }
return len(chunks), nil return len(chunks), nil
@@ -244,10 +261,29 @@ func (e *Engine) Search(ctx context.Context, kb, query string, topK int) ([]Hit,
return cand, nil 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() { func (e *Engine) Close() {
if e.mv != nil { if e.mv != nil {
e.mv.close() e.mv.close()
} }
if e.bleve != nil {
e.bleve.close() // 落盘版释放锁并刷盘
}
e.graph.close(context.Background()) e.graph.close(context.Background())
} }