Add mmr support to redis retriever (#10556)

This commit is contained in:
Bagatur
2023-09-14 08:43:50 -07:00
committed by GitHub
parent ccf71e23e8
commit 7f3f6097e7
2 changed files with 44 additions and 8 deletions

View File

@@ -1425,6 +1425,7 @@ class RedisVectorStoreRetriever(VectorStoreRetriever):
"similarity",
"similarity_distance_threshold",
"similarity_score_threshold",
"mmr",
]
"""Allowed search types."""
@@ -1438,7 +1439,6 @@ class RedisVectorStoreRetriever(VectorStoreRetriever):
) -> List[Document]:
if self.search_type == "similarity":
docs = self.vectorstore.similarity_search(query, **self.search_kwargs)
elif self.search_type == "similarity_distance_threshold":
if self.search_kwargs["distance_threshold"] is None:
raise ValueError(
@@ -1454,6 +1454,10 @@ class RedisVectorStoreRetriever(VectorStoreRetriever):
)
)
docs = [doc for doc, _ in docs_and_similarities]
elif self.search_type == "mmr":
docs = self.vectorstore.max_marginal_relevance_search(
query, **self.search_kwargs
)
else:
raise ValueError(f"search_type of {self.search_type} not allowed.")
return docs