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
This commit is contained in:
Saurabh Misra 2023-03-28 15:19:48 -07:00 committed by GitHub
parent a5bf8c9b9d
commit c33e055f17
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -41,27 +41,23 @@ class ConversationKGMemory(BaseChatMemory, BaseModel):
def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]: def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
"""Return history buffer.""" """Return history buffer."""
entities = self._get_current_entities(inputs) entities = self._get_current_entities(inputs)
summaries = {}
summary_strings = []
for entity in entities: for entity in entities:
knowledge = self.kg.get_entity_knowledge(entity) knowledge = self.kg.get_entity_knowledge(entity)
if knowledge: if knowledge:
summaries[entity] = ". ".join(knowledge) + "." summary = f"On {entity}: {'. '.join(knowledge)}."
summary_strings.append(summary)
context: Union[str, List] context: Union[str, List]
if summaries: if not summary_strings:
summary_strings = [ context = [] if self.return_messages else ""
f"On {entity}: {summary}" for entity, summary in summaries.items() elif self.return_messages:
]
if self.return_messages:
context = [ context = [
self.summary_message_cls(content=text) for text in summary_strings self.summary_message_cls(content=text) for text in summary_strings
] ]
else: else:
context = "\n".join(summary_strings) context = "\n".join(summary_strings)
else:
if self.return_messages:
context = []
else:
context = ""
return {self.memory_key: context} return {self.memory_key: context}
@property @property