Document Loaders: PDF, Word, HTML, Markdown

# core/document.py
from dataclasses import dataclass, field
from typing import Any


@dataclass
class Document:
    """
    Унифицированное представление загруженного документа.
    page_content — очищенный текст.
    metadata    — всё, что поможет при поиске и ответе: источник,
                  страница, заголовок, дата, автор.
    """
    page_content: str
    metadata: dict[str, Any] = field(default_factory=dict)

    def __repr__(self) -> str:
        preview = self.page_content[:80].replace("\n", " ")
        return f"Document(chars={len(self.page_content)}, meta={list(self.metadata)}, preview='{preview}…')"


# Базовый интерфейс загрузчика
from abc import ABC, abstractmethod
from pathlib import Path


class BaseLoader(ABC):
    """Все загрузчики реализуют этот интерфейс."""

    @abstractmethod
    def load(self, source: str | Path) -> list[Document]:
        """Загрузить документ и вернуть список Document (один per страница или весь файл)."""
        ...

    def load_and_split(self, source: str | Path, splitter) -> list[Document]:
        """Загрузить и сразу разбить на чанки."""
        docs = self.load(source)
        return [chunk for doc in docs for chunk in splitter.split(doc)]

import fitz  # PyMuPDF
from pathlib import Path
from core.document import Document


class PDFLoader:
    """
    Загружает PDF через PyMuPDF.
    Возвращает один Document на страницу — это удобно для
    сохранения номера страницы в metadata (нужно для цитат).
    """

    def __init__(self, per_page: bool = True):
        self.per_page = per_page

    def load(self, source: str | Path) -> list[Document]:
        source = Path(source)
        doc = fitz.open(source)
        documents = []

        # Метаданные из свойств файла
        file_meta = doc.metadata  # dict: title, author, creationDate, …

        if self.per_page:
            for page_num, page in enumerate(doc, start=1):
                text = page.get_text("text")          # простое извлечение
                text = self._clean(text)
                if not text.strip():
                    continue                           # пропускаем пустые страницы
                documents.append(Document(
                    page_content=text,
                    metadata={
                        "source":    str(source),
                        "filename":  source.name,
                        "page":      page_num,
                        "total_pages": len(doc),
                        "title":     file_meta.get("title", ""),
                        "author":    file_meta.get("author", ""),
                    }
                ))
        else:
            full_text = "\n\n".join(
                self._clean(page.get_text("text")) for page in doc
            )
            documents.append(Document(
                page_content=full_text,
                metadata={"source": str(source), **file_meta},
            ))

        doc.close()
        return documents

    @staticmethod
    def _clean(text: str) -> str:
        """Убираем мусор: лишние переносы, дефисы в конце строки (перенос слов)."""
        import re
        # Убираем перенос слов в конце строки (дефис + \n → слово слитно)
        text = re.sub(r"-\n(\w)", r"\1", text)
        # Несколько пробелов → один
        text = re.sub(r"[ \t]{2,}", " ", text)
        # Больше двух переносов строки → два
        text = re.sub(r"\n{3,}", "\n\n", text)
        return text.strip()

import fitz
from pathlib import Path


def extract_structured(pdf_path: str | Path) -> list[dict]:
    """
    Возвращает список блоков с типом (heading/text/footer).
    Эвристика: блок — заголовок, если размер шрифта >= median * 1.2.
    """
    import statistics

    doc = fitz.open(pdf_path)
    all_blocks = []

    for page_num, page in enumerate(doc, start=1):
        # get_text("dict") возвращает детальную структуру
        page_dict = page.get_text("dict")
        for block in page_dict["blocks"]:
            if block["type"] != 0:  # 0 = text, 1 = image
                continue
            for line in block["lines"]:
                for span in line["spans"]:
                    all_blocks.append({
                        "text":    span["text"].strip(),
                        "size":    span["size"],    # размер шрифта
                        "bold":    bool(span["flags"] & 2**4),
                        "page":    page_num,
                        "bbox":    span["bbox"],    # (x0, y0, x1, y1)
                    })

    # Определяем медианный размер шрифта для всего документа
    sizes = [b["size"] for b in all_blocks if b["text"]]
    median_size = statistics.median(sizes) if sizes else 12

    result = []
    for block in all_blocks:
        if not block["text"]:
            continue
        # Эвристика: заголовок = крупный шрифт или жирный + чуть крупнее обычного
        is_heading = (
            block["size"] >= median_size * 1.2 or
            (block["bold"] and block["size"] >= median_size * 1.05)
        )
        result.append({
            **block,
            "type": "heading" if is_heading else "text",
        })

    doc.close()
    return result

import pdfplumber
from pathlib import Path
from core.document import Document


class PDFPlumberLoader:
    """
    Загрузчик на pdfplumber.
    Таблицы преобразуются в Markdown-представление,
    чтобы LLM мог их понять в контексте RAG.
    """

    def load(self, source: str | Path) -> list[Document]:
        source = Path(source)
        documents = []

        with pdfplumber.open(source) as pdf:
            for page_num, page in enumerate(pdf.pages, start=1):
                parts = []

                # 1. Обычный текст (без таблиц)
                # Вычитаем области таблиц из страницы, чтобы текст не дублировался
                tables = page.find_tables()
                page_without_tables = page
                for table in tables:
                    page_without_tables = page_without_tables.outside_bbox(
                        table.bbox
                    )
                plain_text = page_without_tables.extract_text() or ""
                if plain_text.strip():
                    parts.append(plain_text.strip())

                # 2. Таблицы → Markdown
                for table in tables:
                    rows = table.extract()
                    if not rows:
                        continue
                    md_table = _rows_to_markdown(rows)
                    parts.append(md_table)

                if not parts:
                    continue

                documents.append(Document(
                    page_content="\n\n".join(parts),
                    metadata={
                        "source":   str(source),
                        "filename": source.name,
                        "page":     page_num,
                        "has_tables": len(tables) > 0,
                    }
                ))

        return documents


def _rows_to_markdown(rows: list[list]) -> str:
    """Преобразует список строк таблицы в Markdown-формат."""
    if not rows:
        return ""
    # Заменяем None на пустую строку
    clean = [[str(cell or "") for cell in row] for row in rows]
    header = "| " + " | ".join(clean[0]) + " |"
    sep    = "| " + " | ".join("---" for _ in clean[0]) + " |"
    body   = "\n".join("| " + " | ".join(row) + " |" for row in clean[1:])
    return "\n".join([header, sep, body])

from docx import Document as DocxDocument
from docx.oxml.ns import qn
from pathlib import Path
from core.document import Document
import re


class DocxLoader:
    """
    Загружает .docx.
    Заголовки (Heading 1/2/3) сохраняются с markdown-разметкой
    и в metadata как outline — структура документа.
    """

    HEADING_STYLES = {
        "Heading 1": "#",
        "Heading 2": "##",
        "Heading 3": "###",
        "Заголовок 1": "#",   # русские имена стилей
        "Заголовок 2": "##",
        "Заголовок 3": "###",
    }

    def load(self, source: str | Path) -> list[Document]:
        source = Path(source)
        docx = DocxDocument(source)

        parts: list[str] = []
        outline: list[dict] = []   # [{level, title}] — оглавление

        for para in docx.paragraphs:
            style_name = para.style.name
            text = para.text.strip()
            if not text:
                continue

            if style_name in self.HEADING_STYLES:
                prefix = self.HEADING_STYLES[style_name]
                parts.append(f"{prefix} {text}")
                outline.append({"level": prefix.count("#"), "title": text})
            else:
                parts.append(text)

        # Таблицы
        for table in docx.tables:
            md_rows = []
            for i, row in enumerate(table.rows):
                cells = [cell.text.strip() for cell in row.cells]
                md_rows.append("| " + " | ".join(cells) + " |")
                if i == 0:
                    md_rows.append("| " + " | ".join("---" for _ in cells) + " |")
            if md_rows:
                parts.append("\n".join(md_rows))

        # Метаданные из core properties
        props = docx.core_properties
        return [Document(
            page_content="\n\n".join(parts),
            metadata={
                "source":   str(source),
                "filename": source.name,
                "title":    props.title or source.stem,
                "author":   props.author or "",
                "created":  str(props.created or ""),
                "outline":  outline,          # оглавление — полезно для chunk-заголовков
            }
        )]

from bs4 import BeautifulSoup, Tag
from pathlib import Path
from core.document import Document
import re


# Теги, которые никогда не содержат полезный контент
NOISE_TAGS = {
    "script", "style", "nav", "header", "footer",
    "aside", "form", "button", "noscript", "iframe",
    "svg", "meta", "link", "head",
}

# Атрибуты class/id, указывающие на boilerplate
NOISE_PATTERNS = re.compile(
    r"(nav|menu|footer|header|sidebar|cookie|banner|ad-|advertisement|"
    r"breadcrumb|social|share|comment|related|popup|modal)",
    re.IGNORECASE,
)


class HTMLLoader:
    """
    Загружает HTML-файл (или строку), убирает boilerplate,
    сохраняет заголовки с markdown-разметкой.
    """

    def load(self, source: str | Path) -> list[Document]:
        source = Path(source)
        html = source.read_text(encoding="utf-8", errors="replace")
        return [self._parse(html, str(source))]

    def load_html(self, html: str, url: str = "") -> Document:
        """Загрузка из строки — удобно для web-скрапинга."""
        return self._parse(html, url)

    def _parse(self, html: str, source: str) -> Document:
        soup = BeautifulSoup(html, "lxml")

        # Метаданные из 
        title = ""
        if soup.title:
            title = soup.title.get_text(strip=True)
        description = ""
        meta_desc = soup.find("meta", attrs={"name": "description"})
        if meta_desc:
            description = meta_desc.get("content", "")

        # Удаляем шумовые теги целиком
        for tag in soup.find_all(NOISE_TAGS):
            tag.decompose()

        # Удаляем элементы с шумовыми class/id
        for tag in soup.find_all(True):
            classes = " ".join(tag.get("class", []))
            tag_id  = tag.get("id", "")
            if NOISE_PATTERNS.search(classes) or NOISE_PATTERNS.search(tag_id):
                tag.decompose()

        # Пытаемся найти основной контент
        main = (
            soup.find("article") or
            soup.find("main") or
            soup.find(id=re.compile(r"content|article|main", re.I)) or
            soup.find(class_=re.compile(r"content|article|post-body", re.I)) or
            soup.body
        )
        if not main:
            main = soup

        text = self._to_text(main)
        text = _normalize_whitespace(text)

        return Document(
            page_content=text,
            metadata={
                "source":      source,
                "title":       title,
                "description": description,
            }
        )

    def _to_text(self, element: Tag) -> str:
        """Рекурсивно конвертируем HTML в текст с markdown-заголовками."""
        parts = []
        for child in element.children:
            if hasattr(child, "name"):
                if child.name in ("h1",):
                    parts.append(f"# {child.get_text(strip=True)}")
                elif child.name in ("h2",):
                    parts.append(f"## {child.get_text(strip=True)}")
                elif child.name in ("h3", "h4"):
                    parts.append(f"### {child.get_text(strip=True)}")
                elif child.name in ("p", "div", "section"):
                    inner = self._to_text(child).strip()
                    if inner:
                        parts.append(inner)
                elif child.name in ("li",):
                    parts.append(f"- {child.get_text(strip=True)}")
                elif child.name in ("code", "pre"):
                    parts.append(f"`{child.get_text()}`")
                else:
                    parts.append(child.get_text(" ", strip=True))
            else:
                text = str(child).strip()
                if text:
                    parts.append(text)
        return "\n\n".join(p for p in parts if p.strip())


def _normalize_whitespace(text: str) -> str:
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()

import frontmatter   # pip install python-frontmatter
import re
from pathlib import Path
from core.document import Document


class MarkdownLoader:
    """
    Загружает Markdown с поддержкой YAML frontmatter.
    Заголовки сохраняются в metadata как outline.
    """

    def load(self, source: str | Path) -> list[Document]:
        source = Path(source)
        raw = source.read_text(encoding="utf-8")

        # python-frontmatter разделяет метаданные и тело
        post = frontmatter.loads(raw)
        meta = dict(post.metadata)   # title, date, tags, author, …
        body = post.content

        # Строим outline из заголовков
        outline = []
        for match in re.finditer(r"^(#{1,3})\s+(.+)$", body, re.MULTILINE):
            outline.append({
                "level": len(match.group(1)),
                "title": match.group(2).strip(),
            })

        return [Document(
            page_content=body,
            metadata={
                "source":   str(source),
                "filename": source.name,
                "title":    meta.get("title", source.stem),
                "date":     str(meta.get("date", "")),
                "tags":     meta.get("tags", []),
                "author":   meta.get("author", ""),
                "outline":  outline,
                **{k: v for k, v in meta.items()
                   if k not in ("title", "date", "tags", "author")},
            }
        )]

from pathlib import Path
from core.document import Document
from loaders.pdf_loader import PDFLoader
from loaders.pdf_plumber_loader import PDFPlumberLoader
from loaders.docx_loader import DocxLoader
from loaders.html_loader import HTMLLoader
from loaders.markdown_loader import MarkdownLoader


class DocumentDispatcher:
    """
    Определяет загрузчик по расширению файла.
    Легко расширяется: добавь новый класс и регистрацию.
    """

    def __init__(self, use_plumber_for_tables: bool = False):
        pdf_cls = PDFPlumberLoader if use_plumber_for_tables else PDFLoader
        self._loaders: dict[str, object] = {
            ".pdf":      pdf_cls(),
            ".docx":     DocxLoader(),
            ".doc":      DocxLoader(),    # python-docx читает оба формата
            ".html":     HTMLLoader(),
            ".htm":      HTMLLoader(),
            ".md":       MarkdownLoader(),
            ".markdown": MarkdownLoader(),
            ".txt":      PlainTextLoader(),
        }

    def load(self, source: str | Path) -> list[Document]:
        source = Path(source)
        ext = source.suffix.lower()
        loader = self._loaders.get(ext)
        if loader is None:
            raise ValueError(
                f"No loader for extension '{ext}'. "
                f"Supported: {list(self._loaders)}"
            )
        return loader.load(source)

    def load_directory(
        self,
        directory: str | Path,
        glob: str = "**/*",
        recursive: bool = True,
    ) -> list[Document]:
        """Загружает все поддерживаемые файлы из директории."""
        directory = Path(directory)
        pattern = glob if recursive else "*"
        documents = []
        supported = set(self._loaders)

        for path in sorted(directory.glob(pattern)):
            if path.is_file() and path.suffix.lower() in supported:
                try:
                    docs = self.load(path)
                    documents.extend(docs)
                    print(f"  ✓ {path.name}: {len(docs)} document(s)")
                except Exception as e:
                    print(f"  ✗ {path.name}: {e}")

        return documents


class PlainTextLoader:
    def load(self, source: Path) -> list[Document]:
        import chardet
        raw = source.read_bytes()
        encoding = chardet.detect(raw)["encoding"] or "utf-8"
        text = raw.decode(encoding, errors="replace")
        return [Document(
            page_content=text,
            metadata={"source": str(source), "filename": source.name}
        )]


# Использование
if __name__ == "__main__":
    dispatcher = DocumentDispatcher(use_plumber_for_tables=True)

    # Один файл
    docs = dispatcher.load("report.pdf")
    for d in docs:
        print(d)

    # Вся папка
    all_docs = dispatcher.load_directory("./knowledge_base")
    print(f"\nTotal: {len(all_docs)} documents loaded")

import os
import hashlib
from datetime import datetime
from pathlib import Path
from core.document import Document


def enrich_metadata(doc: Document, path: Path) -> Document:
    """
    Добавляет в метаданные:
    - file_size: размер файла в байтах
    - modified_at: дата последней правки
    - content_hash: SHA256 от текста (для дедупликации)
    - word_count: количество слов
    - language: определяем по первым 200 символам
    """
    stat = path.stat()
    text = doc.page_content

    # Хэш — для дедупликации перед индексированием
    content_hash = hashlib.sha256(text.encode()).hexdigest()[:16]

    # Количество слов (грубо, но полезно для фильтрации)
    word_count = len(text.split())

    # Простая эвристика языка по частотным словам
    sample = text[:200].lower()
    if any(w in sample for w in ("the ", "and ", "for ", "this ", "with ")):
        lang = "en"
    elif any(w in sample for w in ("и ", "в ", "на ", "не ", "что ")):
        lang = "ru"
    else:
        lang = "unknown"

    doc.metadata.update({
        "file_size":    stat.st_size,
        "modified_at":  datetime.fromtimestamp(stat.st_mtime).isoformat(),
        "content_hash": content_hash,
        "word_count":   word_count,
        "language":     lang,
    })
    return doc


# Пайплайн с обогащением
def load_with_enrichment(
    source: str | Path,
    dispatcher: "DocumentDispatcher",
) -> list[Document]:
    source = Path(source)
    docs = dispatcher.load(source)
    return [enrich_metadata(doc, source) for doc in docs]