02_preprocessing.py

"""
02_preprocessing.py
--------------------
Stage 2 of the RAG pipeline: PREPROCESSING / CLEANING.

Takes the raw documents produced by 01_documents.py and cleans the text so
downstream chunking and embedding steps work on consistent, noise-free input.

Cleaning steps applied:
    - Normalize line endings
    - Collapse repeated whitespace/newlines
    - Strip non-printable / control characters
    - Trim leading/trailing whitespace

Run directly to see a demo of this stage:
    python 02_preprocessing.py
"""

from __future__ import annotations

import re
import importlib.util
import sys
import os
from dataclasses import dataclass, field
from typing import List

_THIS_DIR = os.path.dirname(os.path.abspath(__file__))


def _import_module_from_file(module_name: str, filename: str):
    """Import a numbered pipeline file (e.g. '01_documents.py') as a module.

    Regular `import` statements cannot start with a digit, so numbered
    pipeline stage files are loaded dynamically using importlib.
    """
    path = os.path.join(_THIS_DIR, filename)
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module  # required so dataclasses can resolve cls.__module__
    spec.loader.exec_module(module)
    return module


_documents_module = _import_module_from_file("doc_stage", "01_documents.py")
RawDocument = _documents_module.RawDocument


@dataclass
class CleanDocument:
    """A document after preprocessing/cleaning."""

    doc_id: str
    source: str
    text: str
    metadata: dict = field(default_factory=dict)


_CONTROL_CHARS_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
_MULTI_WHITESPACE_RE = re.compile(r"[ \t]+")
_MULTI_NEWLINE_RE = re.compile(r"\n{3,}")


def clean_text(text: str) -> str:
    """Apply all cleaning steps to a single piece of text."""
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    text = _CONTROL_CHARS_RE.sub("", text)
    text = _MULTI_WHITESPACE_RE.sub(" ", text)
    text = _MULTI_NEWLINE_RE.sub("\n\n", text)
    text = "\n".join(line.strip() for line in text.split("\n"))
    return text.strip()


def preprocess_documents(raw_documents: List[RawDocument]) -> List[CleanDocument]:
    """Clean a list of RawDocument objects into CleanDocument objects."""
    cleaned: List[CleanDocument] = []

    for doc in raw_documents:
        cleaned_text = clean_text(doc.text)
        if not cleaned_text:
            continue
        cleaned.append(
            CleanDocument(
                doc_id=doc.doc_id,
                source=doc.source,
                text=cleaned_text,
                metadata=doc.metadata,
            )
        )

    return cleaned


if __name__ == "__main__":
    raw_docs = _documents_module.load_documents()
    cleaned_docs = preprocess_documents(raw_docs)
    print(f"Preprocessed {len(cleaned_docs)} document(s):")
    for doc in cleaned_docs:
        preview = doc.text[:120].replace("\n", " ")
        print(f"  - {doc.doc_id} ({len(doc.text)} chars after cleaning): {preview}...")