feat: RAG分块策略增强
This commit is contained in:
@@ -0,0 +1,80 @@
|
||||
package com.easyagents.rag.ingestion;
|
||||
|
||||
import com.easyagents.rag.core.RagChunk;
|
||||
import com.easyagents.rag.core.RagChunkTypes;
|
||||
import com.easyagents.rag.core.RagStrategyCodes;
|
||||
import com.easyagents.rag.ingestion.analysis.DocumentStructureAnalyzer;
|
||||
import com.easyagents.rag.ingestion.chunk.RagSplitStrategyRegistry;
|
||||
import com.easyagents.rag.ingestion.model.AnalysisResult;
|
||||
import com.easyagents.rag.ingestion.model.StrategyConfig;
|
||||
import com.easyagents.rag.ingestion.recommend.SplitStrategyRecommender;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
public class RagIngestionPipelineTest {
|
||||
|
||||
private final DocumentStructureAnalyzer analyzer = new DocumentStructureAnalyzer();
|
||||
private final SplitStrategyRecommender recommender = new SplitStrategyRecommender();
|
||||
private final RagSplitStrategyRegistry registry = new RagSplitStrategyRegistry();
|
||||
|
||||
@Test
|
||||
public void shouldRecommendMarkdownStrategy() {
|
||||
String markdown = "# Quick Start\n"
|
||||
+ "Welcome\n\n"
|
||||
+ "## Install\n"
|
||||
+ "Run npm install\n\n"
|
||||
+ "## Usage\n"
|
||||
+ "Run pnpm dev";
|
||||
|
||||
AnalysisResult analysis = recommender.recommend(analyzer.analyze(markdown, "md"));
|
||||
|
||||
Assert.assertEquals(RagStrategyCodes.MARKDOWN_SECTION, analysis.getRecommendedStrategyCode());
|
||||
Assert.assertTrue(analysis.getConfidence().doubleValue() > 0.4D);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void shouldRecommendQaStrategyForEnglishAndChinese() {
|
||||
String qa = "Q: How to reset password?\n"
|
||||
+ "A: Open admin page and click reset.\n\n"
|
||||
+ "问:默认密码是什么?\n"
|
||||
+ "答:由系统配置统一决定。";
|
||||
|
||||
AnalysisResult analysis = recommender.recommend(analyzer.analyze(qa, "txt"));
|
||||
|
||||
Assert.assertEquals(RagStrategyCodes.QA_PAIR, analysis.getRecommendedStrategyCode());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void shouldSplitOutlineDocumentByHeadingPath() {
|
||||
String outline = "第1章 总则\n适用范围说明。\n\n1.1 目标\n定义系统目标。\n\n1.2 范围\n定义系统范围。";
|
||||
AnalysisResult analysis = recommender.recommend(analyzer.analyze(outline, "docx"));
|
||||
StrategyConfig config = StrategyConfig.defaults();
|
||||
config.setStrategyCode(RagStrategyCodes.OUTLINE_SECTION);
|
||||
|
||||
List<RagChunk> chunks = registry.split(analysis, config);
|
||||
|
||||
Assert.assertEquals(3, chunks.size());
|
||||
Assert.assertEquals("第1章 总则", chunks.get(0).getSourceLabel());
|
||||
Assert.assertEquals(2, chunks.get(1).getHeadingPath().size());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void shouldSplitQaDocumentByPair() {
|
||||
String qa = "Q: How to reset password?\n"
|
||||
+ "A: Open admin page and click reset.\n\n"
|
||||
+ "问:默认密码是什么?\n"
|
||||
+ "答:由系统配置统一决定。";
|
||||
AnalysisResult analysis = recommender.recommend(analyzer.analyze(qa, "txt"));
|
||||
StrategyConfig config = StrategyConfig.defaults();
|
||||
config.setStrategyCode(RagStrategyCodes.QA_PAIR);
|
||||
|
||||
List<RagChunk> chunks = registry.split(analysis, config);
|
||||
|
||||
Assert.assertEquals(2, chunks.size());
|
||||
Assert.assertEquals(RagChunkTypes.QA_PAIR, chunks.get(0).getChunkType());
|
||||
Assert.assertTrue(chunks.get(0).getContent().contains("问题"));
|
||||
Assert.assertTrue(chunks.get(1).getAnswer().contains("系统配置"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user