diff --git a/reme2/component/embedding/base_embedding_model.py b/reme2/component/embedding/base_embedding_model.py index 35f4ce97..21453bed 100644 --- a/reme2/component/embedding/base_embedding_model.py +++ b/reme2/component/embedding/base_embedding_model.py @@ -187,6 +187,7 @@ class BaseEmbeddingModel(BaseComponent): async def get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]: """Get embeddings for multiple texts with cache and batching.""" + # TODO change to bytes instead of str truncated_texts = [t[:self.max_input_length] for t in input_text] results: list[list[float] | None] = [None] * len(truncated_texts) texts_to_compute: list[tuple[int, str]] = [] diff --git a/reme2/component/file_parser/base_file_parser.py b/reme2/component/file_parser/base_file_parser.py index fe310a29..580ee445 100644 --- a/reme2/component/file_parser/base_file_parser.py +++ b/reme2/component/file_parser/base_file_parser.py @@ -3,39 +3,16 @@ from abc import abstractmethod from ..base_component import BaseComponent -from ...enumeration import ComponentEnum +from ...enumeration import ComponentEnum, FileSuffixEnum from ...schema import FileChunk, FileMetadata class BaseFileParser(BaseComponent): - """Abstract base class for file format parsers. - - Each parser declares which file suffixes it handles and implements - the parse method to produce FileMetadata and FileChunks. - """ + """Parser that declares handled suffixes and produces FileChunks.""" component_type = ComponentEnum.FILE_PARSER - - suffixes: list[str] = [] - - def __init__(self, chunk_tokens: int = 400, chunk_overlap: int = 80, **kwargs): - super().__init__(**kwargs) - self.chunk_tokens = chunk_tokens - self.chunk_overlap = chunk_overlap - - async def _start(self): - pass - - async def _close(self): - pass + suffixes: list[FileSuffixEnum] = [] @abstractmethod async def parse(self, path: str) -> tuple[FileMetadata, list[FileChunk]]: - """Parse a file into metadata and chunks. - - Args: - path: Absolute path to the file. - - Returns: - Tuple of (FileMetadata, list of FileChunks). - """ + """Parse a file into metadata and chunks.""" diff --git a/reme2/component/file_parser/default_file_parser.py b/reme2/component/file_parser/default_file_parser.py deleted file mode 100644 index e262c569..00000000 --- a/reme2/component/file_parser/default_file_parser.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Default file parser for unknown file types.""" - -import asyncio -from pathlib import Path - -from .base_file_parser import BaseFileParser -from ..component_registry import R -from ...schema import FileChunk, FileMetadata -from ...utils import chunk_markdown - - -@R.register("default") -class DefaultFileParser(BaseFileParser): - """Fallback parser for unknown file types. - - Attempts to read as text and chunk. If the file is binary, - stores metadata only with no chunks. - """ - - suffixes = [] - - def __init__(self, encoding: str = "utf-8", **kwargs): - super().__init__(**kwargs) - self.encoding = encoding - - async def parse(self, path: str) -> tuple[FileMetadata, list[FileChunk]]: - file_path = Path(path) - - def _read_file(): - stat = file_path.stat() - raw = file_path.read_bytes() - try: - return stat, raw.decode(self.encoding) - except (UnicodeDecodeError, ValueError): - return stat, None - - stat, content = await asyncio.to_thread(_read_file) - - file_meta = FileMetadata( - modified_time=stat.st_mtime, - path=str(file_path.absolute()), - ) - - chunks: list[FileChunk] = [] - if content: - chunks = ( - chunk_markdown( - content, - file_meta.path, - self.chunk_tokens, - self.chunk_overlap, - ) - or [] - ) - - return file_meta, chunks diff --git a/reme2/component/file_parser/md_file_parser.py b/reme2/component/file_parser/md_file_parser.py index 4142f658..fdb60373 100644 --- a/reme2/component/file_parser/md_file_parser.py +++ b/reme2/component/file_parser/md_file_parser.py @@ -7,6 +7,7 @@ import frontmatter from .base_file_parser import BaseFileParser from ..component_registry import R +from ...enumeration import FileSuffixEnum from ...schema import FileChunk, FileMetadata from ...utils import chunk_markdown @@ -15,11 +16,13 @@ from ...utils import chunk_markdown class MdFileParser(BaseFileParser): """Parser for Markdown files with YAML frontmatter support.""" - suffixes = [".md", ".markdown"] + suffixes = [FileSuffixEnum.MD, FileSuffixEnum.MARKDOWN] - def __init__(self, encoding: str = "utf-8", **kwargs): + def __init__(self, encoding: str = "utf-8", chunk_tokens: int = 400, chunk_overlap: int = 80, **kwargs): super().__init__(**kwargs) self.encoding = encoding + self.chunk_tokens = chunk_tokens + self.chunk_overlap = chunk_overlap async def parse(self, path: str) -> tuple[FileMetadata, list[FileChunk]]: file_path = Path(path) diff --git a/reme2/component/file_parser/text_file_parser.py b/reme2/component/file_parser/text_file_parser.py new file mode 100644 index 00000000..490714e9 --- /dev/null +++ b/reme2/component/file_parser/text_file_parser.py @@ -0,0 +1,107 @@ +"""Text file parser with built-in chunking.""" + +import hashlib +from pathlib import Path + +import aiofiles + +from .base_file_parser import BaseFileParser +from ..component_registry import R +from ...enumeration import FileSuffixEnum +from ...schema import FileChunk, FileMetadata + + +@R.register("text") +class TextFileParser(BaseFileParser): + """Parser for text files with built-in chunking support.""" + + suffixes = [FileSuffixEnum.TXT] + + def __init__(self, encoding: str = "utf-8", chunk_tokens: int = 400, chunk_overlap: int = 80, **kwargs): + super().__init__(**kwargs) + self.encoding = encoding + self.chunk_size = max(32, chunk_tokens * 4) + self.overlap_size = max(0, chunk_overlap * 4) + + async def parse(self, path: str) -> tuple[FileMetadata, list[FileChunk]]: + file_path = Path(path) + stat = file_path.stat() + + try: + async with aiofiles.open(file_path, encoding=self.encoding) as f: + content = await f.read() + except UnicodeDecodeError: + async with aiofiles.open(file_path, encoding=self.encoding, errors="ignore") as f: + content = await f.read() + except Exception: + content = None + + file_meta = FileMetadata( + file=file_path.stem, + path=str(file_path.absolute()), + st_mtime=stat.st_mtime, + ) + + chunks = self._chunk(content, file_meta.path) if content else [] + return file_meta, chunks + + def _chunk(self, text: str, path: str) -> list[FileChunk]: + """Split text into chunks with overlap.""" + if not text.strip(): + return [] + + lines = text.split("\n") + chunks: list[FileChunk] = [] + buf: list[tuple[str, int]] = [] # (line_content, line_no) + buf_chars = 0 + + for line_no, line in enumerate(lines, 1): + # Split long lines into segments + for start in range(0, max(1, len(line)), self.chunk_size): + seg = line[start:start + self.chunk_size] + seg_chars = len(seg) + 1 # +1 for newline + + # Flush when buffer would exceed limit + if buf and buf_chars + seg_chars > self.chunk_size: + self._flush_chunk(chunks, buf, path) + buf, buf_chars = self._carry_overlap(buf) + + buf.append((seg, line_no)) + buf_chars += seg_chars + + if buf: + self._flush_chunk(chunks, buf, path) + + return [c for c in chunks if c.text.strip()] + + @staticmethod + def _flush_chunk(chunks: list[FileChunk], buf: list[tuple[str, int]], path: str): + """Create a chunk from buffer and append to chunks list.""" + chunk_text = "\n".join(content for content, _ in buf) + start_line, end_line = buf[0][1], buf[-1][1] + h = hashlib.sha256(chunk_text.encode()).hexdigest() + chunk_id = hashlib.sha256(f"{path}:{start_line}:{end_line}:{h}:{len(chunks)}".encode()).hexdigest() + + chunks.append(FileChunk( + id=chunk_id, + path=path, + start_line=start_line, + end_line=end_line, + text=chunk_text, + hash=h, + )) + + def _carry_overlap(self, buf: list[tuple[str, int]]) -> tuple[list[tuple[str, int]], int]: + """Keep overlapping lines from the end of buffer.""" + if self.overlap_size <= 0 or not buf: + return [], 0 + + acc, kept = 0, [] + for content, line_no in reversed(buf): + acc += len(content) + 1 + kept.insert(0, (content, line_no)) + if acc >= self.overlap_size: + break + + buf_chars = sum(len(c) + 1 for c, _ in kept) + return kept, buf_chars diff --git a/reme2/enumeration/__init__.py b/reme2/enumeration/__init__.py index 9887ddee..ef66736c 100644 --- a/reme2/enumeration/__init__.py +++ b/reme2/enumeration/__init__.py @@ -2,8 +2,10 @@ from .chunk_enum import ChunkEnum from .component_enum import ComponentEnum +from .file_suffix_enum import FileSuffixEnum __all__ = [ "ChunkEnum", "ComponentEnum", + "FileSuffixEnum", ] diff --git a/reme2/enumeration/file_suffix_enum.py b/reme2/enumeration/file_suffix_enum.py new file mode 100644 index 00000000..7732e194 --- /dev/null +++ b/reme2/enumeration/file_suffix_enum.py @@ -0,0 +1,19 @@ +"""File suffix enumeration module. + +Defines the file suffixes that can be handled by file parsers. +""" + +from enum import Enum + + +class FileSuffixEnum(str, Enum): + """Enumeration of supported file suffixes. + + Each value represents a file extension that a file parser can handle. + """ + + MD = ".md" + + MARKDOWN = ".markdown" + + TXT = ".txt"