443 lines
14 KiB
Go
443 lines
14 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"regexp"
|
||
"strings"
|
||
"unicode/utf8"
|
||
|
||
"rag-local/common"
|
||
"rag-local/kb/consts"
|
||
"rag-local/kb/dao"
|
||
"rag-local/kb/model/domain"
|
||
"rag-local/kb/model/entity"
|
||
|
||
"github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive"
|
||
"github.com/cloudwego/eino-ext/components/document/transformer/splitter/semantic"
|
||
"github.com/cloudwego/eino/components/document"
|
||
eembedding "github.com/cloudwego/eino/components/embedding"
|
||
"github.com/cloudwego/eino/schema"
|
||
"github.com/gogf/gf/v2/errors/gerror"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
)
|
||
|
||
// 中文分块分隔符：先段落/换行，再句读标点；runeLen 保证按字计数（与标题策略一致）
|
||
var chunkSeparators = []string{"\n\n", "\n", "。", "！", "？", "；", "，", " ", ""}
|
||
var runeLen = utf8.RuneCountInString
|
||
|
||
// SplitAuto 内置组合分块管线（无需用户选择策略），按文档内容自动分派：
|
||
// 1. 结构识别（条文/条款/章回/编号等，行首命中达标）→ 结构化切分，带上下文前缀
|
||
// 2. 标题感知（# / 标题 / 第X章 行首）→ 标题并入段落切分
|
||
// 3. 语义分块（无结构无标题，需 embedder）→ 按句子语义切分，超限块递归兜底
|
||
// 4. 递归兜底（embedder 不可用或语义失败）→ 中文分隔符递归切分
|
||
//
|
||
// 返回的 unitPattern/ctxPattern 仅结构识别路径有值（自动识别出的结构模式，供落库展示）。
|
||
func (s *chunkService) SplitAuto(ctx context.Context, text string, chunkSize, overlap int, embedder eembedding.Embedder) (chunks []string, unitPattern, ctxPattern string, err error) {
|
||
if unit, ctxPat, ok := s.DetectStructure(text); ok {
|
||
g.Log().Infof(ctx, "分块管线：结构识别命中（%s / %s），按结构单元切分", unit, ctxPat)
|
||
return s.splitStructured(text, unit, ctxPat, chunkSize, overlap), unit, ctxPat, nil
|
||
}
|
||
if hasHeadings(text) {
|
||
g.Log().Infof(ctx, "分块管线：未识别结构，含标题行，走标题感知切分")
|
||
return s.SplitText(text, chunkSize, overlap), "", "", nil
|
||
}
|
||
if embedder != nil {
|
||
g.Log().Infof(ctx, "分块管线：无结构无标题，走语义切分")
|
||
sp, spErr := semantic.NewSplitter(ctx, &semantic.Config{
|
||
Embedding: embedder,
|
||
BufferSize: 1,
|
||
MinChunkSize: chunkSize / 2,
|
||
Separators: chunkSeparators[:len(chunkSeparators)-2],
|
||
LenFunc: runeLen,
|
||
})
|
||
if spErr != nil {
|
||
g.Log().Warningf(ctx, "构建语义分块器失败，回退递归分块: %v", spErr)
|
||
} else {
|
||
semanticChunks, tErr := s.transformText(ctx, sp, text)
|
||
if tErr != nil {
|
||
g.Log().Warningf(ctx, "语义分块失败，回退递归分块: %v", tErr)
|
||
} else if len(semanticChunks) > 0 {
|
||
// chunkSize 作为安全上限：语义分块可能产出巨块（长连续语义段），超限块递归二次切分兜底（不重叠）
|
||
chunks, err = s.capChunks(ctx, semanticChunks, chunkSize)
|
||
return chunks, "", "", err
|
||
}
|
||
}
|
||
}
|
||
g.Log().Infof(ctx, "分块管线：无结构无标题，语义分块不可用，走递归切分")
|
||
chunks, err = s.splitRecursive(ctx, text, chunkSize, overlap)
|
||
return chunks, "", "", err
|
||
}
|
||
|
||
// hasHeadings 文本是否含标题行（Markdown # / 标题 / 第X章），标题感知分块保留标题结构
|
||
func hasHeadings(text string) bool {
|
||
for _, line := range strings.Split(text, "\n") {
|
||
if isHeading(strings.TrimSpace(line)) {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// 结构标记候选：单元模式按优先级排列（条文 > 章回 > 小节 > 章节 > 篇部 > 数字编号 > 括号序号）
|
||
var (
|
||
unitPatternCandidates = []string{
|
||
`第[零一二三四五六七八九十百千]+条`,
|
||
`第[零一二三四五六七八九十百千]+回`,
|
||
`第[零一二三四五六七八九十百千]+节`,
|
||
`第[零一二三四五六七八九十百千]+章`,
|
||
`第[零一二三四五六七八九十百千]+[篇部]`,
|
||
`[一二三四五六七八九十]{1,3}[、.]`,
|
||
`\d{1,3}[、.．]`,
|
||
`[（(][一二三四五六七八九十\d]{1,3}[)）]`,
|
||
}
|
||
ctxPatternCandidates = []string{
|
||
`第[零一二三四五六七八九十百千]+[篇部]`,
|
||
`第[零一二三四五六七八九十百千]+章`,
|
||
`第[零一二三四五六七八九十百千]+节`,
|
||
}
|
||
minUnitHits = 5 // 行首命中至少 5 次才认定是结构单元（排除正文偶然提及）
|
||
minCtxHits = 3 // 上下文标记命中下限
|
||
)
|
||
|
||
// DetectStructure 自动识别文本的结构单元与上下文标记：统计各候选模式在行首的出现次数，
|
||
// 命中数达到阈值且单元优先顺序靠前者胜出；上下文取单元之外层级最高的候选。未识别出单元时 ok=false。
|
||
func (s *chunkService) DetectStructure(text string) (unit, ctx string, ok bool) {
|
||
count := func(pat string) int {
|
||
return len(regexp.MustCompile(`(?m)^\s*`+pat).FindAllStringIndex(text, -1))
|
||
}
|
||
hits := make(map[string]int, len(unitPatternCandidates)+len(ctxPatternCandidates))
|
||
for _, p := range unitPatternCandidates {
|
||
hits[p] = count(p)
|
||
}
|
||
for _, p := range ctxPatternCandidates {
|
||
hits[p] = count(p)
|
||
}
|
||
for _, p := range unitPatternCandidates {
|
||
if hits[p] >= minUnitHits {
|
||
unit = p
|
||
break
|
||
}
|
||
}
|
||
if unit == "" {
|
||
return "", "", false
|
||
}
|
||
for _, p := range ctxPatternCandidates {
|
||
if p == unit {
|
||
continue
|
||
}
|
||
if hits[p] >= minCtxHits {
|
||
ctx = p
|
||
break
|
||
}
|
||
}
|
||
return unit, ctx, true
|
||
}
|
||
|
||
// splitStructured 结构单元切分：按单元模式切分（标记保留在块内），上下文模式命中的整行作为前缀；
|
||
// 相邻短单元合并到 chunkSize/2（上限 300 字），超长单元按句兜底切分（后续片段补「（标记 续）」前缀）。
|
||
func (s *chunkService) splitStructured(text, unitPattern, ctxPattern string, chunkSize, overlap int) []string {
|
||
matches := regexp.MustCompile(`(?m)^\s*(`+unitPattern+`)`).FindAllStringIndex(text, -1)
|
||
if len(matches) == 0 {
|
||
return s.SplitText(text, chunkSize, overlap)
|
||
}
|
||
mergeMin := chunkSize / 2
|
||
if mergeMin > 300 {
|
||
mergeMin = 300
|
||
}
|
||
if mergeMin < 1 {
|
||
mergeMin = 1
|
||
}
|
||
var ctxMatches [][]int
|
||
if ctxPattern != "" {
|
||
ctxMatches = regexp.MustCompile(`(?m)^\s*(`+ctxPattern+`)`).FindAllStringIndex(text, -1)
|
||
}
|
||
// ctxBefore 返回 pos 之前最近的上下文标记整行（如「第三章 劳动合同的解除」）
|
||
ctxBefore := func(pos int) string {
|
||
cur := ""
|
||
for _, m := range ctxMatches {
|
||
if m[0] >= pos {
|
||
break
|
||
}
|
||
lineEnd := strings.IndexByte(text[m[0]:], '\n')
|
||
if lineEnd < 0 {
|
||
lineEnd = len(text) - m[0]
|
||
}
|
||
cur = strings.TrimSpace(text[m[0] : m[0]+lineEnd])
|
||
}
|
||
return cur
|
||
}
|
||
|
||
out := make([]string, 0, 16)
|
||
var buf []string
|
||
bufLen := 0
|
||
bufCtx := ""
|
||
flush := func() {
|
||
if len(buf) == 0 {
|
||
return
|
||
}
|
||
body := strings.Join(buf, "\n")
|
||
if bufCtx != "" {
|
||
body = "【" + bufCtx + "】" + body
|
||
}
|
||
out = append(out, body)
|
||
buf, bufLen, bufCtx = nil, 0, ""
|
||
}
|
||
for i, m := range matches {
|
||
start := m[1]
|
||
end := len(text)
|
||
if i+1 < len(matches) {
|
||
end = matches[i+1][0]
|
||
}
|
||
marker := strings.TrimSpace(text[m[0]:m[1]])
|
||
content := strings.TrimSpace(text[start:end])
|
||
if content == "" {
|
||
continue
|
||
}
|
||
item := marker + " " + content
|
||
ctx := ctxBefore(m[0])
|
||
if bufLen > 0 && (ctx != bufCtx || bufLen+runeLen(item) > chunkSize) {
|
||
flush()
|
||
}
|
||
if runeLen(item) > chunkSize {
|
||
pieces := s.SplitText(item, chunkSize, overlap)
|
||
if len(pieces) > 0 {
|
||
if ctx != "" {
|
||
pieces[0] = "【" + ctx + "】" + pieces[0]
|
||
}
|
||
out = append(out, pieces[0])
|
||
for _, p := range pieces[1:] {
|
||
out = append(out, "（"+marker+" 续）"+p)
|
||
}
|
||
}
|
||
continue
|
||
}
|
||
buf = append(buf, item)
|
||
bufLen += runeLen(item)
|
||
bufCtx = ctx
|
||
if bufLen >= mergeMin {
|
||
flush()
|
||
}
|
||
}
|
||
flush()
|
||
return out
|
||
}
|
||
|
||
// splitRecursive 递归字符分块：按分隔符列表递归切分到目标大小，保留 overlap
|
||
func (s *chunkService) splitRecursive(ctx context.Context, text string, chunkSize, overlap int) ([]string, error) {
|
||
sp, err := recursive.NewSplitter(ctx, &recursive.Config{
|
||
ChunkSize: chunkSize,
|
||
OverlapSize: overlap,
|
||
Separators: chunkSeparators,
|
||
LenFunc: runeLen,
|
||
KeepType: recursive.KeepTypeEnd,
|
||
})
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "构建递归分块器失败")
|
||
}
|
||
return s.transformText(ctx, sp, text)
|
||
}
|
||
|
||
// capChunks 超过 maxSize 的块用递归分块器二次切分，防止语义分块产出巨块
|
||
func (s *chunkService) capChunks(ctx context.Context, chunks []string, maxSize int) ([]string, error) {
|
||
if maxSize <= 0 {
|
||
return chunks, nil
|
||
}
|
||
out := make([]string, 0, len(chunks))
|
||
for _, c := range chunks {
|
||
if runeLen(c) <= maxSize {
|
||
out = append(out, c)
|
||
continue
|
||
}
|
||
sub, err := s.splitRecursive(ctx, c, maxSize, 0)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, sub...)
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
func (s *chunkService) transformText(ctx context.Context, sp document.Transformer, text string) ([]string, error) {
|
||
docs, err := sp.Transform(ctx, []*schema.Document{{ID: "0", Content: text}})
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "分块失败")
|
||
}
|
||
out := make([]string, 0, len(docs))
|
||
for _, d := range docs {
|
||
if c := strings.TrimSpace(d.Content); c != "" {
|
||
out = append(out, c)
|
||
}
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
var ChunkService = &chunkService{}
|
||
|
||
type chunkService struct{}
|
||
|
||
// SplitText 文本分块：标题感知 + 固定大小回退，超长段落按句号/换行切分并保留重叠
|
||
func (s *chunkService) SplitText(text string, chunkSize, overlap int) []string {
|
||
if chunkSize <= 0 {
|
||
chunkSize = consts.DefaultChunkSize
|
||
}
|
||
if overlap < 0 {
|
||
overlap = consts.DefaultChunkOverlap
|
||
}
|
||
text = strings.ReplaceAll(text, "\r\n", "\n")
|
||
text = strings.ReplaceAll(text, "\r", "\n")
|
||
|
||
paragraphs := splitParagraphs(text)
|
||
var merged []string
|
||
var cur strings.Builder
|
||
curRunes := 0
|
||
for _, p := range paragraphs {
|
||
pRunes := utf8.RuneCountInString(p)
|
||
if cur.Len() > 0 && curRunes+pRunes > chunkSize {
|
||
merged = append(merged, cur.String())
|
||
cur.Reset()
|
||
curRunes = 0
|
||
}
|
||
cur.WriteString(p)
|
||
cur.WriteString("\n\n")
|
||
curRunes += pRunes + 2
|
||
}
|
||
if cur.Len() > 0 {
|
||
merged = append(merged, cur.String())
|
||
}
|
||
|
||
var chunks []string
|
||
for _, c := range merged {
|
||
if len(c) <= chunkSize {
|
||
chunks = append(chunks, strings.TrimSpace(c))
|
||
continue
|
||
}
|
||
chunks = append(chunks, forceSplit(c, chunkSize, overlap)...)
|
||
}
|
||
return chunks
|
||
}
|
||
|
||
// splitParagraphs 按空行/标题切段；# 标题行并入其后的段落（标题保留在段首，供检索回显）
|
||
func splitParagraphs(text string) []string {
|
||
lines := strings.Split(text, "\n")
|
||
var paras []string
|
||
var cur []string
|
||
flush := func() {
|
||
if len(cur) > 0 {
|
||
paras = append(paras, strings.Join(cur, "\n"))
|
||
cur = nil
|
||
}
|
||
}
|
||
for _, line := range lines {
|
||
trimmed := strings.TrimSpace(line)
|
||
if trimmed == "" {
|
||
flush()
|
||
continue
|
||
}
|
||
if isHeading(trimmed) && len(cur) > 0 {
|
||
flush()
|
||
}
|
||
cur = append(cur, trimmed)
|
||
}
|
||
flush()
|
||
return paras
|
||
}
|
||
|
||
func isHeading(line string) bool {
|
||
return strings.HasPrefix(line, "#") ||
|
||
strings.HasPrefix(line, "标题") || strings.HasPrefix(line, "第") && strings.Contains(line, "章")
|
||
}
|
||
|
||
// forceSplit 超长文本按可读位置切分（rune 安全，避免切在 UTF-8 中间），重叠 overlap 字
|
||
func forceSplit(text string, chunkSize, overlap int) []string {
|
||
runes := []rune(text)
|
||
var result []string
|
||
for len(runes) > chunkSize {
|
||
limit := min(chunkSize, len(runes))
|
||
cut := lastCutPoint(runes[:limit])
|
||
if cut < chunkSize/2 {
|
||
cut = chunkSize
|
||
}
|
||
if chunk := strings.TrimSpace(string(runes[:cut])); chunk != "" {
|
||
result = append(result, chunk)
|
||
}
|
||
runes = runes[max(0, cut-overlap):]
|
||
}
|
||
if chunk := strings.TrimSpace(string(runes)); chunk != "" {
|
||
result = append(result, chunk)
|
||
}
|
||
return result
|
||
}
|
||
|
||
// lastCutPoint 在 limit 内找最后一个可读切点（句号/问号/感叹号/分号/换行/英文标点），找不到返回 -1
|
||
func lastCutPoint(runes []rune) int {
|
||
cut := -1
|
||
for i, r := range runes {
|
||
switch r {
|
||
case '。', '！', '？', '；', '\n', '.', '!', '?', ';':
|
||
cut = i + 1
|
||
}
|
||
}
|
||
return cut
|
||
}
|
||
|
||
// List 分块列表
|
||
func (s *chunkService) List(ctx context.Context, documentId int64, page, pageSize int) ([]*entity.Chunk, int, error) {
|
||
return dao.Chunk.ListByDocument(ctx, documentId, page, pageSize)
|
||
}
|
||
|
||
// InsertAll 写入文档的全部分块；embedder 非空时批量向量化写入 vec0，否则仅写 FTS5
|
||
func (s *chunkService) InsertAll(ctx context.Context, datasetId, documentId int64, chunks []string, embedder eembedding.Embedder) error {
|
||
for start := 0; start < len(chunks); start += consts.EmbedBatchSize {
|
||
end := min(start+consts.EmbedBatchSize, len(chunks))
|
||
var vecs [][]float64
|
||
if embedder != nil {
|
||
v, err := embedder.EmbedStrings(ctx, chunks[start:end])
|
||
if err != nil {
|
||
return gerror.Wrap(err, "分块向量化失败")
|
||
}
|
||
vecs = v
|
||
}
|
||
for j, content := range chunks[start:end] {
|
||
vecJson := ""
|
||
if len(vecs) > 0 && j < len(vecs) {
|
||
vecJson = domain.VecJsonF64(vecs[j])
|
||
}
|
||
if _, err := dao.Chunk.InsertWithVec(ctx, datasetId, documentId, start+j+1, content, "", "", vecJson, 0); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
}
|
||
return dao.Document.UpdateFields(ctx, documentId, g.Map{
|
||
"chunk_count": len(chunks),
|
||
})
|
||
}
|
||
|
||
// Update 编辑分块文本：重新分词 FTS 索引，数据集绑定 embedding 配置时同步重新向量化
|
||
func (s *chunkService) Update(ctx context.Context, id int64, content string) error {
|
||
chunk, err := dao.Chunk.GetOne(ctx, id)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if chunk == nil {
|
||
return gerror.New("分块不存在")
|
||
}
|
||
vecJson := ""
|
||
if cfgId, err := dao.Dataset.GetEmbeddingCfgId(ctx, chunk.DatasetId); err == nil && cfgId > 0 {
|
||
if em, err := BuildEmbedder(ctx, cfgId); err == nil {
|
||
if vecs, err := em.EmbedStrings(ctx, []string{content}); err != nil {
|
||
g.Log().Warningf(ctx, "re-embed chunk %d failed: %v", id, err)
|
||
} else if len(vecs) > 0 {
|
||
vecJson = domain.VecJsonF64(vecs[0])
|
||
}
|
||
} else {
|
||
g.Log().Warningf(ctx, "build embedder failed: %v", err)
|
||
}
|
||
}
|
||
return dao.Chunk.UpdateContent(ctx, id, content, vecJson, common.Tokenize(content))
|
||
}
|
||
|
||
// DeleteByDocument 删除文档全部数据（chunk + vec + fts）
|
||
func (s *chunkService) DeleteByDocument(ctx context.Context, documentId int64) error {
|
||
return dao.Chunk.DeleteByDocument(ctx, documentId)
|
||
}
|