fix: split Markdown lines like markdown-it does

src/raglite/_split_chunklets.py

@@ -1,5 +1,6 @@
 """Split a document into chunklets."""
 
+import re
 from collections.abc import Callable
 
 import numpy as np
@@ -14,8 +15,10 @@ def markdown_chunklet_boundaries(sentences: list[str]) -> FloatVector:
     doc = "".join(sentences)
     md = MarkdownIt()
     tokens = md.parse(doc)
-    # Identify the character index of each line in the document.
-    lines = doc.splitlines(keepends=True)
+    # Identify the character index of each line in the document. Split lines like markdown-it does
+    # (on LF, CRLF, and CR only) rather than with `str.splitlines`, which also splits on form feeds,
+    # U+2028, etc., so that the line numbers in `token.map` stay aligned with the character indexes.
+    lines = re.findall(r"[^\r\n]*(?:\r\n|\r|\n)|[^\r\n]+", doc)
     line_start_char = [0]
     for line in lines[:-1]:
         line_start_char.append(line_start_char[-1] + len(line))

src/raglite/_split_sentences.py

@@ -1,5 +1,6 @@
 """Split a document into sentences."""
 
+import re
 from collections import deque
 from collections.abc import Callable
 from functools import cache
@@ -28,7 +29,10 @@ def get_markdown_heading_indexes(doc: str) -> list[tuple[int, int]]:
         md = MarkdownIt()
         tokens = md.parse(doc)
         headings = []
-        lines = doc.splitlines(keepends=True)
+        # Split lines like markdown-it does (on LF, CRLF, and CR only) rather than with
+        # `str.splitlines`, which also splits on form feeds, U+2028, etc., so that the line numbers
+        # in `token.map` stay aligned with the character indexes.
+        lines = re.findall(r"[^\r\n]*(?:\r\n|\r|\n)|[^\r\n]+", doc)
         line_start_char = [0]
         for line in lines:
             line_start_char.append(line_start_char[-1] + len(line))

tests/test_split_chunklets.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from raglite._split_chunklets import split_chunklets
+from raglite._split_chunklets import markdown_chunklet_boundaries, split_chunklets
 
 
 @pytest.mark.parametrize(
@@ -102,3 +102,28 @@ def test_split_chunklets(sentences_splits: tuple[list[str], list[int]]) -> None:
         chunklet == expected_chunklet
         for chunklet, expected_chunklet in zip(chunklets, expected_chunklets, strict=True)
     )
+
+
+@pytest.mark.parametrize(
+    "line_separator",
+    [
+        pytest.param("\n", id="lf"),
+        pytest.param("\r\n", id="crlf"),
+        pytest.param("\r", id="cr"),
+        pytest.param("\x0b", id="vertical_tab"),
+        pytest.param("\x0c", id="form_feed"),
+        pytest.param("\x1c", id="file_separator"),
+        pytest.param("\x1d", id="group_separator"),
+        pytest.param("\x1e", id="record_separator"),
+        pytest.param("\x85", id="next_line"),
+        pytest.param("\u2028", id="line_separator"),
+        pytest.param("\u2029", id="paragraph_separator"),
+    ],
+)
+def test_markdown_chunklet_boundaries_line_separators(line_separator: str) -> None:
+    """Test that boundaries stay aligned when a document has exotic line separators."""
+    # Only LF, CRLF, and CR are line endings to markdown-it, so the other separators must not shift
+    # the line numbers that markdown-it reports for the paragraphs and the heading.
+    sentences = [f"Intro one.{line_separator}", "Intro two.\n\n", "# Real Heading\n\n", "Body.\n"]
+    boundary_probas = markdown_chunklet_boundaries(sentences)
+    assert boundary_probas.tolist() == [0.5, 0.0, 1.0, 0.0]

tests/test_split_sentences.py

@@ -2,10 +2,11 @@
 
 from pathlib import Path
 
+import numpy as np
 import pytest
 
 from raglite._markdown import document_to_markdown
-from raglite._split_sentences import split_sentences
+from raglite._split_sentences import markdown_sentence_boundaries, split_sentences
 
 
 def test_split_sentences() -> None:
@@ -77,3 +78,32 @@ def test_split_sentences_edge_cases(case: tuple[str, list[str], tuple[int, int |
             sentences[: len(expected_sentences)], expected_sentences, strict=True
         )
     )
+
+
+@pytest.mark.parametrize(
+    "line_separator",
+    [
+        pytest.param("\n", id="lf"),
+        pytest.param("\r\n", id="crlf"),
+        pytest.param("\r", id="cr"),
+        pytest.param("\x0b", id="vertical_tab"),
+        pytest.param("\x0c", id="form_feed"),
+        pytest.param("\x1c", id="file_separator"),
+        pytest.param("\x1d", id="group_separator"),
+        pytest.param("\x1e", id="record_separator"),
+        pytest.param("\x85", id="next_line"),
+        pytest.param("\u2028", id="line_separator"),
+        pytest.param("\u2029", id="paragraph_separator"),
+    ],
+)
+def test_markdown_sentence_boundaries_line_separators(line_separator: str) -> None:
+    """Test that heading boundaries stay aligned when a document has exotic line separators."""
+    # Only LF, CRLF, and CR are line endings to markdown-it, so the other separators must not shift
+    # the line numbers that markdown-it reports for the heading.
+    heading = "# Real Heading\n"
+    doc = f"Intro one.{line_separator}Intro two.\n\n{heading}\nBody.\n"
+    heading_start = doc.index(heading)
+    boundary_probas = markdown_sentence_boundaries(doc)
+    marked_indexes = np.flatnonzero(~np.isnan(boundary_probas))
+    assert marked_indexes.min() == heading_start - 1
+    assert marked_indexes.max() == heading_start + len(heading)