rag/document_loader.py

"""
Document loader for UTM/U-space reference materials.
Supports PDF, DOCX, and plain text files from data/raw/.
"""

import os
from pathlib import Path
from typing import List

from langchain_core.documents import Document
from langchain_community.document_loaders import (
    PyPDFLoader,
    Docx2txtLoader,
    TextLoader,
    DirectoryLoader,
)


SUPPORTED_EXTENSIONS = {
    ".pdf": PyPDFLoader,
    ".docx": Docx2txtLoader,
    ".txt": TextLoader,
    ".md": TextLoader,
}


def load_documents(data_dir: str = "data/raw") -> List[Document]:
    """
    Load all supported documents from data_dir.

    Args:
        data_dir: Path to folder containing UTM reference documents.

    Returns:
        List of LangChain Document objects.
    """
    data_path = Path(data_dir)
    if not data_path.exists():
        raise FileNotFoundError(f"Data directory not found: {data_dir}")

    documents = []
    loaded_files = []

    for file_path in sorted(data_path.rglob("*")):
        ext = file_path.suffix.lower()
        loader_cls = SUPPORTED_EXTENSIONS.get(ext)
        if loader_cls is None:
            continue

        try:
            loader = loader_cls(str(file_path))
            docs = loader.load()
            # Tag each doc with its source file
            for doc in docs:
                doc.metadata["source_file"] = file_path.name
            documents.extend(docs)
            loaded_files.append(file_path.name)
        except Exception as e:
            print(f"Warning: Could not load {file_path.name}: {e}")

    print(f"Loaded {len(documents)} document chunks from {len(loaded_files)} files:")
    for f in loaded_files:
        print(f"  - {f}")

    return documents