mirror of
https://github.com/Mintplex-Labs/langchain-python.git
synced 2026-07-24 20:15:34 -04:00
Ankush Gola
b82cbd1be0
`combine_docs` does not go through the standard chain call path which means that chain callbacks won't be triggered, meaning QA chains won't be traced properly, this fixes that. Also fix several errors in the chat_vector_db notebook
75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
"""Map-reduce chain.
|
|
|
|
Splits up a document, sends the smaller parts to the LLM with one prompt,
|
|
then combines the results with another one.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Dict, List
|
|
|
|
from pydantic import Extra
|
|
|
|
from langchain.chains.base import Chain
|
|
from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
|
|
from langchain.chains.combine_documents.map_reduce import MapReduceDocumentsChain
|
|
from langchain.chains.combine_documents.stuff import StuffDocumentsChain
|
|
from langchain.chains.llm import LLMChain
|
|
from langchain.docstore.document import Document
|
|
from langchain.llms.base import BaseLLM
|
|
from langchain.prompts.base import BasePromptTemplate
|
|
from langchain.text_splitter import TextSplitter
|
|
|
|
|
|
class MapReduceChain(Chain):
|
|
"""Map-reduce chain."""
|
|
|
|
combine_documents_chain: BaseCombineDocumentsChain
|
|
"""Chain to use to combine documents."""
|
|
text_splitter: TextSplitter
|
|
"""Text splitter to use."""
|
|
input_key: str = "input_text" #: :meta private:
|
|
output_key: str = "output_text" #: :meta private:
|
|
|
|
@classmethod
|
|
def from_params(
|
|
cls, llm: BaseLLM, prompt: BasePromptTemplate, text_splitter: TextSplitter
|
|
) -> MapReduceChain:
|
|
"""Construct a map-reduce chain that uses the chain for map and reduce."""
|
|
llm_chain = LLMChain(llm=llm, prompt=prompt)
|
|
reduce_chain = StuffDocumentsChain(llm_chain=llm_chain)
|
|
combine_documents_chain = MapReduceDocumentsChain(
|
|
llm_chain=llm_chain, combine_document_chain=reduce_chain
|
|
)
|
|
return cls(
|
|
combine_documents_chain=combine_documents_chain, text_splitter=text_splitter
|
|
)
|
|
|
|
class Config:
|
|
"""Configuration for this pydantic object."""
|
|
|
|
extra = Extra.forbid
|
|
arbitrary_types_allowed = True
|
|
|
|
@property
|
|
def input_keys(self) -> List[str]:
|
|
"""Expect input key.
|
|
|
|
:meta private:
|
|
"""
|
|
return [self.input_key]
|
|
|
|
@property
|
|
def output_keys(self) -> List[str]:
|
|
"""Return output key.
|
|
|
|
:meta private:
|
|
"""
|
|
return [self.output_key]
|
|
|
|
def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
|
|
# Split the larger text into smaller chunks.
|
|
texts = self.text_splitter.split_text(inputs[self.input_key])
|
|
docs = [Document(page_content=text) for text in texts]
|
|
outputs = self.combine_documents_chain.run(input_documents=docs)
|
|
return {self.output_key: outputs}
|