weichen
4fdbfbc9a2
Details (original message): Add the vector RAG pipeline: paragraph chunker, OpenAI-compatible embeddings and rerank clients, dual-tier vector store with cosine search, and RRF hybrid retrieval
88 lines
2.0 KiB
Go
88 lines
2.0 KiB
Go
// Package rag 实现检索管线（设计 4.12）：分块 →（可选）Embedding → 全文+向量混合
|
|
// 检索（RRF 融合）→（可选）Rerank → TopK。embedding 未配置时自动退化为全文档。
|
|
package rag
|
|
|
|
import (
|
|
"strings"
|
|
|
|
core "agent-system-core"
|
|
)
|
|
|
|
// FixedChunker 段落感知定长分块：优先沿空行段落边界聚包，段落超长按窗口硬切，
|
|
// 相邻块保留 Overlap 个字符上下文（runes 计）。
|
|
type FixedChunker struct {
|
|
Target int // 目标块长（runes）；<=0 用默认 800
|
|
Overlap int // 硬切重叠（runes）；<0 用默认 120
|
|
}
|
|
|
|
func (c FixedChunker) norm() (int, int) {
|
|
t, o := c.Target, c.Overlap
|
|
if t <= 0 {
|
|
t = 800
|
|
}
|
|
if o < 0 {
|
|
o = 120
|
|
}
|
|
if o >= t {
|
|
o = t / 4
|
|
}
|
|
return t, o
|
|
}
|
|
|
|
func (c FixedChunker) Chunk(doc core.ParsedDocument) []core.Chunk {
|
|
target, overlap := c.norm()
|
|
return chunkText(doc.PlainText(), target, overlap)
|
|
}
|
|
|
|
func chunkText(text string, target, overlap int) []core.Chunk {
|
|
paras := splitParas(text)
|
|
var chunks []core.Chunk
|
|
cur := []rune{}
|
|
flush := func() {
|
|
if len(cur) > 0 {
|
|
chunks = append(chunks, core.Chunk{Text: string(cur), Index: len(chunks)})
|
|
cur = nil
|
|
}
|
|
}
|
|
for _, p := range paras {
|
|
pr := []rune(p)
|
|
if len(pr) > target {
|
|
// 超长段落：先封当前块，再按窗口硬切（带重叠）。
|
|
flush()
|
|
for start := 0; start < len(pr); start += target - overlap {
|
|
end := start + target
|
|
if end > len(pr) {
|
|
end = len(pr)
|
|
}
|
|
chunks = append(chunks, core.Chunk{Text: string(pr[start:end]), Index: len(chunks)})
|
|
if end == len(pr) {
|
|
break
|
|
}
|
|
}
|
|
continue
|
|
}
|
|
if len(cur)+len(pr)+1 > target {
|
|
flush()
|
|
}
|
|
if len(cur) > 0 {
|
|
cur = append(cur, '\n')
|
|
}
|
|
cur = append(cur, pr...)
|
|
}
|
|
flush()
|
|
return chunks
|
|
}
|
|
|
|
func splitParas(text string) []string {
|
|
text = strings.ReplaceAll(text, "\r\n", "\n")
|
|
raw := strings.Split(text, "\n\n")
|
|
var out []string
|
|
for _, p := range raw {
|
|
p = strings.TrimSpace(p)
|
|
if p != "" {
|
|
out = append(out, p)
|
|
}
|
|
}
|
|
return out
|
|
}
|