community[patch]: implement qdrant _aembed_query and use it in other async funcs (#19155)

`amax_marginal_relevance_search ` and `asimilarity_search_with_score `
should use an async version of `_embed_query `.
This commit is contained in:
kaijietti 2024-03-16 05:20:12 +08:00 committed by GitHub
parent 527676a753
commit c20aeef79a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -417,8 +417,9 @@ class Qdrant(VectorStore):
Returns: Returns:
List of documents most similar to the query text and distance for each. List of documents most similar to the query text and distance for each.
""" """
query_embedding = await self._aembed_query(query)
return await self.asimilarity_search_with_score_by_vector( return await self.asimilarity_search_with_score_by_vector(
self._embed_query(query), query_embedding,
k, k,
filter=filter, filter=filter,
search_params=search_params, search_params=search_params,
@ -841,7 +842,7 @@ class Qdrant(VectorStore):
Returns: Returns:
List of Documents selected by maximal marginal relevance. List of Documents selected by maximal marginal relevance.
""" """
query_embedding = self._embed_query(query) query_embedding = await self._aembed_query(query)
return await self.amax_marginal_relevance_search_by_vector( return await self.amax_marginal_relevance_search_by_vector(
query_embedding, query_embedding,
k=k, k=k,
@ -2022,6 +2023,26 @@ class Qdrant(VectorStore):
raise ValueError("Neither of embeddings or embedding_function is set") raise ValueError("Neither of embeddings or embedding_function is set")
return embedding.tolist() if hasattr(embedding, "tolist") else embedding return embedding.tolist() if hasattr(embedding, "tolist") else embedding
async def _aembed_query(self, query: str) -> List[float]:
"""Embed query text asynchronously.
Used to provide backward compatibility with `embedding_function` argument.
Args:
query: Query text.
Returns:
List of floats representing the query embedding.
"""
if self.embeddings is not None:
embedding = await self.embeddings.aembed_query(query)
else:
if self._embeddings_function is not None:
embedding = self._embeddings_function(query)
else:
raise ValueError("Neither of embeddings or embedding_function is set")
return embedding.tolist() if hasattr(embedding, "tolist") else embedding
def _embed_texts(self, texts: Iterable[str]) -> List[List[float]]: def _embed_texts(self, texts: Iterable[str]) -> List[List[float]]:
"""Embed search texts. """Embed search texts.