document-management/backend/dms/services/document_content_service.py

"""不依赖AI或Office进程的PDF/DOCX正文提取与检索文本构建。"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

import fitz
from docx import Document as WordDocument

from dms.common.enums import ContentExtractStatus

logger = logging.getLogger(__name__)


@dataclass(frozen=True, slots=True)
class ContentExtraction:
    text: str | None
    status: str
    extracted_at: datetime | None


def build_search_text(
    name: str,
    summary: str | None,
    tags: list[str] | None,
    content_text: str | None,
) -> str:
    return " ".join(
        part for part in [name, summary or "", *(tags or []), content_text or ""]
        if part
    ).strip()


def _word_text(path: Path) -> str:
    document = WordDocument(path)
    values: list[str] = []

    def add_paragraphs(paragraphs) -> None:
        values.extend(
            paragraph.text.strip()
            for paragraph in paragraphs
            if paragraph.text.strip()
        )

    def add_tables(tables) -> None:
        for table in tables:
            for row in table.rows:
                for cell in row.cells:
                    add_paragraphs(cell.paragraphs)
                    add_tables(cell.tables)

    add_paragraphs(document.paragraphs)
    add_tables(document.tables)
    for section in document.sections:
        add_paragraphs(section.header.paragraphs)
        add_tables(section.header.tables)
        add_paragraphs(section.footer.paragraphs)
        add_tables(section.footer.tables)
    return "\n".join(values).strip()


def _pdf_text(path: Path) -> str:
    with fitz.open(stream=path.read_bytes(), filetype="pdf") as document:
        return "\n".join(page.get_text("text") for page in document).strip()


def extract_document_content(path: Path, extension: str) -> ContentExtraction:
    normalized = extension.lower().lstrip(".")
    if normalized not in {"pdf", "docx"}:
        return ContentExtraction(
            text=None,
            status=ContentExtractStatus.UNSUPPORTED.value,
            extracted_at=None,
        )
    extracted_at = datetime.now(timezone.utc).replace(tzinfo=None)
    try:
        text = _pdf_text(path) if normalized == "pdf" else _word_text(path)
    except Exception:
        logger.exception("文档正文提取失败：extension=%s", normalized)
        return ContentExtraction(
            text=None,
            status=ContentExtractStatus.FAILED.value,
            extracted_at=extracted_at,
        )
    return ContentExtraction(
        text=text or None,
        status=(
            ContentExtractStatus.SUCCESS.value
            if text
            else ContentExtractStatus.EMPTY.value
        ),
        extracted_at=extracted_at,
    )


__all__ = ["ContentExtraction", "build_search_text", "extract_document_content"]