From ecb11a4a3274cbf85ae842eb12d71d968240de36 Mon Sep 17 00:00:00 2001 From: xiaohuanshu Date: Fri, 29 Mar 2024 07:38:08 +0800 Subject: [PATCH] langchain[patch]: fix BaseChatMemory get output data error with extra key (#18117) **Description:** At times, BaseChatMemory._get_input_output may acquire some extra keys such as 'intermediate_steps' (agent_executor with return_intermediate_steps set to True) and 'messages' (agent_executor.iter with memory). In these instances, _get_input_output can raise an error due to the presence of multiple keys. The 'output' field should be used as the default field in these cases. **Issue:** #16791 --- .../langchain/langchain/memory/chat_memory.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/libs/langchain/langchain/memory/chat_memory.py b/libs/langchain/langchain/memory/chat_memory.py index ad030c3f71a..671edf9f31b 100644 --- a/libs/langchain/langchain/memory/chat_memory.py +++ b/libs/langchain/langchain/memory/chat_memory.py @@ -1,3 +1,4 @@ +import warnings from abc import ABC from typing import Any, Dict, Optional, Tuple @@ -26,9 +27,21 @@ class BaseChatMemory(BaseMemory, ABC): else: prompt_input_key = self.input_key if self.output_key is None: - if len(outputs) != 1: - raise ValueError(f"One output key expected, got {outputs.keys()}") - output_key = list(outputs.keys())[0] + if len(outputs) == 1: + output_key = list(outputs.keys())[0] + elif "output" in outputs: + output_key = "output" + warnings.warn( + f"'{self.__class__.__name__}' got multiple output keys:" + f" {outputs.keys()}. The default 'output' key is being used." + f" If this is not desired, please manually set 'output_key'." + ) + else: + raise ValueError( + f"Got multiple output keys: {outputs.keys()}, cannot " + f"determine which to store in memory. Please set the " + f"'output_key' explicitly." + ) else: output_key = self.output_key return inputs[prompt_input_key], outputs[output_key]