blob42/langchain
1
0
Fork 0
forked from Archives/langchain
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Harrison/fix save context (#2082)

Browse Source 
Co-authored-by: Saurabh Misra <misra.saurabh1@gmail.com>
This commit is contained in:
Harrison Chase 2023-03-27 23:10:46 -07:00 committed by GitHub
parent b7f392fdd6
commit 365669a7fd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
langchain/memory/entity.py
View File

@ -74,25 +74,27 @@ class ConversationEntityMemory(BaseChatMemory, BaseModel):
def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
"""Save context from this conversation to buffer."""
super().save_context(inputs, outputs)
if self.input_key is None:
prompt_input_key = get_prompt_input_key(inputs, self.memory_variables)
else:
prompt_input_key = self.input_key
for entity in self.entity_cache:
chain = LLMChain(llm=self.llm, prompt=self.entity_summarization_prompt)
# key value store for entity
existing_summary = self.store.get(entity, "")
buffer_string = get_buffer_string(
self.buffer[-self.k * 2 :],
human_prefix=self.human_prefix,
ai_prefix=self.ai_prefix,
)
buffer_string = get_buffer_string(
self.buffer[-self.k * 2 :],
human_prefix=self.human_prefix,
ai_prefix=self.ai_prefix,
)
input_data = inputs[prompt_input_key]
chain = LLMChain(llm=self.llm, prompt=self.entity_summarization_prompt)
for entity in self.entity_cache:
existing_summary = self.store.get(entity, "")
output = chain.predict(
summary=existing_summary,
history=buffer_string,
input=inputs[prompt_input_key],
entity=entity,
history=buffer_string,
input=input_data,
)
self.store[entity] = output.strip()