diff --git a/libs/community/langchain_community/vectorstores/milvus.py b/libs/community/langchain_community/vectorstores/milvus.py index 635f165ed2fa5..ddeefd07d6107 100644 --- a/libs/community/langchain_community/vectorstores/milvus.py +++ b/libs/community/langchain_community/vectorstores/milvus.py @@ -465,6 +465,7 @@ def _load( from pymilvus import Collection, utility from pymilvus.client.types import LoadState + timeout = self.timeout or timeout if ( isinstance(self.col, Collection) and self._get_index() is not None @@ -481,7 +482,7 @@ def add_texts( self, texts: Iterable[str], metadatas: Optional[List[dict]] = None, - timeout: Optional[int] = None, + timeout: Optional[float] = None, batch_size: int = 1000, *, ids: Optional[List[str]] = None, @@ -502,7 +503,7 @@ def add_texts( metadatas (Optional[List[dict]]): Metadata dicts attached to each of the texts. Defaults to None. 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. batch_size (int, optional): Batch size to use for insertion. Defaults to 1000. @@ -590,6 +591,7 @@ def add_texts( # Insert into the collection. try: res: Collection + timeout = self.timeout or timeout res = self.col.insert(insert_list, timeout=timeout, **kwargs) pks.extend(res.primary_keys) except MilvusException as e: @@ -606,7 +608,7 @@ def similarity_search( k: int = 4, param: Optional[dict] = None, expr: Optional[str] = None, - timeout: Optional[int] = None, + timeout: Optional[float] = None, **kwargs: Any, ) -> List[Document]: """Perform a similarity search against the query string. @@ -627,6 +629,7 @@ def similarity_search( if self.col is None: logger.debug("No existing collection to search.") return [] + timeout = self.timeout or timeout res = self.similarity_search_with_score( query=query, k=k, param=param, expr=expr, timeout=timeout, **kwargs ) @@ -638,7 +641,7 @@ def similarity_search_by_vector( k: int = 4, param: Optional[dict] = None, expr: Optional[str] = None, - timeout: Optional[int] = None, + timeout: Optional[float] = None, **kwargs: Any, ) -> List[Document]: """Perform a similarity search against the query string. @@ -659,6 +662,7 @@ def similarity_search_by_vector( if self.col is None: logger.debug("No existing collection to search.") return [] + timeout = self.timeout or timeout res = self.similarity_search_with_score_by_vector( embedding=embedding, k=k, param=param, expr=expr, timeout=timeout, **kwargs ) @@ -670,7 +674,7 @@ def similarity_search_with_score( k: int = 4, param: Optional[dict] = None, expr: Optional[str] = None, - timeout: Optional[int] = None, + timeout: Optional[float] = None, **kwargs: Any, ) -> List[Tuple[Document, float]]: """Perform a search on a query string and return results with score. @@ -685,7 +689,7 @@ def similarity_search_with_score( param (dict): The search params for the specified index. 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. kwargs: Collection.search() keyword arguments. @@ -698,7 +702,7 @@ def similarity_search_with_score( # Embed the query text. embedding = self.embedding_func.embed_query(query) - + timeout = self.timeout or timeout res = self.similarity_search_with_score_by_vector( embedding=embedding, k=k, param=param, expr=expr, timeout=timeout, **kwargs ) @@ -710,7 +714,7 @@ def similarity_search_with_score_by_vector( k: int = 4, param: Optional[dict] = None, expr: Optional[str] = None, - timeout: Optional[int] = None, + timeout: Optional[float] = None, **kwargs: Any, ) -> List[Tuple[Document, float]]: """Perform a search on a query string and return results with score. @@ -725,7 +729,7 @@ def similarity_search_with_score_by_vector( param (dict): The search params for the specified index. 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. kwargs: Collection.search() keyword arguments. @@ -742,7 +746,7 @@ def similarity_search_with_score_by_vector( # Determine result metadata fields with PK. output_fields = self.fields[:] output_fields.remove(self._vector_field) - + timeout = self.timeout or timeout # Perform the search. res = self.col.search( data=[embedding], @@ -772,7 +776,7 @@ def max_marginal_relevance_search( lambda_mult: float = 0.5, param: Optional[dict] = None, expr: Optional[str] = None, - timeout: Optional[int] = None, + timeout: Optional[float] = None, **kwargs: Any, ) -> List[Document]: """Perform a search and return results that are reordered by MMR. @@ -789,7 +793,7 @@ def max_marginal_relevance_search( param (dict, optional): The search params for the specified index. 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. kwargs: Collection.search() keyword arguments. @@ -802,7 +806,7 @@ def max_marginal_relevance_search( return [] embedding = self.embedding_func.embed_query(query) - + timeout = self.timeout or timeout return self.max_marginal_relevance_search_by_vector( embedding=embedding, k=k, @@ -822,7 +826,7 @@ def max_marginal_relevance_search_by_vector( lambda_mult: float = 0.5, param: Optional[dict] = None, expr: Optional[str] = None, - timeout: Optional[int] = None, + timeout: Optional[float] = None, **kwargs: Any, ) -> List[Document]: """Perform a search and return results that are reordered by MMR. @@ -839,7 +843,7 @@ def max_marginal_relevance_search_by_vector( param (dict, optional): The search params for the specified index. 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. kwargs: Collection.search() keyword arguments. @@ -856,7 +860,7 @@ def max_marginal_relevance_search_by_vector( # Determine result metadata fields. output_fields = self.fields[:] output_fields.remove(self._vector_field) - + timeout = self.timeout or timeout # Perform the search. res = self.col.search( data=[embedding],