diff --git a/libs/community/langchain_community/retrievers/milvus.py b/libs/community/langchain_community/retrievers/milvus.py index 1d759155701..fc81c465c08 100644 --- a/libs/community/langchain_community/retrievers/milvus.py +++ b/libs/community/langchain_community/retrievers/milvus.py @@ -18,6 +18,7 @@ class MilvusRetriever(BaseRetriever): embedding_function: Embeddings collection_name: str = "LangChainCollection" + collection_properties: Optional[Dict[str, Any]] = None connection_args: Optional[Dict[str, Any]] = None consistency_level: str = "Session" search_params: Optional[dict] = None @@ -31,6 +32,7 @@ class MilvusRetriever(BaseRetriever): values["store"] = Milvus( values["embedding_function"], values["collection_name"], + values["collection_properties"], values["connection_args"], values["consistency_level"], ) diff --git a/libs/community/langchain_community/vectorstores/milvus.py b/libs/community/langchain_community/vectorstores/milvus.py index 692f2084d5f..a1095a2d519 100644 --- a/libs/community/langchain_community/vectorstores/milvus.py +++ b/libs/community/langchain_community/vectorstores/milvus.py @@ -42,6 +42,10 @@ class Milvus(VectorStore): "LangChainCollection". collection_description (str): The description of the collection. Defaults to "". + collection_properties (Optional[dict[str, any]]): The collection properties. + Defaults to None. + If set, will override collection existing properties. + For example: {"collection.ttl.seconds": 60}. connection_args (Optional[dict[str, any]]): The connection args used for this class comes in the form of a dict. consistency_level (str): The consistency level to use for a collection. @@ -109,6 +113,7 @@ class Milvus(VectorStore): embedding_function: Embeddings, collection_name: str = "LangChainCollection", collection_description: str = "", + collection_properties: Optional[dict[str, Any]] = None, connection_args: Optional[dict[str, Any]] = None, consistency_level: str = "Session", index_params: Optional[dict] = None, @@ -149,6 +154,7 @@ class Milvus(VectorStore): self.embedding_func = embedding_function self.collection_name = collection_name self.collection_description = collection_description + self.collection_properties = collection_properties self.index_params = index_params self.search_params = search_params self.consistency_level = consistency_level @@ -177,6 +183,8 @@ class Milvus(VectorStore): self.collection_name, using=self.alias, ) + if self.collection_properties is not None: + self.col.set_properties(self.collection_properties) # If need to drop old, drop it if drop_old and isinstance(self.col, Collection): self.col.drop() @@ -332,6 +340,9 @@ class Milvus(VectorStore): consistency_level=self.consistency_level, using=self.alias, ) + # Set the collection properties if they exist + if self.collection_properties is not None: + self.col.set_properties(self.collection_properties) except MilvusException as e: logger.error( "Failed to create collection: %s error: %s", self.collection_name, e