From 1c08eaa559fe85b85ca3ff9f2d3fe6ea1c2b42ef Mon Sep 17 00:00:00 2001 From: Sen Huang <48879559+ployts@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:03:58 +0800 Subject: [PATCH] fix: enforce markdown chunk byte limits (#370) --- .../file_chunker/markdown_file_chunker.py | 175 ++++++++---------- reme/config/default.yaml | 1 - reme/config/jinli_lme.yaml | 1 - tests/unit/test_config_parser.py | 1 - tests/unit/test_markdown_file_chunker.py | 132 +++++++------ 5 files changed, 153 insertions(+), 157 deletions(-) diff --git a/reme/components/file_chunker/markdown_file_chunker.py b/reme/components/file_chunker/markdown_file_chunker.py index 373af3c6..1365020c 100644 --- a/reme/components/file_chunker/markdown_file_chunker.py +++ b/reme/components/file_chunker/markdown_file_chunker.py @@ -25,7 +25,6 @@ import yaml from pydantic import ValidationError -from .base_file_chunker import BaseFileChunker from .default_file_chunker import DefaultFileChunker from ..component_registry import R from ...schema import ( @@ -99,26 +98,22 @@ def _toc_join(*parts: str) -> str: @R.register("markdown") -class MarkdownFileChunker(BaseFileChunker): +class MarkdownFileChunker(DefaultFileChunker): """Markdown chunker with breadcrumb context and adjacent-section packing.""" def __init__( self, encoding: str = "utf-8", - chunk_chars: int = 10000, + chunk_byte_size: int = 10000, embed_toc: bool = True, max_ast_sections: int | None = 100, - default_chunker: str = "default", include_frontmatter_in_metadata: bool = False, include_frontmatter_keys_in_metadata: list[str] | None = None, **kwargs, ): - super().__init__(**kwargs) - self.encoding = encoding - self.chunk_chars = max(100, chunk_chars) + super().__init__(encoding=encoding, chunk_byte_size=chunk_byte_size, **kwargs) self.embed_toc = embed_toc self.max_ast_sections = max(0, max_ast_sections) if max_ast_sections is not None else None - self.default_chunker = self.bind(default_chunker, BaseFileChunker, optional=False) self.include_frontmatter_in_metadata = include_frontmatter_in_metadata self.include_frontmatter_keys_in_metadata = list(include_frontmatter_keys_in_metadata or []) @@ -211,10 +206,8 @@ class MarkdownFileChunker(BaseFileChunker): return count def _chunk_plain_text(self, content: str, path: str, line_offset: int) -> list[FileChunk]: - """Use the default byte chunker while preserving original Markdown lines.""" - if not isinstance(self.default_chunker, DefaultFileChunker): - raise RuntimeError("DefaultFileChunker dependency is unavailable; call start() before chunk().") - chunks = self.default_chunker.chunk_content(content, path, parse_links=True) + """Use inherited byte chunking while preserving original Markdown lines.""" + chunks = self.chunk_content(content, path, parse_links=True) if line_offset: for chunk in chunks: chunk.start_line += line_offset @@ -329,12 +322,11 @@ class MarkdownFileChunker(BaseFileChunker): heading = f"{'#' * node.level} {node.heading or ''}" subtree_text = _toc_join(heading, node.text) - if subtree_text and len(subtree_text) <= self.chunk_chars: + if subtree_text and self._byte_len(self._compose_text(prefix, subtree_text)) <= self.chunk_byte_size: return [ self._make_chunk( prefix, subtree_text, - "", node.start_line, node.end_line, path, @@ -342,7 +334,7 @@ class MarkdownFileChunker(BaseFileChunker): ] if node.kind == "body": - return self._split_leaf(node, prefix, "", path, renderer) + return self._split_leaf(node, prefix, path, renderer) child_ancestors = ancestors if node.kind == "section": @@ -365,28 +357,32 @@ class MarkdownFileChunker(BaseFileChunker): return if cache is not None: candidate = _toc_join(cache.text, text) - if len(candidate) <= self.chunk_chars: + candidate_size = self._byte_len(candidate) + if candidate_size <= self.chunk_byte_size: cache.text = candidate cache.end_line = end_line - if len(candidate) == self.chunk_chars: + if candidate_size == self.chunk_byte_size: flush_cache() return flush_cache() - cache_text = _toc_join(new_prefix, text) if self.embed_toc else text + cache_text = self._compose_text(new_prefix, text) cache = FileChunk( path=path, start_line=start_line, end_line=end_line, text=cache_text, ) - if len(cache_text) >= self.chunk_chars: + if self._byte_len(cache_text) >= self.chunk_byte_size: flush_cache() + if heading: + append_to_cache(heading, node.start_line, node.start_line, prefix) + for child in node.children: child_heading = f"{'#' * child.level} {child.heading or ''}" if child.kind == "section" else "" child_text = _toc_join(child_heading, child.text) if child_heading else child.text - if len(child_text) <= self.chunk_chars: + if self._byte_len(child_text) <= self.chunk_byte_size: append_to_cache(child_text, child.start_line, child.end_line, child_prefix) continue @@ -394,45 +390,14 @@ class MarkdownFileChunker(BaseFileChunker): chunks.extend(self._chunk_node(child, child_ancestors, path, renderer)) flush_cache() - if node.kind == "section": - if not chunks: - return [ - self._make_chunk( - prefix, - heading, - "", - node.start_line, - node.start_line, - path, - ), - ] - first_content = self._without_prefix(chunks[0].text, child_prefix) - chunks[0] = self._make_chunk( - prefix, - _toc_join(heading, first_content), - "", - node.start_line, - chunks[0].end_line, - path, - ) return chunks - def _without_prefix(self, text: str, prefix: str) -> str: - """Remove a breadcrumb that this chunker prepended to ``text``.""" - if not self.embed_toc or not prefix: - return text - if text == prefix: - return "" - marker = f"{prefix}\n\n" - return text[len(marker) :] if text.startswith(marker) else text - # -- Leaf splitters: build (text, start, end) units, hand off to packer def _split_leaf( self, body: MdNode, - before: str, - after: str, + breadcrumb: str, path: str, renderer, ) -> list[FileChunk]: @@ -444,18 +409,17 @@ class MarkdownFileChunker(BaseFileChunker): block = body.block if isinstance(block, Table): - return self._split_table(body, before, after, path) + return self._split_table(body, breadcrumb, path) if isinstance(block, CodeFence): - return self._split_code(body, before, after, path) + return self._split_code(body, breadcrumb, path) if isinstance(block, List): - return self._split_list(body, before, after, path, renderer) - return self._split_lines(body, before, after, path) + return self._split_list(body, breadcrumb, path, renderer) + return self._split_lines(body, breadcrumb, path) def _split_table( self, body: MdNode, - before: str, - after: str, + breadcrumb: str, path: str, ) -> list[FileChunk]: """Repeat header + separator on every chunk.""" @@ -473,8 +437,7 @@ class MarkdownFileChunker(BaseFileChunker): units = [(text, line_of(i), line_of(i)) for i, text in enumerate(data)] return self._emit_packed( units, - before, - after, + breadcrumb, path, joiner="\n", wrap=f"{header}\n{{inner}}", @@ -483,8 +446,7 @@ class MarkdownFileChunker(BaseFileChunker): def _split_code( self, body: MdNode, - before: str, - after: str, + breadcrumb: str, path: str, ) -> list[FileChunk]: """Repeat fence opener + closer on every chunk.""" @@ -499,8 +461,7 @@ class MarkdownFileChunker(BaseFileChunker): units = [(indent + ln, start + i, start + i) for i, ln in enumerate(raw.split("\n"))] return self._emit_packed( units, - before, - after, + breadcrumb, path, joiner="\n", wrap=f"{opener}\n{{inner}}\n{fence}", @@ -510,8 +471,7 @@ class MarkdownFileChunker(BaseFileChunker): def _split_list( self, body: MdNode, - before: str, - after: str, + breadcrumb: str, path: str, renderer, ) -> list[FileChunk]: @@ -520,7 +480,7 @@ class MarkdownFileChunker(BaseFileChunker): items = [c for c in (body.block.children or []) if isinstance(c, ListItem)] if not items: - return self._split_lines(body, before, after, path) + return self._split_lines(body, breadcrumb, path) units: list[tuple[str, int, int]] = [] line_offset = body.start_line - (getattr(body.block, "line_number", None) or body.start_line) for it in items: @@ -531,8 +491,7 @@ class MarkdownFileChunker(BaseFileChunker): units.append((text, line, line + text.count("\n"))) return self._emit_packed( units, - before, - after, + breadcrumb, path, joiner="\n", wrap="{inner}", @@ -541,8 +500,7 @@ class MarkdownFileChunker(BaseFileChunker): def _split_lines( self, body: MdNode, - before: str, - after: str, + breadcrumb: str, path: str, ) -> list[FileChunk]: """Last-resort line-greedy split for paragraphs / quotes / html.""" @@ -550,8 +508,7 @@ class MarkdownFileChunker(BaseFileChunker): units = [(line, start + i, start + i) for i, line in enumerate(body.text.split("\n"))] return self._emit_packed( units, - before, - after, + breadcrumb, path, joiner="\n", wrap="{inner}", @@ -560,8 +517,7 @@ class MarkdownFileChunker(BaseFileChunker): def _emit_packed( self, units: list[tuple[str, int, int]], - before: str, - after: str, + breadcrumb: str, path: str, joiner: str, wrap: str, @@ -569,49 +525,52 @@ class MarkdownFileChunker(BaseFileChunker): ) -> list[FileChunk]: """Greedy-pack units into ``wrap`` envelopes; emit each piece. - Envelope (table header, code fence) counts against ``chunk_chars``; - TOC (when on) is additive prefix/suffix downstream. Oversized - units overflow rather than truncate. Multi-piece outputs get - ``[Part X/N]`` markers; single pieces don't. + Envelope (table header, code fence), breadcrumb, separators and the + largest possible ``[Part X/N]`` marker all count against + ``chunk_byte_size``. Oversized atomic units overflow rather than + truncate. Multi-piece outputs get part markers; single pieces don't. """ - envelope = len(wrap.replace("{inner}", "")) - budget = max(64, self.chunk_chars - envelope) - sep_len = len(joiner) + envelope = self._byte_len(wrap.replace("{inner}", "")) + breadcrumb_overhead = self._breadcrumb_overhead(breadcrumb) + part_marker = len(units) > 1 + marker_overhead = self._byte_len(f"[Part {len(units)}/{len(units)}]\n\n") if part_marker else 0 + budget = self.chunk_byte_size - envelope - breadcrumb_overhead - marker_overhead + sep_len = self._byte_len(joiner) parts: list[tuple[str, int, int]] = [] bucket: list[tuple[str, int, int]] = [] - bucket_chars = 0 + bucket_bytes = 0 def flush() -> None: - nonlocal bucket, bucket_chars + nonlocal bucket, bucket_bytes if not bucket: return inner = joiner.join(t for t, _, _ in bucket) parts.append((inner, bucket[0][1], bucket[-1][2])) bucket = [] - bucket_chars = 0 + bucket_bytes = 0 for text, s, e in units: if not text and not allow_empty: continue sep = sep_len if bucket else 0 - if bucket_chars + sep + len(text) > budget: + text_size = self._byte_len(text) + if bucket_bytes + sep + text_size > budget: flush() sep = 0 bucket.append((text, s, e)) - bucket_chars += sep + len(text) + bucket_bytes += sep + text_size flush() total = len(parts) return [ self._make_chunk( - before, + breadcrumb, ( f"[Part {idx}/{total}]\n\n{wrap.replace('{inner}', inner)}" if total > 1 else wrap.replace("{inner}", inner) ), - after, s, e, path, @@ -623,19 +582,47 @@ class MarkdownFileChunker(BaseFileChunker): def _make_chunk( self, - before: str, + breadcrumb: str, content: str, - after: str, start_line: int, end_line: int, path: str, ) -> FileChunk: - """Build one ``FileChunk`` — text is ``before + content + after`` - when ``embed_toc``, otherwise just ``content``.""" - text = _toc_join(before, content, after) if self.embed_toc else content + """Build one chunk with an optional ancestor-heading breadcrumb.""" + text = self._compose_text(breadcrumb, content) return FileChunk( path=path, start_line=start_line, end_line=end_line, text=text, ).set_hash_id() + + def _breadcrumb_overhead(self, breadcrumb: str) -> int: + """Return the breadcrumb and separator cost before one content string.""" + if not self.embed_toc or not breadcrumb: + return 0 + return self._byte_len(_toc_join(breadcrumb, "x")) - 1 + + def _compose_text(self, breadcrumb: str, content: str) -> str: + """Compose a chunk and trim breadcrumbs when they consume its budget.""" + if not self.embed_toc or not breadcrumb: + return content + full = _toc_join(breadcrumb, content) + if self._byte_len(full) <= self.chunk_byte_size: + return full + + available = self.chunk_byte_size - self._byte_len(content) - 2 + if available <= 0: + return content + + breadcrumb_parts = breadcrumb.split("\n\n") + for start in range(len(breadcrumb_parts)): + retained = "\n\n".join(breadcrumb_parts[start:]) + candidate = _toc_join(retained, content) + if self._byte_len(candidate) <= self.chunk_byte_size: + return candidate + return content + + def _byte_len(self, text: str) -> int: + """Return encoded size used by the shared byte chunking contract.""" + return len(text.encode(self.encoding)) diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 6ce0e287..6720d330 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -707,7 +707,6 @@ components: supported_extensions: [ "md" ] embed_toc: true max_ast_sections: 100 - default_chunker: default include_frontmatter_in_metadata: false include_frontmatter_keys_in_metadata: [] # empty = all non-empty frontmatter keys json: diff --git a/reme/config/jinli_lme.yaml b/reme/config/jinli_lme.yaml index 8f3708e0..5b4a6f41 100644 --- a/reme/config/jinli_lme.yaml +++ b/reme/config/jinli_lme.yaml @@ -399,7 +399,6 @@ components: supported_extensions: [ "md" ] embed_toc: true max_ast_sections: 100 - default_chunker: default default: backend: default supported_extensions: [ "json", "jsonl" ] diff --git a/tests/unit/test_config_parser.py b/tests/unit/test_config_parser.py index 7a6207ee..0b7c0472 100644 --- a/tests/unit/test_config_parser.py +++ b/tests/unit/test_config_parser.py @@ -56,7 +56,6 @@ def test_default_config_keeps_frontmatter_chunk_metadata_opt_in(): markdown = cfg["components"]["file_chunker"]["markdown"] assert markdown["embed_toc"] is True assert markdown["max_ast_sections"] == 100 - assert markdown["default_chunker"] == "default" assert markdown["include_frontmatter_in_metadata"] is False # Allow-list defaults to empty; combined with the False above, chunk metadata stays empty. assert markdown["include_frontmatter_keys_in_metadata"] == [] or markdown.get( diff --git a/tests/unit/test_markdown_file_chunker.py b/tests/unit/test_markdown_file_chunker.py index f60950f8..1d7bdcb6 100644 --- a/tests/unit/test_markdown_file_chunker.py +++ b/tests/unit/test_markdown_file_chunker.py @@ -13,9 +13,7 @@ import os import tempfile from unittest.mock import patch -from reme.components import ApplicationContext from reme.components.file_chunker import DefaultFileChunker, MarkdownFileChunker -from reme.enumeration import ComponentEnum class temp_chdir: @@ -89,13 +87,13 @@ def test_parse_frontmatter_metadata_is_opt_in(): "Jon said he lost his job today.\n" ) path = _write_md(tmp, "daily/2023-01-19/locomo-event.md", body) - chunker = MarkdownFileChunker(chunk_chars=500) + chunker = MarkdownFileChunker(chunk_byte_size=500) _, chunks = await chunker.chunk(path) assert len(chunks) == 1 assert chunks[0].metadata == {} - chunker = MarkdownFileChunker(chunk_chars=500, include_frontmatter_in_metadata=True) + chunker = MarkdownFileChunker(chunk_byte_size=500, include_frontmatter_in_metadata=True) _, chunks = await chunker.chunk(path) assert len(chunks) == 1 @@ -126,7 +124,7 @@ def test_parse_frontmatter_metadata_keys_allowlist(): # Allow-list restricted to a single key. chunker = MarkdownFileChunker( - chunk_chars=500, + chunk_byte_size=500, include_frontmatter_in_metadata=True, include_frontmatter_keys_in_metadata=["conversation_date"], ) @@ -136,7 +134,7 @@ def test_parse_frontmatter_metadata_keys_allowlist(): # Allow-list with a key not present in frontmatter is a no-op for that key. chunker = MarkdownFileChunker( - chunk_chars=500, + chunk_byte_size=500, include_frontmatter_in_metadata=True, include_frontmatter_keys_in_metadata=["conversation_date", "absent"], ) @@ -145,7 +143,7 @@ def test_parse_frontmatter_metadata_keys_allowlist(): # Empty allow-list (not None) keeps the legacy "all non-empty keys" behavior. chunker = MarkdownFileChunker( - chunk_chars=500, + chunk_byte_size=500, include_frontmatter_in_metadata=True, include_frontmatter_keys_in_metadata=[], ) @@ -158,7 +156,7 @@ def test_parse_frontmatter_metadata_keys_allowlist(): # Allow-list is ignored when the master toggle is off (back-compat default). chunker = MarkdownFileChunker( - chunk_chars=500, + chunk_byte_size=500, include_frontmatter_in_metadata=False, include_frontmatter_keys_in_metadata=["conversation_date"], ) @@ -170,13 +168,13 @@ def test_parse_frontmatter_metadata_keys_allowlist(): def test_parse_small_body_one_chunk(): - """A body shorter than chunk_chars produces exactly one chunk that contains the body.""" + """A body shorter than chunk_byte_size produces exactly one chunk that contains the body.""" async def run(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): body = "# Hello\n\nthis is a small body." path = _write_md(tmp, "small.md", body) - chunker = MarkdownFileChunker(chunk_chars=500) + chunker = MarkdownFileChunker(chunk_byte_size=500) node, chunks = await chunker.chunk(path) assert len(chunks) == 1 assert "this is a small body" in chunks[0].text @@ -193,7 +191,7 @@ def test_parse_small_children_are_cached_without_recursive_calls(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): sections = "\n\n".join(f"# Section-{i}\n\n{'x' * 30}" for i in range(6)) path = _write_md(tmp, "small-children.md", sections) - chunker = MarkdownFileChunker(chunk_chars=100, embed_toc=False) + chunker = MarkdownFileChunker(chunk_byte_size=100, embed_toc=False) with patch.object( chunker, "_chunk_node", @@ -203,22 +201,22 @@ def test_parse_small_children_are_cached_without_recursive_calls(): assert chunk_node.call_count == 1 assert 1 < len(chunks) < 6 - assert all(len(chunk.text) <= chunker.chunk_chars for chunk in chunks) + assert all(len(chunk.text.encode("utf-8")) <= chunker.chunk_byte_size for chunk in chunks) asyncio.run(run()) def test_parse_child_cache_flushes_at_exact_limit(): - """A cache reaching ``chunk_chars`` is finalized before the next child.""" + """A cache reaching ``chunk_byte_size`` is finalized before the next child.""" async def run(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): path = _write_md(tmp, "exact-cache.md", f"# A\n\n{'x' * 95}\n\n# B\n\ny") - chunker = MarkdownFileChunker(chunk_chars=100) + chunker = MarkdownFileChunker(chunk_byte_size=100) _, chunks = await chunker.chunk(path) assert len(chunks) == 2 - assert len(chunks[0].text) == chunker.chunk_chars + assert len(chunks[0].text) == chunker.chunk_byte_size assert chunks[0].text.startswith("# A") assert chunks[1].text == "# B\n\ny" @@ -233,28 +231,29 @@ def test_parse_oversized_child_flushes_parent_cache(): leaves = "\n\n".join(f"## L{i}\n\n{'x' * 10}" for i in range(6)) body = f"# A\n\na\n\n# Large\n\n{leaves}\n\n# C\n\nc" path = _write_md(tmp, "recursive-boundary.md", body) - chunker = MarkdownFileChunker(chunk_chars=100, embed_toc=False) + chunker = MarkdownFileChunker(chunk_byte_size=100, embed_toc=False) _, chunks = await chunker.chunk(path) assert len(chunks) == 4 assert chunks[0].text == "# A\n\na" - assert chunks[-2].text == f"## L5\n\n{'x' * 10}" + assert chunks[-2].text == f"## L4\n\n{'x' * 10}\n\n## L5\n\n{'x' * 10}" assert chunks[-1].text == "# C\n\nc" asyncio.run(run()) def test_parse_oversized_body_splits(): - """A body exceeding chunk_chars triggers multiple chunks.""" + """A body exceeding chunk_byte_size triggers multiple chunks.""" async def run(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): paras = "\n\n".join(f"paragraph {i} with some content text here." for i in range(50)) body = "# H\n\n" + paras path = _write_md(tmp, "big.md", body) - chunker = MarkdownFileChunker(chunk_chars=200) + chunker = MarkdownFileChunker(chunk_byte_size=200) _, chunks = await chunker.chunk(path) assert len(chunks) > 1 + assert all(len(chunk.text.encode("utf-8")) <= chunker.chunk_byte_size for chunk in chunks) print("✓ test_parse_oversized_body_splits passed") asyncio.run(run()) @@ -268,7 +267,7 @@ def test_parse_chunk_ids_match_node_chunk_ids(): paras = "\n\n".join(f"para {i} body content here." for i in range(40)) body = "# H\n\n" + paras path = _write_md(tmp, "p.md", body) - chunker = MarkdownFileChunker(chunk_chars=200) + chunker = MarkdownFileChunker(chunk_byte_size=200) node, chunks = await chunker.chunk(path) assert node.chunk_ids == [c.id for c in chunks] print("✓ test_parse_chunk_ids_match_node_chunk_ids passed") @@ -350,11 +349,11 @@ def test_parse_links_deduped(): asyncio.run(run()) -def test_parse_min_chunk_chars_clamped(): - """chunk_chars below 100 should be clamped to 100.""" - chunker = MarkdownFileChunker(chunk_chars=10) - assert chunker.chunk_chars == 100 - print("✓ test_parse_min_chunk_chars_clamped passed") +def test_parse_min_chunk_byte_size_clamped(): + """chunk_byte_size below 100 should be clamped to 100.""" + chunker = MarkdownFileChunker(chunk_byte_size=10) + assert chunker.chunk_byte_size == 100 + print("✓ test_parse_min_chunk_byte_size_clamped passed") def test_parse_embed_toc_prefixes_chunk_text(): @@ -364,7 +363,7 @@ def test_parse_embed_toc_prefixes_chunk_text(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): body = "# Top\n\n## Sub\n\nbody-content" path = _write_md(tmp, "toc.md", body) - chunker = MarkdownFileChunker(chunk_chars=200, embed_toc=True) + chunker = MarkdownFileChunker(chunk_byte_size=200, embed_toc=True) _, chunks = await chunker.chunk(path) # Single small section fits; check that the heading appears in text. assert any("Top" in c.text for c in chunks) @@ -388,30 +387,19 @@ def test_count_sections_ignores_fenced_headings_and_supports_setext(): def test_parse_excessive_sections_uses_plain_text_without_ast(): - """Too many sections preserve Markdown metadata while bypassing mistletoe.""" + """Inherited fallback is standalone, bounded, and bypasses mistletoe.""" async def run(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - app_context = ApplicationContext(workspace_dir=tmp) - default_chunker = DefaultFileChunker( - name="default", - app_context=app_context, - chunk_byte_size=160, - overlap_byte_size=20, - ) - app_context.components[ComponentEnum.FILE_CHUNKER] = {"default": default_chunker} - sections = "\n\n".join(f"# Section-{i}\n\n{'x' * 40} [[target.md]]" for i in range(4)) + sections = "\n\n".join(f"# Section-{i}\n\n{'x' * 480} [[target.md]]" for i in range(101)) body = f"---\nname: fallback\n---\n{sections}" - _write_md(tmp, "fallback.md", body) - path = os.path.join(tmp, "fallback.md") + path = _write_md(tmp, "fallback.md", body) chunker = MarkdownFileChunker( - app_context=app_context, - chunk_chars=100, + chunk_byte_size=10000, embed_toc=True, - max_ast_sections=2, + max_ast_sections=100, include_frontmatter_in_metadata=True, ) - await default_chunker.start() await chunker.start() try: with ( @@ -420,19 +408,19 @@ def test_parse_excessive_sections_uses_plain_text_without_ast(): side_effect=AssertionError("fallback must not construct an AST"), ), patch.object( - default_chunker, + chunker, "chunk_content", - wraps=default_chunker.chunk_content, + wraps=chunker.chunk_content, ) as chunk_content, ): node, chunks = await chunker.chunk(path) - assert chunker.default_chunker is default_chunker assert chunk_content.call_count == 1 finally: await chunker.close() - await default_chunker.close() + assert isinstance(chunker, DefaultFileChunker) assert len(chunks) > 1 + assert max(len(chunk.text.encode("utf-8")) for chunk in chunks) <= chunker.chunk_byte_size assert chunks[0].start_line == 4 assert node.front_matter.name == "fallback" assert node.chunk_ids == [chunk.id for chunk in chunks] @@ -440,7 +428,7 @@ def test_parse_excessive_sections_uses_plain_text_without_ast(): ("target.md", "fallback.md"), } assert all(chunk.metadata == {"name": "fallback"} for chunk in chunks) - for i in range(4): + for i in (0, 100): assert any(f"# Section-{i}" in chunk.text for chunk in chunks) asyncio.run(run()) @@ -469,7 +457,7 @@ def test_parse_small_sections_are_merged_and_headings_preserved(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): sections = "\n\n".join(f"## Section-{i:03d}\n\nfact-{i:03d}" for i in range(40)) path = _write_md(tmp, "sections.md", f"# Root\n\n{sections}") - chunker = MarkdownFileChunker(chunk_chars=200, embed_toc=False) + chunker = MarkdownFileChunker(chunk_byte_size=200, embed_toc=False) _, chunks = await chunker.chunk(path) assert 1 < len(chunks) < 40 @@ -487,7 +475,7 @@ def test_parse_embed_toc_uses_breadcrumbs_without_sibling_duplication(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): sections = "\n\n".join(f"## Parallel-{i:03d}\n\nfact-{i:03d}" for i in range(40)) path = _write_md(tmp, "breadcrumbs.md", f"# Root\n\n{sections}") - chunker = MarkdownFileChunker(chunk_chars=200, embed_toc=True) + chunker = MarkdownFileChunker(chunk_byte_size=200, embed_toc=True) _, chunks = await chunker.chunk(path) assert len(chunks) > 1 @@ -507,7 +495,7 @@ def test_parse_heading_heavy_output_grows_linearly(): sections = "\n\n".join(f"## Observation-{i:04d}\n\nsynthetic fact" for i in range(1000)) body = f"# Root\n\n{sections}" path = _write_md(tmp, "linear.md", body) - chunker = MarkdownFileChunker(chunk_chars=10000, embed_toc=True, max_ast_sections=None) + chunker = MarkdownFileChunker(chunk_byte_size=10000, embed_toc=True, max_ast_sections=None) _, chunks = await chunker.chunk(path) assert len(chunks) < 10 @@ -527,14 +515,37 @@ def test_parse_nested_sections_merge_within_recursive_context(): branches.append(f"## {branch}\n\n{leaves}") branch_text = "\n\n".join(branches) path = _write_md(tmp, "nested.md", f"# Root\n\n{branch_text}") - chunker = MarkdownFileChunker(chunk_chars=150, embed_toc=True) + chunker = MarkdownFileChunker(chunk_byte_size=150, embed_toc=True) _, chunks = await chunker.chunk(path) - assert len(chunks) == 4 - assert all(len(chunk.text) <= chunker.chunk_chars for chunk in chunks) - assert chunks[1].text.startswith("# Root\n\n## A\n\n### A2") - assert chunks[2].text.startswith("# Root\n\n## B\n\n### B0") - assert "## B" not in chunks[1].text + assert len(chunks) == 5 + assert all(len(chunk.text.encode("utf-8")) <= chunker.chunk_byte_size for chunk in chunks) + assert chunks[0].text == "# Root" + assert chunks[2].text.startswith("# Root\n\n## A\n\n### A2") + assert chunks[3].text.startswith("# Root\n\n## B\n\n### B0") + assert "## B" not in chunks[2].text + + asyncio.run(run()) + + +def test_parse_breadcrumbs_and_part_markers_share_byte_budget(): + """Breadcrumbs and part labels cannot push multi-byte chunks over budget.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + content_lines = "\n".join(f"内容-{i}-" * 6 for i in range(4)) + body = f"# {'根' * 15}\n\n## {'枝' * 15}\n\n### {'叶' * 15}\n\n{content_lines}" + path = _write_md(tmp, "breadcrumb-budget.md", body) + chunker = MarkdownFileChunker( + chunk_byte_size=100, + embed_toc=True, + max_ast_sections=None, + ) + _, chunks = await chunker.chunk(path) + + assert len(chunks) > 1 + assert any(chunk.text.startswith("[Part ") for chunk in chunks) + assert all(len(chunk.text.encode("utf-8")) <= chunker.chunk_byte_size for chunk in chunks) asyncio.run(run()) @@ -546,7 +557,7 @@ def test_parse_frontmatter_preserves_original_line_numbers(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): body = "---\nname: t\n---\n# H\nline 1\nline 2\n" path = _write_md(tmp, "front-lines.md", body) - chunker = MarkdownFileChunker(chunk_chars=500) + chunker = MarkdownFileChunker(chunk_byte_size=500) _, chunks = await chunker.chunk(path) assert len(chunks) == 1 assert chunks[0].start_line == 4 @@ -564,7 +575,7 @@ def test_parse_frontmatter_offsets_split_table_rows(): rows = "".join(f"| {i} | {i} |\n" for i in range(12)) body = "---\nname: t\n---\n| A | B |\n|---|---|\n" + rows path = _write_md(tmp, "front-table.md", body) - chunker = MarkdownFileChunker(chunk_chars=100) + chunker = MarkdownFileChunker(chunk_byte_size=100) _, chunks = await chunker.chunk(path) assert len(chunks) > 1 assert chunks[0].start_line == 6 @@ -581,7 +592,7 @@ def test_parse_bad_frontmatter_does_not_abort_chunking(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): body = "---\nname: [\n---\n# H\nbody\n" path = _write_md(tmp, "bad-frontmatter.md", body) - chunker = MarkdownFileChunker(chunk_chars=500) + chunker = MarkdownFileChunker(chunk_byte_size=500) node, chunks = await chunker.chunk(path) assert node.front_matter.name == "" assert len(chunks) == 1 @@ -607,7 +618,7 @@ if __name__ == "__main__": test_parse_links_short_and_no_ext_kept_literally() test_parse_links_predicate_inline_and_line() test_parse_links_deduped() - test_parse_min_chunk_chars_clamped() + test_parse_min_chunk_byte_size_clamped() test_parse_embed_toc_prefixes_chunk_text() test_parse_embed_toc_is_enabled_by_default() test_count_sections_ignores_fenced_headings_and_supports_setext() @@ -617,6 +628,7 @@ if __name__ == "__main__": test_parse_embed_toc_uses_breadcrumbs_without_sibling_duplication() test_parse_heading_heavy_output_grows_linearly() test_parse_nested_sections_merge_within_recursive_context() + test_parse_breadcrumbs_and_part_markers_share_byte_budget() test_parse_frontmatter_preserves_original_line_numbers() test_parse_frontmatter_offsets_split_table_rows() test_parse_bad_frontmatter_does_not_abort_chunking()