From f46f1d28afed95eb0cdda26ef14d2e40ab1e9f28 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Mon, 30 Jan 2023 14:48:12 -0800 Subject: [PATCH] expose memory key name (#808) --- langchain/chains/conversation/memory.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/langchain/chains/conversation/memory.py b/langchain/chains/conversation/memory.py index 28f3130ba37..cf174436f4e 100644 --- a/langchain/chains/conversation/memory.py +++ b/langchain/chains/conversation/memory.py @@ -230,12 +230,12 @@ class ConversationEntityMemory(Memory, BaseModel): llm: BaseLLM entity_extraction_prompt: BasePromptTemplate = ENTITY_EXTRACTION_PROMPT entity_summarization_prompt: BasePromptTemplate = ENTITY_SUMMARIZATION_PROMPT - memory_keys: List[str] = ["entities", "history"] #: :meta private: output_key: Optional[str] = None input_key: Optional[str] = None store: Dict[str, Optional[str]] = {} entity_cache: List[str] = [] k: int = 3 + chat_history_key: str = "history" @property def memory_variables(self) -> List[str]: @@ -243,7 +243,7 @@ class ConversationEntityMemory(Memory, BaseModel): :meta private: """ - return ["entities", "history"] + return ["entities", self.chat_history_key] def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]: """Return history buffer.""" @@ -265,7 +265,7 @@ class ConversationEntityMemory(Memory, BaseModel): entity_summaries[entity] = self.store.get(entity, "") self.entity_cache = entities return { - "history": "\n".join(self.buffer[-self.k :]), + self.chat_history_key: "\n".join(self.buffer[-self.k :]), "entities": entity_summaries, }