src/database/loaders.py - master - gurgutan/tutor

src/database/loaders.py
144 строки6 KB
Ivan Slepovichev
initial
07 фев 2025, 18:17
07 фев 2025, 18:17a823a91
"""
Реализация классов для загрузки markdown-файла с рекурсивным разбиением текста по заголовкам.

При загрузке файла текст рекурсивно разбивается по заголовкам от уровня 1 до headers_level.
Для каждого фрагмента формируется заголовок (title) как конкатенация всех заголовков, найденных на пути.

Использование:
    loader = MarkdownLoader(headers_level=2)
    file_path = Path("data/tprg-lections-full.md")
    documents = loader.load(file_path)
    for doc in documents:
        print(doc.title)
"""

from collections import deque
import re
from pathlib import Path
from typing import List, Tuple, Union
from uuid import uuid4

from .models import Document


class MarkdownSplitter:
    """Класс для разбиения markdown-текста на фрагменты по заголовкам определенного уровня.

    Использование:
        splitter = MarkdownSplitter(header_level=2)
        text = "# Title\nSome text\n## Subtitle\nMore text"
        fragments = splitter.split(text=text)
        print(f"Фрагменты:\n{fragments}")
    """

    def __init__(self, header_level: int) -> None:
        """Инициализация MarkdownSplitter с заданием уровня заголовков."""
        self.header_level = header_level

    def split(self, text: str, level: int = 1) -> List[str]:
        """Разбивает markdown-текст на фрагменты по заголовкам заданного уровня.

        Args:
            text: Строка с текстом документа в формате Markdown.

        Returns:
            Список фрагментов, каждый фрагмент начинается с заголовка указанного уровня.

        Example:
            splitter = MarkdownSplitter(header_level=2)
            md_text = "## Header1\nContent1\n## Header2\nContent2"
            result = splitter.split(md_text)
            print(result)  # ['## Header1\nContent1', '## Header2\nContent2']
        """
        if not text:
            return []

        header_token = f'{"#" * level} '

        fragments = []
        current_fragments = []

        lines = text.splitlines()
        for line in lines:
            if line.startswith(header_token):
                if current_fragments:
                    fragments.append("\n".join(current_fragments))
                    current_fragments = []
            current_fragments.append(line)
        if current_fragments:
            fragments.append("\n".join(current_fragments))
        return fragments

    def recursive_split(self, text: str) -> List[Tuple[str, str]]:
        """Рекурсивно разбивает список строк по заголовкам указанного уровня."""

        queue = deque()
        queue.append((0, text, ""))  # (level, text, title)
        fragments = []
        while queue:
            level, current_text, title = queue.popleft()
            if level == self.header_level:
                fragments.append((current_text, title))
            elif level < self.header_level:
                chunks = self.split(current_text, level + 1)
                for chunk in chunks:
                    chunk_lines = chunk.split("\n")
                    if not chunk_lines:
                        continue
                    prefix = title + ". " if title else ""
                    chunk_first_line = chunk_lines[0][:256].strip("# ")
                    chunk_title = prefix + (
                        chunk_first_line if title != chunk_first_line else ""
                    )
                    queue.append((level + 1, chunk.strip(), chunk_title.strip()))
            else:
                raise ValueError(
                    f"Header level {level} is greater than {self.header_level}"
                )

        return fragments


class MarkdownLoader:
    def __init__(self, headers_level: int = 1, min_length: int = 256):
        self.headers_level = headers_level
        self.min_length = min_length

    def load(self, file_path: str | Path) -> List[Document]:
        self.file_path = file_path if isinstance(file_path, Path) else Path(file_path)
        if not self.file_path.exists():
            raise FileNotFoundError(f"File '{self.file_path}' not found.")
        text = self.file_path.read_text(encoding="utf-8")
        splitter = MarkdownSplitter(header_level=self.headers_level)
        fragments = splitter.recursive_split(text=text)
        documents = []
        for uid, (content, title) in enumerate(fragments):
            if len(content) < self.min_length:
                continue
            documents.append(
                Document(
                    id=uuid4().int % 2**32,
                    content=content,
                    metadata={
                        "title": title,
                        "source": str(self.file_path),
                        "tags": "",
                    },
                )
            )
        return documents


# Пример использования класса при запуске скрипта напрямую
if __name__ == "__main__":
    loader = MarkdownLoader(headers_level=2)
    file_path = Path("data/tprg-lections-full.md")
    text = file_path.read_text(encoding="utf-8")

    if file_path.exists():
        documents = loader.load(file_path)
    else:
        raise FileNotFoundError(f"File '{file_path}' not found.")
    print(len(documents))
    for doc in documents:
        print(doc.title)