From ee1adaacaa7d2416bb387911d18f57bbf3bcc235 Mon Sep 17 00:00:00 2001 From: Wickes Wong <27881570+wickes1@users.noreply.github.com> Date: Fri, 3 May 2024 17:25:09 -0400 Subject: [PATCH] langchain[patch]: Fix summary buffer memory with return message flag (#21115) ## Description Memory return could be set as `str` or `message` by `return_messages` flag as mentioned in https://python.langchain.com/docs/modules/memory/#whether-memory-is-a-string-or-a-list-of-messages, where `langchain.chains.conversation.memory.ConversationSummaryBufferMemory` did not implement that. This commit added `buffer_as_str` and `buffer_as_messages` function, and `buffer` now affected by `return_messages` flag. ## Example Test Code and Output ```python # Fix: ConversationSummaryBufferMemory with return_messages flag function # Test code from langchain.chains.conversation.memory import ConversationSummaryBufferMemory from langchain_community.llms.ollama import Ollama llm = Ollama() # Create an instance of ConversationSummaryBufferMemory with return_messages set to True memory = ConversationSummaryBufferMemory(return_messages=True, llm=llm) # Add user and AI messages to the chat memory memory.chat_memory.add_user_message("hi!") memory.chat_memory.add_ai_message("what's up?") # Print the buffer print("Buffer:") print(*map(type, memory.buffer), sep="\n") print(memory.buffer, "\n") # Print the buffer as a string print("Buffer as String:") print(type(memory.buffer_as_str)) print(memory.buffer_as_str, "\n") # Print the buffer as messages print("Buffer as Messages:") print(*map(type, memory.buffer_as_messages), sep="\n") print(memory.buffer_as_messages, "\n") # Print the buffer after setting return_messages to False memory.return_messages = False print("Buffer after setting return_messages to False:") print(type(memory.buffer)) print(memory.buffer, "\n") ``` ```plaintext Buffer: [HumanMessage(content='hi!'), AIMessage(content="what's up?")] Buffer as String: Human: hi! AI: what's up? Buffer as Messages: [HumanMessage(content='hi!'), AIMessage(content="what's up?")] Buffer after setting return_messages to False: Human: hi! AI: what's up? ``` --------- Co-authored-by: Bagatur <22008038+baskaryan@users.noreply.github.com> --- libs/langchain/langchain/memory/summary_buffer.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/libs/langchain/langchain/memory/summary_buffer.py b/libs/langchain/langchain/memory/summary_buffer.py index 23d050f62dd..389da7e4215 100644 --- a/libs/langchain/langchain/memory/summary_buffer.py +++ b/libs/langchain/langchain/memory/summary_buffer.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List +from typing import Any, Dict, List, Union from langchain_core.messages import BaseMessage, get_buffer_string from langchain_core.pydantic_v1 import root_validator @@ -15,8 +15,9 @@ class ConversationSummaryBufferMemory(BaseChatMemory, SummarizerMixin): memory_key: str = "history" @property - def buffer(self) -> List[BaseMessage]: - return self.chat_memory.messages + def buffer(self) -> Union[str, List[BaseMessage]]: + """String buffer of memory.""" + return self.load_memory_variables({})[self.memory_key] @property def memory_variables(self) -> List[str]: @@ -28,7 +29,7 @@ class ConversationSummaryBufferMemory(BaseChatMemory, SummarizerMixin): def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]: """Return history buffer.""" - buffer = self.buffer + buffer = self.chat_memory.messages if self.moving_summary_buffer != "": first_messages: List[BaseMessage] = [ self.summary_message_cls(content=self.moving_summary_buffer)