From 72be3ef31815e49fd314d268cdca72c03127c527 Mon Sep 17 00:00:00 2001 From: Ashley Dunfield Date: Wed, 18 Mar 2026 01:35:07 -0600 Subject: [PATCH 1/3] fix(kg-pipeline): always create entity-to-chunk relationships regardless of create_lexical_graph MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When LLMEntityRelationExtractor was run with create_lexical_graph=False and no lexical_graph_config, lexical_graph_builder was left as None. This caused post_process_chunk() to skip process_chunk_extracted_entities(), so FROM_CHUNK relationships between extracted entities and their source chunks were never created, silently breaking downstream graph queries that traverse entity→chunk. Fix: always instantiate LexicalGraphBuilder (with default or provided config) before the conditional block. The create_lexical_graph flag continues to control whether document/chunk nodes and NEXT_CHUNK/FROM_DOCUMENT relationships are added — it no longer suppresses entity→chunk relationship creation. Also updates the existing test that was asserting the (now-fixed) buggy behavior of relationships == [], and adds three new regression tests. Closes #470 --- .../components/entity_relation_extractor.py | 7 +- .../test_entity_relation_extractor.py | 91 ++++++++++++++++++- 2 files changed, 92 insertions(+), 6 deletions(-) diff --git a/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py b/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py index d71553451..197ddfe8f 100644 --- a/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py +++ b/src/neo4j_graphrag/experimental/components/entity_relation_extractor.py @@ -358,17 +358,14 @@ async def run( schema (GraphSchema | None): Definition of the schema to guide the LLM in its extraction. examples (str): Examples for few-shot learning in the prompt. """ - lexical_graph_builder = None + config = lexical_graph_config or LexicalGraphConfig() + lexical_graph_builder = LexicalGraphBuilder(config=config) lexical_graph = None if self.create_lexical_graph: - config = lexical_graph_config or LexicalGraphConfig() - lexical_graph_builder = LexicalGraphBuilder(config=config) lexical_graph_result = await lexical_graph_builder.run( text_chunks=chunks, document_info=document_info ) lexical_graph = lexical_graph_result.graph - elif lexical_graph_config: - lexical_graph_builder = LexicalGraphBuilder(config=lexical_graph_config) schema = schema or GraphSchema( node_types=(), ) diff --git a/tests/unit/experimental/components/test_entity_relation_extractor.py b/tests/unit/experimental/components/test_entity_relation_extractor.py index 4d9f6f298..fe1942efd 100644 --- a/tests/unit/experimental/components/test_entity_relation_extractor.py +++ b/tests/unit/experimental/components/test_entity_relation_extractor.py @@ -215,7 +215,10 @@ async def test_extractor_llm_badly_formatted_json_gets_fixed() -> None: assert len(res.nodes) == 1 assert res.nodes[0].label == "Person" assert res.nodes[0].embedding_properties == {} - assert res.relationships == [] + # With create_lexical_graph=False a FROM_CHUNK relationship is still created + # so that extracted entities remain linked to the chunk they came from. + from_chunk_rels = [r for r in res.relationships if r.type == "FROM_CHUNK"] + assert len(from_chunk_rels) == 1 @pytest.mark.asyncio @@ -490,3 +493,89 @@ async def test_extractor_structured_output_false_uses_v1() -> None: call_args = llm.ainvoke.call_args assert isinstance(call_args[0][0], str) # First arg is string prompt assert "response_format" not in call_args[1] # No response_format kwarg + + +# --------------------------------------------------------------------------- +# Tests for entity-to-chunk relationship creation (issue #470) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_entity_to_chunk_rels_created_when_lexical_graph_disabled() -> None: + """ + Regression test for #470. + + When create_lexical_graph=False and no lexical_graph_config is provided, + entity-to-chunk (FROM_CHUNK) relationships must still be created. + Previously lexical_graph_builder was left as None in this path, so + process_chunk_extracted_entities() was never called and FROM_CHUNK + relationships were silently missing. + """ + llm = MagicMock(spec=LLMInterface) + llm.ainvoke.return_value = LLMResponse( + content='{"nodes": [{"id": "p1", "label": "Person", "properties": {"name": "Alice"}}, ' + '{"id": "o1", "label": "Organization", "properties": {"name": "Acme"}}], ' + '"relationships": [{"start_node_id": "p1", "end_node_id": "o1", "type": "WORKS_FOR"}]}' + ) + extractor = LLMEntityRelationExtractor( + llm=llm, + on_error=OnError.RAISE, + create_lexical_graph=False, + ) + chunks = TextChunks(chunks=[TextChunk(text="Alice works at Acme.", index=0)]) + res = await extractor.run(chunks=chunks) + + from_chunk_rels = [r for r in res.relationships if r.type == "FROM_CHUNK"] + assert len(from_chunk_rels) == 2, ( + f"Expected 2 FROM_CHUNK relationships (one per entity), got {len(from_chunk_rels)}. " + "This is the regression for issue #470." + ) + + +@pytest.mark.asyncio +async def test_no_chunk_nodes_when_lexical_graph_disabled() -> None: + """ + When create_lexical_graph=False, Chunk and Document nodes must NOT appear + in the graph — only extracted entities and their FROM_CHUNK links. + """ + llm = MagicMock(spec=LLMInterface) + llm.ainvoke.return_value = LLMResponse( + content='{"nodes": [{"id": "p1", "label": "Person", "properties": {}}], "relationships": []}' + ) + extractor = LLMEntityRelationExtractor( + llm=llm, + on_error=OnError.RAISE, + create_lexical_graph=False, + ) + chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)]) + res = await extractor.run(chunks=chunks) + + chunk_nodes = [n for n in res.nodes if n.label == "Chunk"] + doc_nodes = [n for n in res.nodes if n.label == "Document"] + assert chunk_nodes == [], "Chunk nodes must not appear when create_lexical_graph=False" + assert doc_nodes == [], "Document nodes must not appear when create_lexical_graph=False" + + +@pytest.mark.asyncio +async def test_entity_to_chunk_rels_created_with_custom_lexical_graph_config() -> None: + """ + When a custom lexical_graph_config is supplied, FROM_CHUNK relationships + must use the custom relationship type. + """ + from neo4j_graphrag.experimental.components.types import LexicalGraphConfig + + llm = MagicMock(spec=LLMInterface) + llm.ainvoke.return_value = LLMResponse( + content='{"nodes": [{"id": "p1", "label": "Person", "properties": {}}], "relationships": []}' + ) + extractor = LLMEntityRelationExtractor( + llm=llm, + on_error=OnError.RAISE, + create_lexical_graph=False, + ) + config = LexicalGraphConfig(node_to_chunk_relationship_type="EXTRACTED_FROM") + chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)]) + res = await extractor.run(chunks=chunks, lexical_graph_config=config) + + custom_rels = [r for r in res.relationships if r.type == "EXTRACTED_FROM"] + assert len(custom_rels) == 1 From 58fbabdeb4dddece3147422c389ba2519146979a Mon Sep 17 00:00:00 2001 From: AmirLayegh Date: Tue, 30 Jun 2026 16:53:59 +0200 Subject: [PATCH 2/3] Fix failing tests --- .../components/test_entity_relation_extractor.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/unit/experimental/components/test_entity_relation_extractor.py b/tests/unit/experimental/components/test_entity_relation_extractor.py index fe1942efd..8745b761a 100644 --- a/tests/unit/experimental/components/test_entity_relation_extractor.py +++ b/tests/unit/experimental/components/test_entity_relation_extractor.py @@ -514,8 +514,8 @@ async def test_entity_to_chunk_rels_created_when_lexical_graph_disabled() -> Non llm = MagicMock(spec=LLMInterface) llm.ainvoke.return_value = LLMResponse( content='{"nodes": [{"id": "p1", "label": "Person", "properties": {"name": "Alice"}}, ' - '{"id": "o1", "label": "Organization", "properties": {"name": "Acme"}}], ' - '"relationships": [{"start_node_id": "p1", "end_node_id": "o1", "type": "WORKS_FOR"}]}' + '{"id": "o1", "label": "Organization", "properties": {"name": "Acme"}}], ' + '"relationships": [{"start_node_id": "p1", "end_node_id": "o1", "type": "WORKS_FOR"}]}' ) extractor = LLMEntityRelationExtractor( llm=llm, @@ -552,8 +552,12 @@ async def test_no_chunk_nodes_when_lexical_graph_disabled() -> None: chunk_nodes = [n for n in res.nodes if n.label == "Chunk"] doc_nodes = [n for n in res.nodes if n.label == "Document"] - assert chunk_nodes == [], "Chunk nodes must not appear when create_lexical_graph=False" - assert doc_nodes == [], "Document nodes must not appear when create_lexical_graph=False" + assert ( + chunk_nodes == [] + ), "Chunk nodes must not appear when create_lexical_graph=False" + assert ( + doc_nodes == [] + ), "Document nodes must not appear when create_lexical_graph=False" @pytest.mark.asyncio From a3ea7969bf3ff26c79b8e3b0688db15ad9a25b4e Mon Sep 17 00:00:00 2001 From: AmirLayegh Date: Wed, 1 Jul 2026 09:56:16 +0200 Subject: [PATCH 3/3] Fix e2e tests --- tests/e2e/experimental/test_neo4j_reader_component_e2e.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/e2e/experimental/test_neo4j_reader_component_e2e.py b/tests/e2e/experimental/test_neo4j_reader_component_e2e.py index e33a5c0e0..f1fe2cfbe 100644 --- a/tests/e2e/experimental/test_neo4j_reader_component_e2e.py +++ b/tests/e2e/experimental/test_neo4j_reader_component_e2e.py @@ -74,8 +74,11 @@ async def test_neo4j_reader_in_pipeline(driver: neo4j.Driver, llm: MagicMock) -> pipeline_output = await pipeline.run({}) created_graph = pipeline_output.result["extractor"] assert len(created_graph["nodes"]) == 1 - # no lexical graph, so no relationship to the chunk - assert len(created_graph["relationships"]) == 0 + # no lexical graph nodes, but FROM_CHUNK relationship is still created + from_chunk_rels = [ + r for r in created_graph["relationships"] if r["type"] == "FROM_CHUNK" + ] + assert len(from_chunk_rels) == 1 @pytest.mark.asyncio