From 8ba868d3b06662e5bb75ee0971f86cc9d597ded7 Mon Sep 17 00:00:00 2001 From: Christophe Bornet Date: Tue, 4 Jun 2024 22:43:55 +0200 Subject: [PATCH] core[patch]: Add similarity_score_threshold to VectorStore search types (#22477) --- libs/core/langchain_core/vectorstores.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/libs/core/langchain_core/vectorstores.py b/libs/core/langchain_core/vectorstores.py index 96160cedf2..aed6bdd6ca 100644 --- a/libs/core/langchain_core/vectorstores.py +++ b/libs/core/langchain_core/vectorstores.py @@ -156,12 +156,17 @@ class VectorStore(ABC): """Return docs most similar to query using specified search type.""" if search_type == "similarity": return self.similarity_search(query, **kwargs) + elif search_type == "similarity_score_threshold": + docs_and_similarities = self.similarity_search_with_relevance_scores( + query, **kwargs + ) + return [doc for doc, _ in docs_and_similarities] elif search_type == "mmr": return self.max_marginal_relevance_search(query, **kwargs) else: raise ValueError( f"search_type of {search_type} not allowed. Expected " - "search_type to be 'similarity' or 'mmr'." + "search_type to be 'similarity', 'similarity_score_threshold' or 'mmr'." ) async def asearch( @@ -170,12 +175,17 @@ class VectorStore(ABC): """Return docs most similar to query using specified search type.""" if search_type == "similarity": return await self.asimilarity_search(query, **kwargs) + elif search_type == "similarity_score_threshold": + docs_and_similarities = await self.asimilarity_search_with_relevance_scores( + query, **kwargs + ) + return [doc for doc, _ in docs_and_similarities] elif search_type == "mmr": return await self.amax_marginal_relevance_search(query, **kwargs) else: raise ValueError( f"search_type of {search_type} not allowed. Expected " - "search_type to be 'similarity' or 'mmr'." + "search_type to be 'similarity', 'similarity_score_threshold' or 'mmr'." ) @abstractmethod