ingest.py

"""
ingest.py — Build a FAISS vector index from a folder of PDFs.

Usage:
    python ingest.py [--pdf-dir pdfs] [--index-dir faiss_index] [--chunk-size 800] [--chunk-overlap 100]

Environment variables:
    EMBED_MODEL   — Embedding model (default: BAAI/bge-small-en-v1.5)
    PDF_DIR       — Folder containing PDF files (default: pdfs)
    INDEX_DIR     — Where to save the FAISS index (default: faiss_index)
"""

import argparse
import json
import os
from pathlib import Path

import fitz  # PyMuPDF
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from tqdm import tqdm

EMBED_MODEL = os.getenv("EMBED_MODEL", "BAAI/bge-small-en-v1.5")
PDF_DIR     = Path(os.getenv("PDF_DIR", "pdfs"))
INDEX_DIR   = Path(os.getenv("INDEX_DIR", "faiss_index"))
META_FILE   = Path("metadata.json")


def extract_text(pdf_path: Path) -> str:
    doc = fitz.open(str(pdf_path))
    return "\n".join(page.get_text() for page in doc)


def load_pdfs(pdf_dir: Path) -> list[Document]:
    pdfs = sorted(pdf_dir.glob("**/*.pdf"))
    if not pdfs:
        raise FileNotFoundError(f"No PDF files found in {pdf_dir}")
    print(f"Found {len(pdfs)} PDF(s) in {pdf_dir}")

    docs = []
    metadata_records = []
    for pdf_path in tqdm(pdfs, desc="Reading PDFs"):
        text = extract_text(pdf_path)
        if not text.strip():
            print(f"  [WARN] No text extracted from {pdf_path.name}, skipping")
            continue
        doc_id = pdf_path.stem
        docs.append(Document(
            page_content=text,
            metadata={"doc_id": doc_id, "filename": pdf_path.name, "source": str(pdf_path)},
        ))
        metadata_records.append({"doc_id": doc_id, "filename": pdf_path.name})

    with open(META_FILE, "w", encoding="utf-8") as f:
        json.dump(metadata_records, f, ensure_ascii=False, indent=2)
    print(f"Saved metadata for {len(metadata_records)} documents → {META_FILE}")
    return docs


def chunk_documents(docs: list[Document], chunk_size: int, chunk_overlap: int) -> list[Document]:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", ". ", " ", ""],
    )
    chunks = splitter.split_documents(docs)
    print(f"Split into {len(chunks)} chunks (chunk_size={chunk_size}, overlap={chunk_overlap})")
    return chunks


def build_index(chunks: list[Document], embeddings: HuggingFaceEmbeddings, index_dir: Path) -> None:
    print(f"Building FAISS index with {EMBED_MODEL}...")
    vectorstore = FAISS.from_documents(chunks, embeddings)
    index_dir.mkdir(parents=True, exist_ok=True)
    vectorstore.save_local(str(index_dir))
    print(f"FAISS index saved → {index_dir}/")


def main():
    parser = argparse.ArgumentParser(description="Ingest PDFs into a FAISS vector index")
    parser.add_argument("--pdf-dir",      type=Path, default=PDF_DIR)
    parser.add_argument("--index-dir",    type=Path, default=INDEX_DIR)
    parser.add_argument("--chunk-size",   type=int,  default=1500)
    parser.add_argument("--chunk-overlap",type=int,  default=200)
    args = parser.parse_args()

    print(f"Loading embedding model: {EMBED_MODEL}")
    embeddings = HuggingFaceEmbeddings(
        model_name=EMBED_MODEL,
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True},
    )

    docs   = load_pdfs(args.pdf_dir)
    chunks = chunk_documents(docs, args.chunk_size, args.chunk_overlap)
    build_index(chunks, embeddings, args.index_dir)
    print("\nDone! You can now run: python app.py")


if __name__ == "__main__":
    main()