diff --git a/src/core/splitters/__init__.py b/src/core/splitters/__init__.py new file mode 100644 index 0000000..4fc8eed --- /dev/null +++ b/src/core/splitters/__init__.py @@ -0,0 +1,16 @@ +"""文档分块器包 — 支持 Markdown / 纯文本 / PDF / HTML.""" + +from src.core.splitters.base import Splitter, BaseTextSplitter +from src.core.splitters.markdown import MarkdownSplitter +from src.core.splitters.text import TextSplitter +from src.core.splitters.registry import get_splitter, register_splitter, SUPPORTED_SUFFIXES + +__all__ = [ + "Splitter", + "BaseTextSplitter", + "MarkdownSplitter", + "TextSplitter", + "get_splitter", + "register_splitter", + "SUPPORTED_SUFFIXES", +] diff --git a/src/core/splitters/base.py b/src/core/splitters/base.py new file mode 100644 index 0000000..a9f9697 --- /dev/null +++ b/src/core/splitters/base.py @@ -0,0 +1,85 @@ +"""Splitter Protocol 和文本切分基类.""" +import re +from abc import ABC, abstractmethod +from typing import Protocol + + +class Splitter(Protocol): + """文档分块器接口 — 将文本拆分为带元数据的 chunk 列表. + + 每个 chunk 为 dict: {"content": str, "section_title": str, "heading_level": int, ...} + """ + + def split(self, text: str, source_file: str = "") -> list[dict]: ... + + +class BaseTextSplitter(ABC): + """文本切分基类 — 提供段落切分和硬切逻辑,子类实现 split().""" + + def __init__(self, max_size: int = 1000, overlap: int = 100): + self.max_size = max_size + self.overlap = overlap + + @abstractmethod + def split(self, text: str, source_file: str = "") -> list[dict]: ... + + def _split_by_paragraphs( + self, text: str, section_title: str = "", heading_level: int = 0 + ) -> list[dict]: + """按段落边界拆分超长文本,若单段仍超长则硬切.""" + paragraphs = re.split(r"\n\n+", text) + chunks = [] + current = "" + + for para in paragraphs: + if len(para) > self.max_size: + if current.strip(): + chunks.append(self._make_chunk(current, section_title, heading_level)) + current = "" + for sub in self._split_single_paragraph(para): + chunks.append(self._make_chunk(sub, section_title, heading_level)) + continue + + if len(current) + len(para) > self.max_size and current: + chunks.append(self._make_chunk(current, section_title, heading_level)) + if self.overlap > 0 and len(current) > self.overlap: + current = current[-self.overlap:] + "\n\n" + para + else: + current = para + else: + current = f"{current}\n\n{para}" if current else para + + if current.strip(): + chunks.append(self._make_chunk(current, section_title, heading_level)) + + return chunks + + def _split_single_paragraph(self, text: str) -> list[str]: + """按字符边界硬切单个超长段落(带 overlap).""" + parts = [] + start = 0 + while start < len(text): + end = start + self.max_size + if end >= len(text): + parts.append(text[start:].strip()) + break + break_point = end + for sep in ("。", "!", "?", "\n", ". ", " "): + pos = text.rfind(sep, start, end) + if pos > start: + break_point = pos + len(sep) + break + part = text[start:break_point].strip() + if part: + parts.append(part) + next_start = break_point - self.overlap if self.overlap > 0 else break_point + start = max(start + 1, next_start) + return parts + + @staticmethod + def _make_chunk(content: str, section_title: str, heading_level: int) -> dict: + return { + "content": content.strip(), + "section_title": section_title, + "heading_level": heading_level, + }