diff --git a/easy-agents-embedding/easy-agents-embedding-openai/src/main/java/com/easyagents/embedding/openai/OpenAIEmbeddingModel.java b/easy-agents-embedding/easy-agents-embedding-openai/src/main/java/com/easyagents/embedding/openai/OpenAIEmbeddingModel.java index afbc8c2..cd8e0ed 100644 --- a/easy-agents-embedding/easy-agents-embedding-openai/src/main/java/com/easyagents/embedding/openai/OpenAIEmbeddingModel.java +++ b/easy-agents-embedding/easy-agents-embedding-openai/src/main/java/com/easyagents/embedding/openai/OpenAIEmbeddingModel.java @@ -11,13 +11,19 @@ import com.easyagents.core.store.VectorData; import com.easyagents.core.util.JSONUtil; import com.easyagents.core.util.Maps; import com.easyagents.core.util.StringUtil; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.Locale; import java.util.Map; public class OpenAIEmbeddingModel extends BaseEmbeddingModel { + private static final Logger LOG = LoggerFactory.getLogger(OpenAIEmbeddingModel.class); + private static final String REDACTED_HEADER_VALUE = "[REDACTED]"; + private HttpClient httpClient = new HttpClient(); public OpenAIEmbeddingModel(OpenAIEmbeddingConfig config) { @@ -40,27 +46,45 @@ public class OpenAIEmbeddingModel extends BaseEmbeddingModel sanitizeHeadersForLogging(Map headers) { + Map sanitizedHeaders = new LinkedHashMap<>(); + if (headers == null) { + return sanitizedHeaders; + } + headers.forEach((name, value) -> sanitizedHeaders.put( + name, + "Authorization".equalsIgnoreCase(name) ? REDACTED_HEADER_VALUE : value + )); + return sanitizedHeaders; + } } diff --git a/easy-agents-embedding/easy-agents-embedding-openai/src/test/java/com/easyagents/embedding/openai/OpenAIEmbeddingModelTest.java b/easy-agents-embedding/easy-agents-embedding-openai/src/test/java/com/easyagents/embedding/openai/OpenAIEmbeddingModelTest.java index 1017c08..edf3a5f 100644 --- a/easy-agents-embedding/easy-agents-embedding-openai/src/test/java/com/easyagents/embedding/openai/OpenAIEmbeddingModelTest.java +++ b/easy-agents-embedding/easy-agents-embedding-openai/src/test/java/com/easyagents/embedding/openai/OpenAIEmbeddingModelTest.java @@ -7,6 +7,7 @@ import com.easyagents.core.model.exception.ModelException; import org.junit.Assert; import org.junit.Test; +import java.util.LinkedHashMap; import java.util.Map; /** @@ -68,4 +69,46 @@ public class OpenAIEmbeddingModelTest { Assert.assertTrue(exception.getMessage().contains("data[0].embedding")); } + + /** + * Verifies that invalid JSON responses are reported as embedding response parsing failures. + */ + @Test + public void shouldThrowModelExceptionWhenResponseIsInvalidJson() { + OpenAIEmbeddingConfig config = new OpenAIEmbeddingConfig(); + config.setProvider("test-provider"); + config.setModel("BAAI/bge-m3"); + config.setApiKey("test-key"); + OpenAIEmbeddingModel model = new OpenAIEmbeddingModel(config); + model.setHttpClient(new HttpClient() { + @Override + public String post(String url, Map headers, String payload) { + return "not-json"; + } + }); + + ModelException exception = Assert.assertThrows( + ModelException.class, + () -> model.embed(Document.of("hello")) + ); + + Assert.assertEquals("Failed to parse embedding response.", exception.getMessage()); + Assert.assertNotNull(exception.getCause()); + } + + /** + * Verifies that diagnostic headers retain ordinary values without exposing credentials. + */ + @Test + public void shouldRedactAuthorizationHeaderForFailureLogging() { + Map headers = new LinkedHashMap<>(); + headers.put("Content-Type", "application/json"); + headers.put("Authorization", "Bearer test-key"); + + Map sanitizedHeaders = OpenAIEmbeddingModel.sanitizeHeadersForLogging(headers); + + Assert.assertEquals("application/json", sanitizedHeaders.get("Content-Type")); + Assert.assertEquals("[REDACTED]", sanitizedHeaders.get("Authorization")); + Assert.assertEquals("Bearer test-key", headers.get("Authorization")); + } } diff --git a/easy-agents-rag/easy-agents-rag-core/src/main/java/com/easyagents/rag/core/BgeM3ChunkSafety.java b/easy-agents-rag/easy-agents-rag-core/src/main/java/com/easyagents/rag/core/BgeM3ChunkSafety.java new file mode 100644 index 0000000..4865ae5 --- /dev/null +++ b/easy-agents-rag/easy-agents-rag-core/src/main/java/com/easyagents/rag/core/BgeM3ChunkSafety.java @@ -0,0 +1,230 @@ +package com.easyagents.rag.core; + +import java.text.Normalizer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +/** + * Provides conservative BGE-M3 token estimation and hard-limit text splitting. + * + *

BGE-M3 uses an XLM-R SentencePiece tokenizer. Loading the model vocabulary + * during ingestion would introduce runtime model downloads, so this guard counts + * NFKC-normalized Unicode code points. The estimate is intentionally conservative + * for ordinary Chinese and mixed-language knowledge documents.

+ */ +public final class BgeM3ChunkSafety { + + private static final double PREFERRED_BOUNDARY_MIN_RATIO = 0.8D; + + private BgeM3ChunkSafety() { + } + + /** + * Estimates the content token count after tokenizer-compatible normalization. + * + * @param content chunk content + * @return conservative token estimate, or zero for empty content + */ + public static int estimateContentTokens(String content) { + if (content == null || content.isEmpty()) { + return 0; + } + String normalized = Normalizer.normalize(content, Normalizer.Form.NFKC); + return normalized.codePointCount(0, normalized.length()); + } + + /** + * Checks whether content stays within the shared BGE-M3 embedding budget. + * + * @param content chunk content + * @return true when the content can be sent to BGE-M3 safely + */ + public static boolean isWithinHardLimit(String content) { + return estimateContentTokens(content) + <= RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT; + } + + /** + * Splits content using the shared BGE-M3 hard token limit. + * + * @param content source content + * @return ordered, non-overlapping source ranges + */ + public static List splitToHardLimit(String content) { + return splitToTokenLimit(content, RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT); + } + + /** + * Splits content to an explicit conservative token limit. + * + * @param content source content + * @param maxTokens maximum estimated tokens per part + * @return ordered, non-overlapping source ranges + * @throws IllegalArgumentException when maxTokens is not positive + */ + public static List splitToTokenLimit(String content, int maxTokens) { + if (maxTokens <= 0) { + throw new IllegalArgumentException("maxTokens must be positive"); + } + if (content == null || content.isEmpty()) { + return Collections.emptyList(); + } + if (estimateContentTokens(content) <= maxTokens) { + return Collections.singletonList(new ChunkRange(0, content.length())); + } + + List parts = new ArrayList(); + int start = 0; + while (start < content.length()) { + int end = findSplitEnd(content, start, maxTokens); + if (end <= start) { + throw new IllegalStateException("BGE-M3 hard split did not advance the text cursor"); + } + parts.add(new ChunkRange(start, end)); + start = end; + } + return parts; + } + + /** + * Finds a safe source boundary while preferring nearby sentence endings. + * + * @param content complete source content + * @param start current part start + * @param maxTokens maximum estimated tokens + * @return exclusive end offset + */ + private static int findSplitEnd(String content, int start, int maxTokens) { + int cursor = start; + int currentTokens = 0; + int preferredEnd = -1; + int preferredBoundaryThreshold = Math.max( + 1, (int) Math.floor(maxTokens * PREFERRED_BOUNDARY_MIN_RATIO)); + while (cursor < content.length()) { + int codePoint = content.codePointAt(cursor); + int next = cursor + Character.charCount(codePoint); + int nextTokens = currentTokens + estimateCodePointTokens(codePoint); + if (nextTokens > maxTokens) { + break; + } + cursor = next; + currentTokens = nextTokens; + if (currentTokens >= preferredBoundaryThreshold + && isPreferredBoundary(codePoint)) { + preferredEnd = cursor; + } + } + int candidateEnd = preferredEnd > start ? preferredEnd : cursor; + if (candidateEnd <= start) { + candidateEnd = start + Character.charCount(content.codePointAt(start)); + } + return fitWithinTokenLimit(content, start, candidateEnd, maxTokens); + } + + /** + * Shrinks a candidate when Unicode normalization expands its token estimate. + * + * @param content complete source content + * @param start part start + * @param candidateEnd proposed exclusive end + * @param maxTokens maximum estimated tokens + * @return exclusive end within the requested token limit + */ + private static int fitWithinTokenLimit( + String content, + int start, + int candidateEnd, + int maxTokens) { + if (estimateContentTokens(content.substring(start, candidateEnd)) + <= maxTokens) { + return candidateEnd; + } + int low = 1; + int high = content.codePointCount(start, candidateEnd); + while (low < high) { + int middle = low + (high - low + 1) / 2; + int end = content.offsetByCodePoints(start, middle); + if (estimateContentTokens(content.substring(start, end)) <= maxTokens) { + low = middle; + } else { + high = middle - 1; + } + } + return content.offsetByCodePoints(start, low); + } + + /** + * Estimates the normalized token cost of one source code point. + * + * @param codePoint source Unicode code point + * @return conservative token estimate for the code point + */ + private static int estimateCodePointTokens(int codePoint) { + if (codePoint >= 0 && codePoint <= 0x7F) { + return 1; + } + String source = new String(Character.toChars(codePoint)); + String normalized = Normalizer.normalize(source, Normalizer.Form.NFKC); + return Math.max(1, normalized.codePointCount(0, normalized.length())); + } + + /** + * Determines whether a code point is a suitable semantic boundary. + * + * @param codePoint current Unicode code point + * @return true for sentence, paragraph, or whitespace boundaries + */ + private static boolean isPreferredBoundary(int codePoint) { + return Character.isWhitespace(codePoint) + || codePoint == '。' + || codePoint == '！' + || codePoint == '？' + || codePoint == '；' + || codePoint == '.' + || codePoint == '!' + || codePoint == '?' + || codePoint == ';'; + } + + /** + * Immutable UTF-16 source range returned by the hard splitter. + */ + public static final class ChunkRange { + + private final int start; + private final int end; + + /** + * Creates a source range. + * + * @param start inclusive UTF-16 offset + * @param end exclusive UTF-16 offset + */ + public ChunkRange(int start, int end) { + if (start < 0 || end < start) { + throw new IllegalArgumentException("Invalid chunk range"); + } + this.start = start; + this.end = end; + } + + /** + * Returns the inclusive start offset. + * + * @return inclusive UTF-16 offset + */ + public int getStart() { + return start; + } + + /** + * Returns the exclusive end offset. + * + * @return exclusive UTF-16 offset + */ + public int getEnd() { + return end; + } + } +} diff --git a/easy-agents-rag/easy-agents-rag-core/src/main/java/com/easyagents/rag/core/RagDefaults.java b/easy-agents-rag/easy-agents-rag-core/src/main/java/com/easyagents/rag/core/RagDefaults.java index a78c377..7a2aaba 100644 --- a/easy-agents-rag/easy-agents-rag-core/src/main/java/com/easyagents/rag/core/RagDefaults.java +++ b/easy-agents-rag/easy-agents-rag-core/src/main/java/com/easyagents/rag/core/RagDefaults.java @@ -1,12 +1,26 @@ package com.easyagents.rag.core; +/** + * RAG ingestion default values and safety limits. + */ public final class RagDefaults { private RagDefaults() { } + /** Default strategy chunk size. */ public static final int CHUNK_SIZE = 512; + /** Default overlap size. */ public static final int OVERLAP_SIZE = 128; + /** Default Markdown heading level. */ public static final int MD_SPLITTER_LEVEL = 2; + /** Default spreadsheet rows per chunk. */ public static final int ROWS_PER_CHUNK = 1; + /** BGE-M3 maximum sequence length. */ + public static final int BGE_M3_MAX_SEQUENCE_TOKENS = 8192; + /** Reserved budget for model special tokens and tokenizer estimation differences. */ + public static final int BGE_M3_RESERVED_TOKENS = 128; + /** Hard upper token estimate for chunk content sent to BGE-M3. */ + public static final int BGE_M3_HARD_CHUNK_TOKEN_LIMIT = + BGE_M3_MAX_SEQUENCE_TOKENS - BGE_M3_RESERVED_TOKENS; } diff --git a/easy-agents-rag/easy-agents-rag-ingestion/src/main/java/com/easyagents/rag/ingestion/chunk/RagSplitStrategyRegistry.java b/easy-agents-rag/easy-agents-rag-ingestion/src/main/java/com/easyagents/rag/ingestion/chunk/RagSplitStrategyRegistry.java index d915857..b46bcc6 100644 --- a/easy-agents-rag/easy-agents-rag-ingestion/src/main/java/com/easyagents/rag/ingestion/chunk/RagSplitStrategyRegistry.java +++ b/easy-agents-rag/easy-agents-rag-ingestion/src/main/java/com/easyagents/rag/ingestion/chunk/RagSplitStrategyRegistry.java @@ -21,19 +21,19 @@ public class RagSplitStrategyRegistry { strategyCode = analysisResult.getRecommendedStrategyCode(); } String normalizedContent = analysisResult.getNormalizedContent(); + List chunks; if (RagStrategyCodes.MARKDOWN_SECTION.equals(strategyCode)) { - return buildMarkdownChunks(normalizedContent, strategyConfig); + chunks = buildMarkdownChunks(normalizedContent, strategyConfig); + } else if (RagStrategyCodes.OUTLINE_SECTION.equals(strategyCode)) { + chunks = buildOutlineChunks(normalizedContent, strategyConfig); + } else if (RagStrategyCodes.QA_PAIR.equals(strategyCode)) { + chunks = buildQaChunks(normalizedContent, strategyConfig); + } else if (RagStrategyCodes.CUSTOM_REGEX.equals(strategyCode)) { + chunks = buildRegexChunks(normalizedContent, strategyConfig); + } else { + chunks = buildParagraphChunks(normalizedContent, strategyConfig); } - if (RagStrategyCodes.OUTLINE_SECTION.equals(strategyCode)) { - return buildOutlineChunks(normalizedContent, strategyConfig); - } - if (RagStrategyCodes.QA_PAIR.equals(strategyCode)) { - return buildQaChunks(normalizedContent, strategyConfig); - } - if (RagStrategyCodes.CUSTOM_REGEX.equals(strategyCode)) { - return buildRegexChunks(normalizedContent, strategyConfig); - } - return buildParagraphChunks(normalizedContent, strategyConfig); + return postProcess(enforceBgeM3HardLimit(chunks)); } private List buildMarkdownChunks(String content, StrategyConfig strategyConfig) { @@ -143,7 +143,7 @@ public class RagSplitStrategyRegistry { )); } } - return postProcess(result); + return result; } private List buildQaChunks(String content, StrategyConfig strategyConfig) { @@ -162,7 +162,14 @@ public class RagSplitStrategyRegistry { Matcher questionMatcher = QUESTION_PREFIX.matcher(line); Matcher answerMatcher = ANSWER_PREFIX.matcher(line); if (questionMatcher.matches()) { - qaIndex = flushQaChunk(result, currentQuestion, questionSlices, answerSlices, qaIndex, strategyConfig); + qaIndex = flushQaChunk( + result, + content, + currentQuestion, + questionSlices, + answerSlices, + qaIndex, + strategyConfig); currentQuestion = questionMatcher.group(2).trim(); questionSlices = new ArrayList(); answerSlices = new ArrayList(); @@ -187,11 +194,19 @@ public class RagSplitStrategyRegistry { questionSlices.add(lineSlice); } } - flushQaChunk(result, currentQuestion, questionSlices, answerSlices, qaIndex, strategyConfig); - return postProcess(result); + flushQaChunk( + result, + content, + currentQuestion, + questionSlices, + answerSlices, + qaIndex, + strategyConfig); + return result; } private int flushQaChunk(List result, + String content, String currentQuestion, List questionSlices, List answerSlices, @@ -203,27 +218,59 @@ public class RagSplitStrategyRegistry { if (answerSlices == null || answerSlices.isEmpty()) { return qaIndex; } - String question = joinLineSlices(questionSlices); - String answer = joinLineSlices(answerSlices); - String baseContent = "问题：" + question + "\n答案：" + answer; - List subContents = baseContent.length() > safeChunkSize(strategyConfig) - ? splitLongContent(baseContent, strategyConfig.getChunkSize()) - : Collections.singletonList(baseContent); - int total = subContents.size(); - List sourceRanges = buildQaSourceRanges(questionSlices, answerSlices); - for (int i = 0; i < subContents.size(); i++) { + TextRange questionRange = mergeLineSlices(questionSlices); + TextRange rawAnswerRange = mergeLineSlices(answerSlices); + TextRange answerRange = rawAnswerRange == null + ? null + : trimRange(content, rawAnswerRange.start, rawAnswerRange.end); + if (questionRange == null || answerRange == null) { + return qaIndex; + } + + String question = currentQuestion.trim(); + String answer = content.substring(answerRange.start, answerRange.end); + String qaPrefix = "问题：" + question + "\n答案："; + int requestedTokens = Math.min( + safeChunkSize(strategyConfig), + RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT); + int prefixTokens = BgeM3ChunkSafety.estimateContentTokens(qaPrefix); + boolean includePrefix = prefixTokens + < RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT; + int answerTokenLimit = includePrefix + ? Math.max(1, Math.max(requestedTokens, prefixTokens + 1) - prefixTokens) + : RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT; + List answerPartRanges = new ArrayList(); + for (BgeM3ChunkSafety.ChunkRange answerPart + : BgeM3ChunkSafety.splitToTokenLimit(answer, answerTokenLimit)) { + TextRange absoluteAnswerRange = trimRange( + content, + answerRange.start + answerPart.getStart(), + answerRange.start + answerPart.getEnd()); + if (absoluteAnswerRange != null) { + answerPartRanges.add(absoluteAnswerRange); + } + } + int total = answerPartRanges.size(); + for (int i = 0; i < answerPartRanges.size(); i++) { + TextRange absoluteAnswerRange = answerPartRanges.get(i); + String answerFragment = content.substring( + absoluteAnswerRange.start, + absoluteAnswerRange.end); + List sourceRanges = new ArrayList(); + sourceRanges.add(questionRange); + sourceRanges.add(absoluteAnswerRange); RagChunk chunk = createChunk( RagChunkTypes.QA_PAIR, "Q" + qaIndex + " " + question, Collections.emptyList(), - subContents.get(i), + includePrefix ? qaPrefix + answerFragment : answerFragment, result.size() + 1, i + 1, total, sourceRanges ); chunk.setQuestion(question); - chunk.setAnswer(answer); + chunk.setAnswer(answerFragment); chunk.getOptions().put(RagMetadataKeys.QA_GROUP_ID, "qa-" + qaIndex); result.add(chunk); } @@ -254,7 +301,7 @@ public class RagSplitStrategyRegistry { )); index++; } - return postProcess(result); + return result; } private List buildRegexChunks(String content, StrategyConfig strategyConfig) { @@ -264,42 +311,129 @@ public class RagSplitStrategyRegistry { Pattern pattern = Pattern.compile(regex); Matcher matcher = pattern.matcher(content); int segmentStart = 0; + boolean retainRegexMatch = Boolean.TRUE.equals( + strategyConfig.getRetainRegexMatch()); while (matcher.find()) { index = addRegexChunk(result, content, segmentStart, matcher.start(), index); - segmentStart = matcher.end(); + // 保留时将匹配内容归入下一分块；默认继续从匹配结束位置开始。 + segmentStart = retainRegexMatch ? matcher.start() : matcher.end(); } addRegexChunk(result, content, segmentStart, content.length(), index); - return postProcess(result); + return result; } private List splitLongContent(String content, Integer chunkSize) { - int size = chunkSize == null || chunkSize.intValue() <= 0 ? RagDefaults.CHUNK_SIZE : chunkSize.intValue(); - String[] paragraphs = content.split("\\n\\s*\\n"); + int size = chunkSize == null || chunkSize.intValue() <= 0 + ? RagDefaults.CHUNK_SIZE + : Math.min(chunkSize.intValue(), RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT); List parts = new ArrayList(); - StringBuilder current = new StringBuilder(); - for (String paragraph : paragraphs) { - String text = paragraph.trim(); - if (!StringUtil.hasText(text)) { - continue; + for (BgeM3ChunkSafety.ChunkRange range + : BgeM3ChunkSafety.splitToTokenLimit(content, size)) { + String part = content.substring(range.getStart(), range.getEnd()).trim(); + if (StringUtil.hasText(part)) { + parts.add(part); } - if (current.length() > 0 && current.length() + text.length() + 2 > size) { - parts.add(current.toString().trim()); - current = new StringBuilder(); - } - if (current.length() > 0) { - current.append("\n\n"); - } - current.append(text); - } - if (current.length() > 0) { - parts.add(current.toString().trim()); - } - if (parts.isEmpty()) { - parts.add(content); } return parts; } + /** + * Applies the shared BGE-M3 context safety guard after strategy-specific splitting. + * + * @param chunks chunks produced by a semantic or length strategy + * @return chunks whose conservative token estimates stay within the hard limit + */ + private List enforceBgeM3HardLimit(List chunks) { + List result = new ArrayList(); + for (RagChunk chunk : chunks) { + String content = chunk.getContent(); + if (BgeM3ChunkSafety.isWithinHardLimit(content)) { + result.add(chunk); + continue; + } + List parts = + BgeM3ChunkSafety.splitToHardLimit(content); + for (int i = 0; i < parts.size(); i++) { + BgeM3ChunkSafety.ChunkRange range = parts.get(i); + result.add(copyHardSplitChunk( + chunk, + content, + new TextRange(range.getStart(), range.getEnd()), + i + 1, + parts.size())); + } + } + return result; + } + + /** + * Copies a strategy chunk while replacing only content and hard-split metadata. + * + * @param source original strategy chunk + * @param sourceContent original chunk content + * @param range selected content range + * @param partNo forced part number + * @param partTotal forced part count + * @return copied hard-split chunk + */ + private RagChunk copyHardSplitChunk(RagChunk source, + String sourceContent, + TextRange range, + int partNo, + int partTotal) { + RagChunk copy = new RagChunk(); + copy.setChunkType(source.getChunkType()); + copy.setSourceLabel(source.getSourceLabel()); + copy.setHeadingPath(new ArrayList(source.getHeadingPath())); + copy.setContent(sourceContent.substring(range.start, range.end)); + copy.setQuestion(source.getQuestion()); + copy.setAnswer(RagChunkTypes.QA_PAIR.equals(source.getChunkType()) + ? copy.getContent() + : source.getAnswer()); + copy.setPartNo(Integer.valueOf(partNo)); + copy.setPartTotal(Integer.valueOf(partTotal)); + copy.setWarnings(new ArrayList(source.getWarnings())); + copy.setOptions(new LinkedHashMap(source.getOptions())); + copy.getOptions().put("hardSplit", Boolean.TRUE); + copy.getOptions().put("hardSplitTokenLimit", + Integer.valueOf(RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT)); + adjustSingleSourceRange(copy, sourceContent.length(), range); + return copy; + } + + /** + * Narrows an exact single source range to the forced subrange. + * + * @param chunk copied hard-split chunk + * @param sourceLength original chunk content length + * @param relativeRange forced range relative to the original chunk + */ + @SuppressWarnings("unchecked") + private void adjustSingleSourceRange(RagChunk chunk, int sourceLength, TextRange relativeRange) { + Object rawRanges = chunk.getOptions().get(RagMetadataKeys.SOURCE_RANGES); + if (!(rawRanges instanceof List) || ((List) rawRanges).size() != 1) { + return; + } + Object rawRange = ((List) rawRanges).get(0); + if (!(rawRange instanceof Map)) { + return; + } + Map sourceRange = (Map) rawRange; + Object rawStart = sourceRange.get("start"); + Object rawEnd = sourceRange.get("end"); + if (!(rawStart instanceof Number) || !(rawEnd instanceof Number)) { + return; + } + int sourceStart = ((Number) rawStart).intValue(); + int sourceEnd = ((Number) rawEnd).intValue(); + if (sourceEnd - sourceStart != sourceLength) { + return; + } + chunk.getOptions().put(RagMetadataKeys.SOURCE_RANGES, + toSourceRangeMaps(Collections.singletonList( + new TextRange(sourceStart + relativeRange.start, sourceStart + relativeRange.end)))); + } + private List postProcess(List chunks) { List result = new ArrayList(); Set dedup = new HashSet(); @@ -318,7 +452,8 @@ public class RagSplitStrategyRegistry { } chunk.setChunkId("chunk-" + index); chunk.setCharCount(Integer.valueOf(content.length())); - chunk.setTokenEstimate(Integer.valueOf(Math.max(1, content.length() / 4))); + chunk.setTokenEstimate(Integer.valueOf( + Math.max(1, BgeM3ChunkSafety.estimateContentTokens(content)))); result.add(chunk); index++; } @@ -460,7 +595,14 @@ public class RagSplitStrategyRegistry { private int safeOverlap(StrategyConfig strategyConfig) { Integer overlapSize = strategyConfig.getOverlapSize(); - return overlapSize == null || overlapSize.intValue() < 0 ? RagDefaults.OVERLAP_SIZE : overlapSize.intValue(); + int overlap = overlapSize == null || overlapSize.intValue() < 0 + ? RagDefaults.OVERLAP_SIZE + : overlapSize.intValue(); + int chunkSize = safeChunkSize(strategyConfig); + if (overlap >= chunkSize) { + throw new IllegalArgumentException("overlapSize must be smaller than chunkSize"); + } + return overlap; } private String joinAndTrim(List lines) { diff --git a/easy-agents-rag/easy-agents-rag-ingestion/src/main/java/com/easyagents/rag/ingestion/model/StrategyConfig.java b/easy-agents-rag/easy-agents-rag-ingestion/src/main/java/com/easyagents/rag/ingestion/model/StrategyConfig.java index b4d4da2..07971f0 100644 --- a/easy-agents-rag/easy-agents-rag-ingestion/src/main/java/com/easyagents/rag/ingestion/model/StrategyConfig.java +++ b/easy-agents-rag/easy-agents-rag-ingestion/src/main/java/com/easyagents/rag/ingestion/model/StrategyConfig.java @@ -11,6 +11,10 @@ public class StrategyConfig implements Serializable { private Integer chunkSize = RagDefaults.CHUNK_SIZE; private Integer overlapSize = RagDefaults.OVERLAP_SIZE; private String regex; + /** + * Whether a custom-regex match is retained at the beginning of the next chunk. + */ + private Boolean retainRegexMatch; private Integer rowsPerChunk = RagDefaults.ROWS_PER_CHUNK; private Integer mdSplitterLevel = RagDefaults.MD_SPLITTER_LEVEL; @@ -24,6 +28,7 @@ public class StrategyConfig implements Serializable { copy.setChunkSize(this.chunkSize); copy.setOverlapSize(this.overlapSize); copy.setRegex(this.regex); + copy.setRetainRegexMatch(this.retainRegexMatch); copy.setRowsPerChunk(this.rowsPerChunk); copy.setMdSplitterLevel(this.mdSplitterLevel); return copy; @@ -61,6 +66,25 @@ public class StrategyConfig implements Serializable { this.regex = regex; } + /** + * Returns whether custom-regex matches are retained in the next chunk. + * + * @return {@code true} to retain each match at the beginning of the next chunk; + * otherwise {@code false} or {@code null} + */ + public Boolean getRetainRegexMatch() { + return retainRegexMatch; + } + + /** + * Sets whether custom-regex matches are retained in the next chunk. + * + * @param retainRegexMatch {@code true} to retain each match; {@code false} to discard it + */ + public void setRetainRegexMatch(Boolean retainRegexMatch) { + this.retainRegexMatch = retainRegexMatch; + } + public Integer getRowsPerChunk() { return rowsPerChunk; } diff --git a/easy-agents-rag/easy-agents-rag-ingestion/src/test/java/com/easyagents/rag/ingestion/RagIngestionPipelineTest.java b/easy-agents-rag/easy-agents-rag-ingestion/src/test/java/com/easyagents/rag/ingestion/RagIngestionPipelineTest.java index d1d94f3..669b024 100644 --- a/easy-agents-rag/easy-agents-rag-ingestion/src/test/java/com/easyagents/rag/ingestion/RagIngestionPipelineTest.java +++ b/easy-agents-rag/easy-agents-rag-ingestion/src/test/java/com/easyagents/rag/ingestion/RagIngestionPipelineTest.java @@ -2,6 +2,7 @@ package com.easyagents.rag.ingestion; import com.easyagents.rag.core.RagChunk; import com.easyagents.rag.core.RagChunkTypes; +import com.easyagents.rag.core.RagDefaults; import com.easyagents.rag.core.RagStrategyCodes; import com.easyagents.rag.ingestion.analysis.DocumentStructureAnalyzer; import com.easyagents.rag.ingestion.chunk.RagSplitStrategyRegistry; @@ -97,6 +98,157 @@ public class RagIngestionPipelineTest { assertHasValidSourceRanges(analysis, chunks.get(0)); } + @Test + public void shouldEnforceBgeM3HardLimitForAutoQaTxt() { + StringBuilder qa = new StringBuilder("问：自动导入为什么失败？\n答："); + int i = 0; + while (qa.length() < 77946) { + qa.append((char) ('一' + i % 20)); + if (i > 0 && i % 700 == 0) { + qa.append('。'); + } + i++; + } + qa.setLength(77946); + AnalysisResult analysis = recommender.recommend(analyzer.analyze(qa.toString(), "txt")); + StrategyConfig config = StrategyConfig.defaults(); + config.setStrategyCode(RagStrategyCodes.AUTO); + + List chunks = registry.split(analysis, config); + + Assert.assertEquals(RagStrategyCodes.QA_PAIR, analysis.getRecommendedStrategyCode()); + Assert.assertEquals(77946, analysis.getNormalizedContent().length()); + Assert.assertTrue(chunks.size() > 1); + assertWithinBgeM3HardLimit(chunks); + Assert.assertTrue(chunks.stream().allMatch( + chunk -> chunk.getTokenEstimate().intValue() <= RagDefaults.CHUNK_SIZE)); + Assert.assertTrue(chunks.stream().allMatch( + chunk -> chunk.getAnswer() != null + && chunk.getAnswer().length() < analysis.getNormalizedContent().length())); + assertHasValidSourceRanges(analysis, chunks.get(0)); + } + + @Test + public void shouldEnforceBgeM3HardLimitForCustomRegexChunk() { + StringBuilder content = new StringBuilder(); + for (int i = 0; i < 20000; i++) { + content.append((char) ('甲' + i % 16)); + if (i > 0 && i % 997 == 0) { + content.append('；'); + } + } + AnalysisResult analysis = recommender.recommend(analyzer.analyze(content.toString(), "txt")); + StrategyConfig config = StrategyConfig.defaults(); + config.setStrategyCode(RagStrategyCodes.CUSTOM_REGEX); + config.setRegex("\\|NEVER_MATCH\\|"); + + List chunks = registry.split(analysis, config); + + Assert.assertTrue(chunks.size() > 1); + assertWithinBgeM3HardLimit(chunks); + assertHasValidSourceRanges(analysis, chunks.get(0)); + } + + /** + * Verifies that custom-regex matches remain discarded when the new option is omitted. + */ + @Test + public void shouldDiscardCustomRegexMatchByDefault() { + AnalysisResult analysis = new AnalysisResult(); + analysis.setNormalizedContent( + "第一个问题及其答案内容。第二个问题及其答案内容。"); + StrategyConfig config = StrategyConfig.defaults(); + config.setStrategyCode(RagStrategyCodes.CUSTOM_REGEX); + config.setRegex(""); + + List chunks = registry.split(analysis, config); + + Assert.assertEquals(2, chunks.size()); + Assert.assertEquals("第一个问题及其答案内容。", chunks.get(0).getContent()); + Assert.assertEquals("第二个问题及其答案内容。", chunks.get(1).getContent()); + } + + /** + * Verifies that an enabled option retains each match at the beginning of the next chunk. + */ + @Test + public void shouldRetainCustomRegexMatchInNextChunk() { + AnalysisResult analysis = new AnalysisResult(); + analysis.setNormalizedContent( + "第一个问题及其答案内容。第二个问题及其答案内容。"); + StrategyConfig config = StrategyConfig.defaults(); + config.setStrategyCode(RagStrategyCodes.CUSTOM_REGEX); + config.setRegex(""); + config.setRetainRegexMatch(Boolean.TRUE); + + List chunks = registry.split(analysis, config); + + Assert.assertEquals(2, chunks.size()); + Assert.assertEquals( + "第一个问题及其答案内容。", + chunks.get(0).getContent()); + Assert.assertEquals( + "第二个问题及其答案内容。", + chunks.get(1).getContent()); + assertHasValidSourceRanges(analysis, chunks.get(1)); + } + + @Test + public void shouldUseChineseAwareTokenEstimate() { + String content = "知识库自动导入向量化失败，需要检查分块大小和模型上下文窗口。"; + AnalysisResult analysis = recommender.recommend(analyzer.analyze(content, "txt")); + StrategyConfig config = StrategyConfig.defaults(); + config.setStrategyCode(RagStrategyCodes.PARAGRAPH_LENGTH); + + List chunks = registry.split(analysis, config); + + Assert.assertEquals(Integer.valueOf(content.codePointCount(0, content.length())), + chunks.get(0).getTokenEstimate()); + } + + @Test + public void shouldRejectOverlapEqualToChunkSize() { + AnalysisResult analysis = recommender.recommend( + analyzer.analyze("知识库分块参数需要保证游标持续向前推进。", "txt")); + StrategyConfig config = StrategyConfig.defaults(); + config.setStrategyCode(RagStrategyCodes.PARAGRAPH_LENGTH); + config.setChunkSize(128); + config.setOverlapSize(128); + + try { + registry.split(analysis, config); + Assert.fail("overlapSize 等于 chunkSize 时应拒绝分块"); + } catch (IllegalArgumentException expected) { + Assert.assertTrue(expected.getMessage().contains("overlapSize")); + } + } + + @Test + public void shouldRejectOverlapGreaterThanChunkSize() { + AnalysisResult analysis = recommender.recommend( + analyzer.analyze("知识库分块参数需要保证游标持续向前推进。", "txt")); + StrategyConfig config = StrategyConfig.defaults(); + config.setStrategyCode(RagStrategyCodes.PARAGRAPH_LENGTH); + config.setChunkSize(128); + config.setOverlapSize(512); + + try { + registry.split(analysis, config); + Assert.fail("overlapSize 大于 chunkSize 时应拒绝分块"); + } catch (IllegalArgumentException expected) { + Assert.assertTrue(expected.getMessage().contains("overlapSize")); + } + } + + private void assertWithinBgeM3HardLimit(List chunks) { + for (RagChunk chunk : chunks) { + Assert.assertNotNull(chunk.getTokenEstimate()); + Assert.assertTrue( + "chunk token estimate exceeds BGE-M3 hard limit: " + chunk.getTokenEstimate(), + chunk.getTokenEstimate().intValue() <= RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT); + } + } + @SuppressWarnings("unchecked") private void assertHasValidSourceRanges(AnalysisResult analysis, RagChunk chunk) { Object rawRanges = chunk.getOptions().get("sourceRanges"); diff --git a/easy-agents-search-engine/easy-agents-search-engine-es/src/main/java/com/easyagents/engine/es/ElasticSearcher.java b/easy-agents-search-engine/easy-agents-search-engine-es/src/main/java/com/easyagents/engine/es/ElasticSearcher.java index 7f7f07a..f3a8190 100644 --- a/easy-agents-search-engine/easy-agents-search-engine-es/src/main/java/com/easyagents/engine/es/ElasticSearcher.java +++ b/easy-agents-search-engine/easy-agents-search-engine-es/src/main/java/com/easyagents/engine/es/ElasticSearcher.java @@ -2,8 +2,6 @@ package com.easyagents.engine.es; import co.elastic.clients.elasticsearch.ElasticsearchClient; import co.elastic.clients.elasticsearch.core.*; -import co.elastic.clients.elasticsearch.core.bulk.BulkOperation; -import co.elastic.clients.elasticsearch.core.bulk.IndexOperation; import co.elastic.clients.elasticsearch.core.search.SourceConfig; import co.elastic.clients.json.JsonData; import co.elastic.clients.json.jackson.JacksonJsonpMapper; @@ -25,19 +23,46 @@ import org.slf4j.LoggerFactory; import javax.net.ssl.SSLContext; import javax.net.ssl.TrustManager; import javax.net.ssl.X509TrustManager; +import java.io.IOException; import java.security.KeyManagementException; import java.security.NoSuchAlgorithmException; import java.security.cert.X509Certificate; import java.util.*; -public class ElasticSearcher implements DocumentSearcher { +/** + * 基于 Elasticsearch 的关键词搜索器。 + * + *

底层客户端线程安全并与搜索器实例共享生命周期，避免逐文档重复创建网络连接。

+ */ +public class ElasticSearcher implements DocumentSearcher, AutoCloseable { private static final Logger LOG = LoggerFactory.getLogger(ElasticSearcher.class); + static final int MAX_BULK_OPERATIONS = 200; + static final long MAX_BULK_ESTIMATED_BYTES = 5L * 1024L * 1024L; + private static final long BULK_OPERATION_METADATA_BYTES = 128L; private final ESConfig esConfig; + private final JacksonJsonpMapper jsonpMapper = new JacksonJsonpMapper(); + private final ElasticsearchTransport transport; + private final ElasticsearchClient client; + /** + * 创建 Elasticsearch 搜索器及其共享客户端。 + * + * @param esConfig Elasticsearch 配置 + * @throws IllegalStateException 客户端初始化失败时抛出 + */ public ElasticSearcher(ESConfig esConfig) { - this.esConfig = esConfig; + this.esConfig = Objects.requireNonNull(esConfig, "ESConfig 不能为 null"); + RestClient openedRestClient = null; + try { + openedRestClient = buildRestClient(); + this.transport = new RestClientTransport(openedRestClient, jsonpMapper); + this.client = new ElasticsearchClient(transport); + } catch (Exception e) { + closeQuietly(openedRestClient); + throw new IllegalStateException("初始化 Elasticsearch 客户端失败", e); + } } // 忽略 SSL 的 client 构建逻辑 @@ -76,52 +101,128 @@ public class ElasticSearcher implements DocumentSearcher { /** - * 添加文档到Elasticsearch + * 添加文档到 Elasticsearch。 + * + * @param document 待写入文档 + * @return 写入成功时返回 {@code true} */ @Override public boolean addDocument(Document document) { - if (document == null || document.getContent() == null) { + return addDocuments(Collections.singletonList(document)); + } + + /** + * 使用有界 Bulk 请求批量写入文档。 + * + * @param documents 待写入文档 + * @return 全部文档写入成功时返回 {@code true} + */ + @Override + public boolean addDocuments(List documents) { + if (!areValidDocuments(documents)) { return false; } - - RestClient restClient = null; - ElasticsearchTransport transport = null; try { - restClient = buildRestClient(); - transport = new RestClientTransport(restClient, new JacksonJsonpMapper()); - ElasticsearchClient client = new ElasticsearchClient(transport); - - Map source = buildSource(document); - String documentId = document.getId().toString(); - IndexOperation indexOp = IndexOperation.of(i -> i - .index(esConfig.getIndexName()) - .id(documentId) - .document(JsonData.of(source)) - ); - - BulkOperation bulkOp = BulkOperation.of(b -> b.index(indexOp)); - BulkRequest request = BulkRequest.of(b -> b.operations(Collections.singletonList(bulkOp))); - BulkResponse response = client.bulk(request); - return !response.errors(); - + List> batches = partitionDocuments(documents); + for (int batchIndex = 0; batchIndex < batches.size(); batchIndex++) { + List batch = batches.get(batchIndex); + BulkResponse response = client.bulk(buildBulkRequest(batch)); + if (!isBulkSuccessful(response, "写入", batchIndex, batch.size())) { + return false; + } + } + return true; } catch (Exception e) { - LOG.error(e.getMessage(), e); + LOG.error("Elasticsearch 批量写入异常: count={}", documents.size(), e); return false; - } finally { - closeResources(transport, restClient); } } + /** + * 按操作数量和序列化后的估算字节数拆分 Bulk 批次。 + * + * @param documents 已完成校验的文档 + * @return 有界文档批次 + * @throws IOException 文档序列化失败时抛出 + * @throws IllegalArgumentException 单个文档超过批次字节限制时抛出 + */ + List> partitionDocuments(List documents) throws IOException { + List> batches = new ArrayList<>(); + List currentBatch = new ArrayList<>(Math.min(documents.size(), MAX_BULK_OPERATIONS)); + long currentBatchBytes = 0L; + for (Document document : documents) { + long documentBytes = estimateBulkOperationBytes(document); + if (documentBytes > MAX_BULK_ESTIMATED_BYTES) { + throw new IllegalArgumentException( + "单个 Elasticsearch 索引文档超过 Bulk 字节限制: id=" + document.getId()); + } + if (!currentBatch.isEmpty() + && (currentBatch.size() >= MAX_BULK_OPERATIONS + || currentBatchBytes + documentBytes > MAX_BULK_ESTIMATED_BYTES)) { + batches.add(currentBatch); + currentBatch = new ArrayList<>(Math.min(documents.size(), MAX_BULK_OPERATIONS)); + currentBatchBytes = 0L; + } + currentBatch.add(document); + currentBatchBytes += documentBytes; + } + if (!currentBatch.isEmpty()) { + batches.add(currentBatch); + } + return batches; + } + + private long estimateBulkOperationBytes(Document document) throws IOException { + long sourceBytes = jsonpMapper.objectMapper().writeValueAsBytes(buildSource(document)).length; + return sourceBytes + + utf8Length(esConfig.getIndexName()) + + utf8Length(document.getId().toString()) + + BULK_OPERATION_METADATA_BYTES; + } + + private int utf8Length(String value) { + return value == null ? 0 : value.getBytes(java.nio.charset.StandardCharsets.UTF_8).length; + } + + /** + * 构建包含全部文档的 Bulk 请求。 + * + * @param documents 已完成校验的文档 + * @return Bulk 请求 + */ + BulkRequest buildBulkRequest(List documents) { + BulkRequest.Builder builder = new BulkRequest.Builder(); + for (Document document : documents) { + builder.operations(operation -> operation.index(index -> index + .index(esConfig.getIndexName()) + .id(document.getId().toString()) + .document(JsonData.of(buildSource(document))) + )); + } + return builder.build(); + } + + private boolean areValidDocuments(List documents) { + if (documents == null || documents.isEmpty()) { + return false; + } + for (Document document : documents) { + if (document == null || document.getId() == null || document.getContent() == null) { + return false; + } + } + return true; + } + + /** + * 按请求条件搜索文档。 + * + * @param request 搜索请求 + * @return 命中文档 + */ @Override public List searchDocuments(KeywordSearchRequest request) { - RestClient restClient = null; - ElasticsearchTransport transport = null; - try { - restClient = buildRestClient(); - transport = new RestClientTransport(restClient, new JacksonJsonpMapper()); - ElasticsearchClient client = new ElasticsearchClient(transport); - SearchResponse response = client.search(buildSearchRequest(request), Map.class); List results = new ArrayList<>(); response.hits().hits().forEach(hit -> { @@ -137,54 +238,114 @@ public class ElasticSearcher implements DocumentSearcher { } catch (Exception e) { LOG.error(e.getMessage(), e); return Collections.emptyList(); - } finally { - closeResources(transport, restClient); } } + /** + * 删除指定文档。 + * + * @param id 文档 ID + * @return 删除成功时返回 {@code true} + */ @Override public boolean deleteDocument(Object id) { - if (id == null) { + return deleteDocuments(Collections.singletonList(id)); + } + + /** + * 使用有界 Bulk 请求批量删除文档。 + * + *

删除不存在的文档按幂等成功处理；仅 Bulk 条目包含实际错误时返回失败。

+ * + * @param ids 文档 ID 集合 + * @return 全部删除操作成功时返回 {@code true} + */ + @Override + public boolean deleteDocuments(Collection ids) { + if (ids == null || ids.isEmpty()) { return false; } - - RestClient restClient = null; - ElasticsearchTransport transport = null; + List documentIds = new ArrayList<>(ids.size()); + for (Object id : ids) { + if (id == null) { + return false; + } + documentIds.add(id.toString()); + } try { - restClient = buildRestClient(); - transport = new RestClientTransport(restClient, new JacksonJsonpMapper()); - ElasticsearchClient client = new ElasticsearchClient(transport); - - DeleteRequest request = DeleteRequest.of(d -> d - .index(esConfig.getIndexName()) - .id(id.toString()) - ); - - DeleteResponse response = client.delete(request); - return response.result() == co.elastic.clients.elasticsearch._types.Result.Deleted; - + for (int start = 0, batchIndex = 0; + start < documentIds.size(); + start += MAX_BULK_OPERATIONS, batchIndex++) { + int end = Math.min(start + MAX_BULK_OPERATIONS, documentIds.size()); + List batch = documentIds.subList(start, end); + BulkResponse response = client.bulk(buildDeleteBulkRequest(batch)); + if (!isBulkSuccessful(response, "删除", batchIndex, batch.size())) { + return false; + } + } + return true; } catch (Exception e) { - LOG.error("Error deleting document with id: " + id, e); + LOG.error("Elasticsearch 批量删除异常: count={}", documentIds.size(), e); return false; - } finally { - closeResources(transport, restClient); } } + /** + * 构建批量删除请求。 + * + * @param ids 文档 ID + * @return Bulk 删除请求 + */ + BulkRequest buildDeleteBulkRequest(List ids) { + BulkRequest.Builder builder = new BulkRequest.Builder(); + for (String id : ids) { + builder.operations(operation -> operation.delete(delete -> delete + .index(esConfig.getIndexName()) + .id(id) + )); + } + return builder.build(); + } + + private boolean isBulkSuccessful(BulkResponse response, + String operation, + int batchIndex, + int batchSize) { + if (response != null && !response.errors()) { + return true; + } + if (response == null) { + LOG.error("Elasticsearch 批量{}未返回结果: batchIndex={}, batchSize={}", + operation, batchIndex, batchSize); + return false; + } + response.items().stream() + .filter(item -> item.error() != null) + .forEach(item -> LOG.error( + "Elasticsearch 批量{}失败: batchIndex={}, index={}, id={}, status={}, reason={}", + operation, + batchIndex, + item.index(), + item.id(), + item.status(), + item.error().reason() + )); + return false; + } + + /** + * 更新指定文档。 + * + * @param document 待更新文档 + * @return 更新成功时返回 {@code true} + */ @Override public boolean updateDocument(Document document) { if (document == null || document.getId() == null) { return false; } - RestClient restClient = null; - ElasticsearchTransport transport = null; - try { - restClient = buildRestClient(); - transport = new RestClientTransport(restClient, new JacksonJsonpMapper()); - ElasticsearchClient client = new ElasticsearchClient(transport); - UpdateRequest, Map> request = UpdateRequest.of(u -> u .index(esConfig.getIndexName()) .id(document.getId().toString()) @@ -199,20 +360,17 @@ public class ElasticSearcher implements DocumentSearcher { } catch (Exception e) { LOG.error("Error updating document with id: " + document.getId(), e); return false; - } finally { - closeResources(transport, restClient); } } - - private void closeResources(AutoCloseable... closeables) { - for (AutoCloseable closeable : closeables) { - try { - if (closeable != null) - closeable.close(); - } catch (Exception e) { - LOG.error("Error closing resource", e); - } + private static void closeQuietly(AutoCloseable closeable) { + if (closeable == null) { + return; + } + try { + closeable.close(); + } catch (Exception ignored) { + // 初始化失败时仅执行尽力清理，原始异常由构造方法继续抛出。 } } @@ -284,19 +442,27 @@ public class ElasticSearcher implements DocumentSearcher { ); } + /** + * 检查 Elasticsearch 服务是否可用。 + * + * @return 服务可访问时返回 {@code true} + */ public boolean checkAvailable() { - RestClient restClient = null; - ElasticsearchTransport transport = null; try { - restClient = buildRestClient(); - transport = new RestClientTransport(restClient, new JacksonJsonpMapper()); - ElasticsearchClient client = new ElasticsearchClient(transport); return client.info() != null; } catch (Exception e) { LOG.error("Elasticsearch availability check failed", e); return false; - } finally { - closeResources(transport, restClient); } } + + /** + * 关闭共享传输层及其底层 RestClient。 + * + * @throws IOException 关闭客户端资源失败时抛出 + */ + @Override + public void close() throws IOException { + transport.close(); + } } diff --git a/easy-agents-search-engine/easy-agents-search-engine-es/src/test/java/com/easyagents/engine/es/ElasticSearcherQueryBuilderTest.java b/easy-agents-search-engine/easy-agents-search-engine-es/src/test/java/com/easyagents/engine/es/ElasticSearcherQueryBuilderTest.java index 31b929d..a3c1ecc 100644 --- a/easy-agents-search-engine/easy-agents-search-engine-es/src/test/java/com/easyagents/engine/es/ElasticSearcherQueryBuilderTest.java +++ b/easy-agents-search-engine/easy-agents-search-engine-es/src/test/java/com/easyagents/engine/es/ElasticSearcherQueryBuilderTest.java @@ -1,48 +1,185 @@ package com.easyagents.engine.es; import co.elastic.clients.elasticsearch.core.SearchRequest; +import co.elastic.clients.elasticsearch.core.BulkRequest; import com.easyagents.core.document.Document; import com.easyagents.search.engine.service.KeywordSearchMetadataKeys; import com.easyagents.search.engine.service.KeywordSearchRequest; import org.junit.Assert; import org.junit.Test; +import java.util.ArrayList; +import java.util.List; import java.util.Map; +/** + * {@link ElasticSearcher} 请求构建回归测试。 + */ public class ElasticSearcherQueryBuilderTest { + /** + * 验证搜索请求同时包含多字段检索和知识库过滤。 + * + * @throws Exception Elasticsearch 客户端关闭失败时抛出 + */ @Test - public void shouldBuildSearchRequestWithMultiMatchAndKnowledgeFilter() { - ElasticSearcher searcher = new ElasticSearcher(config()); - KeywordSearchRequest request = KeywordSearchRequest.of("客服", 5); - request.setKnowledgeId("100"); + public void shouldBuildSearchRequestWithMultiMatchAndKnowledgeFilter() throws Exception { + try (ElasticSearcher searcher = new ElasticSearcher(config())) { + KeywordSearchRequest request = KeywordSearchRequest.of("客服", 5); + request.setKnowledgeId("100"); - SearchRequest searchRequest = searcher.buildSearchRequest(request); + SearchRequest searchRequest = searcher.buildSearchRequest(request); - Assert.assertEquals(5, searchRequest.size().intValue()); - Assert.assertNotNull(searchRequest.query().bool()); - Assert.assertEquals(1, searchRequest.query().bool().must().size()); - Assert.assertNotNull(searchRequest.query().bool().must().get(0).multiMatch()); - Assert.assertEquals(2, searchRequest.query().bool().must().get(0).multiMatch().fields().size()); - Assert.assertEquals(1, searchRequest.query().bool().filter().size()); - Assert.assertEquals("knowledgeId", searchRequest.query().bool().filter().get(0).term().field()); + Assert.assertEquals(5, searchRequest.size().intValue()); + Assert.assertNotNull(searchRequest.query().bool()); + Assert.assertEquals(1, searchRequest.query().bool().must().size()); + Assert.assertNotNull(searchRequest.query().bool().must().get(0).multiMatch()); + Assert.assertEquals(2, searchRequest.query().bool().must().get(0).multiMatch().fields().size()); + Assert.assertEquals(1, searchRequest.query().bool().filter().size()); + Assert.assertEquals("knowledgeId", searchRequest.query().bool().filter().get(0).term().field()); + } } + /** + * 验证知识库 ID 会写入顶层字段。 + * + * @throws Exception Elasticsearch 客户端关闭失败时抛出 + */ @Test - public void shouldExtractKnowledgeIdToTopLevelSource() { - ElasticSearcher searcher = new ElasticSearcher(config()); + public void shouldExtractKnowledgeIdToTopLevelSource() throws Exception { + try (ElasticSearcher searcher = new ElasticSearcher(config())) { + Document document = new Document(); + document.setId("1"); + document.setTitle("title"); + document.setContent("content"); + document.addMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID, "100"); + + Map source = searcher.buildSource(document); + + Assert.assertEquals("100", source.get(KeywordSearchMetadataKeys.KNOWLEDGE_ID)); + Assert.assertTrue(source.get("metadataMap") instanceof Map); + } + } + + /** + * 验证多个文档会进入同一个 Bulk 请求并保留稳定文档 ID。 + * + * @throws Exception Elasticsearch 客户端关闭失败时抛出 + */ + @Test + public void shouldBuildSingleBulkRequestForMultipleDocuments() throws Exception { + try (ElasticSearcher searcher = new ElasticSearcher(config())) { + Document first = new Document(); + first.setId("first"); + first.setContent("first content"); + Document second = new Document(); + second.setId("second"); + second.setContent("second content"); + + BulkRequest request = searcher.buildBulkRequest(List.of(first, second)); + + Assert.assertEquals(2, request.operations().size()); + Assert.assertEquals("first", request.operations().get(0).index().id()); + Assert.assertEquals("second", request.operations().get(1).index().id()); + } + } + + /** + * 验证 Bulk 写入会按最大操作数拆分，避免单次请求无界增长。 + * + * @throws Exception Elasticsearch 客户端关闭或文档序列化失败时抛出 + */ + @Test + public void shouldPartitionBulkRequestsByOperationCount() throws Exception { + try (ElasticSearcher searcher = new ElasticSearcher(config())) { + List documents = new ArrayList<>(); + for (int index = 0; index <= ElasticSearcher.MAX_BULK_OPERATIONS; index++) { + documents.add(document("count-" + index, "content")); + } + + List> batches = searcher.partitionDocuments(documents); + + Assert.assertEquals(2, batches.size()); + Assert.assertEquals(ElasticSearcher.MAX_BULK_OPERATIONS, batches.get(0).size()); + Assert.assertEquals(1, batches.get(1).size()); + } + } + + /** + * 验证 Bulk 写入会按序列化字节数拆分。 + * + * @throws Exception Elasticsearch 客户端关闭或文档序列化失败时抛出 + */ + @Test + public void shouldPartitionBulkRequestsByEstimatedBytes() throws Exception { + try (ElasticSearcher searcher = new ElasticSearcher(config())) { + int contentLength = (int) (ElasticSearcher.MAX_BULK_ESTIMATED_BYTES / 2L); + Document first = document("bytes-first", "a".repeat(contentLength)); + Document second = document("bytes-second", "b".repeat(contentLength)); + + List> batches = searcher.partitionDocuments(List.of(first, second)); + + Assert.assertEquals(2, batches.size()); + Assert.assertEquals(1, batches.get(0).size()); + Assert.assertEquals(1, batches.get(1).size()); + } + } + + /** + * 验证单个超限文档会在发送请求前失败。 + * + * @throws Exception Elasticsearch 客户端关闭或文档序列化失败时抛出 + */ + @Test + public void shouldRejectSingleDocumentOverByteLimit() throws Exception { + try (ElasticSearcher searcher = new ElasticSearcher(config())) { + int contentLength = (int) ElasticSearcher.MAX_BULK_ESTIMATED_BYTES; + Document oversized = document("oversized", "a".repeat(contentLength)); + + try { + searcher.partitionDocuments(List.of(oversized)); + Assert.fail("单个超限文档应拒绝构建 Bulk 请求"); + } catch (IllegalArgumentException expected) { + Assert.assertTrue(expected.getMessage().contains("oversized")); + } + } + } + + /** + * 验证批量删除请求包含全部稳定文档 ID。 + * + * @throws Exception Elasticsearch 客户端关闭失败时抛出 + */ + @Test + public void shouldBuildBulkDeleteRequest() throws Exception { + try (ElasticSearcher searcher = new ElasticSearcher(config())) { + BulkRequest request = searcher.buildDeleteBulkRequest(List.of("first", "second")); + + Assert.assertEquals(2, request.operations().size()); + Assert.assertEquals("first", request.operations().get(0).delete().id()); + Assert.assertEquals("second", request.operations().get(1).delete().id()); + } + } + + /** + * 创建测试文档。 + * + * @param id 文档 ID + * @param content 文档内容 + * @return 文档 + */ + private Document document(String id, String content) { Document document = new Document(); - document.setId("1"); - document.setTitle("title"); - document.setContent("content"); - document.addMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID, "100"); - - Map source = searcher.buildSource(document); - - Assert.assertEquals("100", source.get(KeywordSearchMetadataKeys.KNOWLEDGE_ID)); - Assert.assertTrue(source.get("metadataMap") instanceof Map); + document.setId(id); + document.setContent(content); + return document; } + /** + * 创建测试用 Elasticsearch 配置。 + * + * @return Elasticsearch 配置 + */ private ESConfig config() { ESConfig config = new ESConfig(); config.setHost("http://127.0.0.1:9200"); diff --git a/easy-agents-search-engine/easy-agents-search-engine-lucene/src/main/java/com/easyagents/search/engine/lucene/LuceneSearcher.java b/easy-agents-search-engine/easy-agents-search-engine-lucene/src/main/java/com/easyagents/search/engine/lucene/LuceneSearcher.java index 5edd49e..84c4a7c 100644 --- a/easy-agents-search-engine/easy-agents-search-engine-lucene/src/main/java/com/easyagents/search/engine/lucene/LuceneSearcher.java +++ b/easy-agents-search-engine/easy-agents-search-engine-lucene/src/main/java/com/easyagents/search/engine/lucene/LuceneSearcher.java @@ -29,7 +29,7 @@ import org.apache.lucene.queryparser.classic.QueryParser; import org.apache.lucene.search.*; import org.apache.lucene.store.Directory; import org.apache.lucene.store.FSDirectory; -import org.jetbrains.annotations.NotNull; +import org.apache.lucene.util.IOUtils; import org.lionsoul.jcseg.ISegment; import org.lionsoul.jcseg.analyzer.JcsegAnalyzer; import org.lionsoul.jcseg.dic.DictionaryFactory; @@ -40,17 +40,38 @@ import org.slf4j.LoggerFactory; import java.io.File; import java.io.IOException; import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.LinkedHashSet; import java.util.List; import java.util.Objects; +import java.util.Set; -public class LuceneSearcher implements DocumentSearcher { +/** + * 基于 Lucene 的本地关键词搜索器。 + * + *

每个实例长期复用一个线程安全的 {@link IndexWriter}，避免并发任务重复抢占 + * 同一索引目录的 {@code write.lock}。

+ */ +public class LuceneSearcher implements DocumentSearcher, AutoCloseable { private static final Logger LOG = LoggerFactory.getLogger(LuceneSearcher.class); - private Directory directory; + private final Directory directory; + private final Analyzer analyzer; + private final IndexWriter indexWriter; + /** + * 创建 Lucene 搜索器并打开索引写入器。 + * + * @param config Lucene 配置 + * @throws IllegalStateException 索引目录或写入器初始化失败时抛出 + */ public LuceneSearcher(LuceneConfig config) { Objects.requireNonNull(config, "LuceneConfig 不能为 null"); + Directory openedDirectory = null; + Analyzer openedAnalyzer = null; + IndexWriter openedWriter = null; try { String indexDirPath = config.getIndexDirPath(); // 索引目录路径 File indexDir = new File(indexDirPath); @@ -58,89 +79,123 @@ public class LuceneSearcher implements DocumentSearcher { throw new IllegalStateException("can not mkdirs for path: " + indexDirPath); } - this.directory = FSDirectory.open(indexDir.toPath()); - } catch (IOException e) { + openedDirectory = FSDirectory.open(indexDir.toPath()); + openedAnalyzer = createAnalyzer(); + openedWriter = new IndexWriter(openedDirectory, new IndexWriterConfig(openedAnalyzer)); + } catch (IOException | RuntimeException e) { + IOUtils.closeWhileHandlingException(openedWriter, openedAnalyzer, openedDirectory); LOG.error("初始化 Lucene 索引失败", e); - throw new RuntimeException(e); + throw new IllegalStateException("初始化 Lucene 索引失败", e); } + this.directory = openedDirectory; + this.analyzer = openedAnalyzer; + this.indexWriter = openedWriter; } + /** + * 以文档 ID 为键写入或覆盖单个文档。 + * + * @param document 待写入文档 + * @return 写入成功时返回 {@code true} + */ @Override public boolean addDocument(Document document) { - if (document == null || document.getContent() == null) return false; + return addDocuments(Collections.singletonList(document)); + } - IndexWriter indexWriter = null; - try { - indexWriter = createIndexWriter(); - - org.apache.lucene.document.Document luceneDoc = new org.apache.lucene.document.Document(); - luceneDoc.add(new StringField("id", document.getId().toString(), Field.Store.YES)); - luceneDoc.add(new TextField("content", document.getContent(), Field.Store.YES)); - - if (document.getTitle() != null) { - luceneDoc.add(new TextField("title", document.getTitle(), Field.Store.YES)); + /** + * 批量写入或覆盖文档，并在批次结束后统一提交。 + * + * @param documents 待写入文档 + * @return 全部文档写入成功时返回 {@code true} + */ + @Override + public boolean addDocuments(List documents) { + if (documents == null || documents.isEmpty()) { + return false; + } + List luceneDocuments = new ArrayList<>(documents.size()); + for (Document document : documents) { + if (document == null || document.getId() == null || document.getContent() == null) { + return false; + } + luceneDocuments.add(toLuceneDocument(document)); + } + try { + for (org.apache.lucene.document.Document luceneDocument : luceneDocuments) { + String documentId = luceneDocument.get("id"); + // updateDocument 以 ID 执行 upsert，确保失败重试不会生成重复索引。 + indexWriter.updateDocument(new Term("id", documentId), luceneDocument); } - appendKnowledgeId(document, luceneDoc); - - indexWriter.addDocument(luceneDoc); indexWriter.commit(); return true; } catch (Exception e) { - LOG.error("添加文档失败", e); + LOG.error("批量添加文档失败: count={}", documents.size(), e); return false; - } finally { - close(indexWriter); } } - + /** + * 删除指定文档。 + * + * @param id 文档 ID + * @return 删除成功时返回 {@code true} + */ @Override public boolean deleteDocument(Object id) { - if (id == null) return false; + return deleteDocuments(Collections.singletonList(id)); + } - IndexWriter indexWriter = null; + /** + * 批量删除指定文档，并在批次结束后统一提交。 + * + * @param ids 文档 ID 集合 + * @return 全部文档删除成功时返回 {@code true} + */ + @Override + public boolean deleteDocuments(Collection ids) { + if (ids == null || ids.isEmpty()) { + return false; + } + Set uniqueIds = new LinkedHashSet<>(); + for (Object id : ids) { + if (id == null) { + return false; + } + uniqueIds.add(id.toString()); + } try { - indexWriter = createIndexWriter(); - Term term = new Term("id", id.toString()); - indexWriter.deleteDocuments(term); + Term[] terms = new Term[uniqueIds.size()]; + int index = 0; + for (String id : uniqueIds) { + terms[index++] = new Term("id", id); + } + indexWriter.deleteDocuments(terms); indexWriter.commit(); return true; } catch (IOException e) { - LOG.error("删除文档失败", e); + LOG.error("批量删除文档失败: count={}", uniqueIds.size(), e); return false; - } finally { - close(indexWriter); } } + /** + * 以文档 ID 为键更新文档。 + * + * @param document 待更新文档 + * @return 更新成功时返回 {@code true} + */ @Override public boolean updateDocument(Document document) { - if (document == null || document.getId() == null) return false; - - IndexWriter indexWriter = null; - try { - indexWriter = createIndexWriter(); - Term term = new Term("id", document.getId().toString()); - - org.apache.lucene.document.Document luceneDoc = new org.apache.lucene.document.Document(); - luceneDoc.add(new StringField("id", document.getId().toString(), Field.Store.YES)); - luceneDoc.add(new TextField("content", document.getContent(), Field.Store.YES)); - - if (document.getTitle() != null) { - luceneDoc.add(new TextField("title", document.getTitle(), Field.Store.YES)); - } - appendKnowledgeId(document, luceneDoc); - indexWriter.updateDocument(term, luceneDoc); - indexWriter.commit(); - return true; - } catch (IOException e) { - LOG.error("更新文档失败", e); - return false; - } finally { - close(indexWriter); - } + return addDocument(document); } + /** + * 按请求条件搜索文档。 + * + * @param request 搜索请求 + * @return 命中文档 + */ @Override public List searchDocuments(KeywordSearchRequest request) { List results = new ArrayList<>(); @@ -171,7 +226,6 @@ public class LuceneSearcher implements DocumentSearcher { Query buildQuery(KeywordSearchRequest request) { try { - Analyzer analyzer = createAnalyzer(); String keyword = request == null ? null : request.getKeyword(); QueryParser titleQueryParser = new QueryParser("title", analyzer); @@ -195,20 +249,22 @@ public class LuceneSearcher implements DocumentSearcher { return null; } - - @NotNull - private IndexWriter createIndexWriter() throws IOException { - Analyzer analyzer = createAnalyzer(); - IndexWriterConfig indexWriterConfig = new IndexWriterConfig(analyzer); - return new IndexWriter(directory, indexWriterConfig); - } - - private static Analyzer createAnalyzer() { SegmenterConfig config = new SegmenterConfig(true); return new JcsegAnalyzer(ISegment.Type.NLP, config, DictionaryFactory.createSingletonDictionary(config)); } + private org.apache.lucene.document.Document toLuceneDocument(Document document) { + org.apache.lucene.document.Document luceneDoc = new org.apache.lucene.document.Document(); + luceneDoc.add(new StringField("id", document.getId().toString(), Field.Store.YES)); + luceneDoc.add(new TextField("content", document.getContent(), Field.Store.YES)); + if (document.getTitle() != null) { + luceneDoc.add(new TextField("title", document.getTitle(), Field.Store.YES)); + } + appendKnowledgeId(document, luceneDoc); + return luceneDoc; + } + private void appendKnowledgeId(Document document, org.apache.lucene.document.Document luceneDoc) { if (document == null || document.getMetadataMap() == null) { return; @@ -219,13 +275,13 @@ public class LuceneSearcher implements DocumentSearcher { } } - public void close(IndexWriter indexWriter) { - try { - if (indexWriter != null) { - indexWriter.close(); - } - } catch (IOException e) { - LOG.error("关闭 Lucene 失败", e); - } + /** + * 关闭写入器、分词器和索引目录。 + * + * @throws IOException 任一 Lucene 资源关闭失败时抛出 + */ + @Override + public void close() throws IOException { + IOUtils.close(indexWriter, analyzer, directory); } } diff --git a/easy-agents-search-engine/easy-agents-search-engine-lucene/src/test/java/com/easyagents/search/engine/lucene/LuceneSearcherTest.java b/easy-agents-search-engine/easy-agents-search-engine-lucene/src/test/java/com/easyagents/search/engine/lucene/LuceneSearcherTest.java index 4f21723..e7804dc 100644 --- a/easy-agents-search-engine/easy-agents-search-engine-lucene/src/test/java/com/easyagents/search/engine/lucene/LuceneSearcherTest.java +++ b/easy-agents-search-engine/easy-agents-search-engine-lucene/src/test/java/com/easyagents/search/engine/lucene/LuceneSearcherTest.java @@ -8,38 +8,145 @@ import org.junit.Test; import java.nio.file.Files; import java.nio.file.Path; +import java.util.ArrayList; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +/** + * {@link LuceneSearcher} 回归测试。 + */ public class LuceneSearcherTest { + /** + * 验证知识库过滤同时覆盖标题和正文。 + * + * @throws Exception 临时目录或 Lucene 资源操作失败时抛出 + */ @Test public void shouldFilterByKnowledgeIdAndSearchTitleAndContent() throws Exception { Path tempDir = Files.createTempDirectory("lucene-searcher-test"); LuceneConfig config = new LuceneConfig(); config.setIndexDirPath(tempDir.toString()); - LuceneSearcher searcher = new LuceneSearcher(config); + try (LuceneSearcher searcher = new LuceneSearcher(config)) { + Document first = new Document(); + first.setId("1"); + first.setTitle("客服标题"); + first.setContent("这里没有关键字"); + first.addMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID, "100"); - Document first = new Document(); - first.setId("1"); - first.setTitle("客服标题"); - first.setContent("这里没有关键字"); - first.addMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID, "100"); + Document second = new Document(); + second.setId("2"); + second.setTitle("别的知识库"); + second.setContent("客服内容"); + second.addMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID, "200"); - Document second = new Document(); - second.setId("2"); - second.setTitle("别的知识库"); - second.setContent("客服内容"); - second.addMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID, "200"); + Assert.assertTrue(searcher.addDocument(first)); + Assert.assertTrue(searcher.addDocument(second)); - Assert.assertTrue(searcher.addDocument(first)); - Assert.assertTrue(searcher.addDocument(second)); + KeywordSearchRequest request = KeywordSearchRequest.of("客服", 10); + request.setKnowledgeId("100"); + List results = searcher.searchDocuments(request); - KeywordSearchRequest request = KeywordSearchRequest.of("客服", 10); - request.setKnowledgeId("100"); - List results = searcher.searchDocuments(request); + Assert.assertEquals(1, results.size()); + Assert.assertEquals("1", String.valueOf(results.get(0).getId())); + Assert.assertEquals("100", String.valueOf( + results.get(0).getMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID))); + } + } - Assert.assertEquals(1, results.size()); - Assert.assertEquals("1", String.valueOf(results.get(0).getId())); - Assert.assertEquals("100", String.valueOf(results.get(0).getMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID))); + /** + * 验证多个导入线程可共享同一个 IndexWriter 完成批量写入。 + * + * @throws Exception 并发执行或 Lucene 资源操作失败时抛出 + */ + @Test + public void shouldWriteConcurrentBatchesWithoutLockConflict() throws Exception { + Path tempDir = Files.createTempDirectory("lucene-searcher-concurrent-test"); + LuceneConfig config = new LuceneConfig(); + config.setIndexDirPath(tempDir.toString()); + int workerCount = 4; + int documentsPerWorker = 20; + ExecutorService executor = Executors.newFixedThreadPool(workerCount); + try (LuceneSearcher searcher = new LuceneSearcher(config)) { + CountDownLatch startSignal = new CountDownLatch(1); + List> futures = new ArrayList<>(); + for (int worker = 0; worker < workerCount; worker++) { + int currentWorker = worker; + futures.add(executor.submit(() -> { + startSignal.await(); + List documents = new ArrayList<>(); + for (int offset = 0; offset < documentsPerWorker; offset++) { + Document document = new Document(); + document.setId(currentWorker + "-" + offset); + document.setContent("concurrent indexing content"); + document.addMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID, "concurrent-kb"); + documents.add(document); + } + return searcher.addDocuments(documents); + })); + } + startSignal.countDown(); + for (Future future : futures) { + Assert.assertTrue(future.get()); + } + + int expectedCount = workerCount * documentsPerWorker; + KeywordSearchRequest request = KeywordSearchRequest.of("concurrent", expectedCount); + request.setKnowledgeId("concurrent-kb"); + Assert.assertEquals(expectedCount, searcher.searchDocuments(request).size()); + } finally { + executor.shutdownNow(); + } + } + + /** + * 验证相同文档 ID 重试写入时执行覆盖而不会保留重复旧索引。 + * + * @throws Exception 临时目录或 Lucene 资源操作失败时抛出 + */ + @Test + public void shouldUpsertDocumentByIdOnRetry() throws Exception { + Path tempDir = Files.createTempDirectory("lucene-searcher-upsert-test"); + LuceneConfig config = new LuceneConfig(); + config.setIndexDirPath(tempDir.toString()); + try (LuceneSearcher searcher = new LuceneSearcher(config)) { + Document document = new Document(); + document.setId("retry-id"); + document.setContent("oldkeyword"); + Assert.assertTrue(searcher.addDocument(document)); + + document.setContent("newkeyword"); + Assert.assertTrue(searcher.addDocument(document)); + + Assert.assertTrue(searcher.searchDocuments("oldkeyword", 10).isEmpty()); + Assert.assertEquals(1, searcher.searchDocuments("newkeyword", 10).size()); + } + } + + /** + * 验证多个文档可在一次提交中批量删除。 + * + * @throws Exception 临时目录或 Lucene 资源操作失败时抛出 + */ + @Test + public void shouldDeleteDocumentsInSingleBatch() throws Exception { + Path tempDir = Files.createTempDirectory("lucene-searcher-delete-test"); + LuceneConfig config = new LuceneConfig(); + config.setIndexDirPath(tempDir.toString()); + try (LuceneSearcher searcher = new LuceneSearcher(config)) { + Document first = new Document(); + first.setId("delete-first"); + first.setContent("batchdelete"); + Document second = new Document(); + second.setId("delete-second"); + second.setContent("batchdelete"); + + Assert.assertTrue(searcher.addDocuments(List.of(first, second))); + Assert.assertTrue(searcher.deleteDocuments(List.of(first.getId(), second.getId()))); + Assert.assertTrue(searcher.searchDocuments("batchdelete", 10).isEmpty()); + } } } diff --git a/easy-agents-search-engine/easy-agents-search-engine-service/src/main/java/com/easyagents/search/engine/service/DocumentSearcher.java b/easy-agents-search-engine/easy-agents-search-engine-service/src/main/java/com/easyagents/search/engine/service/DocumentSearcher.java index 00e84c6..ab248fc 100644 --- a/easy-agents-search-engine/easy-agents-search-engine-service/src/main/java/com/easyagents/search/engine/service/DocumentSearcher.java +++ b/easy-agents-search-engine/easy-agents-search-engine-service/src/main/java/com/easyagents/search/engine/service/DocumentSearcher.java @@ -17,23 +17,106 @@ package com.easyagents.search.engine.service; import com.easyagents.core.document.Document; +import java.util.Collection; import java.util.List; +/** + * 关键词搜索引擎统一接口。 + */ public interface DocumentSearcher { + /** + * 写入单个文档。 + * + * @param document 待写入文档 + * @return 写入成功时返回 {@code true} + */ boolean addDocument(Document document); + /** + * 批量写入文档。 + * + *

默认实现保持现有搜索引擎兼容性；支持原生批量写入的实现应覆盖该方法。

+ * + * @param documents 待写入文档 + * @return 全部文档写入成功时返回 {@code true} + */ + default boolean addDocuments(List documents) { + if (documents == null || documents.isEmpty()) { + return false; + } + for (Document document : documents) { + if (!addDocument(document)) { + return false; + } + } + return true; + } + + /** + * 删除指定文档。 + * + * @param id 文档 ID + * @return 删除成功时返回 {@code true} + */ boolean deleteDocument(Object id); + /** + * 批量删除指定文档。 + * + *

默认实现保持现有搜索引擎兼容性，并确保所有文档都会尝试删除； + * 支持原生批量删除的实现应覆盖该方法。

+ * + * @param ids 文档 ID 集合 + * @return 全部文档删除成功时返回 {@code true} + */ + default boolean deleteDocuments(Collection ids) { + if (ids == null || ids.isEmpty()) { + return false; + } + boolean success = true; + for (Object id : ids) { + if (!deleteDocument(id)) { + success = false; + } + } + return success; + } + + /** + * 更新指定文档。 + * + * @param document 待更新文档 + * @return 更新成功时返回 {@code true} + */ boolean updateDocument(Document document); + /** + * 使用默认返回数量搜索文档。 + * + * @param keyword 搜索关键词 + * @return 命中文档 + */ default List searchDocuments(String keyword) { return searchDocuments(KeywordSearchRequest.of(keyword, 10)); } + /** + * 搜索指定数量的文档。 + * + * @param keyword 搜索关键词 + * @param count 最大返回数量 + * @return 命中文档 + */ default List searchDocuments(String keyword, int count) { return searchDocuments(KeywordSearchRequest.of(keyword, count)); } + /** + * 按请求条件搜索文档。 + * + * @param request 搜索请求 + * @return 命中文档 + */ List searchDocuments(KeywordSearchRequest request); }