rag.py

from pathlib import Path
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings

BASE_DIR = Path(__file__).resolve().parent
KB_DIR = BASE_DIR / "data" / "kb"
PERSIST_DIR = BASE_DIR / ".faiss"

EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

_embeddings = None
_vectorstore = None

def get_embeddings():
    global _embeddings
    if _embeddings is None:
        print("[RAG] Loading embedding model...")
        _embeddings = HuggingFaceEmbeddings(
            model_name=EMBEDDING_MODEL,
            model_kwargs={"device": "cpu"},
            encode_kwargs={"normalize_embeddings": True}
        )
        print("[RAG] Embedding model loaded.")
    return _embeddings

def build_default_vectorstore():
    global _vectorstore
    if _vectorstore is not None:
        return _vectorstore
    
    if not PERSIST_DIR.exists():
        print(f"[RAG] Vectorstore not found: {PERSIST_DIR}")
        return None
    
    try:
        print("[RAG] Loading FAISS database...")
        _vectorstore = FAISS.load_local(
            folder_path=str(PERSIST_DIR),
            embeddings=get_embeddings(),
            allow_dangerous_deserialization=True,
        )
        print("[RAG] Vectorstore loaded successfully.")
        return _vectorstore
    except Exception as e:
        print(f"[RAG] Failed to load vectorstore: {e}")
        return None

def get_retriever(k: int = 4):
    vectorstore = build_default_vectorstore()
    if vectorstore is None:
        return None
    return vectorstore.as_retriever(search_kwargs={"k": k})