src/chains.py

from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from langchain_classic.chains.retrieval import create_retrieval_chain

from src.llm import get_llm
from src.vectorstore import get_vectorstore
from src.prompts import QA_PROMPT, DOCUMENT_PROMPT
from src.config import TOP_K

def build_chain():

    llm = get_llm()

    retriever = get_vectorstore().as_retriever(
        search_type="mmr",
        search_kwargs={"k": TOP_K}
    )

    combine = create_stuff_documents_chain(
        llm,
        QA_PROMPT,
        document_prompt=DOCUMENT_PROMPT
    )

    return create_retrieval_chain(
        retriever,
        combine
    )