1"""Load question answering chains."""23from collections.abc import Mapping4from typing import Any, Protocol56from langchain_core._api import deprecated7from langchain_core.callbacks import BaseCallbackManager, Callbacks8from langchain_core.language_models import BaseLanguageModel9from langchain_core.prompts import BasePromptTemplate1011from langchain_classic.chains import ReduceDocumentsChain12from langchain_classic.chains.combine_documents.base import BaseCombineDocumentsChain13from langchain_classic.chains.combine_documents.map_reduce import (14    MapReduceDocumentsChain,15)16from langchain_classic.chains.combine_documents.map_rerank import (17    MapRerankDocumentsChain,18)19from langchain_classic.chains.combine_documents.refine import RefineDocumentsChain20from langchain_classic.chains.combine_documents.stuff import StuffDocumentsChain21from langchain_classic.chains.llm import LLMChain22from langchain_classic.chains.question_answering import (23    map_reduce_prompt,24    refine_prompts,25    stuff_prompt,26)27from langchain_classic.chains.question_answering.map_rerank_prompt import (28    PROMPT as MAP_RERANK_PROMPT,29)303132class LoadingCallable(Protocol):33    """Interface for loading the combine documents chain."""3435    def __call__(36        self,37        llm: BaseLanguageModel,38        **kwargs: Any,39    ) -> BaseCombineDocumentsChain:40        """Callable to load the combine documents chain."""414243def _load_map_rerank_chain(44    llm: BaseLanguageModel,45    *,46    prompt: BasePromptTemplate = MAP_RERANK_PROMPT,47    verbose: bool = False,48    document_variable_name: str = "context",49    rank_key: str = "score",50    answer_key: str = "answer",51    callback_manager: BaseCallbackManager | None = None,52    callbacks: Callbacks = None,53    **kwargs: Any,54) -> MapRerankDocumentsChain:55    llm_chain = LLMChain(56        llm=llm,57        prompt=prompt,58        verbose=verbose,59        callback_manager=callback_manager,60        callbacks=callbacks,61    )62    return MapRerankDocumentsChain(63        llm_chain=llm_chain,64        rank_key=rank_key,65        answer_key=answer_key,66        document_variable_name=document_variable_name,67        verbose=verbose,68        callback_manager=callback_manager,69        **kwargs,70    )717273def _load_stuff_chain(74    llm: BaseLanguageModel,75    *,76    prompt: BasePromptTemplate | None = None,77    document_variable_name: str = "context",78    verbose: bool | None = None,79    callback_manager: BaseCallbackManager | None = None,80    callbacks: Callbacks = None,81    **kwargs: Any,82) -> StuffDocumentsChain:83    _prompt = prompt or stuff_prompt.PROMPT_SELECTOR.get_prompt(llm)84    llm_chain = LLMChain(85        llm=llm,86        prompt=_prompt,87        verbose=verbose,88        callback_manager=callback_manager,89        callbacks=callbacks,90    )91    # TODO: document prompt92    return StuffDocumentsChain(93        llm_chain=llm_chain,94        document_variable_name=document_variable_name,95        verbose=verbose,96        callback_manager=callback_manager,97        callbacks=callbacks,98        **kwargs,99    )100101102def _load_map_reduce_chain(103    llm: BaseLanguageModel,104    *,105    question_prompt: BasePromptTemplate | None = None,106    combine_prompt: BasePromptTemplate | None = None,107    combine_document_variable_name: str = "summaries",108    map_reduce_document_variable_name: str = "context",109    collapse_prompt: BasePromptTemplate | None = None,110    reduce_llm: BaseLanguageModel | None = None,111    collapse_llm: BaseLanguageModel | None = None,112    verbose: bool | None = None,113    callback_manager: BaseCallbackManager | None = None,114    callbacks: Callbacks = None,115    token_max: int = 3000,116    **kwargs: Any,117) -> MapReduceDocumentsChain:118    _question_prompt = (119        question_prompt or map_reduce_prompt.QUESTION_PROMPT_SELECTOR.get_prompt(llm)120    )121    _combine_prompt = (122        combine_prompt or map_reduce_prompt.COMBINE_PROMPT_SELECTOR.get_prompt(llm)123    )124    map_chain = LLMChain(125        llm=llm,126        prompt=_question_prompt,127        verbose=verbose,128        callback_manager=callback_manager,129        callbacks=callbacks,130    )131    _reduce_llm = reduce_llm or llm132    reduce_chain = LLMChain(133        llm=_reduce_llm,134        prompt=_combine_prompt,135        verbose=verbose,136        callback_manager=callback_manager,137        callbacks=callbacks,138    )139    # TODO: document prompt140    combine_documents_chain = StuffDocumentsChain(141        llm_chain=reduce_chain,142        document_variable_name=combine_document_variable_name,143        verbose=verbose,144        callback_manager=callback_manager,145        callbacks=callbacks,146    )147    if collapse_prompt is None:148        collapse_chain = None149        if collapse_llm is not None:150            msg = (151                "collapse_llm provided, but collapse_prompt was not: please "152                "provide one or stop providing collapse_llm."153            )154            raise ValueError(msg)155    else:156        _collapse_llm = collapse_llm or llm157        collapse_chain = StuffDocumentsChain(158            llm_chain=LLMChain(159                llm=_collapse_llm,160                prompt=collapse_prompt,161                verbose=verbose,162                callback_manager=callback_manager,163                callbacks=callbacks,164            ),165            document_variable_name=combine_document_variable_name,166            verbose=verbose,167            callback_manager=callback_manager,168        )169    reduce_documents_chain = ReduceDocumentsChain(170        combine_documents_chain=combine_documents_chain,171        collapse_documents_chain=collapse_chain,172        token_max=token_max,173        verbose=verbose,174    )175    return MapReduceDocumentsChain(176        llm_chain=map_chain,177        document_variable_name=map_reduce_document_variable_name,178        reduce_documents_chain=reduce_documents_chain,179        verbose=verbose,180        callback_manager=callback_manager,181        callbacks=callbacks,182        **kwargs,183    )184185186def _load_refine_chain(187    llm: BaseLanguageModel,188    *,189    question_prompt: BasePromptTemplate | None = None,190    refine_prompt: BasePromptTemplate | None = None,191    document_variable_name: str = "context_str",192    initial_response_name: str = "existing_answer",193    refine_llm: BaseLanguageModel | None = None,194    verbose: bool | None = None,195    callback_manager: BaseCallbackManager | None = None,196    callbacks: Callbacks = None,197    **kwargs: Any,198) -> RefineDocumentsChain:199    _question_prompt = (200        question_prompt or refine_prompts.QUESTION_PROMPT_SELECTOR.get_prompt(llm)201    )202    _refine_prompt = refine_prompt or refine_prompts.REFINE_PROMPT_SELECTOR.get_prompt(203        llm,204    )205    initial_chain = LLMChain(206        llm=llm,207        prompt=_question_prompt,208        verbose=verbose,209        callback_manager=callback_manager,210        callbacks=callbacks,211    )212    _refine_llm = refine_llm or llm213    refine_chain = LLMChain(214        llm=_refine_llm,215        prompt=_refine_prompt,216        verbose=verbose,217        callback_manager=callback_manager,218        callbacks=callbacks,219    )220    return RefineDocumentsChain(221        initial_llm_chain=initial_chain,222        refine_llm_chain=refine_chain,223        document_variable_name=document_variable_name,224        initial_response_name=initial_response_name,225        verbose=verbose,226        callback_manager=callback_manager,227        callbacks=callbacks,228        **kwargs,229    )230231232@deprecated(233    since="0.2.13",234    removal="2.0.0",235    alternative="langchain.agents.create_agent",236    addendum=(237        "Build new RAG flows with `create_agent` and a retrieval tool. See "238        "https://docs.langchain.com/oss/python/langchain/rag"239    ),240)241def load_qa_chain(242    llm: BaseLanguageModel,243    chain_type: str = "stuff",244    verbose: bool | None = None,  # noqa: FBT001245    callback_manager: BaseCallbackManager | None = None,246    **kwargs: Any,247) -> BaseCombineDocumentsChain:248    """Load question answering chain.249250    Args:251        llm: Language Model to use in the chain.252        chain_type: Type of document combining chain to use. Should be one of "stuff",253            "map_reduce", "map_rerank", and "refine".254        verbose: Whether chains should be run in verbose mode or not. Note that this255            applies to all chains that make up the final chain.256        callback_manager: Callback manager to use for the chain.257        **kwargs: Additional keyword arguments.258259    Returns:260        A chain to use for question answering.261    """262    loader_mapping: Mapping[str, LoadingCallable] = {263        "stuff": _load_stuff_chain,264        "map_reduce": _load_map_reduce_chain,265        "refine": _load_refine_chain,266        "map_rerank": _load_map_rerank_chain,267    }268    if chain_type not in loader_mapping:269        msg = (270            f"Got unsupported chain type: {chain_type}. "271            f"Should be one of {loader_mapping.keys()}"272        )273        raise ValueError(msg)274    return loader_mapping[chain_type](275        llm,276        verbose=verbose,277        callback_manager=callback_manager,278        **kwargs,279    )