rag-local/kb/service/chunk_service.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"

	eembedding "github.com/cloudwego/eino/components/embedding"
	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/frame/g"
)

var ChunkService = &chunkService{}

type chunkService struct{}

// SplitText 文本分块：标题感知 + 固定大小回退，超长段落按句号/换行切分并保留重叠
func (s *chunkService) SplitText(text string) []string {
	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 > consts.MaxChunkSize {
			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) <= consts.MaxChunkSize {
			chunks = append(chunks, strings.TrimSpace(c))
			continue
		}
		chunks = append(chunks, forceSplit(c)...)
	}
	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) []string {
	runes := []rune(text)
	var result []string
	for len(runes) > consts.MaxChunkSize {
		limit := min(consts.MaxChunkSize, len(runes))
		cut := lastCutPoint(runes[:limit])
		if cut < consts.MaxChunkSize/2 {
			cut = consts.MaxChunkSize
		}
		if chunk := strings.TrimSpace(string(runes[:cut])); chunk != "" {
			result = append(result, chunk)
		}
		runes = runes[max(0, cut-consts.ChunkOverlap):]
	}
	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)
}