Expose all components in SimpleKGPipeline

Author: stellasia

src/neo4j_graphrag/experimental/pipeline/config/object_config.py

@@ -244,10 +244,10 @@ def parse(self, resolved_data: dict[str, Any] | None = None) -> Embedder:
         return self.root.parse(resolved_data)
 
 
-ComponentT = TypeVar("ComponentT", bound=Component)
+ComponentGeneric = TypeVar("ComponentGeneric")
 
 
-class ComponentConfig(ObjectConfig[ComponentT]):
+class ComponentConfig(ObjectConfig[ComponentGeneric], Generic[ComponentGeneric]):
     """A config model for all components.
 
     In addition to the object config, components can have pre-defined parameters
@@ -259,25 +259,27 @@ class ComponentConfig(ObjectConfig[ComponentT]):
     DEFAULT_MODULE = "neo4j_graphrag.experimental.components"
     INTERFACE = Component
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     def get_run_params(self, resolved_data: dict[str, Any]) -> dict[str, Any]:
         self._global_data = resolved_data
         return self.resolve_params(self.run_params_)
 
 
 class ComponentType(
-    RootModel[Generic[ComponentT]],
-    # Generic[ComponentT]
+    RootModel[Union[ComponentGeneric, ComponentConfig[ComponentGeneric]]],
+    Generic[ComponentGeneric],
 ):
-    root: Union[ComponentT, ComponentConfig[ComponentT]]
+    root: Union[ComponentGeneric, ComponentConfig[ComponentGeneric]]
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    def parse(self, resolved_data: dict[str, Any] | None = None) -> ComponentT:
-        if isinstance(self.root, Component):
-            return self.root  # type: ignore
-        return self.root.parse(resolved_data)
+    def parse(self, resolved_data: dict[str, Any] | None = None) -> ComponentGeneric:
+        if isinstance(self.root, ComponentConfig):
+            return self.root.parse(resolved_data)
+        return self.root
 
     def get_run_params(self, resolved_data: dict[str, Any]) -> dict[str, Any]:
-        if isinstance(self.root, Component):
-            return {}
-        return self.root.get_run_params(resolved_data)
+        if isinstance(self.root, ComponentConfig):
+            return self.root.get_run_params(resolved_data)
+        return {}

src/neo4j_graphrag/experimental/pipeline/config/template_pipeline/simple_kg_builder.py

@@ -87,7 +87,7 @@ class SimpleKGPipelineConfig(TemplatePipelineConfig):
     chunk_embedder: Optional[ComponentType[TextChunkEmbedder]] = None
     extractor: Optional[ComponentType[EntityRelationExtractor]] = None
     kg_writer: Optional[ComponentType[KGWriter]] = None
-    resolver: Optional[list[ComponentType[EntityResolver]]] = None
+    resolver: Optional[ComponentType[EntityResolver]] = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -158,6 +158,8 @@ def _get_run_params_for_writer(self) -> dict[str, Any]:
     def _get_resolver(self) -> Optional[EntityResolver]:
         if not self.perform_entity_resolution:
             return None
+        if self.resolver:
+            return self.resolver.parse(self._global_data)
         return SinglePropertyExactMatchResolver(
             driver=self.get_default_neo4j_driver(),
             neo4j_database=self.neo4j_database,

src/neo4j_graphrag/experimental/pipeline/kg_builder.py

@@ -94,7 +94,7 @@ def __init__(
         chunk_embedder: Optional[TextChunkEmbedder] = None,
         extractor: Optional[EntityRelationExtractor] = None,
         kg_writer: Optional[KGWriter] = None,
-        resolver: Optional[list[EntityResolver]] = None,
+        resolver: Optional[EntityResolver] = None,
         on_error: str = "IGNORE",
         prompt_template: Union[ERExtractionTemplate, str] = ERExtractionTemplate(),
         perform_entity_resolution: bool = True,
@@ -112,12 +112,16 @@ def __init__(
                 potential_schema=potential_schema,
                 from_pdf=from_pdf,
                 pdf_loader=ComponentType(pdf_loader) if pdf_loader else None,
-                schema_builder=ComponentType(schema_builder) if schema_builder else None,
+                schema_builder=ComponentType(schema_builder)
+                if schema_builder
+                else None,
                 text_splitter=ComponentType(text_splitter) if text_splitter else None,
-                chunk_embedder=ComponentType(chunk_embedder) if chunk_embedder else None,
+                chunk_embedder=ComponentType(chunk_embedder)
+                if chunk_embedder
+                else None,
                 extractor=ComponentType(extractor) if extractor else None,
                 kg_writer=ComponentType(kg_writer) if kg_writer else None,
-                resolver=[ComponentType(r) for r in resolver] if resolver else None,
+                resolver=ComponentType(resolver) if resolver else None,
                 on_error=OnError(on_error),
                 prompt_template=prompt_template,
                 perform_entity_resolution=perform_entity_resolution,

tests/unit/experimental/pipeline/config/test_pipeline_config.py

@@ -345,7 +345,7 @@ def test_abstract_pipeline_config_resolve_component_definition_no_run_params(
 ) -> None:
     mock_component_parse.return_value = component
     config = AbstractPipelineConfig()
-    component_type = ComponentType(component)
+    component_type: ComponentType[Component] = ComponentType(component)
     component_definition = config._resolve_component_definition("name", component_type)
     assert isinstance(component_definition, ComponentDefinition)
     mock_component_parse.assert_called_once_with({})