rag/vector_store.py

"""
Vector store builder and retriever for UTM RAG pipeline.
Uses FAISS locally with sentence-transformers embeddings.
No GPU required.
"""

import os
from pathlib import Path
from typing import List, Tuple

from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings


# Multilingual embedding model — supports English + Chinese
EMBEDDING_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
VECTOR_STORE_PATH = "data/processed/faiss_index"


def build_vector_store(
    documents: List[Document],
    chunk_size: int = None,
    chunk_overlap: int = None,
    save_path: str = VECTOR_STORE_PATH,
) -> FAISS:
    """
    Split documents and build a FAISS vector store.

    Chunk strategy tuned for lecture slides + regulatory PDFs:
    - chunk_size=800: large enough to capture a full slide concept
    - chunk_overlap=100: preserve context across slide boundaries
    - separators prioritise paragraph/slide breaks before sentence breaks
    """
    chunk_size = chunk_size or int(os.getenv("CHUNK_SIZE", 800))
    chunk_overlap = chunk_overlap or int(os.getenv("CHUNK_OVERLAP", 100))

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", ". ", "! ", "? ", ", ", " ", ""],
    )
    chunks = splitter.split_documents(documents)
    # Filter out empty or whitespace-only chunks
    chunks = [c for c in chunks if c.page_content and c.page_content.strip()]
    print(f"Split into {len(chunks)} chunks (size={chunk_size}, overlap={chunk_overlap})")

    print(f"Loading embedding model: {EMBEDDING_MODEL}")
    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)

    print("Building FAISS index...")
    vectorstore = FAISS.from_documents(chunks, embeddings)

    Path(save_path).parent.mkdir(parents=True, exist_ok=True)
    vectorstore.save_local(save_path)
    print(f"Vector store saved to: {save_path}")

    return vectorstore


def load_vector_store(load_path: str = VECTOR_STORE_PATH) -> FAISS:
    """Load a previously saved FAISS vector store."""
    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
    vectorstore = FAISS.load_local(
        load_path, embeddings, allow_dangerous_deserialization=True
    )
    print(f"Vector store loaded from: {load_path}")
    return vectorstore


def retrieve_context(
    query: str,
    vectorstore: FAISS,
    top_k: int = None,
) -> Tuple[str, List[Document]]:
    """
    Retrieve the most relevant document chunks for a query.

    Returns:
        Tuple of (combined context string, list of source Documents).
    """
    top_k = top_k or int(os.getenv("TOP_K_RETRIEVAL", 6))
    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
    docs = retriever.invoke(query)

    context = "\n\n---\n\n".join(
        f"[Source: {d.metadata.get('source_file', 'unknown')}, "
        f"Page: {d.metadata.get('page', 'N/A')}]\n{d.page_content}"
        for d in docs
    )
    return context, docs