Address review findings: from_stages() factory, chunking fix, observability, NER error handling

- Replace ExtractionConfig(stages=...) with ExtractionConfig.from_stages() factory
  method to prevent misuse of mutually exclusive parameters
- Move chunking above custom stages early return so it applies to both paths
- Add INFO log showing custom pipeline composition
- Use Pydantic PrivateAttr for NER model caching instead of fragile hasattr
- Add logging and error handling for NER model loading
- Simplify from_stages() to call cls() for future-proof attribute inheritance
- Add integration tests for custom pipeline path (chunking + schema injection)
- Update docs and notebook to use from_stages() API

Author: mykola-pereyma

docs-site/src/content/docs/lexical-graph/composable-extraction-pipeline-usage.md

@@ -18,12 +18,12 @@ builder.add(LLMTopicExtractionStage())
 transforms = builder.build()
 ```
 
-Or pass stages directly to `ExtractionConfig`:
+Or use `ExtractionConfig.from_stages()` for a custom composable pipeline:
 
 ```python
 from graphrag_toolkit.lexical_graph.lexical_graph_index import ExtractionConfig, LexicalGraphIndex
 
-config = ExtractionConfig(
+config = ExtractionConfig.from_stages(
     stages=[
         LLMPropositionStage(),
         LLMTopicExtractionStage(),
@@ -94,7 +94,7 @@ schema = ExtractionSchema(
 Add the filter after topic extraction:
 
 ```python
-config = ExtractionConfig(
+config = ExtractionConfig.from_stages(
     stages=[
         LLMPropositionStage(),
         LLMTopicExtractionStage(),
@@ -105,7 +105,7 @@ config = ExtractionConfig(
 
 When `strict=False` (the default), the filter stage passes everything through unchanged — useful for soft constraints where you want the schema for prompt guidance but not hard filtering.
 
-When `schema` is provided to `ExtractionConfig`, it is automatically injected into the `LLMTopicExtractionStage` prompt — guiding the LLM to focus on your entity types and relationships during extraction, not just during filtering.
+When `schema` is provided to `ExtractionConfig.from_stages()`, it is automatically injected into the `LLMTopicExtractionStage` prompt — guiding the LLM to focus on your entity types and relationships during extraction, not just during filtering.
 
 ## Combining NER with LLM Extraction
 
@@ -119,7 +119,7 @@ from graphrag_toolkit.lexical_graph.indexing.extract import (
     EntityMergeStage,
 )
 
-config = ExtractionConfig(
+config = ExtractionConfig.from_stages(
     stages=[
         LLMPropositionStage(),
         NERExtractionStage(
@@ -164,7 +164,7 @@ from graphrag_toolkit.lexical_graph.indexing.extract import (
     LLMTopicExtractionStage,
 )
 
-config = ExtractionConfig(
+config = ExtractionConfig.from_stages(
     stages=[
         LocalPropositionStage(),  # Runs locally, no LLM cost
         LLMTopicExtractionStage(),
@@ -196,7 +196,7 @@ schema = ExtractionSchema(
     strict=True,
 )
 
-config = ExtractionConfig(
+config = ExtractionConfig.from_stages(
     stages=[
         LLMPropositionStage(),
         NERExtractionStage(entity_labels=['Person', 'Organization', 'Technology']),
@@ -243,4 +243,4 @@ The `PipelineBuilder` validates that `input_keys()` are available (from initial
 
 ## Backward Compatibility
 
-When `stages` is not provided to `ExtractionConfig`, the existing default pipeline (LLM propositions → topic extraction) is used unchanged. All existing code continues to work without modification.
+When `ExtractionConfig.from_stages()` is not used, the existing default pipeline (LLM propositions → topic extraction) is used unchanged. All existing code continues to work without modification.

docs-site/src/content/docs/lexical-graph/composable-extraction-pipeline-validation.md

@@ -151,7 +151,7 @@ assert config.stages is None
 assert config.schema is None
 
 # Custom stages
-config = ExtractionConfig(
+config = ExtractionConfig.from_stages(
     stages=[LLMPropositionStage(), LLMTopicExtractionStage()],
     schema=ExtractionSchema(
         entity_types={'Person': EntityTypeConfig()},

examples/lexical-graph-local-dev/notebooks/06-Composable-Extraction-Pipeline.ipynb

@@ -75,22 +75,22 @@
    "outputs": [],
    "source": [
     "# Pattern A: Default \u2014 same as built-in behavior, but explicit\n",
-    "default_pipeline = ExtractionConfig(\n",
+    "default_pipeline = ExtractionConfig.from_stages(\n",
     "    stages=[\n",
     "        LLMPropositionStage(),\n",
     "        LLMTopicExtractionStage(),\n",
     "    ]\n",
     ")\n",
     "\n",
     "# Pattern B: Skip propositions \u2014 extract topics directly from text (faster, cheaper)\n",
-    "direct_pipeline = ExtractionConfig(\n",
+    "direct_pipeline = ExtractionConfig.from_stages(\n",
     "    stages=[\n",
     "        LLMTopicExtractionStage(use_propositions=False),\n",
     "    ]\n",
     ")\n",
     "\n",
     "# Pattern C: Local propositions \u2014 use a local model instead of LLM (free, no API calls)\n",
-    "local_pipeline = ExtractionConfig(\n",
+    "local_pipeline = ExtractionConfig.from_stages(\n",
     "    stages=[\n",
     "        LocalPropositionStage(),\n",
     "        LLMTopicExtractionStage(),\n",
@@ -139,7 +139,7 @@
     "print(schema.format_as_prompt_constraint())\n",
     "\n",
     "# Use it in a pipeline\n",
-    "extraction_config = ExtractionConfig(\n",
+    "extraction_config = ExtractionConfig.from_stages(\n",
     "    stages=[\n",
     "        LLMPropositionStage(),\n",
     "        LLMTopicExtractionStage(),\n",
@@ -208,7 +208,7 @@
     "    strict=True,\n",
     ")\n",
     "\n",
-    "extraction_config = ExtractionConfig(\n",
+    "extraction_config = ExtractionConfig.from_stages(\n",
     "    stages=[\n",
     "        LLMPropositionStage(),                                              # 1. Break text into propositions\n",
     "        NERExtractionStage(                                                 # 2. Find entities locally (free)\n",
@@ -319,7 +319,7 @@
     "    strict=True,\n",
     ")\n",
     "\n",
-    "extraction_config = ExtractionConfig(\n",
+    "extraction_config = ExtractionConfig.from_stages(\n",
     "    stages=[\n",
     "        LLMPropositionStage(),\n",
     "        NERExtractionStage(\n",

lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/ner_extraction_stage.py

@@ -2,10 +2,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-from typing import List, Optional, Sequence
+from typing import Any, List, Optional, Sequence
 
 from llama_index.core.schema import BaseNode, TransformComponent
 from llama_index.core.bridge.pydantic import Field
+from pydantic import PrivateAttr
 
 from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage
 
@@ -20,6 +21,7 @@ class NERTransform(TransformComponent):
     model_name: str = Field(default='urchade/gliner_base')
     entity_labels: List[str] = Field(default_factory=list)
     threshold: float = Field(default=0.5)
+    _model: Any = PrivateAttr(default=None)
 
     def __call__(self, nodes: Sequence[BaseNode], **kwargs) -> Sequence[BaseNode]:
         try:
@@ -30,8 +32,16 @@ def __call__(self, nodes: Sequence[BaseNode], **kwargs) -> Sequence[BaseNode]:
                 'Install it with: pip install gliner'
             )
 
-        if not hasattr(self, '_model'):
-            self._model = GLiNER.from_pretrained(self.model_name)
+        if self._model is None:
+            logger.info('Loading NER model: %s', self.model_name)
+            try:
+                self._model = GLiNER.from_pretrained(self.model_name)
+            except Exception as e:
+                raise RuntimeError(
+                    f'Failed to load NER model "{self.model_name}". '
+                    f'Ensure the model name is correct and you have internet access. '
+                    f'Error: {e}'
+                ) from e
 
         for node in nodes:
             text = node.get_content()

lexical-graph/src/graphrag_toolkit/lexical_graph/lexical_graph_index.py

@@ -57,6 +57,9 @@ class ExtractionConfig():
     and topic or metadata filtering. It provides flexibility in customizing
     the extraction behavior and prompts to tailor it to specific use cases.
 
+    For custom composable pipelines, use the ExtractionConfig.from_stages()
+    class method instead of the constructor.
+
     Attributes:
         enable_proposition_extraction (bool): Determines whether proposition
         extraction is enabled. Defaults to True.
@@ -88,23 +91,40 @@ def __init__(self,
                  extract_propositions_prompt_template: Optional[str] = None,
                  extract_topics_prompt_template: Optional[str] = None,
                  extraction_filters: Optional[MetadataFiltersType] = None,
-                 extraction_llm: Optional[ExtractionLLMType] = None,
-                 stages: Optional[List[ExtractionStage]] = None,
-                 schema: Optional[ExtractionSchema] = None):
+                 extraction_llm: Optional[ExtractionLLMType] = None):
         self.enable_proposition_extraction = enable_proposition_extraction
         self.preferred_entity_classifications = preferred_entity_classifications if preferred_entity_classifications is not None else []
         self.preferred_topics = preferred_topics if preferred_topics is not None else []
         self.infer_entity_classifications = infer_entity_classifications
         self.extract_propositions_prompt_template = extract_propositions_prompt_template
         self.extract_topics_prompt_template = extract_topics_prompt_template
         self.extraction_filters = FilterConfig(extraction_filters)
-        self.stages = stages
-        self.schema = schema
+        self.stages = None
+        self.schema = None
         if extraction_llm is not None:
             self.extraction_llm = extraction_llm if isinstance(extraction_llm, LLMCache) else GraphRAGConfig.to_llm(extraction_llm)
         else:
             self.extraction_llm = None
 
+    @classmethod
+    def from_stages(cls, stages: List['ExtractionStage'], schema: Optional['ExtractionSchema'] = None):
+        """Create config with a custom composable extraction pipeline.
+
+        When using custom stages, configure extraction behavior (LLM, prompts,
+        classifications) directly on the stage instances.
+
+        Args:
+            stages: Ordered list of ExtractionStage instances defining the pipeline.
+            schema: Optional ExtractionSchema for type constraints and filtering.
+
+        Returns:
+            ExtractionConfig configured for custom pipeline execution.
+        """
+        config = cls()
+        config.stages = stages
+        config.schema = schema
+        return config
+
 
 class BuildConfig():
     """
@@ -346,6 +366,10 @@ def _configure_extraction_pipeline(self, config: IndexingConfig):
         pre_processors = []
         components = []
 
+        if config.chunking:
+            for c in config.chunking:
+                components.append(c)
+
         # If custom stages are provided, use PipelineBuilder
         if config.extraction.stages:
             schema = config.extraction.schema
@@ -355,13 +379,13 @@ def _configure_extraction_pipeline(self, config: IndexingConfig):
                 if schema and isinstance(stage, LLMTopicExtractionStage) and stage._schema is None:
                     stage._schema = schema
                 builder.add(stage)
-            components = builder.build()
+            components.extend(builder.build())
+            logger.info(
+                'Custom extraction pipeline: %s',
+                ' → '.join(type(s).__name__ for s in config.extraction.stages)
+            )
             return (pre_processors, components)
 
-        if config.chunking:
-            for c in config.chunking:
-                components.append(c)
-
         if config.extraction.enable_proposition_extraction:
             if config.batch_config:
                 components.append(BatchLLMPropositionExtractorSync(

lexical-graph/tests/unit/test_lexical_graph_index.py

@@ -12,6 +12,10 @@
 from graphrag_toolkit.lexical_graph import ExtractionConfig
 from graphrag_toolkit.lexical_graph.lexical_graph_index import LexicalGraphIndex
 from graphrag_toolkit.lexical_graph.utils.llm_cache import LLMCache
+from graphrag_toolkit.lexical_graph.indexing.extract import (
+    ExtractionStage, ExtractionSchema, EntityTypeConfig,
+    LLMPropositionStage, LLMTopicExtractionStage
+)
 
 class TestExtractionConfig:
 
@@ -46,6 +50,127 @@ def test_extraction_lmm_configured_with_llm_cache_returns_llm_cache(self):
 
         assert extraction_config.extraction_llm == llm_cache
 
+    def test_default_config_has_no_stages(self):
+        config = ExtractionConfig()
+        assert config.stages is None
+        assert config.schema is None
+
+    def test_from_stages_creates_config_with_stages(self):
+        stages = [LLMPropositionStage(), LLMTopicExtractionStage()]
+        config = ExtractionConfig.from_stages(stages=stages)
+        assert config.stages == stages
+        assert config.schema is None
+        assert config.extraction_llm is None
+
+    def test_from_stages_with_schema(self):
+        stages = [LLMPropositionStage(), LLMTopicExtractionStage()]
+        schema = ExtractionSchema(
+            entity_types={'Person': EntityTypeConfig(description='A person')},
+            strict=True,
+        )
+        config = ExtractionConfig.from_stages(stages=stages, schema=schema)
+        assert config.stages == stages
+        assert config.schema == schema
+        assert config.schema.strict is True
+
+class TestCustomPipelineIntegration:
+
+    def test_from_stages_produces_chunking_plus_stage_transforms(self):
+        """Verify custom stages path includes chunking and stage transforms."""
+        from graphrag_toolkit.lexical_graph.indexing.extract import SchemaFilterStage
+        from graphrag_toolkit.lexical_graph.lexical_graph_index import IndexingConfig
+        from llama_index.core.node_parser import SentenceSplitter
+
+        schema = ExtractionSchema(
+            entity_types={'Person': EntityTypeConfig(description='A person')},
+            strict=True,
+        )
+        extraction_config = ExtractionConfig.from_stages(
+            stages=[LLMPropositionStage(), LLMTopicExtractionStage(), SchemaFilterStage(schema=schema)],
+            schema=schema,
+        )
+        indexing_config = IndexingConfig(extraction=extraction_config)
+
+        graph_store = Mock()
+        vector_store = Mock()
+
+        with (
+            patch(
+                'graphrag_toolkit.lexical_graph.lexical_graph_index.GraphStoreFactory.for_graph_store',
+                return_value=graph_store,
+            ),
+            patch(
+                'graphrag_toolkit.lexical_graph.lexical_graph_index.MultiTenantGraphStore.wrap',
+                return_value=graph_store,
+            ),
+            patch(
+                'graphrag_toolkit.lexical_graph.lexical_graph_index.VectorStoreFactory.for_vector_store',
+                return_value=vector_store,
+            ),
+            patch(
+                'graphrag_toolkit.lexical_graph.lexical_graph_index.MultiTenantVectorStore.wrap',
+                return_value=vector_store,
+            ),
+            patch.object(LexicalGraphIndex, '_configure_extraction_pipeline', return_value=([], [])),
+        ):
+            index = LexicalGraphIndex(graph_store, vector_store)
+
+        # Call the real method directly
+        (pre_processors, components) = LexicalGraphIndex._configure_extraction_pipeline(index, indexing_config)
+
+        # Chunking should be first (default SentenceSplitter)
+        assert isinstance(components[0], SentenceSplitter)
+        # 3 stage transforms after chunking
+        assert len(components) == 4
+        assert len(pre_processors) == 0
+
+    def test_from_stages_schema_injected_into_topic_stage(self):
+        """Verify schema is auto-injected into LLMTopicExtractionStage."""
+        from graphrag_toolkit.lexical_graph.lexical_graph_index import IndexingConfig
+
+        schema = ExtractionSchema(
+            entity_types={'Person': EntityTypeConfig(description='A human being')},
+        )
+        topic_stage = LLMTopicExtractionStage()
+        assert topic_stage._schema is None
+
+        extraction_config = ExtractionConfig.from_stages(
+            stages=[LLMPropositionStage(), topic_stage],
+            schema=schema,
+        )
+        indexing_config = IndexingConfig(extraction=extraction_config)
+
+        graph_store = Mock()
+        vector_store = Mock()
+
+        with (
+            patch(
+                'graphrag_toolkit.lexical_graph.lexical_graph_index.GraphStoreFactory.for_graph_store',
+                return_value=graph_store,
+            ),
+            patch(
+                'graphrag_toolkit.lexical_graph.lexical_graph_index.MultiTenantGraphStore.wrap',
+                return_value=graph_store,
+            ),
+            patch(
+                'graphrag_toolkit.lexical_graph.lexical_graph_index.VectorStoreFactory.for_vector_store',
+                return_value=vector_store,
+            ),
+            patch(
+                'graphrag_toolkit.lexical_graph.lexical_graph_index.MultiTenantVectorStore.wrap',
+                return_value=vector_store,
+            ),
+            patch.object(LexicalGraphIndex, '_configure_extraction_pipeline', return_value=([], [])),
+        ):
+            index = LexicalGraphIndex(graph_store, vector_store)
+
+        # Call the real method directly
+        LexicalGraphIndex._configure_extraction_pipeline(index, indexing_config)
+
+        # Schema should be injected into the topic stage
+        assert topic_stage._schema == schema
+
+
 class TestLexicalGraphIndex:
 
     def test_init_invokes_graph_store_init_hook(self):