fix: bound automatic schema extraction input

CHANGELOG.md

@@ -7,6 +7,10 @@
 - Experimental: the schema-from-text extraction prompt now instructs the LLM to define each relationship type once and reuse it across patterns, using distinct type names only when patterns genuinely need different properties or constraints.
 - Experimental: LLM-auto-generated schemas now reconcile duplicate `relationship_types` (entries sharing the same label) by merging them into a single type that carries the union of their properties, emitting a warning log. This reflects that Neo4j relationship types are global per name.
 
+### Fixed
+
+- Experimental: `SimpleKGPipeline` automatic schema extraction now uses splitter chunks instead of passing the full file or inline text directly to the schema LLM prompt.
+
 ## 1.17.0
 
 ### Added

src/neo4j_graphrag/experimental/components/schema.py

@@ -54,6 +54,7 @@
     SchemaExtractionError,
     SchemaValidationError,
 )
+from neo4j_graphrag.experimental.components.types import TextChunk
 from neo4j_graphrag.experimental.pipeline.component import Component, DataModel
 from neo4j_graphrag.experimental.pipeline.types.schema import (
     EntityInputType,
@@ -1726,17 +1727,21 @@ async def _run_with_prompt_based_extraction(self, prompt: str) -> GraphSchema:
         return validate_extraction_dict_to_graph_schema(extracted_schema)
 
     @validate_call
-    async def run(self, text: str, examples: str = "", **kwargs: Any) -> GraphSchema:
+    async def run(
+        self, text: Union[str, List[TextChunk]], examples: str = "", **kwargs: Any
+    ) -> GraphSchema:
         """
         Asynchronously extracts the schema from text and returns a GraphSchema object.
 
         Args:
-            text (str): the text from which the schema will be inferred.
+            text (str | List[TextChunk]): the text from which the schema will be inferred.
             examples (str): examples to guide schema extraction.
         Returns:
             GraphSchema: A configured schema object, extracted automatically and
             constructed asynchronously.
         """
+        if isinstance(text, list):
+            text = text[0].text if text else ""
         prompt: str = self._prompt_template.format(text=text, examples=examples)
 
         if self.use_structured_output:

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

@@ -337,9 +337,9 @@ def _get_connections(self) -> list[ConnectionDefinition]:
             if not self.has_user_provided_schema():
                 connections.append(
                     ConnectionDefinition(
-                        start="file_loader",
+                        start="splitter",
                         end="schema",
-                        input_config={"text": "file_loader.text"},
+                        input_config={"text": "splitter.chunks"},
                     )
                 )
 
@@ -354,6 +354,14 @@ def _get_connections(self) -> list[ConnectionDefinition]:
                 )
             )
         else:
+            if not self.has_user_provided_schema():
+                connections.append(
+                    ConnectionDefinition(
+                        start="splitter",
+                        end="schema",
+                        input_config={"text": "splitter.chunks"},
+                    )
+                )
             connections.append(
                 ConnectionDefinition(
                     start="schema",
@@ -436,9 +444,6 @@ def get_run_params(self, user_input: dict[str, Any]) -> dict[str, Any]:
                     "Expected 'text' argument when 'from_file' is False."
                 )
             run_params["splitter"]["text"] = text
-            # Add full text to schema component for automatic schema extraction
-            if not self.has_user_provided_schema():
-                run_params["schema"]["text"] = text
             run_params["extractor"]["document_info"] = dict(
                 path=user_input.get(
                     "file_path",

tests/unit/experimental/components/test_schema.py

@@ -37,6 +37,7 @@
     validate_extraction_dict_to_graph_schema,
     _merge_duplicate_relationship_types,
 )
+from neo4j_graphrag.experimental.components.types import TextChunk
 import os
 import tempfile
 import yaml
@@ -1313,6 +1314,26 @@ async def test_schema_from_text_run_valid_response(
     assert schema_config.patterns[0] == ("Person", "WORKS_FOR", "Organization")
 
 
+@pytest.mark.asyncio
+async def test_schema_from_text_run_uses_first_chunk(
+    schema_from_text: SchemaFromTextExtractor,
+    mock_llm: AsyncMock,
+    valid_schema_json: str,
+) -> None:
+    mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json)
+
+    await schema_from_text.run(
+        text=[
+            TextChunk(text="first chunk for schema", index=0),
+            TextChunk(text="second chunk should not be used", index=1),
+        ]
+    )
+
+    prompt_arg = mock_llm.ainvoke.call_args[0][0]
+    assert "first chunk for schema" in prompt_arg
+    assert "second chunk should not be used" not in prompt_arg
+
+
 @pytest.mark.asyncio
 async def test_schema_from_text_run_invalid_json(
     schema_from_text: SchemaFromTextExtractor,

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

@@ -313,7 +313,7 @@ def test_simple_kg_pipeline_config_connections_from_file() -> None:
     assert len(connections) == 7
     expected_connections = [
         ("file_loader", "splitter"),
-        ("file_loader", "schema"),
+        ("splitter", "schema"),
         ("schema", "extractor"),
         ("splitter", "chunk_embedder"),
         ("chunk_embedder", "extractor"),
@@ -324,14 +324,31 @@ def test_simple_kg_pipeline_config_connections_from_file() -> None:
         assert (actual.start, actual.end) == expected
 
 
+def test_simple_kg_pipeline_config_automatic_schema_uses_splitter_from_file() -> None:
+    config = SimpleKGPipelineConfig(
+        from_file=True,
+        text_splitter=FixedSizeSplitter(chunk_size=25, chunk_overlap=0),
+    )
+    connections = config._get_connections()
+    schema_edges = [
+        c
+        for c in connections
+        if c.end == "schema" and c.input_config.get("text") is not None
+    ]
+    assert len(schema_edges) == 1
+    assert schema_edges[0].start == "splitter"
+    assert schema_edges[0].input_config == {"text": "splitter.chunks"}
+
+
 def test_simple_kg_pipeline_config_connections_from_text() -> None:
     config = SimpleKGPipelineConfig(
         from_file=False,
         perform_entity_resolution=False,
     )
     connections = config._get_connections()
-    assert len(connections) == 5
+    assert len(connections) == 6
     expected_connections = [
+        ("splitter", "schema"),
         ("schema", "extractor"),
         ("splitter", "chunk_embedder"),
         ("chunk_embedder", "extractor"),
@@ -351,7 +368,7 @@ def test_simple_kg_pipeline_config_connections_with_er() -> None:
     assert len(connections) == 8
     expected_connections = [
         ("file_loader", "splitter"),
-        ("file_loader", "schema"),
+        ("splitter", "schema"),
         ("schema", "extractor"),
         ("splitter", "chunk_embedder"),
         ("chunk_embedder", "extractor"),
@@ -374,7 +391,7 @@ def test_simple_kg_pipeline_config_run_params_from_text_text() -> None:
     config = SimpleKGPipelineConfig(from_file=False)
     run_params = config.get_run_params({"text": "my text"})
     assert run_params["splitter"] == {"text": "my text"}
-    assert run_params["schema"] == {"text": "my text"}
+    assert "schema" not in run_params
     assert run_params["extractor"]["document_info"]["path"] == "document.txt"