001/*
002 *  Copyright (c) 2023-2026, Agents-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.agentsflex.core.document.splitter;
017
018import com.agentsflex.core.document.Document;
019import com.agentsflex.core.document.DocumentSplitter;
020import com.agentsflex.core.document.id.DocumentIdGenerator;
021import com.agentsflex.core.model.chat.ChatModel;
022import com.agentsflex.core.model.chat.ChatOptions;
023import org.slf4j.Logger;
024import org.slf4j.LoggerFactory;
025
026import java.util.ArrayList;
027import java.util.Collections;
028import java.util.List;
029import java.util.stream.Collectors;
030
031/**
032 * AIDocumentSplitter：基于大模型（AI/LLM）的语义文档拆分器。
033 * 使用 "---" 作为段落分隔符，避免 JSON 解析风险。
034 * 支持注入 fallback 拆分器以提高鲁棒性。
035 */
036public class AIDocumentSplitter implements DocumentSplitter {
037
038    private static final Logger log = LoggerFactory.getLogger(AIDocumentSplitter.class);
039
040    private static final String DEFAULT_SPLIT_PROMPT_TEMPLATE =
041        "你是一个专业的文档处理助手，请将以下长文档按语义拆分为多个逻辑连贯的段落块。\n" +
042            "要求：\n" +
043            "1. 每个块应保持主题/语义完整性，避免在句子中间切断。\n" +
044            "2. 每个块长度建议在 200-500 字之间（可根据内容灵活调整）。\n" +
045            "3. **不要添加任何解释、编号、前缀或后缀**。\n" +
046            "4. **仅用三连短横线 \"---\" 作为块之间的分隔符**，格式如下：\n" +
047            "\n" +
048            "块1内容\n" +
049            "---\n" +
050            "块2内容\n" +
051            "---\n" +
052            "块3内容\n" +
053            "\n" +
054            "注意：开头不要有 ---，结尾也不要有多余的 ---。\n" +
055            "\n" +
056            "文档内容如下：\n" +
057            "{document}";
058
059    private static final String CHUNK_SEPARATOR = "---";
060
061    private final ChatModel chatModel;
062    private String splitPromptTemplate = DEFAULT_SPLIT_PROMPT_TEMPLATE;
063    private ChatOptions chatOptions = new ChatOptions.Builder().temperature(0.2f).build();
064    private int maxChunks = 20;
065    private int maxTotalLength = 10000;
066
067    // 可配置的 fallback 拆分器
068    private DocumentSplitter fallbackSplitter;
069
070    public AIDocumentSplitter(ChatModel chatModel) {
071        this.chatModel = chatModel;
072    }
073
074    @Override
075    public List<Document> split(Document document, DocumentIdGenerator idGenerator) {
076        if (document == null || document.getContent() == null || document.getContent().trim().isEmpty()) {
077            return Collections.emptyList();
078        }
079
080        String content = document.getContent().trim();
081        if (content.length() > maxTotalLength) {
082            log.warn("文档过长（{} 字符），已截断至 {} 字符", content.length(), maxTotalLength);
083            content = content.substring(0, maxTotalLength);
084        }
085
086        List<String> chunks;
087        try {
088            String prompt = splitPromptTemplate.replace("{document}", content);
089            String llmOutput = chatModel.chat(prompt, chatOptions);
090
091            chunks = parseChunksBySeparator(llmOutput, CHUNK_SEPARATOR);
092        } catch (Exception e) {
093            log.error("AI 拆分失败，使用 fallback 拆分器", e);
094            if (fallbackSplitter == null) {
095                log.error("没有可用的 fallback 拆分器，请检查配置");
096                return Collections.emptyList();
097            }
098            List<Document> fallbackDocs = fallbackSplitter.split(document, idGenerator);
099            if (fallbackDocs.size() > maxChunks) {
100                return new ArrayList<>(fallbackDocs.subList(0, maxChunks));
101            }
102            return fallbackDocs;
103        }
104
105        List<String> validChunks = chunks.stream()
106            .map(String::trim)
107            .filter(s -> !s.isEmpty())
108            .limit(maxChunks)
109            .collect(Collectors.toList());
110
111        List<Document> result = new ArrayList<>();
112        for (String chunk : validChunks) {
113            Document doc = new Document();
114            doc.setContent(chunk);
115            doc.setTitle(document.getTitle());
116            if (idGenerator != null) {
117                doc.setId(idGenerator.generateId(doc));
118            }
119            result.add(doc);
120        }
121
122        return result;
123    }
124
125    private List<String> parseChunksBySeparator(String text, String separator) {
126        if (text == null || text.trim().isEmpty()) {
127            return Collections.emptyList();
128        }
129
130        String[] parts = text.split(separator, -1);
131        List<String> chunks = new ArrayList<>();
132        for (String part : parts) {
133            String trimmed = part.trim();
134            if (!trimmed.isEmpty()) {
135                chunks.add(trimmed);
136            }
137        }
138
139        if (chunks.size() == 1 && text.contains(separator)) {
140            return tryAlternativeSplit(text, separator);
141        }
142
143        return chunks;
144    }
145
146    private List<String> tryAlternativeSplit(String text, String separator) {
147        String normalized = text.replaceAll("\\s*---\\s*", "---");
148        return parseChunksBySeparator(normalized, separator);
149    }
150
151    // ===== Getters & Setters =====
152
153    public void setFallbackSplitter(DocumentSplitter fallbackSplitter) {
154        this.fallbackSplitter = fallbackSplitter;
155    }
156
157    public void setSplitPromptTemplate(String splitPromptTemplate) {
158        if (splitPromptTemplate != null && !splitPromptTemplate.trim().isEmpty()) {
159            this.splitPromptTemplate = splitPromptTemplate;
160        }
161    }
162
163    public void setChatOptions(ChatOptions chatOptions) {
164        if (chatOptions != null) {
165            this.chatOptions = chatOptions;
166        }
167    }
168
169    public void setMaxChunks(int maxChunks) {
170        this.maxChunks = Math.max(1, maxChunks);
171    }
172
173    public void setMaxTotalLength(int maxTotalLength) {
174        this.maxTotalLength = Math.max(100, maxTotalLength);
175    }
176}