From 22d844dc0795e7e53a4cc499bf4974cb83df490d Mon Sep 17 00:00:00 2001 From: "Jiaping(JP) Zhang" Date: Fri, 19 May 2023 13:05:24 -0700 Subject: [PATCH] Add async search with relevance score (#4558) Add the async version for the search with relevance score Co-authored-by: Dev 2049 --- langchain/vectorstores/base.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/langchain/vectorstores/base.py b/langchain/vectorstores/base.py index 11a75885..6b689cd1 100644 --- a/langchain/vectorstores/base.py +++ b/langchain/vectorstores/base.py @@ -166,6 +166,17 @@ class VectorStore(ABC): """ raise NotImplementedError + async def asimilarity_search_with_relevance_scores( + self, query: str, k: int = 4, **kwargs: Any + ) -> List[Tuple[Document, float]]: + """Return docs most similar to query.""" + + # This is a temporary workaround to make the similarity search + # asynchronous. The proper solution is to make the similarity search + # asynchronous in the vector store implementations. + func = partial(self.similarity_search_with_relevance_scores, query, k, **kwargs) + return await asyncio.get_event_loop().run_in_executor(None, func) + async def asimilarity_search( self, query: str, k: int = 4, **kwargs: Any ) -> List[Document]: @@ -384,6 +395,13 @@ class VectorStoreRetriever(BaseRetriever, BaseModel): docs = await self.vectorstore.asimilarity_search( query, **self.search_kwargs ) + elif self.search_type == "similarity_score_threshold": + docs_and_similarities = ( + await self.vectorstore.asimilarity_search_with_relevance_scores( + query, **self.search_kwargs + ) + ) + docs = [doc for doc, _ in docs_and_similarities] elif self.search_type == "mmr": docs = await self.vectorstore.amax_marginal_relevance_search( query, **self.search_kwargs