update register name

OpenSPG · Dec 27, 2024 · 192f213 · 192f213
1 parent 26fdbc0
commit 192f213
Show file tree

Hide file tree

Showing 8 changed files with 8 additions and 0 deletions.
diff --git a/kag/builder/component/extractor/schema_constraint_extractor.py b/kag/builder/component/extractor/schema_constraint_extractor.py
@@ -30,6 +30,7 @@
 logger = logging.getLogger(__name__)
 
 
+@ExtractorABC.register("schema_constraint")
 @ExtractorABC.register("schema_constraint_extractor")
 class SchemaConstraintExtractor(ExtractorABC):
     """

diff --git a/kag/builder/component/postprocessor/kag_postprocessor.py b/kag/builder/component/postprocessor/kag_postprocessor.py
@@ -24,6 +24,7 @@
 logger = logging.getLogger()
 
 
+@PostProcessorABC.register("base", as_default=True)
 @PostProcessorABC.register("kag_post_processor", as_default=True)
 class KAGPostProcessor(PostProcessorABC):
     """

diff --git a/kag/builder/component/reader/txt_reader.py b/kag/builder/component/reader/txt_reader.py
@@ -19,6 +19,7 @@
 from knext.common.base.runnable import Input, Output
 
 
+@ReaderABC.register("txt")
 @ReaderABC.register("txt_reader")
 class TXTReader(ReaderABC):
     """

diff --git a/kag/builder/component/scanner/directory_scanner.py b/kag/builder/component/scanner/directory_scanner.py
@@ -18,6 +18,7 @@
 from knext.common.base.runnable import Input, Output
 
 
+@ScannerABC.register("dir")
 @ScannerABC.register("dir_file_scanner")
 class DirectoryScanner(ScannerABC):
     """

diff --git a/kag/builder/component/splitter/length_splitter.py b/kag/builder/component/splitter/length_splitter.py
@@ -19,6 +19,7 @@
 from kag.builder.component.splitter.base_table_splitter import BaseTableSplitter
 
 
+@SplitterABC.register("length")
 @SplitterABC.register("length_splitter")
 class LengthSplitter(BaseTableSplitter):
     """

diff --git a/kag/builder/component/vectorizer/batch_vectorizer.py b/kag/builder/component/vectorizer/batch_vectorizer.py
@@ -129,6 +129,7 @@ def batch_generate(self, node_batch, batch_size=32):
         manager.patch()
 
 
+@VectorizerABC.register("batch")
 @VectorizerABC.register("batch_vectorizer")
 class BatchVectorizer(VectorizerABC):
     """

diff --git a/kag/builder/component/writer/kg_writer.py b/kag/builder/component/writer/kg_writer.py
@@ -28,6 +28,7 @@ class AlterOperationEnum(str, Enum):
     Delete = "DELETE"
 
 
+@SinkWriterABC.register("kg", as_default=True)
 @SinkWriterABC.register("kg_writer", as_default=True)
 class KGWriter(SinkWriterABC):
     """

diff --git a/kag/builder/default_chain.py b/kag/builder/default_chain.py
@@ -84,6 +84,7 @@ def build(self, **kwargs):
     #             node.checkpointer.close()
 
 
+@KAGBuilderChain.register("unstructured")
 @KAGBuilderChain.register("unstructured_builder_chain")
 class DefaultUnstructuredBuilderChain(KAGBuilderChain):
     """