S SmartDocs
系列: Ricky python 63 行 · 更新於 2026-04-07

document_processor.py

Ricky/RAG_project/processors/document_processor.py

"""
Document chunking: clean text and split into overlapping chunks.

Uses RecursiveCharacterTextSplitter as the baseline.  Summary documents
(type='summary') are kept as single chunks to avoid fragmenting curated text.
"""
import logging

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

from config.settings import CHUNK_SIZE, CHUNK_OVERLAP
from processors.source_metadata import ensure_source_metadata

logger = logging.getLogger(__name__)


def clean_text(text: str) -> str:
    """Lightweight text cleaning: remove null bytes and collapse whitespace."""
    if not text:
        return ""
    return " ".join(text.replace("\x00", " ").split())


def process_documents(documents: list[Document]) -> list[Document]:
    """
    Clean -> split -> metadata-tag all documents.

    Summary documents (metadata type='summary') are treated as atomic chunks
    so their curated structure is never fragmented across chunk boundaries.
    """
    summaries: list[Document] = []
    regular: list[Document] = []

    for doc in documents:
        content = clean_text(doc.page_content)
        if not content:
            continue
        doc.page_content = content
        doc.metadata = ensure_source_metadata(doc.metadata)
        if doc.metadata.get("type") == "summary":
            summaries.append(doc)
        else:
            regular.append(doc)

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
    )
    regular_chunks = splitter.split_documents(regular)

    all_chunks = regular_chunks + summaries
    logger.info(
        "Document processor: %d regular + %d summary = %d total chunks",
        len(regular_chunks),
        len(summaries),
        len(all_chunks),
    )
    print(
        f"Document processor complete: {len(regular)} regular + {len(summaries)} summary "
        f"-> {len(all_chunks)} chunks"
    )
    return all_chunks

相關文章