From ffe35c396c48e4a20acc827cc92ce72dfc87cbd7 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Tue, 27 Dec 2022 19:53:45 -0500 Subject: [PATCH] unify return types across map-reduce and refine (#442) --- langchain/chains/combine_documents/map_reduce.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/langchain/chains/combine_documents/map_reduce.py b/langchain/chains/combine_documents/map_reduce.py index dd192062..18389edf 100644 --- a/langchain/chains/combine_documents/map_reduce.py +++ b/langchain/chains/combine_documents/map_reduce.py @@ -147,7 +147,8 @@ class MapReduceDocumentsChain(BaseCombineDocumentsChain, BaseModel): result_docs, **kwargs ) if self.return_map_steps: - extra_return_dict = {"map_steps": results} + _results = [r[self.llm_chain.output_key] for r in results] + extra_return_dict = {"map_steps": _results} else: extra_return_dict = {} output, _ = self.combine_document_chain.combine_docs(result_docs, **kwargs)