from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
from langchain.chains import MapReduceDocumentsChain, ReduceDocumentsChain, LLMChain, StuffDocumentsChain
from langchain.prompts import PromptTemplate

def get_map_reduce_chain(pipeline_or_llm,model_type)-> LLMChain:

    if model_type == "openai":
        llm = pipeline_or_llm
        map_template  = """The following is a set of documents
            {docs}
            Based on this list of docs, please identify the main themes.
            Helpful Answer:"""
        map_prompt = PromptTemplate.from_template(map_template)
        reduce_template = """The following is set of summaries:
        {docs}
        Take these and distill into a final, consolidated summary of the main themes.
        Helpful Answer:"""
        reduce_prompt = PromptTemplate.from_template(reduce_template)

    else:
        map_prompt = PromptTemplate.from_template(template="{docs}")
        reduce_prompt = PromptTemplate.from_template(template="{docs}")
        llm = HuggingFacePipeline(pipeline=pipeline_or_llm)


    map_chain = LLMChain(llm = llm, prompt=map_prompt)
    reduce_chain = LLMChain(llm = llm, prompt = reduce_prompt,verbose = True)
    combine_documents_chain = StuffDocumentsChain(llm_chain=reduce_chain, document_variable_name="docs")
    reduce_documents_chain = ReduceDocumentsChain(
        combine_documents_chain=combine_documents_chain,
        collapse_documents_chain=combine_documents_chain,
        token_max=16384,
        verbose = True,
    )
    map_reduce_chain = MapReduceDocumentsChain(
        llm_chain=map_chain,
        reduce_documents_chain=reduce_documents_chain,
        document_variable_name="docs",
        return_intermediate_steps=False,
        verbose = True,
    )
        
    return map_reduce_chain