From d5aa277b94f89167044d8631c3b14f914b82bce1 Mon Sep 17 00:00:00 2001
From: axiangcoding <axiangcoding@gmail.com>
Date: Thu, 11 Jan 2024 12:29:01 +0800
Subject: [PATCH] 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
---
 .../langchain_community/retrievers/milvus.py          |  2 ++
 .../langchain_community/vectorstores/milvus.py        | 11 +++++++++++
 2 files changed, 13 insertions(+)

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