WIP: make all component parameters in simplekgpipeline

Author: stellasia

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

@@ -244,7 +244,10 @@ def parse(self, resolved_data: dict[str, Any] | None = None) -> Embedder:
         return self.root.parse(resolved_data)
 
 
-class ComponentConfig(ObjectConfig[Component]):
+ComponentT = TypeVar("ComponentT", bound=Component)
+
+
+class ComponentConfig(ObjectConfig[ComponentT]):
     """A config model for all components.
 
     In addition to the object config, components can have pre-defined parameters
@@ -261,14 +264,17 @@ def get_run_params(self, resolved_data: dict[str, Any]) -> dict[str, Any]:
         return self.resolve_params(self.run_params_)
 
 
-class ComponentType(RootModel):  # type: ignore[type-arg]
-    root: Union[Component, ComponentConfig]
+class ComponentType(
+    RootModel[Generic[ComponentT]],
+    # Generic[ComponentT]
+):
+    root: Union[ComponentT, ComponentConfig[ComponentT]]
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    def parse(self, resolved_data: dict[str, Any] | None = None) -> Component:
+    def parse(self, resolved_data: dict[str, Any] | None = None) -> ComponentT:
         if isinstance(self.root, Component):
-            return self.root
+            return self.root  # type: ignore
         return self.root.parse(resolved_data)
 
     def get_run_params(self, resolved_data: dict[str, Any]) -> dict[str, Any]:

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

@@ -20,6 +20,7 @@
 from pydantic import field_validator
 
 from neo4j_graphrag.embeddings import Embedder
+from neo4j_graphrag.experimental.pipeline import Component
 from neo4j_graphrag.experimental.pipeline.config.base import AbstractConfig
 from neo4j_graphrag.experimental.pipeline.config.object_config import (
     ComponentType,
@@ -83,7 +84,7 @@ def validate_embedders(
         return embedders
 
     def _resolve_component_definition(
-        self, name: str, config: ComponentType
+        self, name: str, config: ComponentType[Component]
     ) -> ComponentDefinition:
         component = config.parse(self._global_data)
         if hasattr(config.root, "run_params_"):
@@ -188,7 +189,7 @@ class PipelineConfig(AbstractPipelineConfig):
     """Configuration class for raw pipelines.
     Config must contain all components and connections."""
 
-    component_config: dict[str, ComponentType]
+    component_config: dict[str, ComponentType[Component]]
     connection_config: list[ConnectionDefinition]
     template_: Literal[PipelineType.NONE] = PipelineType.NONE
 

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

@@ -23,7 +23,7 @@
     OnError,
 )
 from neo4j_graphrag.experimental.components.kg_writer import KGWriter, Neo4jWriter
-from neo4j_graphrag.experimental.components.pdf_loader import PdfLoader
+from neo4j_graphrag.experimental.components.pdf_loader import PdfLoader, DataLoader
 from neo4j_graphrag.experimental.components.resolver import (
     EntityResolver,
     SinglePropertyExactMatchResolver,
@@ -81,17 +81,21 @@ class SimpleKGPipelineConfig(TemplatePipelineConfig):
     lexical_graph_config: Optional[LexicalGraphConfig] = None
     neo4j_database: Optional[str] = None
 
-    pdf_loader: Optional[ComponentType] = None
-    kg_writer: Optional[ComponentType] = None
-    text_splitter: Optional[ComponentType] = None
+    pdf_loader: Optional[ComponentType[DataLoader]] = None
+    schema_builder: Optional[ComponentType[SchemaBuilder]] = None
+    text_splitter: Optional[ComponentType[TextSplitter]] = None
+    chunk_embedder: Optional[ComponentType[TextChunkEmbedder]] = None
+    extractor: Optional[ComponentType[EntityRelationExtractor]] = None
+    kg_writer: Optional[ComponentType[KGWriter]] = None
+    resolver: Optional[list[ComponentType[EntityResolver]]] = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    def _get_pdf_loader(self) -> Optional[PdfLoader]:
+    def _get_pdf_loader(self) -> Optional[DataLoader]:
         if not self.from_pdf:
             return None
         if self.pdf_loader:
-            return self.pdf_loader.parse(self._global_data)  # type: ignore
+            return self.pdf_loader.parse(self._global_data)
         return PdfLoader()
 
     def _get_run_params_for_pdf_loader(self) -> dict[str, Any]:
@@ -103,7 +107,7 @@ def _get_run_params_for_pdf_loader(self) -> dict[str, Any]:
 
     def _get_splitter(self) -> TextSplitter:
         if self.text_splitter:
-            return self.text_splitter.parse(self._global_data)  # type: ignore
+            return self.text_splitter.parse(self._global_data)
         return FixedSizeSplitter()
 
     def _get_run_params_for_splitter(self) -> dict[str, Any]:
@@ -112,9 +116,13 @@ def _get_run_params_for_splitter(self) -> dict[str, Any]:
         return {}
 
     def _get_chunk_embedder(self) -> TextChunkEmbedder:
+        if self.chunk_embedder:
+            return self.chunk_embedder.parse(self._global_data)
         return TextChunkEmbedder(embedder=self.get_default_embedder())
 
     def _get_schema(self) -> SchemaBuilder:
+        if self.schema_builder:
+            return self.schema_builder.parse(self._global_data)
         return SchemaBuilder()
 
     def _get_run_params_for_schema(self) -> dict[str, Any]:
@@ -125,6 +133,8 @@ def _get_run_params_for_schema(self) -> dict[str, Any]:
         }
 
     def _get_extractor(self) -> EntityRelationExtractor:
+        if self.extractor:
+            return self.extractor.parse(self._global_data)
         return LLMEntityRelationExtractor(
             llm=self.get_default_llm(),
             prompt_template=self.prompt_template,
@@ -134,7 +144,7 @@ def _get_extractor(self) -> EntityRelationExtractor:
 
     def _get_writer(self) -> KGWriter:
         if self.kg_writer:
-            return self.kg_writer.parse(self._global_data)  # type: ignore
+            return self.kg_writer.parse(self._global_data)
         return Neo4jWriter(
             driver=self.get_default_neo4j_driver(),
             neo4j_database=self.neo4j_database,

src/neo4j_graphrag/experimental/pipeline/kg_builder.py

@@ -21,9 +21,15 @@
 from pydantic import ValidationError
 
 from neo4j_graphrag.embeddings import Embedder
-from neo4j_graphrag.experimental.components.entity_relation_extractor import OnError
+from neo4j_graphrag.experimental.components.embedder import TextChunkEmbedder
+from neo4j_graphrag.experimental.components.entity_relation_extractor import (
+    OnError,
+    EntityRelationExtractor,
+)
 from neo4j_graphrag.experimental.components.kg_writer import KGWriter
 from neo4j_graphrag.experimental.components.pdf_loader import DataLoader
+from neo4j_graphrag.experimental.components.resolver import EntityResolver
+from neo4j_graphrag.experimental.components.schema import SchemaBuilder
 from neo4j_graphrag.experimental.components.text_splitters.base import TextSplitter
 from neo4j_graphrag.experimental.components.types import LexicalGraphConfig
 from neo4j_graphrag.experimental.pipeline.config.object_config import ComponentType
@@ -82,9 +88,13 @@ def __init__(
         relations: Optional[Sequence[RelationInputType]] = None,
         potential_schema: Optional[List[tuple[str, str, str]]] = None,
         from_pdf: bool = True,
-        text_splitter: Optional[TextSplitter] = None,
         pdf_loader: Optional[DataLoader] = None,
+        schema_builder: Optional[SchemaBuilder] = None,
+        text_splitter: Optional[TextSplitter] = None,
+        chunk_embedder: Optional[TextChunkEmbedder] = None,
+        extractor: Optional[EntityRelationExtractor] = None,
         kg_writer: Optional[KGWriter] = None,
+        resolver: Optional[list[EntityResolver]] = None,
         on_error: str = "IGNORE",
         prompt_template: Union[ERExtractionTemplate, str] = ERExtractionTemplate(),
         perform_entity_resolution: bool = True,
@@ -102,8 +112,12 @@ def __init__(
                 potential_schema=potential_schema,
                 from_pdf=from_pdf,
                 pdf_loader=ComponentType(pdf_loader) if pdf_loader else None,
-                kg_writer=ComponentType(kg_writer) if kg_writer 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,
+                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,
                 on_error=OnError(on_error),
                 prompt_template=prompt_template,
                 perform_entity_resolution=perform_entity_resolution,

tests/unit/experimental/pipeline/config/template_pipeline/test_simple_kg_builder.py

@@ -71,7 +71,7 @@ def test_simple_kg_pipeline_config_pdf_loader_class_overwrite_but_from_pdf_is_fa
 def test_simple_kg_pipeline_config_pdf_loader_from_pdf_is_true_class_overwrite_from_config(
     mock_component_parse: Mock,
 ) -> None:
-    my_pdf_loader_config = ComponentConfig(
+    my_pdf_loader_config: ComponentConfig[PdfLoader] = ComponentConfig(
         class_="",
     )
     my_pdf_loader = PdfLoader()
@@ -92,7 +92,7 @@ def test_simple_kg_pipeline_config_text_splitter() -> None:
 def test_simple_kg_pipeline_config_text_splitter_overwrite(
     mock_component_parse: Mock,
 ) -> None:
-    my_text_splitter_config = ComponentConfig(
+    my_text_splitter_config: ComponentConfig[FixedSizeSplitter] = ComponentConfig(
         class_="",
     )
     my_text_splitter = FixedSizeSplitter()
@@ -184,7 +184,7 @@ def test_simple_kg_pipeline_config_writer_overwrite(
     _: Mock,
     driver: neo4j.Driver,
 ) -> None:
-    my_writer_config = ComponentConfig(
+    my_writer_config: ComponentConfig[Neo4jWriter] = ComponentConfig(
         class_="",
     )
     my_writer = Neo4jWriter(driver, neo4j_database="my_db")

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

@@ -366,7 +366,7 @@ def test_abstract_pipeline_config_resolve_component_definition_with_run_params(
     mock_component_parse.return_value = component
     mock_resolve_params.return_value = {"param": "resolver param result"}
     config = AbstractPipelineConfig()
-    component_type = ComponentType(
+    component_type: ComponentType[Component] = ComponentType(
         ComponentConfig(class_="", params_={}, run_params_={"param1": "value1"})
     )
     component_definition = config._resolve_component_definition("name", component_type)