4 天以前 83e1b4d0e661f11a407fd6ea86e906b9b87b7180
yudao-module-ai/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/splitter/SemanticTextSplitter.java
@@ -3,7 +3,6 @@
import cn.hutool.core.util.StrUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.transformer.splitter.TextSplitter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -13,289 +12,117 @@
/**
 * 语义化文本切片器
 *
 * <p>功能特点:
 * <ul>
 *   <li>优先在段落边界(双换行)处切分</li>
 *   <li>其次在句子边界(句号、问号、感叹号)处切分</li>
 *   <li>避免在句子中间截断,保持语义完整性</li>
 *   <li>支持中英文标点符号识别</li>
 * </ul>
 *
 * @author runzhen
 * 优先在段落边界处切分,其次在句子边界处切分,保持语义完整性
 */
@Slf4j
public class SemanticTextSplitter extends TextSplitter {
    /**
     * 分段的最大 Token 数
     */
    private final int chunkSize;
    /**
     * 段落重叠大小(用于保持上下文连贯性)
     */
    private final int chunkOverlap;
    /**
     * 段落分隔符(按优先级排序)
     */
    private static final List<String> PARAGRAPH_SEPARATORS = Arrays.asList(
            "\n\n\n",    // 三个换行
            "\n\n",      // 双换行
            "\n"         // 单换行
    );
    /**
     * 句子结束标记(中英文标点)
     */
    private static final Pattern SENTENCE_END_PATTERN = Pattern.compile(
            "[。!?.!?]+[\\s\"'))】\\]]*"
    );
    /**
     * Token 估算器
     */
    private static final List<String> PARAGRAPH_SEPARATORS = Arrays.asList("\n\n\n", "\n\n", "\n");
    private static final Pattern SENTENCE_END_PATTERN = Pattern.compile("[。!?.!?]+[\\s\"'))】\\]]*");
    private final MarkdownQaSplitter.TokenEstimator tokenEstimator;
    public SemanticTextSplitter(int chunkSize, int chunkOverlap) {
        this.chunkSize = chunkSize;
        this.chunkOverlap = Math.min(chunkOverlap, chunkSize / 2); // 重叠不超过一半
        this.chunkOverlap = Math.min(chunkOverlap, chunkSize / 2);
        this.tokenEstimator = new SimpleTokenEstimator();
    }
    public SemanticTextSplitter(int chunkSize) {
        this(chunkSize, 50); // 默认重叠 50 个 Token
    }
    public SemanticTextSplitter(int chunkSize) { this(chunkSize, 50); }
    @Override
    protected List<String> splitText(String text) {
        if (StrUtil.isEmpty(text)) {
            return Collections.emptyList();
        }
        if (StrUtil.isEmpty(text)) return Collections.emptyList();
        return splitTextRecursive(text);
    }
    /**
     * 切分文本(递归策略)
     *
     * @param text 待切分文本
     * @return 切分后的文本块列表
     */
    private List<String> splitTextRecursive(String text) {
        List<String> chunks = new ArrayList<>();
        // 如果文本不超过限制,直接返回
        int textTokens = tokenEstimator.estimate(text);
        if (textTokens <= chunkSize) {
            chunks.add(text.trim());
            return chunks;
        }
        if (textTokens <= chunkSize) return Collections.singletonList(text.trim());
        // 尝试按不同分隔符切分
        List<String> splits = null;
        String usedSeparator = null;
        for (String separator : PARAGRAPH_SEPARATORS) {
            if (text.contains(separator)) {
                splits = Arrays.asList(text.split(Pattern.quote(separator)));
                usedSeparator = separator;
                break;
            }
            if (text.contains(separator)) { splits = Arrays.asList(text.split(Pattern.quote(separator))); usedSeparator = separator; break; }
        }
        // 如果没有找到段落分隔符,按句子切分
        if (splits == null || splits.size() == 1) {
            splits = splitBySentences(text);
            usedSeparator = ""; // 句子切分不需要分隔符
        }
        // 合并小片段
        chunks = mergeSplits(splits, usedSeparator);
        return chunks;
        if (splits == null || splits.size() == 1) { splits = splitBySentences(text); usedSeparator = ""; }
        return mergeSplits(splits, usedSeparator);
    }
    /**
     * 按句子切分
     *
     * @param text 待切分文本
     * @return 句子列表
     */
    private List<String> splitBySentences(String text) {
        // 使用正则表达式匹配句子结束位置
        List<String> sentences = new ArrayList<>();
        int lastEnd = 0;
        Matcher matcher = SENTENCE_END_PATTERN.matcher(text);
        while (matcher.find()) {
            String sentence = text.substring(lastEnd, matcher.end()).trim();
            if (StrUtil.isNotEmpty(sentence)) {
                sentences.add(sentence);
            }
            if (StrUtil.isNotEmpty(sentence)) sentences.add(sentence);
            lastEnd = matcher.end();
        }
        // 添加剩余部分
        if (lastEnd < text.length()) {
            String remaining = text.substring(lastEnd).trim();
            if (StrUtil.isNotEmpty(remaining)) {
                sentences.add(remaining);
            }
        }
        if (lastEnd < text.length()) { String remaining = text.substring(lastEnd).trim(); if (StrUtil.isNotEmpty(remaining)) sentences.add(remaining); }
        return sentences.isEmpty() ? Collections.singletonList(text) : sentences;
    }
    /**
     * 合并切分后的小片段
     *
     * @param splits 切分后的片段列表
     * @param separator 片段间的分隔符
     * @return 合并后的文本块列表
     */
    private List<String> mergeSplits(List<String> splits, String separator) {
        List<String> chunks = new ArrayList<>();
        List<String> currentChunks = new ArrayList<>();
        int currentLength = 0;
        for (String split : splits) {
            if (StrUtil.isEmpty(split)) {
                continue;
            }
            if (StrUtil.isEmpty(split)) continue;
            int splitTokens = tokenEstimator.estimate(split);
            // 如果单个片段就超过限制,进一步递归切分
            if (splitTokens > chunkSize) {
                // 先保存当前累积的块
                if (!currentChunks.isEmpty()) {
                    String chunkText = String.join(separator, currentChunks);
                    chunks.add(chunkText.trim());
                    currentChunks.clear();
                    currentLength = 0;
                }
                // 递归切分大片段
                if (!separator.isEmpty()) {
                    // 如果是段落分隔符,尝试按句子切分
                    chunks.addAll(splitTextRecursive(split));
                } else {
                    // 如果已经是句子级别,强制按字符切分
                    chunks.addAll(forceSplitLongText(split));
                }
                if (!currentChunks.isEmpty()) { chunks.add(String.join(separator, currentChunks).trim()); currentChunks.clear(); currentLength = 0; }
                if (!separator.isEmpty()) { chunks.addAll(splitTextRecursive(split)); } else { chunks.addAll(forceSplitLongText(split)); }
                continue;
            }
            // 计算加上分隔符的 Token 数
            int separatorTokens = StrUtil.isEmpty(separator) ? 0 : tokenEstimator.estimate(separator);
            // 如果加上这个片段会超过限制
            if (!currentChunks.isEmpty() && currentLength + splitTokens + separatorTokens > chunkSize) {
                // 保存当前块
                String chunkText = String.join(separator, currentChunks);
                chunks.add(chunkText.trim());
                // 处理重叠:保留最后几个片段
                chunks.add(String.join(separator, currentChunks).trim());
                currentChunks = getOverlappingChunks(currentChunks, separator);
                currentLength = estimateTokens(currentChunks, separator);
            }
            // 添加当前片段
            currentChunks.add(split);
            currentLength += splitTokens + separatorTokens;
        }
        // 添加最后一块
        if (!currentChunks.isEmpty()) {
            String chunkText = String.join(separator, currentChunks);
            chunks.add(chunkText.trim());
        }
        if (!currentChunks.isEmpty()) chunks.add(String.join(separator, currentChunks).trim());
        return chunks;
    }
    /**
     * 获取重叠的片段(用于保持上下文)
     *
     * @param chunks 当前片段列表
     * @param separator 片段间的分隔符
     * @return 重叠的片段列表
     */
    private List<String> getOverlappingChunks(List<String> chunks, String separator) {
        if (chunkOverlap == 0 || chunks.isEmpty()) {
            return new ArrayList<>();
        }
        // 从后往前取片段,直到达到重叠大小
        if (chunkOverlap == 0 || chunks.isEmpty()) return new ArrayList<>();
        List<String> overlapping = new ArrayList<>();
        int tokens = 0;
        for (int i = chunks.size() - 1; i >= 0; i--) {
            String chunk = chunks.get(i);
            int chunkTokens = tokenEstimator.estimate(chunk);
            if (tokens + chunkTokens > chunkOverlap) {
                break;
            }
            // 添加到重叠列表前端
            if (tokens + chunkTokens > chunkOverlap) break;
            overlapping.add(0, chunk);
            tokens += chunkTokens + (StrUtil.isEmpty(separator) ? 0 : tokenEstimator.estimate(separator));
        }
        return overlapping;
    }
    /**
     * 估算片段列表的总 Token 数
     *
     * @param chunks 片段列表
     * @param separator 片段间的分隔符
     * @return 总 Token 数
     */
    private int estimateTokens(List<String> chunks, String separator) {
        int total = 0;
        for (int i = 0; i < chunks.size(); i++) {
            total += tokenEstimator.estimate(chunks.get(i));
            if (i < chunks.size() - 1 && StrUtil.isNotEmpty(separator)) {
                total += tokenEstimator.estimate(separator);
            }
        }
        for (int i = 0; i < chunks.size(); i++) { total += tokenEstimator.estimate(chunks.get(i)); if (i < chunks.size() - 1 && StrUtil.isNotEmpty(separator)) total += tokenEstimator.estimate(separator); }
        return total;
    }
    /**
     * 强制切分长文本(当语义切分失败时)
     *
     * @param text 待切分文本
     * @return 切分后的文本块列表
     */
    private List<String> forceSplitLongText(String text) {
        List<String> chunks = new ArrayList<>();
        int charsPerChunk = (int) (chunkSize * 0.8); // 保守估计
        for (int i = 0; i < text.length(); i += charsPerChunk) {
            int end = Math.min(i + charsPerChunk, text.length());
            String chunk = text.substring(i, end);
            chunks.add(chunk.trim());
        }
        log.warn("文本过长,已强制按字符切分,可能影响语义完整性");
        int charsPerChunk = (int)(chunkSize * 0.8);
        for (int i = 0; i < text.length(); i += charsPerChunk) { int end = Math.min(i + charsPerChunk, text.length()); chunks.add(text.substring(i, end).trim()); }
        return chunks;
    }
    /**
     * 简单的 Token 估算器实现
     */
    private static class SimpleTokenEstimator implements MarkdownQaSplitter.TokenEstimator {
        @Override
        public int estimate(String text) {
            if (StrUtil.isEmpty(text)) {
                return 0;
            }
            int chineseChars = 0;
            int englishWords = 0;
            // 简单统计中英文
            for (char c : text.toCharArray()) {
                if (c >= 0x4E00 && c <= 0x9FA5) {
                    chineseChars++;
                }
            }
            // 英文单词估算
            String[] words = text.split("\\s+");
            for (String word : words) {
                if (word.matches(".*[a-zA-Z].*")) {
                    englishWords++;
                }
            }
            return chineseChars + (int) (englishWords * 1.3);
            if (StrUtil.isEmpty(text)) return 0;
            int chineseChars = 0, englishWords = 0;
            for (char c : text.toCharArray()) { if (c >= 0x4E00 && c <= 0x9FA5) chineseChars++; }
            for (String word : text.split("\\s+")) { if (word.matches(".*[a-zA-Z].*")) englishWords++; }
            return chineseChars + (int)(englishWords * 1.3);
        }
    }
}