core: чанкер markdown по заголовкам с перекрытием

2026-10-02 01:14:00 +03:00
parent d62a26f5b1
commit de3cb24edc
4 changed files with 196 additions and 1 deletions
@@ -0,0 +1,78 @@
+package memo.core
+
+data class Chunk(
+    val heading: String,
+    val line: Int,
+    val text: String,
+    val ord: Int,
+)
+
+private const val CHUNK_SIZE = 4000
+private const val CHUNK_OVERLAP = 600
+
+private val HEADING_REGEX = Regex("^(#{1,3}) (.*)$")
+
+fun chunkMarkdown(text: String): List<Chunk> {
+    val sections = mutableListOf<Triple<String, Int, String>>()
+
+    var currentHeading = ""
+    var currentLine = 1
+    val currentBody = StringBuilder()
+
+    val preamble = StringBuilder()
+    var foundHeading = false
+
+    fun flush(heading: String, line: Int, body: String) {
+        val trimmed = body.trim()
+        if (trimmed.isNotEmpty()) {
+            sections.add(Triple(heading, line, trimmed))
+        }
+    }
+
+    val lines = text.split('\n')
+    for ((idx, raw) in lines.withIndex()) {
+        val lineNo = idx + 1
+        val match = HEADING_REGEX.matchEntire(raw)
+        if (match != null) {
+            if (!foundHeading) {
+                flush("", 1, preamble.toString())
+                foundHeading = true
+            } else {
+                flush(currentHeading, currentLine, currentBody.toString())
+            }
+            currentHeading = match.groupValues[2].trim()
+            currentLine = lineNo
+            currentBody.setLength(0)
+        } else {
+            if (foundHeading) {
+                currentBody.append(raw).append('\n')
+            } else {
+                preamble.append(raw).append('\n')
+            }
+        }
+    }
+
+    if (foundHeading) {
+        flush(currentHeading, currentLine, currentBody.toString())
+    } else {
+        flush("", 1, preamble.toString())
+    }
+
+    val result = mutableListOf<Chunk>()
+    var ord = 0
+    for ((heading, line, body) in sections) {
+        if (body.length <= CHUNK_SIZE) {
+            result.add(Chunk(heading, line, body, ord++))
+        } else {
+            val step = CHUNK_SIZE - CHUNK_OVERLAP
+            var i = 0
+            while (i < body.length) {
+                val end = minOf(i + CHUNK_SIZE, body.length)
+                result.add(Chunk(heading, line, body.substring(i, end), ord++))
+                if (end == body.length) break
+                i += step
+            }
+        }
+    }
+    return result
+}
@@ -0,0 +1,60 @@
+package memo.core
+
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+class ChunkerTest {
+
+    @Test
+    fun splitsByHeadingLevels() {
+        val input = "# A\nbodyA\n## B\nbodyB\n### C\nbodyC"
+        val chunks = chunkMarkdown(input)
+        assertEquals(3, chunks.size)
+        assertEquals(listOf("A", "B", "C"), chunks.map { it.heading })
+        assertEquals(listOf(1, 3, 5), chunks.map { it.line })
+        assertEquals(listOf(0, 1, 2), chunks.map { it.ord })
+    }
+
+    @Test
+    fun textBeforeFirstHeadingBecomesChunkWithEmptyHeading() {
+        val input = "преамбула\n# Заголовок\nтело"
+        val chunks = chunkMarkdown(input)
+        assertEquals(2, chunks.size)
+        assertEquals("", chunks[0].heading)
+        assertEquals(1, chunks[0].line)
+        assertEquals("Заголовок", chunks[1].heading)
+    }
+
+    @Test
+    fun fourHashesIsNotAHeading() {
+        val input = "# A\nbody\n#### не заголовок\nmore"
+        val chunks = chunkMarkdown(input)
+        assertEquals(1, chunks.size)
+        assertEquals("A", chunks[0].heading)
+        assertTrue(chunks[0].text.contains("#### не заголовок"))
+        assertTrue(chunks[0].text.contains("more"))
+    }
+
+    @Test
+    fun emptySectionsAreDropped() {
+        val input = "# A\nsome text\n## B\n## C"
+        val chunks = chunkMarkdown(input)
+        assertEquals(1, chunks.size)
+        assertEquals("A", chunks[0].heading)
+    }
+
+    @Test
+    fun longSectionIsSplitWithOverlap() {
+        val body = "а".repeat(6000)
+        val input = "# A\n$body"
+        val chunks = chunkMarkdown(input)
+        assertEquals(2, chunks.size)
+        assertEquals(4000, chunks[0].text.length)
+        assertEquals(2600, chunks[1].text.length)
+        assertEquals(
+            chunks[0].text.substring(4000 - 600),
+            chunks[1].text.substring(0, 600),
+        )
+    }
+}