From df19785aedd95d0e44c6842f6ca97086701db0d0 Mon Sep 17 00:00:00 2001 From: Khanh Le Date: Mon, 15 Jun 2026 20:28:53 -0500 Subject: [PATCH] fix: bound automatic schema extraction input --- CHANGELOG.md | 4 +++ .../experimental/components/schema.py | 9 +++++-- .../template_pipeline/simple_kg_builder.py | 15 +++++++---- .../experimental/components/test_schema.py | 21 ++++++++++++++++ .../test_simple_kg_builder.py | 25 ++++++++++++++++--- 5 files changed, 63 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b53ff38a9..2c032fa6f 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/neo4j_graphrag/experimental/components/schema.py b/src/neo4j_graphrag/experimental/components/schema.py index d6acaef54..136777319 100644 --- a/src/neo4j_graphrag/experimental/components/schema.py +++ b/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: diff --git a/src/neo4j_graphrag/experimental/pipeline/config/template_pipeline/simple_kg_builder.py b/src/neo4j_graphrag/experimental/pipeline/config/template_pipeline/simple_kg_builder.py index ec6a1ffb1..1188ab7fd 100644 --- a/src/neo4j_graphrag/experimental/pipeline/config/template_pipeline/simple_kg_builder.py +++ b/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", diff --git a/tests/unit/experimental/components/test_schema.py b/tests/unit/experimental/components/test_schema.py index 696216d32..2fc754bc2 100644 --- a/tests/unit/experimental/components/test_schema.py +++ b/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, diff --git a/tests/unit/experimental/pipeline/config/template_pipeline/test_simple_kg_builder.py b/tests/unit/experimental/pipeline/config/template_pipeline/test_simple_kg_builder.py index 57f76552f..589fe0df0 100644 --- a/tests/unit/experimental/pipeline/config/template_pipeline/test_simple_kg_builder.py +++ b/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"