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:
@@ -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)}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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() {
|
||||||
|
|||||||
@@ -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)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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("窗口数应封顶 ≤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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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 {
|
||||||
|
|||||||
@@ -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())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user