diff --git a/reme/components/file_chunker/default_file_chunker.py b/reme/components/file_chunker/default_file_chunker.py index 23457f5b..acc545a8 100644 --- a/reme/components/file_chunker/default_file_chunker.py +++ b/reme/components/file_chunker/default_file_chunker.py @@ -2,6 +2,7 @@ from bisect import bisect_right from pathlib import Path +from typing import Literal import aiofiles import yaml @@ -11,14 +12,28 @@ from ..component_registry import R from ...schema import FileChunk, FileFrontMatter, FileNode from ...utils.wikilink_handler import WikilinkHandler +InvalidEncodingPolicy = Literal["replace", "strict"] + @R.register("default") class DefaultFileChunker(BaseFileChunker): """Default chunker that splits files into byte-based overlapping chunks.""" - def __init__(self, encoding: str = "utf-8", chunk_byte_size: int = 10000, overlap_byte_size: int = 100, **kwargs): + def __init__( + self, + encoding: str = "utf-8", + invalid_encoding_policy: InvalidEncodingPolicy = "replace", + chunk_byte_size: int = 10000, + overlap_byte_size: int = 100, + **kwargs, + ): super().__init__(**kwargs) + if invalid_encoding_policy not in {"replace", "strict"}: + raise ValueError( + f"invalid_encoding_policy must be 'replace' or 'strict', got {invalid_encoding_policy!r}", + ) self.encoding = encoding + self.invalid_encoding_policy = invalid_encoding_policy self.chunk_byte_size = max(100, chunk_byte_size) self.overlap_byte_size = max(4, overlap_byte_size) @@ -37,13 +52,28 @@ class DefaultFileChunker(BaseFileChunker): front_matter = FileFrontMatter() return front_matter, text[end_idx + 4 :].lstrip("\n") + async def _read_text_for_indexing(self, file_path: Path) -> str: + """Decode text for a derived index without modifying the source file.""" + async with aiofiles.open(file_path, "rb") as f: + data = await f.read() + try: + return data.decode(self.encoding) + except UnicodeDecodeError as exc: + if self.invalid_encoding_policy == "strict": + raise + invalid_bytes = data[exc.start : exc.end].hex(" ") + self.logger.warning( + f"Invalid {self.encoding} in {file_path} at byte {exc.start} (bytes: {invalid_bytes}); " + "indexed with replacement characters; source file unchanged", + ) + return data.decode(self.encoding, errors="replace") + async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: file_path = Path(path) stat = file_path.stat() rel_path = self.to_workspace_relative(path) - async with aiofiles.open(file_path, encoding=self.encoding) as f: - text = await f.read() + text = await self._read_text_for_indexing(file_path) if not text: return FileNode(path=rel_path, st_mtime=stat.st_mtime), [] diff --git a/reme/components/file_chunker/markdown_file_chunker.py b/reme/components/file_chunker/markdown_file_chunker.py index 7da94ae6..a01e3e82 100644 --- a/reme/components/file_chunker/markdown_file_chunker.py +++ b/reme/components/file_chunker/markdown_file_chunker.py @@ -24,7 +24,7 @@ import yaml from pydantic import ValidationError -from .default_file_chunker import DefaultFileChunker +from .default_file_chunker import DefaultFileChunker, InvalidEncodingPolicy from ..component_registry import R from ...schema import ( FileChunk, @@ -103,6 +103,7 @@ class MarkdownFileChunker(DefaultFileChunker): def __init__( self, encoding: str = "utf-8", + invalid_encoding_policy: InvalidEncodingPolicy = "replace", chunk_byte_size: int = 10000, embed_toc: bool = True, max_ast_sections: int | None = 100, @@ -110,7 +111,12 @@ class MarkdownFileChunker(DefaultFileChunker): include_frontmatter_keys_in_metadata: list[str] | None = None, **kwargs, ): - super().__init__(encoding=encoding, chunk_byte_size=chunk_byte_size, **kwargs) + super().__init__( + encoding=encoding, + invalid_encoding_policy=invalid_encoding_policy, + 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.include_frontmatter_in_metadata = include_frontmatter_in_metadata @@ -119,7 +125,8 @@ class MarkdownFileChunker(DefaultFileChunker): async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: file_path = Path(path) rel_path = self.to_workspace_relative(path) - front_matter, content, line_offset = self._parse_front_matter(file_path.read_text(encoding=self.encoding)) + text = await self._read_text_for_indexing(file_path) + front_matter, content, line_offset = self._parse_front_matter(text) chunks: list[FileChunk] = [] if content and content.strip(): diff --git a/tests/unit/test_default_file_chunker.py b/tests/unit/test_default_file_chunker.py index bf22d77e..d5bbdf23 100644 --- a/tests/unit/test_default_file_chunker.py +++ b/tests/unit/test_default_file_chunker.py @@ -97,6 +97,16 @@ def test_parse_with_custom_encoding(): asyncio.run(run()) +def test_invalid_encoding_policy_is_validated(): + """Reject misspelled policies instead of silently changing decoding behavior.""" + try: + DefaultFileChunker(invalid_encoding_policy="ignore") + except ValueError as exc: + assert "invalid_encoding_policy" in str(exc) + else: + raise AssertionError("invalid encoding policy must be rejected") + + def test_parse_links_bare(): """Bare wikilink: [[target]].""" links = WikilinkHandler.extract_links("see [[note]]", "src.md") diff --git a/tests/unit/test_markdown_file_chunker.py b/tests/unit/test_markdown_file_chunker.py index 118443ca..54676607 100644 --- a/tests/unit/test_markdown_file_chunker.py +++ b/tests/unit/test_markdown_file_chunker.py @@ -57,6 +57,51 @@ def test_parse_empty_file(): asyncio.run(run()) +def test_invalid_utf8_is_replaced_without_modifying_source(): + """Bad source bytes degrade the derived index but remain untouched on disk.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + source = b"# valid\ncontent before \xcd content after\n" + path = os.path.join(tmp, "invalid.md") + with open(path, "wb") as f: + f.write(source) + + chunker = MarkdownFileChunker() + with patch.object(chunker.logger, "warning") as warning: + node, chunks = await chunker.chunk("invalid.md") + + assert node.path == "invalid.md" + assert "content before \ufffd content after" in "\n".join(chunk.text for chunk in chunks) + with open(path, "rb") as f: + assert f.read() == source + warning.assert_called_once_with( + "Invalid utf-8 in invalid.md at byte 23 (bytes: cd); " + "indexed with replacement characters; source file unchanged", + ) + + asyncio.run(run()) + + +def test_invalid_utf8_strict_policy_still_raises(): + """Strict mode remains available when callers require exact decoding.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + with open(os.path.join(tmp, "invalid.md"), "wb") as f: + f.write(b"valid\xcdinvalid") + + chunker = MarkdownFileChunker(invalid_encoding_policy="strict") + try: + await chunker.chunk("invalid.md") + except UnicodeDecodeError as exc: + assert exc.start == 5 + else: + raise AssertionError("strict policy must reject invalid UTF-8") + + asyncio.run(run()) + + def test_parse_frontmatter_only(): """A file with only frontmatter (no body) → no chunks, no links."""