community: Fix Milvus got multiple values for keyword argument 'timeout' (#19232)

- **Description:** Fix Milvus got multiple values for keyword argument
'timeout'
- **Issue:**  fix #18580
- @baskaryan @eyurtsev PTAL
This commit is contained in:
Guangdong Liu 2024-03-19 11:44:25 +08:00 committed by GitHub
parent 95904fe443
commit c3310c5e7f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -465,6 +465,7 @@ class Milvus(VectorStore):
from pymilvus import Collection, utility from pymilvus import Collection, utility
from pymilvus.client.types import LoadState from pymilvus.client.types import LoadState
timeout = self.timeout or timeout
if ( if (
isinstance(self.col, Collection) isinstance(self.col, Collection)
and self._get_index() is not None and self._get_index() is not None
@ -481,7 +482,7 @@ class Milvus(VectorStore):
self, self,
texts: Iterable[str], texts: Iterable[str],
metadatas: Optional[List[dict]] = None, metadatas: Optional[List[dict]] = None,
timeout: Optional[int] = None, timeout: Optional[float] = None,
batch_size: int = 1000, batch_size: int = 1000,
*, *,
ids: Optional[List[str]] = None, ids: Optional[List[str]] = None,
@ -502,7 +503,7 @@ class Milvus(VectorStore):
metadatas (Optional[List[dict]]): Metadata dicts attached to each of metadatas (Optional[List[dict]]): Metadata dicts attached to each of
the texts. Defaults to None. the texts. Defaults to None.
should be less than 65535 bytes. Required and work when auto_id is False. should be less than 65535 bytes. Required and work when auto_id is False.
timeout (Optional[int]): Timeout for each batch insert. Defaults timeout (Optional[float]): Timeout for each batch insert. Defaults
to None. to None.
batch_size (int, optional): Batch size to use for insertion. batch_size (int, optional): Batch size to use for insertion.
Defaults to 1000. Defaults to 1000.
@ -590,6 +591,7 @@ class Milvus(VectorStore):
# Insert into the collection. # Insert into the collection.
try: try:
res: Collection res: Collection
timeout = self.timeout or timeout
res = self.col.insert(insert_list, timeout=timeout, **kwargs) res = self.col.insert(insert_list, timeout=timeout, **kwargs)
pks.extend(res.primary_keys) pks.extend(res.primary_keys)
except MilvusException as e: except MilvusException as e:
@ -606,7 +608,7 @@ class Milvus(VectorStore):
k: int = 4, k: int = 4,
param: Optional[dict] = None, param: Optional[dict] = None,
expr: Optional[str] = None, expr: Optional[str] = None,
timeout: Optional[int] = None, timeout: Optional[float] = None,
**kwargs: Any, **kwargs: Any,
) -> List[Document]: ) -> List[Document]:
"""Perform a similarity search against the query string. """Perform a similarity search against the query string.
@ -627,6 +629,7 @@ class Milvus(VectorStore):
if self.col is None: if self.col is None:
logger.debug("No existing collection to search.") logger.debug("No existing collection to search.")
return [] return []
timeout = self.timeout or timeout
res = self.similarity_search_with_score( res = self.similarity_search_with_score(
query=query, k=k, param=param, expr=expr, timeout=timeout, **kwargs query=query, k=k, param=param, expr=expr, timeout=timeout, **kwargs
) )
@ -638,7 +641,7 @@ class Milvus(VectorStore):
k: int = 4, k: int = 4,
param: Optional[dict] = None, param: Optional[dict] = None,
expr: Optional[str] = None, expr: Optional[str] = None,
timeout: Optional[int] = None, timeout: Optional[float] = None,
**kwargs: Any, **kwargs: Any,
) -> List[Document]: ) -> List[Document]:
"""Perform a similarity search against the query string. """Perform a similarity search against the query string.
@ -659,6 +662,7 @@ class Milvus(VectorStore):
if self.col is None: if self.col is None:
logger.debug("No existing collection to search.") logger.debug("No existing collection to search.")
return [] return []
timeout = self.timeout or timeout
res = self.similarity_search_with_score_by_vector( res = self.similarity_search_with_score_by_vector(
embedding=embedding, k=k, param=param, expr=expr, timeout=timeout, **kwargs embedding=embedding, k=k, param=param, expr=expr, timeout=timeout, **kwargs
) )
@ -670,7 +674,7 @@ class Milvus(VectorStore):
k: int = 4, k: int = 4,
param: Optional[dict] = None, param: Optional[dict] = None,
expr: Optional[str] = None, expr: Optional[str] = None,
timeout: Optional[int] = None, timeout: Optional[float] = None,
**kwargs: Any, **kwargs: Any,
) -> List[Tuple[Document, float]]: ) -> List[Tuple[Document, float]]:
"""Perform a search on a query string and return results with score. """Perform a search on a query string and return results with score.
@ -685,7 +689,7 @@ class Milvus(VectorStore):
param (dict): The search params for the specified index. param (dict): The search params for the specified index.
Defaults to None. Defaults to None.
expr (str, optional): Filtering expression. Defaults to None. expr (str, optional): Filtering expression. Defaults to None.
timeout (int, optional): How long to wait before timeout error. timeout (float, optional): How long to wait before timeout error.
Defaults to None. Defaults to None.
kwargs: Collection.search() keyword arguments. kwargs: Collection.search() keyword arguments.
@ -698,7 +702,7 @@ class Milvus(VectorStore):
# Embed the query text. # Embed the query text.
embedding = self.embedding_func.embed_query(query) embedding = self.embedding_func.embed_query(query)
timeout = self.timeout or timeout
res = self.similarity_search_with_score_by_vector( res = self.similarity_search_with_score_by_vector(
embedding=embedding, k=k, param=param, expr=expr, timeout=timeout, **kwargs embedding=embedding, k=k, param=param, expr=expr, timeout=timeout, **kwargs
) )
@ -710,7 +714,7 @@ class Milvus(VectorStore):
k: int = 4, k: int = 4,
param: Optional[dict] = None, param: Optional[dict] = None,
expr: Optional[str] = None, expr: Optional[str] = None,
timeout: Optional[int] = None, timeout: Optional[float] = None,
**kwargs: Any, **kwargs: Any,
) -> List[Tuple[Document, float]]: ) -> List[Tuple[Document, float]]:
"""Perform a search on a query string and return results with score. """Perform a search on a query string and return results with score.
@ -725,7 +729,7 @@ class Milvus(VectorStore):
param (dict): The search params for the specified index. param (dict): The search params for the specified index.
Defaults to None. Defaults to None.
expr (str, optional): Filtering expression. Defaults to None. expr (str, optional): Filtering expression. Defaults to None.
timeout (int, optional): How long to wait before timeout error. timeout (float, optional): How long to wait before timeout error.
Defaults to None. Defaults to None.
kwargs: Collection.search() keyword arguments. kwargs: Collection.search() keyword arguments.
@ -742,7 +746,7 @@ class Milvus(VectorStore):
# Determine result metadata fields with PK. # Determine result metadata fields with PK.
output_fields = self.fields[:] output_fields = self.fields[:]
output_fields.remove(self._vector_field) output_fields.remove(self._vector_field)
timeout = self.timeout or timeout
# Perform the search. # Perform the search.
res = self.col.search( res = self.col.search(
data=[embedding], data=[embedding],
@ -772,7 +776,7 @@ class Milvus(VectorStore):
lambda_mult: float = 0.5, lambda_mult: float = 0.5,
param: Optional[dict] = None, param: Optional[dict] = None,
expr: Optional[str] = None, expr: Optional[str] = None,
timeout: Optional[int] = None, timeout: Optional[float] = None,
**kwargs: Any, **kwargs: Any,
) -> List[Document]: ) -> List[Document]:
"""Perform a search and return results that are reordered by MMR. """Perform a search and return results that are reordered by MMR.
@ -789,7 +793,7 @@ class Milvus(VectorStore):
param (dict, optional): The search params for the specified index. param (dict, optional): The search params for the specified index.
Defaults to None. Defaults to None.
expr (str, optional): Filtering expression. Defaults to None. expr (str, optional): Filtering expression. Defaults to None.
timeout (int, optional): How long to wait before timeout error. timeout (float, optional): How long to wait before timeout error.
Defaults to None. Defaults to None.
kwargs: Collection.search() keyword arguments. kwargs: Collection.search() keyword arguments.
@ -802,7 +806,7 @@ class Milvus(VectorStore):
return [] return []
embedding = self.embedding_func.embed_query(query) embedding = self.embedding_func.embed_query(query)
timeout = self.timeout or timeout
return self.max_marginal_relevance_search_by_vector( return self.max_marginal_relevance_search_by_vector(
embedding=embedding, embedding=embedding,
k=k, k=k,
@ -822,7 +826,7 @@ class Milvus(VectorStore):
lambda_mult: float = 0.5, lambda_mult: float = 0.5,
param: Optional[dict] = None, param: Optional[dict] = None,
expr: Optional[str] = None, expr: Optional[str] = None,
timeout: Optional[int] = None, timeout: Optional[float] = None,
**kwargs: Any, **kwargs: Any,
) -> List[Document]: ) -> List[Document]:
"""Perform a search and return results that are reordered by MMR. """Perform a search and return results that are reordered by MMR.
@ -839,7 +843,7 @@ class Milvus(VectorStore):
param (dict, optional): The search params for the specified index. param (dict, optional): The search params for the specified index.
Defaults to None. Defaults to None.
expr (str, optional): Filtering expression. Defaults to None. expr (str, optional): Filtering expression. Defaults to None.
timeout (int, optional): How long to wait before timeout error. timeout (float, optional): How long to wait before timeout error.
Defaults to None. Defaults to None.
kwargs: Collection.search() keyword arguments. kwargs: Collection.search() keyword arguments.
@ -856,7 +860,7 @@ class Milvus(VectorStore):
# Determine result metadata fields. # Determine result metadata fields.
output_fields = self.fields[:] output_fields = self.fields[:]
output_fields.remove(self._vector_field) output_fields.remove(self._vector_field)
timeout = self.timeout or timeout
# Perform the search. # Perform the search.
res = self.col.search( res = self.col.search(
data=[embedding], data=[embedding],