diff --git a/python/semantic_kernel/text/text_chunker.py b/python/semantic_kernel/text/text_chunker.py index f79796399b6d..c730bb2a62cb 100644 --- a/python/semantic_kernel/text/text_chunker.py +++ b/python/semantic_kernel/text/text_chunker.py @@ -137,15 +137,12 @@ def _split_text_paragraph(text: list[str], max_tokens: int, token_counter: Calla sec_last_para = paragraphs[-2] if token_counter(last_para) < max_tokens / 4: - last_para_tokens = last_para.split(" ") - sec_last_para_tokens = sec_last_para.split(" ") - last_para_token_count = len(last_para_tokens) - sec_last_para_token_count = len(sec_last_para_tokens) - - if last_para_token_count + sec_last_para_token_count <= max_tokens: - sec_last_para = " ".join(sec_last_para_tokens) + NEWLINE - last_para = " ".join(last_para_tokens) - new_sec_last_para = sec_last_para + last_para + # Compare against `token_counter`, not `len(text.split(" "))`. Word counts and token + # counts are different units, so the old check could pass while the merged paragraph + # was over `max_tokens` -- exactly the limit this function exists to keep. + new_sec_last_para = f"{sec_last_para}{NEWLINE}{last_para}" + + if token_counter(new_sec_last_para) <= max_tokens: paragraphs[-2] = new_sec_last_para.strip() paragraphs.pop() diff --git a/python/tests/unit/text/test_text_chunker.py b/python/tests/unit/text/test_text_chunker.py index 648f84b752c5..5ade4dd9404f 100644 --- a/python/tests/unit/text/test_text_chunker.py +++ b/python/tests/unit/text/test_text_chunker.py @@ -8,6 +8,7 @@ split_plaintext_lines, split_plaintext_paragraph, ) +from semantic_kernel.text.text_chunker import _split_text_paragraph, _token_counter NEWLINE = os.linesep @@ -532,3 +533,25 @@ def test_split_md_on_newlines(): max_token_per_line = 15 split = split_markdown_paragraph(test, max_token_per_line) assert expected == split + + +def test_short_last_paragraph_merge_respects_max_tokens(): + """Folding the last paragraph back must not push it over `max_tokens`.""" + max_tokens = 20 + # The default counter is len(text) // 4, so these are 20 and 4 tokens but one word each. + lines = ["a" * 80, "b" * 16] + + paragraphs = _split_text_paragraph(lines, max_tokens) + + assert all(_token_counter(p) <= max_tokens for p in paragraphs), [_token_counter(p) for p in paragraphs] + + +def test_short_last_paragraph_still_merges_when_it_fits(): + """A trailing paragraph that does fit is still folded back.""" + max_tokens = 20 + lines = ["a" * 20, "b" * 8] + + paragraphs = _split_text_paragraph(lines, max_tokens) + + assert len(paragraphs) == 1 + assert _token_counter(paragraphs[0]) <= max_tokens