Removed register stage decorator

morpheus/service/milvus_vector_db_service.py

@@ -255,7 +255,9 @@ def _collection_insert(self, name: str, data: list[list] | list[dict],
         return result_dict
 
     @with_collection_lock
-    def insert_dataframe(self, name: str, df: cudf.DataFrame | pd.DataFrame,
+    def insert_dataframe(self,
+                         name: str,
+                         df: typing.Union[cudf.DataFrame, pd.DataFrame],
                          **kwargs: dict[str, typing.Any]) -> dict[str, typing.Any]:
         """
         Converts dataframe to rows and insert to a collection in the Milvus vector database.
@@ -264,7 +266,7 @@ def insert_dataframe(self, name: str, df: cudf.DataFrame | pd.DataFrame,
         ----------
         name : str
             Name of the collection to be inserted.
-        df : cudf.DataFrame | pd.DataFrame
+        df : typing.Union[cudf.DataFrame, pd.DataFrame]
             Dataframe to be inserted in the collection.
         **kwargs : dict[str, typing.Any]
             Additional keyword arguments containing collection configuration.

morpheus/service/vector_db_service.py

@@ -52,15 +52,18 @@ def insert(self, name: str, data: list[list] | list[dict], **kwargs: dict[str, t
         pass
 
     @abstractmethod
-    def insert_dataframe(self, name: str, df: cudf.DataFrame | pd.DataFrame, **kwargs: dict[str, typing.Any]) -> dict:
+    def insert_dataframe(self,
+                         name: str,
+                         df: typing.Union[cudf.DataFrame, pd.DataFrame],
+                         **kwargs: dict[str, typing.Any]) -> dict:
         """
         Converts dataframe to rows and insert into the vector database resource.
 
         Parameters
         ----------
         name : str
             Name of the resource to be inserted.
-        df : cudf.DataFrame | pd.DataFrame
+        df : typing.Union[cudf.DataFrame, pd.DataFrame]
             Dataframe to be inserted.
         **kwargs : dict[str, typing.Any]
             Additional keyword arguments containing collection configuration.

morpheus/stages/output/write_to_vector_db.py

@@ -18,7 +18,6 @@
 import mrc
 from mrc.core import operators as ops
 
-from morpheus.cli import register_stage
 from morpheus.config import Config
 from morpheus.messages import ControlMessage
 from morpheus.pipeline.single_port_stage import SinglePortStage
@@ -29,7 +28,6 @@
 logger = logging.getLogger(__name__)
 
 
-@register_stage("to-vector-db")
 class WriteToVectorDBStage(SinglePortStage):
     """
     Writes messages to a Vector Database.