diff --git a/reme2/component/file_parser/default_file_parser.py b/reme2/component/file_parser/default_file_parser.py index b1f30bae..5ea7254c 100644 --- a/reme2/component/file_parser/default_file_parser.py +++ b/reme2/component/file_parser/default_file_parser.py @@ -1,11 +1,12 @@ +from bisect import bisect_right from pathlib import Path import aiofiles +import yaml from .base_file_parser import BaseFileParser from ..component_registry import R -from ...schema import FileChunk, FileNode -from ...utils.common_utils import hash_text +from ...schema import FileChunk, FileNode, FileFrontMatter @R.register("default") @@ -18,6 +19,30 @@ class DefaultFileParser(BaseFileParser): self.chunk_byte_size = max(100, chunk_byte_size) self.overlap_byte_size = max(4, overlap_byte_size) + @staticmethod + def _parse_front_matter(text: str) -> tuple[FileFrontMatter, str]: + """Parse Markdown front_matter and return (front_matter, remaining_content).""" + if not text.startswith("---"): + return FileFrontMatter(), text + + # Find the closing --- (must be at the start of a line) + end_idx = text.find("\n---", 3) + if end_idx == -1: + return FileFrontMatter(), text + + # Parse YAML front_matter + yaml_content = text[3:end_idx].strip() + try: + data = yaml.safe_load(yaml_content) or {} + if not isinstance(data, dict): + data = {} + except yaml.YAMLError: + data = {} + + front_matter = FileFrontMatter(**data) + remaining = text[end_idx + 4:].lstrip("\n") + return front_matter, remaining + async def parse(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: file_path = Path(path) stat = file_path.stat() @@ -29,24 +54,36 @@ class DefaultFileParser(BaseFileParser): if not data: return FileNode(path=rel_path, st_mtime=stat.st_mtime), [] - newline_positions = [i for i, b in enumerate(data) if b == ord(b"\n")] + text = data.decode(self.encoding, errors="ignore") + front_matter, content = self._parse_front_matter(text) + + if not content: + return FileNode(path=rel_path, st_mtime=stat.st_mtime, front_matter=front_matter), [] + + content_bytes = content.encode(self.encoding) + newline_positions = [i for i, b in enumerate(content_bytes) if b == ord(b"\n")] chunks: list[FileChunk] = [] step = self.chunk_byte_size - self.overlap_byte_size start = 0 - while start < len(data): - end = min(start + self.chunk_byte_size, len(data)) - text = data[start:end].decode(self.encoding, errors="ignore") - start_line = sum(1 for p in newline_positions if p < start) + 1 - end_line = sum(1 for p in newline_positions if p < end) + 1 + while start < len(content_bytes): + end = min(start + self.chunk_byte_size, len(content_bytes)) + chunk_text = content_bytes[start:end].decode(self.encoding, errors="ignore") + start_line = bisect_right(newline_positions, start - 1) + 1 + end_line = bisect_right(newline_positions, end - 1) + 1 + if content_bytes[end - 1] == ord(b"\n"): + end_line -= 1 chunks.append(FileChunk( path=rel_path, start_line=start_line, end_line=end_line, - text=text, + text=chunk_text, ).set_hash_id()) - start += step if end < len(data) else len(data) + if end >= len(content_bytes): + break + start += step - return FileNode(path=rel_path, st_mtime=stat.st_mtime, chunk_ids=[c.id for c in chunks]), chunks \ No newline at end of file + return FileNode(path=rel_path, st_mtime=stat.st_mtime, front_matter=front_matter, + chunk_ids=[c.id for c in chunks]), chunks