Python: fix paragraph merge comparing word count against token budget

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()
 

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