data_loader.py

"""
Data loading module: reads .txt/.md/.pdf files and splits them into
paragraph-level LangChain Documents with metadata (source, chunk_id,
section, and page for PDFs).

Requirements:
    pip install langchain-core pypdf
"""

import os
from typing import Any, Dict, List, Optional

from langchain_core.documents import Document

# PDF text extraction
from pypdf import PdfReader

SUPPORTED_EXTENSIONS = (".txt", ".md", ".pdf")

# Set the path to the file you want to load here — change this to
# point at your .txt, .md, or .pdf file. This is the single place
# the file path is configured; hybrid_rag_pipeline.py doesn't set it.
FILE_PATH = "test.txt"


# ---------------------------------------------------
# Shared paragraph -> Document chunking logic
# ---------------------------------------------------

def _paragraphs_to_documents(
    paragraphs: List[str],
    source_name: str,
    start_chunk_id: int = 0,
    extra_metadata: Optional[Dict[str, Any]] = None,
    current_section: str = "unknown",
) -> List[Document]:
    """
    Convert a list of paragraph strings into Documents, carrying forward
    a naive "section" heading guess across paragraphs.
    """
    docs = []
    for offset, para in enumerate(paragraphs):
        first_line = para.splitlines()[0].strip()
        if len(first_line) < 60 and not first_line.endswith((".", "?", "!")):
            current_section = first_line

        metadata = {
            "source": source_name,
            "chunk_id": start_chunk_id + offset,
            "section": current_section,
        }
        if extra_metadata:
            metadata.update(extra_metadata)

        docs.append(Document(page_content=para, metadata=metadata))

    return docs


# ---------------------------------------------------
# Load .txt Documents (with metadata)
# ---------------------------------------------------

def load_text_documents(file_path: str) -> List[Document]:
    """
    Load a .txt/.md file and split it into paragraph-level Documents.

    Each Document gets metadata you can later filter on:
      - source:      the file it came from
      - chunk_id:    its position in the file
      - section:     a naive heading guess (first line-like token),
                      useful as an example metadata filter field
    """
    with open(file_path, "r", encoding="utf-8") as f:
        content = f.read()

    paragraphs = [p.strip() for p in content.split("\n\n") if p.strip()]

    return _paragraphs_to_documents(
        paragraphs,
        source_name=os.path.basename(file_path),
    )


# ---------------------------------------------------
# Load .pdf Documents (with metadata)
# ---------------------------------------------------

def load_pdf_documents(file_path: str) -> List[Document]:
    """
    Load a .pdf file, extract text page by page, and split each page
    into paragraph-level Documents.

    Each Document gets the same metadata fields as load_text_documents,
    plus:
      - page: the 1-indexed PDF page number the chunk came from

    Note: extraction quality depends on the PDF — scanned/image-only
    PDFs will yield little or no text (they'd need OCR first).
    """
    reader = PdfReader(file_path)
    source_name = os.path.basename(file_path)

    all_docs: List[Document] = []
    current_section = "unknown"
    chunk_id = 0

    for page_num, page in enumerate(reader.pages, start=1):
        text = page.extract_text() or ""
        paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]

        if not paragraphs:
            continue

        page_docs = _paragraphs_to_documents(
            paragraphs,
            source_name=source_name,
            start_chunk_id=chunk_id,
            extra_metadata={"page": page_num},
            current_section=current_section,
        )

        # carry section guess forward into the next page
        if page_docs:
            current_section = page_docs[-1].metadata["section"]

        chunk_id += len(page_docs)
        all_docs.extend(page_docs)

    return all_docs


# ---------------------------------------------------
# Generic dispatcher: pick the right loader by extension
# ---------------------------------------------------

def load_documents(file_path: Optional[str] = None) -> List[Document]:
    """
    Load a document file into paragraph-level Documents, dispatching to
    the right loader based on file extension.

    If file_path is omitted, falls back to the FILE_PATH constant
    defined above — so callers (like hybrid_rag_pipeline.py) don't
    need to know or set a path themselves.

    Supported: .txt, .md, .pdf
    """
    path = file_path or FILE_PATH
    ext = os.path.splitext(path)[1].lower()

    if ext == ".pdf":
        return load_pdf_documents(path)
    elif ext in (".txt", ".md"):
        return load_text_documents(path)
    else:
        raise ValueError(
            f"Unsupported file type: '{ext}'. Supported extensions: "
            f"{', '.join(SUPPORTED_EXTENSIONS)}"
        )


if __name__ == "__main__":
    # Quick standalone check: run `python data_loader.py` to verify a
    # file loads and chunks as expected before wiring it into the
    # full pipeline.
    docs = load_documents()
    print(f"Loaded {len(docs)} documents from '{FILE_PATH}'.")
    for doc in docs[:3]:
        print("\n---")
        print(doc.metadata)
        print(doc.page_content[:200])