src/ingest.py

"""Step 3 — Ingestion (the "retrieval" half of RAG).

Pipeline:
    1. Load every PDF in data/
    2. Split the text into overlapping chunks
    3. Embed each chunk into a vector (sentence-transformers)
    4. Store the vectors in a persistent ChromaDB database

Run once after adding/changing documents:
    python -m src.ingest
"""
from __future__ import annotations

import shutil

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma

from src.config import (
    DATA_DIR, CHROMA_DIR, CHUNK_SIZE, CHUNK_OVERLAP,
    EMBED_MODEL, COLLECTION_NAME,
)


def load_documents() -> list:
    """Load all PDFs from data/ into LangChain Document objects (one per page)."""
    pdf_paths = sorted(DATA_DIR.glob("*.pdf"))
    if not pdf_paths:
        raise FileNotFoundError(
            f"No PDFs found in {DATA_DIR}. "
            f"Run `python -m src.download_sources` first, or add your own PDFs."
        )
    docs = []
    for path in pdf_paths:
        print(f"  loading {path.name} ...")
        pages = PyPDFLoader(str(path)).load()
        # Keep a clean source label for citations later.
        for page in pages:
            page.metadata["source"] = path.name
        docs.extend(pages)
    print(f"Loaded {len(docs)} pages from {len(pdf_paths)} PDF(s).")
    return docs


def split_documents(docs: list) -> list:
    """Split pages into overlapping chunks for fine-grained retrieval."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        # Try to split on natural boundaries first (paragraphs, then lines).
        separators=["\n\n", "\n", ". ", " ", ""],
    )
    chunks = splitter.split_documents(docs)
    print(f"Split into {len(chunks)} chunks "
          f"(size={CHUNK_SIZE}, overlap={CHUNK_OVERLAP} chars).")
    return chunks


def build_vector_store(chunks: list, reset: bool = True) -> Chroma:
    """Embed chunks and persist them in ChromaDB."""
    if reset and CHROMA_DIR.exists():
        print("Removing previous vector store ...")
        shutil.rmtree(CHROMA_DIR)

    print(f"Loading embedding model '{EMBED_MODEL}' (first run downloads it) ...")
    embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL)

    print("Embedding chunks and writing to ChromaDB ...")
    store = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        collection_name=COLLECTION_NAME,
        persist_directory=str(CHROMA_DIR),
    )
    print(f"Vector store ready at {CHROMA_DIR}")
    return store


def main() -> None:
    print("=== Ingestion ===")
    docs = load_documents()
    chunks = split_documents(docs)
    build_vector_store(chunks)
    print("Done. You can now run the chatbot.")


if __name__ == "__main__":
    main()