Updated Milvus VDB tests

morpheus/service/milvus_vector_db_service.py

@@ -60,7 +60,7 @@ def wrapper(self, name, *args, **kwargs):
 
 class MilvusVectorDBService(VectorDBService):
     """
-    Service class for Milvus Vector Database implementation. This class provides methods for interacting
+    Service class for Milvus Vector Database implementation. This class provides functions for interacting
     with a Milvus vector database.
 
     Parameters
@@ -219,7 +219,6 @@ def create(self, name: str, overwrite: bool = False, **kwargs: dict[str, typing.
             schema = pymilvus.CollectionSchema(fields=schema_fields, **schema_conf)
 
             if index_conf:
-                # Preserve original index configuration.
                 field_name = index_conf.pop("field_name")
                 metric_type = index_conf.pop("metric_type")
                 index_param = self._client.prepare_index_params(field_name=field_name,
@@ -241,8 +240,7 @@ def create(self, name: str, overwrite: bool = False, **kwargs: dict[str, typing.
                     self._client.create_partition(collection_name=name, partition_name=part["name"], timeout=timeout)
 
     @with_collection_lock
-    def insert(self, name: str, data: typing.Union[list[list], list[dict], dict], **kwargs: dict[str,
-                                                                                                 typing.Any]) -> dict:
+    def insert(self, name: str, data: typing.Any, **kwargs: dict[str, typing.Any]) -> dict:
         """
         Insert a collection specific data in the Milvus vector database.
 
@@ -265,6 +263,7 @@ def insert(self, name: str, data: typing.Union[list[list], list[dict], dict], **
         RuntimeError
             If the collection not exists exists.
         """
+
         return self._collection_insert(name, data, **kwargs)
 
     def _collection_insert(self,
@@ -404,6 +403,11 @@ def update(self, name: str, data: list[dict], **kwargs: dict[str, typing.Any]) -
             Data to be updated in the resource.
         **kwargs : dict[str, typing.Any]
             Extra keyword arguments specific to upsert operation.
+
+        Returns
+        -------
+        dict
+            Returns result of the updated operation stats.
         """
 
         if not isinstance(data, list):
@@ -440,7 +444,7 @@ def delete_by_keys(self, name: str, keys: typing.Union[int, str, list], **kwargs
         Returns
         -------
         typing.Any
-            Returns vectors of the given keys that are delete from the resource.
+            Returns result of the given keys that are delete from the collection.
         """
 
         response = self._client.delete(collection_name=name, pks=keys, **kwargs)
@@ -464,7 +468,7 @@ def delete(self, name: str, expr: str, **kwargs: dict[str, typing.Any]) -> typin
         Returns
         -------
         typing.Any
-            Returns vectors of the given keys that are delete from the resource.
+            Returns result of the given keys that are delete from the collection.
         """
 
         return self._client.delete_by_expr(collection_name=name, expression=expr, **kwargs)
@@ -484,6 +488,11 @@ def retrieve_by_keys(self, name: str, keys: typing.Union[int, str, list], **kwar
             or a list of either.
         **kwargs :  dict[str, typing.Any]
             Additional keyword arguments for the retrieval operation.
+
+        Returns
+        -------
+        list[dict]
+            Returns result rows of the given keys from the collection.
         """
 
         result = None
@@ -509,6 +518,11 @@ def count(self, name: str, **kwargs: dict[str, typing.Any]) -> int:
             Name of the collection.
         **kwargs :  dict[str, typing.Any]
             Additional keyword arguments for the count operation.
+
+        Returns
+        -------
+        int
+            Returns number of entities in the collection.
         """
 
         return self._client.num_entities(collection_name=name, **kwargs)
@@ -575,6 +589,11 @@ def describe(self, name: str, **kwargs: dict[str, typing.Any]) -> dict:
             Name of the collection.
         **kwargs : dict[str, typing.Any]
             Additional keyword arguments specific to the Milvus vector database.
+
+        Returns
+        -------
+        dict
+            Returns collection information.
         """
 
         return self._client.describe_collection(collection_name=name, **kwargs)

morpheus/service/vector_db_service.py

@@ -118,7 +118,7 @@ def drop(self, name: str, **kwargs: dict[str, typing.Any]) -> None:
         pass
 
     @abstractmethod
-    def update(self, name: str, data: typing.Any, **kwargs: dict[str, typing.Any]) -> None:
+    def update(self, name: str, data: typing.Any, **kwargs: dict[str, typing.Any]) -> dict:
         """
         Update data in the vector database.
 
@@ -130,6 +130,11 @@ def update(self, name: str, data: typing.Any, **kwargs: dict[str, typing.Any]) -
             Data to be updated in the resource.
         **kwargs : dict[str, typing.Any]
             Extra keyword arguments specific to the vector database implementation.
+
+        Returns
+        -------
+        dict
+            Returns result of the updated operation stats.
         """
 
         pass
@@ -179,6 +184,11 @@ def describe(self, name: str, **kwargs: dict[str, typing.Any]) -> dict:
             Name of the resource.
         **kwargs : dict[str, typing.Any]
             Extra keyword arguments specific to the vector database implementation.
+
+        Returns
+        -------
+        dict
+            Returns resource information.
         """
 
         pass
@@ -263,7 +273,7 @@ def retrieve_by_keys(self, name: str, keys: typing.Any, **kwargs: dict[str, typi
         Returns
         -------
         typing.Any
-            Returns inserted vectors of the given keys that exists in the resource.
+            Returns rows of the given keys that exists in the resource.
         """
         pass