mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
fix: enforce markdown chunk byte limits (#370)
Some checks failed
Pre-commit / run (ubuntu-latest) (push) Has been cancelled
Tests ReMe / Unit Tests - py3.11 (push) Has been cancelled
Tests ReMe / Unit Tests - py3.12 (push) Has been cancelled
Tests ReMe / Unit Tests - py3.13 (push) Has been cancelled
Windows Smoke / CLI smoke - py3.11 (push) Has been cancelled
Some checks failed
Pre-commit / run (ubuntu-latest) (push) Has been cancelled
Tests ReMe / Unit Tests - py3.11 (push) Has been cancelled
Tests ReMe / Unit Tests - py3.12 (push) Has been cancelled
Tests ReMe / Unit Tests - py3.13 (push) Has been cancelled
Windows Smoke / CLI smoke - py3.11 (push) Has been cancelled
This commit is contained in:
parent
987f275985
commit
1c08eaa559
5 changed files with 153 additions and 157 deletions
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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" ]
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue