From 453d4c3a991b7d1d00f53b24229a21d2fa0c76e1 Mon Sep 17 00:00:00 2001 From: Mohammad Mohtashim <45242107+keenborder786@users.noreply.github.com> Date: Thu, 20 Jul 2023 19:23:27 +0500 Subject: [PATCH] VectorStoreRetrieverMemory exclude additional input keys feature (#7941) - Description: Added a parameter in VectorStoreRetrieverMemory which filters the input given by the key when constructing the buffering the document for Vector. This feature is helpful if you have certain inputs apart from the VectorMemory's own memory_key that needs to be ignored e.g when using combined memory, we might need to filter the memory_key of the other memory, Please see the issue. - Issue: #7695 - Tag maintainer: @rlancemartin, @eyurtsev --------- Co-authored-by: Bagatur --- langchain/memory/vectorstore.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/langchain/memory/vectorstore.py b/langchain/memory/vectorstore.py index d5c40f26b99..60e0ae7acc2 100644 --- a/langchain/memory/vectorstore.py +++ b/langchain/memory/vectorstore.py @@ -1,6 +1,6 @@ """Class for a VectorStore-backed memory object.""" -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Sequence, Union from pydantic import Field @@ -25,6 +25,9 @@ class VectorStoreRetrieverMemory(BaseMemory): return_docs: bool = False """Whether or not to return the result of querying the database directly.""" + exclude_input_keys: Sequence[str] = Field(default_factory=tuple) + """Input keys to exclude in addition to memory key when constructing the document""" + @property def memory_variables(self) -> List[str]: """The list of keys emitted from the load_memory_variables method.""" @@ -55,10 +58,13 @@ class VectorStoreRetrieverMemory(BaseMemory): ) -> List[Document]: """Format context from this conversation to buffer.""" # Each document should only include the current turn, not the chat history - filtered_inputs = {k: v for k, v in inputs.items() if k != self.memory_key} + exclude = set(self.exclude_input_keys) + exclude.add(self.memory_key) + filtered_inputs = {k: v for k, v in inputs.items() if k not in exclude} texts = [ f"{k}: {v}" for k, v in list(filtered_inputs.items()) + list(outputs.items()) + if k not in exclude ] page_content = "\n".join(texts) return [Document(page_content=page_content)]