community: add collection_properties parameter to Milvus (#15788)

- **Description:** add collection_properties parameter to Milvus. See
[pymilvus set_properties()
description](https://milvus.io/api-reference/pymilvus/v2.3.x/Collection/set_properties().md)
  - **Issue:** None
  - **Dependencies:** None
  - **Twitter handle:** None
This commit is contained in:
axiangcoding 2024-01-11 12:29:01 +08:00 committed by GitHub
parent 9e1ed17bfb
commit d5aa277b94
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 13 additions and 0 deletions

View File

@ -18,6 +18,7 @@ class MilvusRetriever(BaseRetriever):
embedding_function: Embeddings embedding_function: Embeddings
collection_name: str = "LangChainCollection" collection_name: str = "LangChainCollection"
collection_properties: Optional[Dict[str, Any]] = None
connection_args: Optional[Dict[str, Any]] = None connection_args: Optional[Dict[str, Any]] = None
consistency_level: str = "Session" consistency_level: str = "Session"
search_params: Optional[dict] = None search_params: Optional[dict] = None
@ -31,6 +32,7 @@ class MilvusRetriever(BaseRetriever):
values["store"] = Milvus( values["store"] = Milvus(
values["embedding_function"], values["embedding_function"],
values["collection_name"], values["collection_name"],
values["collection_properties"],
values["connection_args"], values["connection_args"],
values["consistency_level"], values["consistency_level"],
) )

View File

@ -42,6 +42,10 @@ class Milvus(VectorStore):
"LangChainCollection". "LangChainCollection".
collection_description (str): The description of the collection. Defaults to 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 connection_args (Optional[dict[str, any]]): The connection args used for
this class comes in the form of a dict. this class comes in the form of a dict.
consistency_level (str): The consistency level to use for a collection. consistency_level (str): The consistency level to use for a collection.
@ -109,6 +113,7 @@ class Milvus(VectorStore):
embedding_function: Embeddings, embedding_function: Embeddings,
collection_name: str = "LangChainCollection", collection_name: str = "LangChainCollection",
collection_description: str = "", collection_description: str = "",
collection_properties: Optional[dict[str, Any]] = None,
connection_args: Optional[dict[str, Any]] = None, connection_args: Optional[dict[str, Any]] = None,
consistency_level: str = "Session", consistency_level: str = "Session",
index_params: Optional[dict] = None, index_params: Optional[dict] = None,
@ -149,6 +154,7 @@ class Milvus(VectorStore):
self.embedding_func = embedding_function self.embedding_func = embedding_function
self.collection_name = collection_name self.collection_name = collection_name
self.collection_description = collection_description self.collection_description = collection_description
self.collection_properties = collection_properties
self.index_params = index_params self.index_params = index_params
self.search_params = search_params self.search_params = search_params
self.consistency_level = consistency_level self.consistency_level = consistency_level
@ -177,6 +183,8 @@ class Milvus(VectorStore):
self.collection_name, self.collection_name,
using=self.alias, 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 need to drop old, drop it
if drop_old and isinstance(self.col, Collection): if drop_old and isinstance(self.col, Collection):
self.col.drop() self.col.drop()
@ -332,6 +340,9 @@ class Milvus(VectorStore):
consistency_level=self.consistency_level, consistency_level=self.consistency_level,
using=self.alias, 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: except MilvusException as e:
logger.error( logger.error(
"Failed to create collection: %s error: %s", self.collection_name, e "Failed to create collection: %s error: %s", self.collection_name, e