diff --git a/astrbot/core/knowledge_base/chunking/recursive.py b/astrbot/core/knowledge_base/chunking/recursive.py index e27ffbd1b7..1f2016a80a 100644 --- a/astrbot/core/knowledge_base/chunking/recursive.py +++ b/astrbot/core/knowledge_base/chunking/recursive.py @@ -56,6 +56,14 @@ async def chunk(self, text: str, **kwargs) -> list[str]: overlap = kwargs.get("chunk_overlap", self.chunk_overlap) chunk_size = kwargs.get("chunk_size", self.chunk_size) + # 与 _split_by_character 保持一致的参数防护,保证任何输出块都不超过 + # chunk_size 上限。 + if chunk_size <= 0: + raise ValueError("chunk_size must be greater than 0") + if overlap < 0: + raise ValueError("chunk_overlap must be non-negative") + if overlap >= chunk_size: + raise ValueError("chunk_overlap must be less than chunk_size") text_length = self.length_function(text) if text_length <= chunk_size: @@ -114,10 +122,16 @@ async def chunk(self, text: str, **kwargs) -> list[str]: overlap_start = max(0, len(combined_text) - overlap) if overlap_start > 0: overlap_text = combined_text[overlap_start:] - current_chunk = [overlap_text, split] - current_chunk_length = ( - self.length_function(overlap_text) + split_length - ) + overlap_length = self.length_function(overlap_text) + if overlap_length + split_length > chunk_size: + # 携带 overlap 会超出 chunk_size,改为将 + # overlap 独立成块,避免输出块超限。 + final_chunks.append(overlap_text) + current_chunk = [split] + current_chunk_length = split_length + else: + current_chunk = [overlap_text, split] + current_chunk_length = overlap_length + split_length else: current_chunk = [split] current_chunk_length = split_length diff --git a/tests/test_recursive_chunker.py b/tests/test_recursive_chunker.py new file mode 100644 index 0000000000..19e39a2c3e --- /dev/null +++ b/tests/test_recursive_chunker.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import pytest + +from astrbot.core.knowledge_base.chunking.recursive import RecursiveCharacterChunker + + +def _make_split(chars: int, separator: str = "\n\n") -> str: + return "a" * chars + separator + + +@pytest.mark.asyncio +async def test_overlap_never_pushes_chunks_over_chunk_size() -> None: + # Regression test for #9901: with chunk_size=100 / overlap=30 and 80-char + # splits, the old code built "overlap + split" chunks of 110 chars. + chunker = RecursiveCharacterChunker( + chunk_size=100, + chunk_overlap=30, + separators=["\n\n"], + ) + text = "".join(_make_split(78) for _ in range(3)) + + chunks = await chunker.chunk(text, chunk_size=100, chunk_overlap=30) + + assert chunks + assert all(len(chunk) <= 100 for chunk in chunks), chunks + assert all(chunk for chunk in chunks) + + +@pytest.mark.asyncio +async def test_small_text_returned_as_single_chunk() -> None: + chunker = RecursiveCharacterChunker() + + assert await chunker.chunk("hello", chunk_size=100, chunk_overlap=10) == ["hello"] + assert await chunker.chunk("", chunk_size=100, chunk_overlap=10) == [] + + +@pytest.mark.asyncio +async def test_normal_text_chunks_within_limit() -> None: + chunker = RecursiveCharacterChunker( + chunk_size=50, + chunk_overlap=10, + separators=["\n\n"], + ) + text = "\n\n".join(f"段落{i}" + "字" * 20 for i in range(6)) + + chunks = await chunker.chunk(text, chunk_size=50, chunk_overlap=10) + + assert all(0 < len(chunk) <= 50 for chunk in chunks), chunks + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "text", + [ + "a" * 30 + "\n\n" + "b" * 200, # separator-driven path + "a" * 200, # falls through to _split_by_character + ], +) +async def test_invalid_overlap_raises(text: str) -> None: + chunker = RecursiveCharacterChunker() + + with pytest.raises(ValueError, match="chunk_overlap must be less than chunk_size"): + await chunker.chunk(text, chunk_size=100, chunk_overlap=100)