langchain: Add async methods to VectorStoreRetrieverMemory

libs/langchain/langchain/memory/vectorstore.py

@@ -39,20 +39,34 @@ def _get_prompt_input_key(self, inputs: Dict[str, Any]) -> str:
             return get_prompt_input_key(inputs, self.memory_variables)
         return self.input_key
 
-    def load_memory_variables(
-        self, inputs: Dict[str, Any]
+    def _documents_to_memory_variables(
+        self, docs: List[Document]
     ) -> Dict[str, Union[List[Document], str]]:
-        """Return history buffer."""
-        input_key = self._get_prompt_input_key(inputs)
-        query = inputs[input_key]
-        docs = self.retriever.get_relevant_documents(query)
         result: Union[List[Document], str]
         if not self.return_docs:
             result = "\n".join([doc.page_content for doc in docs])
         else:
             result = docs
         return {self.memory_key: result}
 
+    def load_memory_variables(
+        self, inputs: Dict[str, Any]
+    ) -> Dict[str, Union[List[Document], str]]:
+        """Return history buffer."""
+        input_key = self._get_prompt_input_key(inputs)
+        query = inputs[input_key]
+        docs = self.retriever.get_relevant_documents(query)
+        return self._documents_to_memory_variables(docs)
+
+    async def aload_memory_variables(
+        self, inputs: Dict[str, Any]
+    ) -> Dict[str, Union[List[Document], str]]:
+        """Return history buffer."""
+        input_key = self._get_prompt_input_key(inputs)
+        query = inputs[input_key]
+        docs = await self.retriever.aget_relevant_documents(query)
+        return self._documents_to_memory_variables(docs)
+
     def _form_documents(
         self, inputs: Dict[str, Any], outputs: Dict[str, str]
     ) -> List[Document]:
@@ -73,5 +87,15 @@ def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
         documents = self._form_documents(inputs, outputs)
         self.retriever.add_documents(documents)
 
+    async def asave_context(
+        self, inputs: Dict[str, Any], outputs: Dict[str, str]
+    ) -> None:
+        """Save context from this conversation to buffer."""
+        documents = self._form_documents(inputs, outputs)
+        await self.retriever.aadd_documents(documents)
+
     def clear(self) -> None:
         """Nothing to clear."""
+
+    async def aclear(self) -> None:
+        """Nothing to clear."""