rag/retriever.py

from __future__ import annotations

from functools import lru_cache
from langchain_community.retrievers import BM25Retriever
from langchain_classic.retrievers import EnsembleRetriever
from .config import settings
from .splitter import documents
from .vectorstore import get_vectorstore


@lru_cache(maxsize=1)
def get_retriever() -> EnsembleRetriever:
    bm25_retriever = BM25Retriever.from_documents(documents)
    bm25_retriever.k = settings.top_k

    dense_retriever = get_vectorstore().as_retriever(search_kwargs={"k": settings.top_k})

    return EnsembleRetriever(
        retrievers=[bm25_retriever, dense_retriever],
        weights=[0.5, 0.5],
    )