common/rag/eino/document_semantic.go

package eino

import (
	"context"

	"github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive"
	"github.com/cloudwego/eino-ext/components/document/transformer/splitter/semantic"
	"github.com/cloudwego/eino/schema"
	"github.com/gogf/gf/v2/frame/g"
)

// SemanticSplitDocument 语义分割文档
func SemanticSplitDocument(ctx context.Context, docs []*schema.Document) (res []*schema.Document, err error) {
	// 默认分隔符（支持中英文）
	separators := []string{"\n\n", "\n", "。", "！", "？", "；", ".", "!", "?", ";"}
	// 读取配置，使用合理的默认值
	bufferSize := g.Cfg().MustGet(ctx, "eino.splitter.bufferSize").Int()
	minChunkSize := g.Cfg().MustGet(ctx, "eino.splitter.minChunkSize").Int()
	percentile := g.Cfg().MustGet(ctx, "eino.splitter.percentile").Float64()
	batchSize := g.Cfg().MustGet(ctx, "eino.splitter.batchSize").Int()
	if batchSize <= 0 {
		batchSize = 10 // doubao-embedding-vision 限制每批最多 10 个
	}

	// 使用批量包装器
	var batchEmbedder *BatchEmbedder
	provider := g.Cfg().MustGet(ctx, "eino.embedding.provider").String()
	switch provider {
	case providerArk:
		batchEmbedder = NewBatchEmbedder(EmbedderArk, batchSize)
	case providerOpenai:
		batchEmbedder = NewBatchEmbedder(EmbedderOpenAI, batchSize)
	case providerDashscope:
		batchEmbedder = NewBatchEmbedder(EmbedderDashscope, batchSize)
	}

	splitter, err := semantic.NewSplitter(ctx, &semantic.Config{
		Embedding:    batchEmbedder,
		BufferSize:   bufferSize,
		MinChunkSize: minChunkSize,
		Percentile:   percentile,
		Separators:   separators,
	})
	if err != nil {
		return
	}
	return splitter.Transform(ctx, docs)
}

// RecursiveSplitDocument 递归分割文档
func RecursiveSplitDocument(ctx context.Context, docs []*schema.Document) (res []*schema.Document, err error) {
	// 默认分隔符（支持中英文）
	separators := []string{"\n\n", "\n", "。", "！", "？", "；", ".", "!", "?", ";"}
	splitter, err := recursive.NewSplitter(ctx, &recursive.Config{
		ChunkSize:   512,
		OverlapSize: 100,
		KeepType:    recursive.KeepTypeNone,
		Separators:  separators,
	})
	if err != nil {
		return
	}
	return splitter.Transform(ctx, docs)
}