mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
fix(index): tolerate invalid text encoding
This commit is contained in:
parent
626c850ccb
commit
d4bfa3aaea
4 changed files with 98 additions and 6 deletions
|
|
@ -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), []
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue