Add a Vector Database Service to allow stages to read and write to VDBs

docker/conda/environments/cuda11.8_dev.yml

@@ -111,3 +111,4 @@ dependencies:
         - databricks-connect
         - pytest-kafka==0.6.0
         - pymilvus==2.3.1
+        - milvus

morpheus/service/milvus_client.py

@@ -19,8 +19,24 @@
 
 logger = logging.getLogger(__name__)
 
-
-def handle_exceptions(method_name: str, error_message: str):
+# Milvus data type mapping dictionary
+MILVUS_DATA_TYPE_MAP = {
+    "int8": pymilvus.DataType.INT8,
+    "int16": pymilvus.DataType.INT16,
+    "int32": pymilvus.DataType.INT32,
+    "int64": pymilvus.DataType.INT64,
+    "bool": pymilvus.DataType.BOOL,
+    "float": pymilvus.DataType.FLOAT,
+    "double": pymilvus.DataType.DOUBLE,
+    "binary_vector": pymilvus.DataType.BINARY_VECTOR,
+    "float_vector": pymilvus.DataType.FLOAT_VECTOR,
+    "string": pymilvus.DataType.STRING,
+    "varchar": pymilvus.DataType.VARCHAR,
+    "json": pymilvus.DataType.JSON,
+}
+
+
+def handle_exceptions(method_name: str, error_message: str) -> typing.Callable:
     """
     Decorator function to handle exceptions and log errors.
 
@@ -33,7 +49,7 @@ def handle_exceptions(method_name: str, error_message: str):
 
     Returns
     -------
-    Callable
+    typing.Callable
         Decorated function.
     """
 
@@ -226,7 +242,22 @@ def drop_index(self, collection_name: str, field_name: str, index_name: str) ->
         conn = self._get_connection()
         conn.drop_index(collection_name=collection_name, field_name=field_name, index_name=index_name)
 
-    @handle_exceptions("get_collection", "Error getting collection")
-    def get_collection(self, collection_name: str, **kwargs: dict[str, typing.Any]) -> str:
+    @handle_exceptions("get_collection", "Error getting collection object")
+    def get_collection(self, collection_name: str, **kwargs: dict[str, typing.Any]) -> pymilvus.Collection:
+        """
+        Returns `pymilvus.Collection` object associated with the given collection name.
+
+        Parameters
+        ----------
+        collection_name : str
+            Name of the collection to delete from.
+        **kwargs : dict
+            Additional keyword arguments to get Collection instance.
+
+        Returns
+        -------
+        pymilvus.Collection
+            Returns pymilvus Collection instance.
+        """
         collection = pymilvus.Collection(name=collection_name, using=self._using, **kwargs)
         return collection