From c33e055f17d59e225cc009c49b28d4400d56e709 Mon Sep 17 00:00:00 2001 From: Saurabh Misra Date: Tue, 28 Mar 2023 15:19:48 -0700 Subject: [PATCH] Improve ConversationKGMemory and its function load_memory_variables (#1999) 1. Removed the `summaries` dictionary in favor of directly appending to the summary_strings list, which avoids the unnecessary double-loop. 2. Simplified the logic for populating the `context` variable. Co-created with GPT-4 @agihouse --- langchain/memory/kg.py | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/langchain/memory/kg.py b/langchain/memory/kg.py index b0aaa4c0..ddf7ff33 100644 --- a/langchain/memory/kg.py +++ b/langchain/memory/kg.py @@ -41,27 +41,23 @@ class ConversationKGMemory(BaseChatMemory, BaseModel): def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]: """Return history buffer.""" entities = self._get_current_entities(inputs) - summaries = {} + + summary_strings = [] for entity in entities: knowledge = self.kg.get_entity_knowledge(entity) if knowledge: - summaries[entity] = ". ".join(knowledge) + "." + summary = f"On {entity}: {'. '.join(knowledge)}." + summary_strings.append(summary) context: Union[str, List] - if summaries: - summary_strings = [ - f"On {entity}: {summary}" for entity, summary in summaries.items() + if not summary_strings: + context = [] if self.return_messages else "" + elif self.return_messages: + context = [ + self.summary_message_cls(content=text) for text in summary_strings ] - if self.return_messages: - context = [ - self.summary_message_cls(content=text) for text in summary_strings - ] - else: - context = "\n".join(summary_strings) else: - if self.return_messages: - context = [] - else: - context = "" + context = "\n".join(summary_strings) + return {self.memory_key: context} @property