feat: smart chunking to avoid chunk boundaries mid-sentence

1 files changed, 129 insertions, 0 deletions

diff --git a/app/tests/test_chunking.py b/app/tests/test_chunking.py
index 2904e40..bcae6b2 100644
--- a/app/tests/test_chunking.py
+++ b/app/tests/test_chunking.py
@@ -113,6 +113,135 @@ class SplitIntoChunksTests(unittest.TestCase):
         self.assertNotIn("1, 000", joined)
         self.assertNotIn("12: 30", joined)
 
+    def test_smart_strategy_distributes_sentences(self):
+        # The smart strategy leaves whole sentences intact: a chunk
+        # crossing the target (85% of the limit) ends at the next clean
+        # boundary early; legacy keeps packing until the limit.
+        sentences = " ".join(
+            " ".join(f"Sw{j}" for j in range(n)) + "."
+            for n in (9, 8, 3, 9, 8, 3))
+        smart = split_into_chunks(sentences, max_words=20, smart=True)
+        legacy = split_into_chunks(sentences, max_words=20, smart=False)
+        self.assertEqual([len(c.split()) for c in smart], [17, 20, 3])
+        self.assertEqual([len(c.split()) for c in legacy], [20, 20])
+        self.assertEqual(self._content(smart),
+                         self._content(sentences.split(" ")))
+
+    def test_dialogue_tag_stays_with_its_quote(self):
+        # A closing quote followed by a lowercase attribution ("she
+        # said.") is one sentence unit: the chunk never separates them.
+        text = "“Come here!” she said. “Are you coming?” He did not answer."
+        chunks = split_into_chunks(text, max_words=10, smart=True)
+        self.assertIn("“Come here!” she said.", chunks[0])
+        self._assert_bounded(chunks, 10)
+        self._assert_content(text, chunks)
+
+    def test_abbreviations_and_initials_do_not_split(self):
+        text = ("Dr. Smith met J. K. R. at the U.S. border with No. 3 "
+                "shortly. He continued onward.")
+        chunks = split_into_chunks(text, max_words=50, smart=True)
+        self.assertEqual(len(chunks), 1)
+        self.assertIn("Dr. Smith", chunks[0])
+        self.assertIn("J. K. R.", chunks[0])
+        self.assertIn("U.S.", chunks[0])
+        self._assert_content(text, chunks)
+
+    def test_decimals_are_never_boundaries(self):
+        text = "Pi is 3.14159 and the toll was 1,000,000 miles. Next one."
+        chunks = split_into_chunks(text, max_words=12, smart=True)
+        joined = " ".join(chunks)
+        self.assertIn("3.14159", joined)
+        self.assertIn("1,000,000", joined)
+        self._assert_bounded(chunks, 12)
+        self._assert_content(text, chunks)
+
+    def test_unclosed_quote_still_respects_the_limit(self):
+        # A quote with no closing mark must not grow chunks or silence
+        # later boundaries; units carry the open state and stay bounded.
+        text = "“No closing quote follows. " + "filler words here. " * 20
+        chunks = split_into_chunks(text, max_words=10, smart=True)
+        self._assert_bounded(chunks, 10)
+        self._assert_content(text, chunks)
+
+    def test_paragraph_boundary_resets_an_open_quote(self):
+        # Blank lines reset the quotation state: an unclosed quote in
+        # one paragraph cannot trap later paragraphs inside the quote.
+        text = ("“Still open. filler words inside the quote here.\n\n"
+                "New paragraph. Words outside any quote now, plenty.")
+        chunks = split_into_chunks(text, max_words=10, smart=True)
+        self._assert_bounded(chunks, 10)
+        self._assert_content(text, chunks)
+        second = [c for c in chunks if "New paragraph." in c]
+        self.assertTrue(second)
+        # The paragraph's text starts a fresh chunk: the open-quote
+        # prefix did not swallow it.
+        self.assertTrue(
+            any(chunk.startswith("New paragraph.") or
+                "“Still open." not in chunk
+                for chunk in chunks))
+
+    def test_quote_retreat_keeps_short_quotes_whole(self):
+        # When the next sentence would overflow, the break retreats to
+        # the last boundary that ended outside a quotation — before the
+        # quote — so a short quotation lands in one request.
+        text = ("Plain opening words here to fill. "
+                "“A short quote. With two sentences. Inside.” "
+                "More trailing prose follows this quote. And more.")
+        chunks = split_into_chunks(text, max_words=14, smart=True)
+        self._assert_bounded(chunks, 14)
+        self._assert_content(text, chunks)
+        quote_chunks = [c for c in chunks if "short quote" in c]
+        self.assertTrue(quote_chunks)
+        self.assertIn("Inside.", quote_chunks[0])
+
+    def test_oversized_single_tokens_respect_char_cap(self):
+        # Degenerate text (punctuation-free run): bounded by the hard
+        # character cap, with every character preserved.
+        text = "a" * 5000
+        chunks = split_into_chunks(text, max_words=250, smart=True)
+        self.assertGreater(len(chunks), 1)
+        self.assertTrue(all(len(c) <= 250 * 8 for c in chunks))
+        self.assertEqual("".join(chunks), text)
+
+    def test_non_spaced_script_stays_bounded(self):
+        # CJK text: sentence terminators break units without
+        # whitespace, and each ideograph counts as a word.
+        text = "。".join(["字" * 15 for _ in range(30)]) + "。"
+        chunks = split_into_chunks(text, max_words=40, smart=True)
+        self._assert_bounded(chunks, 40)
+        self._assert_content(text, chunks)
+        joined = " ".join(chunks).split()
+        self.assertTrue(all(len(j) <= 33 for j in joined))
+
+    def test_legacy_and_smart_matching_smart_negation(self):
+        # smart=False preserves the legacy splitting exactly.
+        sentences = " ".join(
+            " ".join(f"Sw{j}" for j in range(9)) + "." for _ in range(10))
+        with patch.object(config, "CHUNK_SIZE", 25):
+            legacy_default = split_into_chunks(sentences, smart=False)
+        self.assertEqual([len(c.split()) for c in legacy_default],
+                         [18] * 5)
+        self.assertEqual(self._content(legacy_default),
+                         self._content(sentences.split()))
+
+    # -- helpers ---------------------------------------------------------
+
+    @staticmethod
+    def _content(chunks_or_words):
+        """Non-whitespace content of the pieces (or of a token list)."""
+        pieces = (chunks_or_words if isinstance(chunks_or_words, str)
+                  else " ".join(chunks_or_words))
+        return "".join(pieces.split())
+
+    def _assert_content(self, text, chunks):
+        self.assertEqual(self._content(chunks), self._content(text))
+
+    def _assert_bounded(self, chunks, max_words):
+        from converter.chunking import _count_words
+        for chunk in chunks:
+            self.assertLessEqual(_count_words(chunk), max_words,
+                                 chunk[:40])
+
 
 if __name__ == "__main__":
     unittest.main()