allow embeddings vector to be used for mmr searching (#2620)

Signed-off-by: rishabh208gupta <rishabhgupta.52pp@gmail.com>

Author: rishabh208gupta

elasticsearch/helpers/vectorstore/_async/vectorstore.py

@@ -232,7 +232,7 @@ async def delete(  # type: ignore[no-untyped-def]
     async def search(
         self,
         *,
-        query: Optional[str],
+        query: Optional[str] = None,
         query_vector: Optional[List[float]] = None,
         k: int = 4,
         num_candidates: int = 50,
@@ -344,8 +344,9 @@ async def _create_index_if_not_exists(self) -> None:
     async def max_marginal_relevance_search(
         self,
         *,
-        embedding_service: AsyncEmbeddingService,
-        query: str,
+        query: Optional[str] = None,
+        query_embedding: Optional[List[float]] = None,
+        embedding_service: Optional[AsyncEmbeddingService] = None,
         vector_field: str,
         k: int = 4,
         num_candidates: int = 20,
@@ -361,6 +362,8 @@ async def max_marginal_relevance_search(
             among selected documents.
 
         :param query (str): Text to look up documents similar to.
+        :param query_embedding: Input embedding vector. If given, input query string is
+            ignored.
         :param k (int): Number of Documents to return. Defaults to 4.
         :param fetch_k (int): Number of Documents to fetch to pass to MMR algorithm.
         :param lambda_mult (float): Number between 0 and 1 that determines the degree
@@ -381,7 +384,10 @@ async def max_marginal_relevance_search(
             remove_vector_query_field_from_metadata = False
 
         # Embed the query
-        query_embedding = await embedding_service.embed_query(query)
+        if self.embedding_service and not query_embedding:
+            if not query:
+                raise ValueError("specify a query or a query_embedding to search")
+            query_embedding = await self.embedding_service.embed_query(query)
 
         # Fetch the initial documents
         got_hits = await self.search(

elasticsearch/helpers/vectorstore/_sync/vectorstore.py

@@ -229,7 +229,7 @@ def delete(  # type: ignore[no-untyped-def]
     def search(
         self,
         *,
-        query: Optional[str],
+        query: Optional[str] = None,
         query_vector: Optional[List[float]] = None,
         k: int = 4,
         num_candidates: int = 50,
@@ -341,8 +341,9 @@ def _create_index_if_not_exists(self) -> None:
     def max_marginal_relevance_search(
         self,
         *,
-        embedding_service: EmbeddingService,
-        query: str,
+        query: Optional[str] = None,
+        query_embedding: Optional[List[float]] = None,
+        embedding_service: Optional[EmbeddingService] = None,
         vector_field: str,
         k: int = 4,
         num_candidates: int = 20,
@@ -358,6 +359,8 @@ def max_marginal_relevance_search(
             among selected documents.
 
         :param query (str): Text to look up documents similar to.
+        :param query_embedding: Input embedding vector. If given, input query string is
+            ignored.
         :param k (int): Number of Documents to return. Defaults to 4.
         :param fetch_k (int): Number of Documents to fetch to pass to MMR algorithm.
         :param lambda_mult (float): Number between 0 and 1 that determines the degree
@@ -378,7 +381,10 @@ def max_marginal_relevance_search(
             remove_vector_query_field_from_metadata = False
 
         # Embed the query
-        query_embedding = embedding_service.embed_query(query)
+        if self.embedding_service and not query_embedding:
+            if not query:
+                raise ValueError("specify a query or a query_embedding to search")
+            query_embedding = self.embedding_service.embed_query(query)
 
         # Fetch the initial documents
         got_hits = self.search(

test_elasticsearch/test_server/test_vectorstore/test_vectorstore.py

@@ -822,6 +822,7 @@ def test_max_marginal_relevance_search(
         texts = ["foo", "bar", "baz"]
         vector_field = "vector_field"
         text_field = "text_field"
+        query_embedding = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0]
         embedding_service = ConsistentFakeEmbeddings()
         store = VectorStore(
             index=index,
@@ -834,7 +835,6 @@ def test_max_marginal_relevance_search(
         store.add_texts(texts)
 
         mmr_output = store.max_marginal_relevance_search(
-            embedding_service=embedding_service,
             query=texts[0],
             vector_field=vector_field,
             k=3,
@@ -843,8 +843,17 @@ def test_max_marginal_relevance_search(
         sim_output = store.search(query=texts[0], k=3)
         assert mmr_output == sim_output
 
+        # search using query embeddings vector instead of query
+        mmr_output = store.max_marginal_relevance_search(
+            query_embedding=query_embedding,
+            vector_field=vector_field,
+            k=3,
+            num_candidates=3,
+        )
+        sim_output = store.search(query_vector=query_embedding, k=3)
+        assert mmr_output == sim_output
+
         mmr_output = store.max_marginal_relevance_search(
-            embedding_service=embedding_service,
             query=texts[0],
             vector_field=vector_field,
             k=2,
@@ -855,7 +864,6 @@ def test_max_marginal_relevance_search(
         assert mmr_output[1]["_source"][text_field] == texts[1]
 
         mmr_output = store.max_marginal_relevance_search(
-            embedding_service=embedding_service,
             query=texts[0],
             vector_field=vector_field,
             k=2,
@@ -868,7 +876,6 @@ def test_max_marginal_relevance_search(
 
         # if fetch_k < k, then the output will be less than k
         mmr_output = store.max_marginal_relevance_search(
-            embedding_service=embedding_service,
             query=texts[0],
             vector_field=vector_field,
             k=3,