246 lines
7.4 KiB
Go
246 lines
7.4 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"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
|
||
|
||
// SplitByStrategy 按数据集分块策略分块；semantic 需要 embedder，其余策略可传 nil
|
||
func (s *chunkService) SplitByStrategy(ctx context.Context, strategy, text string, chunkSize, overlap int, embedder eembedding.Embedder) ([]string, error) {
|
||
switch strategy {
|
||
case consts.ChunkStrategyRecursive:
|
||
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)
|
||
case consts.ChunkStrategySemantic:
|
||
if embedder == nil {
|
||
return nil, gerror.New("语义分块需要向量模型")
|
||
}
|
||
sp, err := semantic.NewSplitter(ctx, &semantic.Config{
|
||
Embedding: embedder,
|
||
BufferSize: 1,
|
||
MinChunkSize: chunkSize / 2,
|
||
Separators: chunkSeparators[:len(chunkSeparators)-2],
|
||
LenFunc: runeLen,
|
||
})
|
||
if err != nil {
|
||
return nil, gerror.Wrap(err, "构建语义分块器失败")
|
||
}
|
||
return s.transformText(ctx, sp, text)
|
||
default:
|
||
return s.SplitText(text, chunkSize, overlap), 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),
|
||
"status": consts.DocumentStatusDone,
|
||
})
|
||
}
|
||
|
||
// 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)
|
||
}
|