Add ann_insert.milvus

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>

__init__.py

@@ -0,0 +1,4 @@
+from .milvus import Milvus
+
+def milvus(*args, **kwargs):
+    return Milvus(*args, **kwargs)

milvus.py

@@ -0,0 +1,38 @@
+import numpy as np
+from towhee import register
+from pymilvus import Collection
+
+
+@register(output_schema=['mr'])
+class Milvus:
+    """
+    Milvus ANN index class.
+    """
+
+    def __init__(self, collection):
+        """
+        Get an existing collection.
+        """
+        if isinstance(collection, str):
+            collection = Collection(collection)
+        self._collection = collection
+
+    def __call__(self, data):
+        """
+        Insert data to Milvus.
+
+        Args:
+        data (`list`):
+            The data to insert into milvus.
+
+        Returns:
+            A MutationResult object contains `insert_count` represents how many and a `primary_keys` of primary keys.
+
+        """
+        vectors = []
+        if isinstance(data, np.ndarray):
+            data = [data]
+        for v in data:
+            vectors.append(v if isinstance(v, list) else [v])
+        mr = self._collection.insert(vectors)
+        return mr

requirements.txt

@@ -0,0 +1,3 @@
+pymilvus
+towhee
+numpy