mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-11 22:51:10 +00:00
feat(parser): add base file parser and concrete implementations
Some checks failed
Pre-commit / run (ubuntu-latest) (push) Has been cancelled
Some checks failed
Pre-commit / run (ubuntu-latest) (push) Has been cancelled
- Introduce BaseFileParser abstract class with component registration - Add MdFileParser implementation for markdown files with YAML frontmatter - Create TextFileParser implementation with built-in chunking support - Implement file suffix enumeration for parser type safety - Add chunking logic with configurable token size and overlap - Support text file parsing with error handling for encoding issues - Include line number tracking and content hashing for file chunks
This commit is contained in:
parent
701669310c
commit
3fc3fd65e8
7 changed files with 138 additions and 85 deletions
|
|
@ -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]] = []
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
107
reme2/component/file_parser/text_file_parser.py
Normal file
107
reme2/component/file_parser/text_file_parser.py
Normal file
|
|
@ -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
|
||||
|
|
@ -2,8 +2,10 @@
|
|||
|
||||
from .chunk_enum import ChunkEnum
|
||||
from .component_enum import ComponentEnum
|
||||
from .file_suffix_enum import FileSuffixEnum
|
||||
|
||||
__all__ = [
|
||||
"ChunkEnum",
|
||||
"ComponentEnum",
|
||||
"FileSuffixEnum",
|
||||
]
|
||||
|
|
|
|||
19
reme2/enumeration/file_suffix_enum.py
Normal file
19
reme2/enumeration/file_suffix_enum.py
Normal file
|
|
@ -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"
|
||||
Loading…
Add table
Reference in a new issue