rag.py

# app/champ/rag.py
import torch

from pathlib import Path

from langchain_community.vectorstores import FAISS as LCFAISS
from langchain_huggingface import HuggingFaceEmbeddings

from constants import BASE_DIR, HF_TOKEN


def load_vector_store(
    base_dir: Path = BASE_DIR,
    hf_token: str = HF_TOKEN,
    rag_relpath: str = "rag_data/FAISS_ALLEN_20260129",
    embedding_model: str = "BAAI/bge-large-en-v1.5",
    device: str = "cuda" if torch.cuda.is_available() else "cpu",
) -> LCFAISS:
    rag_path = base_dir / rag_relpath

    model_embedding_kwargs = {"device": device, "use_auth_token": hf_token}
    encode_kwargs = {"normalize_embeddings": True}

    embeddings = HuggingFaceEmbeddings(
        model_name=embedding_model,
        model_kwargs=model_embedding_kwargs,
        encode_kwargs=encode_kwargs,
    )

    return LCFAISS.load_local(
        str(rag_path),
        embeddings,
        allow_dangerous_deserialization=True,  # safe because you built the files
    )