참고소스 수정본

참고/neo4j-graphrag-python-main/tests/e2e/experimental/pipeline/config/test_pipeline_runner_e2e.py

@@ -0,0 +1,306 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import os
+from typing import Any
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+from neo4j import Driver
+from neo4j_graphrag.experimental.pipeline.config.runner import PipelineRunner
+from neo4j_graphrag.experimental.pipeline.pipeline import PipelineResult
+from neo4j_graphrag.llm import LLMResponse
+
+_SIMPLE_KG_PIPELINE_LLM_RESPONSE = """{
+                            "nodes": [
+                                {
+                                    "id": "0",
+                                    "label": "Person",
+                                    "properties": {
+                                        "name": "Harry Potter"
+                                    }
+                                },
+                                {
+                                    "id": "1",
+                                    "label": "Person",
+                                    "properties": {
+                                        "name": "Alastor Mad-Eye Moody"
+                                    }
+                                },
+                                {
+                                    "id": "2",
+                                    "label": "Organization",
+                                    "properties": {
+                                        "name": "The Order of the Phoenix"
+                                    }
+                                }
+                            ],
+                            "relationships": [
+                                {
+                                    "type": "KNOWS",
+                                    "start_node_id": "0",
+                                    "end_node_id": "1"
+                                },
+                                {
+                                    "type": "LED_BY",
+                                    "start_node_id": "2",
+                                    "end_node_id": "1"
+                                }
+                            ]
+                        }"""
+
+
+@pytest.fixture(scope="function", autouse=True)
+def clear_db(driver: Driver) -> Any:
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    yield
+
+
+@pytest.mark.asyncio
+async def test_pipeline_from_json_config(harry_potter_text: str, driver: Mock) -> None:
+    os.environ["NEO4J_URI"] = "neo4j://localhost:7687"
+    os.environ["NEO4J_USER"] = "neo4j"
+    os.environ["NEO4J_PASSWORD"] = "password"
+
+    runner = PipelineRunner.from_config_file(
+        "tests/e2e/data/config_files/pipeline_config.json"
+    )
+    res = await runner.run({"splitter": {"text": harry_potter_text}})
+    assert isinstance(res, PipelineResult)
+    meta = res.result["writer"]["metadata"]
+    assert "statistics" in meta
+    assert meta["statistics"]["node_count"] == 11
+    assert meta["statistics"]["relationship_count"] == 10
+    assert "nodes_per_label" in meta["statistics"]
+    assert "rel_per_type" in meta["statistics"]
+    assert "input_files_count" in meta["statistics"]
+    assert "input_files_total_size_bytes" in meta["statistics"]
+    nodes = driver.execute_query("MATCH (n) RETURN n")
+    assert len(nodes.records) == 11
+
+
+@pytest.mark.asyncio
+async def test_pipeline_from_yaml_config(harry_potter_text: str, driver: Mock) -> None:
+    os.environ["NEO4J_URI"] = "neo4j://localhost:7687"
+    os.environ["NEO4J_USER"] = "neo4j"
+    os.environ["NEO4J_PASSWORD"] = "password"
+
+    runner = PipelineRunner.from_config_file(
+        "tests/e2e/data/config_files/pipeline_config.yaml"
+    )
+    res = await runner.run({"splitter": {"text": harry_potter_text}})
+    assert isinstance(res, PipelineResult)
+    meta = res.result["writer"]["metadata"]
+    assert "statistics" in meta
+    assert meta["statistics"]["node_count"] == 11
+    assert meta["statistics"]["relationship_count"] == 10
+    assert "nodes_per_label" in meta["statistics"]
+    assert "rel_per_type" in meta["statistics"]
+    assert "input_files_count" in meta["statistics"]
+    assert "input_files_total_size_bytes" in meta["statistics"]
+
+    nodes = driver.execute_query("MATCH (n) RETURN n")
+    assert len(nodes.records) == 11
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.runner.SimpleKGPipelineConfig.get_default_embedder"
+)
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.runner.SimpleKGPipelineConfig.get_default_llm"
+)
+@pytest.mark.asyncio
+async def test_simple_kg_pipeline_from_json_config(
+    mock_llm: Mock, mock_embedder: Mock, harry_potter_text: str, driver: Mock
+) -> None:
+    mock_llm.return_value.ainvoke = AsyncMock(
+        side_effect=[
+            LLMResponse(
+                content=_SIMPLE_KG_PIPELINE_LLM_RESPONSE,
+            ),
+        ]
+    )
+    mock_embedder.return_value.async_embed_query = AsyncMock(
+        side_effect=[
+            [1.0, 2.0],
+        ]
+    )
+
+    os.environ["NEO4J_URI"] = "neo4j://localhost:7687"
+    os.environ["NEO4J_USER"] = "neo4j"
+    os.environ["NEO4J_PASSWORD"] = "password"
+    os.environ["OPENAI_API_KEY"] = "sk-my-secret-key"
+    os.environ["MY_OPENAI_KEY"] = "my-openai-key"
+
+    runner = PipelineRunner.from_config_file(
+        "tests/e2e/data/config_files/simple_kg_pipeline_config.json"
+    )
+
+    # check extras and API keys are handled as expected
+    config = runner.config
+    assert config is not None
+    # extras must be resolved:
+    assert config._global_data["extras"] == {"openai_api_key": "my-openai-key"}
+    # API key for LLM is read from env vars (see config file)
+    default_llm = config._global_data["llm_config"]["default"]
+    assert default_llm.client.api_key == "sk-my-secret-key"
+    # API key for embedder is read from extras (see config file)
+    default_embedder = config._global_data["embedder_config"]["default"]
+    assert default_embedder.client.api_key == "my-openai-key"
+
+    # then run pipeline and check results
+    res = await runner.run({"file_path": "tests/e2e/data/documents/harry_potter.pdf"})
+    assert isinstance(res, PipelineResult)
+    assert res.result["resolver"] == {
+        "number_of_nodes_to_resolve": 3,
+        "number_of_created_nodes": 3,
+    }
+    nodes = driver.execute_query("MATCH (n) RETURN n")
+    # 1 chunk + 1 document + 3 __Entity__ nodes
+    assert len(nodes.records) == 5
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.runner.SimpleKGPipelineConfig.get_default_embedder"
+)
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.runner.SimpleKGPipelineConfig.get_default_llm"
+)
+@pytest.mark.asyncio
+async def test_simple_kg_pipeline_from_json_config_with_markdown(
+    mock_llm: Mock, mock_embedder: Mock, harry_potter_text: str, driver: Mock
+) -> None:
+    mock_llm.return_value.ainvoke = AsyncMock(
+        side_effect=[
+            LLMResponse(
+                content=_SIMPLE_KG_PIPELINE_LLM_RESPONSE,
+            ),
+        ]
+    )
+    mock_embedder.return_value.async_embed_query = AsyncMock(
+        side_effect=[
+            [1.0, 2.0],
+        ]
+    )
+
+    os.environ["NEO4J_URI"] = "neo4j://localhost:7687"
+    os.environ["NEO4J_USER"] = "neo4j"
+    os.environ["NEO4J_PASSWORD"] = "password"
+    os.environ["OPENAI_API_KEY"] = "sk-my-secret-key"
+    os.environ["MY_OPENAI_KEY"] = "my-openai-key"
+
+    runner = PipelineRunner.from_config_file(
+        "tests/e2e/data/config_files/simple_kg_pipeline_config.json"
+    )
+
+    config = runner.config
+    assert config is not None
+    assert config._global_data["extras"] == {"openai_api_key": "my-openai-key"}
+    default_llm = config._global_data["llm_config"]["default"]
+    assert default_llm.client.api_key == "sk-my-secret-key"
+    default_embedder = config._global_data["embedder_config"]["default"]
+    assert default_embedder.client.api_key == "my-openai-key"
+
+    res = await runner.run({"file_path": "tests/e2e/data/documents/harry_potter.md"})
+    assert isinstance(res, PipelineResult)
+    assert res.result["resolver"] == {
+        "number_of_nodes_to_resolve": 3,
+        "number_of_created_nodes": 3,
+    }
+    nodes = driver.execute_query("MATCH (n) RETURN n")
+    assert len(nodes.records) == 5
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.runner.SimpleKGPipelineConfig.get_default_embedder"
+)
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.runner.SimpleKGPipelineConfig.get_default_llm"
+)
+@pytest.mark.asyncio
+async def test_simple_kg_pipeline_from_yaml_config(
+    mock_llm: Mock, mock_embedder: Mock, harry_potter_text: str, driver: Mock
+) -> None:
+    mock_llm.return_value.ainvoke = AsyncMock(
+        side_effect=[
+            LLMResponse(
+                content=_SIMPLE_KG_PIPELINE_LLM_RESPONSE,
+            ),
+        ]
+    )
+    mock_embedder.return_value.async_embed_query = AsyncMock(
+        side_effect=[
+            [1.0, 2.0],
+        ]
+    )
+
+    os.environ["NEO4J_URI"] = "neo4j://localhost:7687"
+    os.environ["NEO4J_USER"] = "neo4j"
+    os.environ["NEO4J_PASSWORD"] = "password"
+    os.environ["OPENAI_API_KEY"] = "sk-my-secret-key"
+
+    runner = PipelineRunner.from_config_file(
+        "tests/e2e/data/config_files/simple_kg_pipeline_config.yaml"
+    )
+    res = await runner.run({"file_path": "tests/e2e/data/documents/harry_potter.pdf"})
+    assert isinstance(res, PipelineResult)
+    # print(await runner.pipeline.store.get_result_for_component(res.run_id, "splitter"))
+    assert res.result["resolver"] == {
+        "number_of_nodes_to_resolve": 3,
+        "number_of_created_nodes": 3,
+    }
+    nodes = driver.execute_query("MATCH (n) RETURN n")
+    # 1 chunk + 1 document + 3 nodes
+    assert len(nodes.records) == 5
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.runner.SimpleKGPipelineConfig.get_default_embedder"
+)
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.runner.SimpleKGPipelineConfig.get_default_llm"
+)
+@pytest.mark.asyncio
+async def test_simple_kg_pipeline_from_yaml_config_with_markdown(
+    mock_llm: Mock, mock_embedder: Mock, harry_potter_text: str, driver: Mock
+) -> None:
+    mock_llm.return_value.ainvoke = AsyncMock(
+        side_effect=[
+            LLMResponse(
+                content=_SIMPLE_KG_PIPELINE_LLM_RESPONSE,
+            ),
+        ]
+    )
+    mock_embedder.return_value.async_embed_query = AsyncMock(
+        side_effect=[
+            [1.0, 2.0],
+        ]
+    )
+
+    os.environ["NEO4J_URI"] = "neo4j://localhost:7687"
+    os.environ["NEO4J_USER"] = "neo4j"
+    os.environ["NEO4J_PASSWORD"] = "password"
+    os.environ["OPENAI_API_KEY"] = "sk-my-secret-key"
+
+    runner = PipelineRunner.from_config_file(
+        "tests/e2e/data/config_files/simple_kg_pipeline_config.yaml"
+    )
+    res = await runner.run({"file_path": "tests/e2e/data/documents/harry_potter.md"})
+    assert isinstance(res, PipelineResult)
+    assert res.result["resolver"] == {
+        "number_of_nodes_to_resolve": 3,
+        "number_of_created_nodes": 3,
+    }
+    nodes = driver.execute_query("MATCH (n) RETURN n")
+    assert len(nodes.records) == 5

참고/neo4j-graphrag-python-main/tests/e2e/experimental/test_entity_resolver_component_e2e.py

@@ -0,0 +1,180 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import neo4j
+import pytest
+from neo4j_graphrag.experimental.components.resolver import (
+    SinglePropertyExactMatchResolver,
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_resolver_single_node(driver: neo4j.Driver) -> None:
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    driver.execute_query(
+        """
+        CREATE (d:Document {id: "0", path: "path"})
+        CREATE (c:Chunk {id: "0:0"})
+        CREATE (c)-[:FROM_DOCUMENT]->(d)
+        CREATE (alice:__Entity__:Person {id: "0:0:1", name: "Alice"})
+        CREATE (alice)-[:FROM_CHUNK]->(c)
+        """
+    )
+    resolver = SinglePropertyExactMatchResolver(driver)
+    res = await resolver.run()
+    # __Entity__ nodes attached to a chunk
+    assert res.number_of_nodes_to_resolve == 1
+    # Alice
+    assert res.number_of_created_nodes == 1
+
+    records, _, _ = driver.execute_query(
+        "MATCH path=(:Person {name: 'Alice'}) RETURN path"
+    )
+    assert len(records) == 1
+    path = records[0].get("path")
+    assert path.start_node.get("name") == "Alice"
+    assert path.start_node.labels == frozenset({"__Entity__", "Person"})
+    assert path.start_node.get("id") == "0:0:1"
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_resolver_two_nodes_and_relationships(driver: neo4j.Driver) -> None:
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    driver.execute_query(
+        """
+        CREATE (d:Document {id: "0", path: "path"})
+        CREATE (c:Chunk {id: "0:0"})
+        CREATE (c)-[:FROM_DOCUMENT]->(d)
+        CREATE (alice1:__Entity__:Person {id: "0:0:1", name: "Alice"})
+        CREATE (alice2:__Entity__:Person {id: "0:0:2", name: "Alice"})
+        CREATE (sweden:__Entity__:Country {id: "0:0:3", name: "Sweden"})
+        CREATE (alice1)-[:LIVES_IN]->(sweden)
+        CREATE (alice1)-[:FROM_CHUNK]->(c)
+        CREATE (alice2)-[:FROM_CHUNK]->(c)
+        CREATE (sweden)-[:FROM_CHUNK]->(c)
+        """
+    )
+    resolver = SinglePropertyExactMatchResolver(driver)
+    res = await resolver.run()
+    # __Entity__ nodes attached to a chunk
+    assert res.number_of_nodes_to_resolve == 3
+    # Alice and Sweden
+    assert res.number_of_created_nodes == 2
+
+    # check the domain graph
+    records, _, _ = driver.execute_query(
+        "MATCH path=(:Person {name: 'Alice'})-[:LIVES_IN]->(:Country {name: 'Sweden'}) RETURN path"
+    )
+    assert len(records) == 1
+    path = records[0].get("path")
+    assert path.start_node.get("name") == "Alice"
+    assert path.start_node.labels == frozenset({"__Entity__", "Person"})
+    assert path.end_node.get("name") == "Sweden"
+    assert path.end_node.labels == frozenset({"__Entity__", "Country"})
+    assert len(path.relationships) == 1
+    assert path.relationships[0].type == "LIVES_IN"
+
+    # check the lexical graph
+    records, _, _ = driver.execute_query(
+        "MATCH path=(:Person {name: 'Alice'})-[:FROM_CHUNK]->(:Chunk) RETURN path"
+    )
+    assert len(records) == 1
+    path = records[0].get("path")
+    assert path.start_node.get("name") == "Alice"
+    assert path.end_node.labels == frozenset({"Chunk"})
+    assert path.end_node.get("id") == "0:0"
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_resolver_same_name_different_labels(driver: neo4j.Driver) -> None:
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    driver.execute_query(
+        """
+        CREATE (d:Document {id: "0", path: "path"})
+        CREATE (c:Chunk {id: "0:0"})
+        CREATE (c)-[:FROM_DOCUMENT]->(d)
+        CREATE (alice1:__Entity__:Person {id: "0:0:1", name: "Alice"})
+        CREATE (alice2:__Entity__:Human {id: "0:0:2", name: "Alice"})
+        CREATE (alice1)-[:FROM_CHUNK]->(c)
+        CREATE (alice2)-[:FROM_CHUNK]->(c)
+        CREATE (sweden)-[:FROM_CHUNK]->(c)
+        """
+    )
+    resolver = SinglePropertyExactMatchResolver(driver)
+    res = await resolver.run()
+    # __Entity__ nodes attached to a chunk
+    assert res.number_of_nodes_to_resolve == 2
+    # Alice Person and Alice Human
+    assert res.number_of_created_nodes == 2
+
+    records, _, _ = driver.execute_query("MATCH (alice {name: 'Alice'}) RETURN alice")
+    assert len(records) == 2
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_resolver_custom_property(driver: neo4j.Driver) -> None:
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    driver.execute_query(
+        """
+        CREATE (d:Document {id: "0", path: "path"})
+        CREATE (c:Chunk {id: "0:0"})
+        CREATE (c)-[:FROM_DOCUMENT]->(d)
+        CREATE (alice:__Entity__:Person {id: "0:0:1", name: "Alice"})
+        CREATE (alicia:__Entity__:Person {id: "0:0:1", name: "Alicia"})
+        CREATE (alice)-[:FROM_CHUNK]->(c)
+        CREATE (alicia)-[:FROM_CHUNK]->(c)
+        """
+    )
+    resolver = SinglePropertyExactMatchResolver(driver, resolve_property="id")
+    res = await resolver.run()
+    # __Entity__ nodes attached to a chunk
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 1
+
+    records, _, _ = driver.execute_query("MATCH (person:Person) RETURN person")
+    assert len(records) == 1
+    assert records[0].get("person").get("name") in ["Alice", "Alicia"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_resolver_custom_filter(driver: neo4j.Driver) -> None:
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    driver.execute_query(
+        """
+        CREATE (d:Document {id: "0", path: "path"})
+        CREATE (c:Chunk {id: "0:0"})
+        CREATE (c)-[:FROM_DOCUMENT]->(d)
+        CREATE (alice1:__Entity__:Person {id: "0:0:1", name: "Alice"})
+        CREATE (alice2:__Entity__:Person {id: "0:0:2", name: "Alice"})
+        CREATE (sweden:__Entity__:Country {id: "0:0:3", name: "Sweden"})
+        CREATE (alice1)-[:LIVES_IN]->(sweden)
+        CREATE (alice1)-[:FROM_CHUNK]->(c)
+        CREATE (alice2)-[:FROM_CHUNK]->(c)
+        CREATE (sweden)-[:FROM_CHUNK]->(c)
+        """
+    )
+    resolver = SinglePropertyExactMatchResolver(
+        driver, filter_query="WHERE not entity:Person"
+    )
+    res = await resolver.run()
+    # __Entity__ nodes attached to a chunk without a Person label
+    # so only Country here
+    assert res.number_of_nodes_to_resolve == 1
+    # Sweden
+    assert res.number_of_created_nodes == 1

참고/neo4j-graphrag-python-main/tests/e2e/experimental/test_graph_pruning_component_e2e.py

@@ -0,0 +1,588 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from typing import Any
+
+import pytest
+
+from neo4j_graphrag.experimental.components.graph_pruning import (
+    GraphPruning,
+    PruningReason,
+)
+from neo4j_graphrag.experimental.components.schema import GraphSchema
+from neo4j_graphrag.experimental.components.types import (
+    Neo4jGraph,
+    Neo4jNode,
+    Neo4jRelationship,
+)
+
+
+@pytest.fixture
+def extracted_graph() -> Neo4jGraph:
+    """This is the graph to be pruned in all the below tests,
+    using different schema configuration.
+    """
+    return Neo4jGraph(
+        nodes=[
+            Neo4jNode(
+                id="1",
+                label="Person",
+                properties={
+                    "name": "John Doe",
+                },
+            ),
+            Neo4jNode(
+                id="2",
+                label="Person",
+                properties={
+                    "height": 180,
+                },
+            ),
+            Neo4jNode(
+                id="3",
+                label="Person",
+                properties={
+                    "name": "Jane Doe",
+                    "weight": 90,
+                },
+            ),
+            Neo4jNode(
+                id="10",
+                label="Organization",
+                properties={
+                    "name": "Azerty Inc.",
+                    "created": 1999,
+                },
+            ),
+        ],
+        relationships=[
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="2",
+                type="KNOWS",
+                properties={"firstMetIn": 2025},
+            ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="3",
+                type="KNOWS",
+            ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="2",
+                type="MANAGES",
+            ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="10",
+                type="MANAGES",
+            ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="10",
+                type="WORKS_FOR",
+            ),
+        ],
+    )
+
+
+async def _test(
+    extracted_graph: Neo4jGraph, schema_dict: dict[str, Any], expected_graph: Neo4jGraph
+) -> None:
+    schema = GraphSchema.model_validate(schema_dict)
+    pruner = GraphPruning()
+    res = await pruner.run(extracted_graph, schema)
+    assert res.graph == expected_graph
+
+
+@pytest.mark.asyncio
+async def test_graph_pruning_loose(extracted_graph: Neo4jGraph) -> None:
+    """Loose schema:
+    - no required properties
+    - all additional* allowed
+
+    => we keep everything from the extracted graph
+    """
+    schema_dict = {
+        "node_types": [
+            {
+                "label": "Person",
+                "properties": [
+                    {"name": "name", "type": "STRING"},
+                    {"name": "height", "type": "INTEGER"},
+                ],
+                "additional_properties": True,
+            }
+        ],
+        "relationship_types": [
+            {
+                "label": "KNOWS",
+            }
+        ],
+        "patterns": [
+            ("Person", "KNOWS", "Person"),
+        ],
+        "additional_node_types": True,
+        "additional_relationship_types": True,
+        "additional_patterns": True,
+    }
+    await _test(extracted_graph, schema_dict, extracted_graph)
+
+
+@pytest.mark.asyncio
+async def test_graph_pruning_missing_required_property(
+    extracted_graph: Neo4jGraph,
+) -> None:
+    """Person node type has a required 'name' property:
+    - extracted nodes without this property are pruned
+    - any relationship tied to this node is also pruned
+    """
+    schema_dict = {
+        "node_types": [
+            {
+                "label": "Person",
+                "properties": [
+                    {
+                        "name": "name",
+                        "type": "STRING",
+                        "required": True,
+                    },
+                    {"name": "height", "type": "INTEGER"},
+                ],
+                "additional_properties": True,
+            }
+        ],
+        "relationship_types": [
+            {
+                "label": "KNOWS",
+            }
+        ],
+        "patterns": [
+            ("Person", "KNOWS", "Person"),
+        ],
+        "additional_node_types": True,
+        "additional_relationship_types": True,
+        "additional_patterns": True,
+    }
+    filtered_graph = Neo4jGraph(
+        nodes=[
+            Neo4jNode(
+                id="1",
+                label="Person",
+                properties={
+                    "name": "John Doe",
+                },
+            ),
+            # do not have the required "name" property
+            # Neo4jNode(
+            #     id="2",
+            #     label="Person",
+            #     properties={
+            #         "height": 180,
+            #     }
+            # ),
+            Neo4jNode(
+                id="3",
+                label="Person",
+                properties={
+                    "name": "Jane Doe",
+                    "weight": 90,
+                },
+            ),
+            Neo4jNode(
+                id="10",
+                label="Organization",
+                properties={
+                    "name": "Azerty Inc.",
+                    "created": 1999,
+                },
+            ),
+        ],
+        relationships=[
+            # node "2" was pruned
+            # Neo4jRelationship(
+            #     start_node_id="1",
+            #     end_node_id="2",
+            #     type="KNOWS",
+            #     properties={"firstMetIn": 2025},
+            # ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="3",
+                type="KNOWS",
+            ),
+            # node "2" was pruned
+            # Neo4jRelationship(
+            #     start_node_id="1",
+            #     end_node_id="2",
+            #     type="MANAGES",
+            # ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="10",
+                type="MANAGES",
+            ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="10",
+                type="WORKS_FOR",
+            ),
+        ],
+    )
+    await _test(extracted_graph, schema_dict, filtered_graph)
+
+
+@pytest.mark.asyncio
+async def test_graph_pruning_existence_constraint_node_property_explicit(
+    extracted_graph: Neo4jGraph,
+) -> None:
+    """Same outcome as ``test_graph_pruning_missing_required_property`` using EXISTENCE only (no ``required``)."""
+    schema_dict = {
+        "node_types": [
+            {
+                "label": "Person",
+                "properties": [
+                    {"name": "name", "type": "STRING"},
+                    {"name": "height", "type": "INTEGER"},
+                ],
+                "additional_properties": True,
+            }
+        ],
+        "relationship_types": [
+            {
+                "label": "KNOWS",
+            }
+        ],
+        "patterns": [
+            ("Person", "KNOWS", "Person"),
+        ],
+        "constraints": [
+            {
+                "type": "EXISTENCE",
+                "node_type": "Person",
+                "property_names": ["name"],
+                "relationship_type": None,
+            }
+        ],
+        "additional_node_types": True,
+        "additional_relationship_types": True,
+        "additional_patterns": True,
+    }
+    schema = GraphSchema.model_validate(schema_dict)
+    assert schema.existence_property_names_for_node("Person") == {"name"}
+
+    filtered_graph = Neo4jGraph(
+        nodes=[
+            Neo4jNode(
+                id="1",
+                label="Person",
+                properties={
+                    "name": "John Doe",
+                },
+            ),
+            Neo4jNode(
+                id="3",
+                label="Person",
+                properties={
+                    "name": "Jane Doe",
+                    "weight": 90,
+                },
+            ),
+            Neo4jNode(
+                id="10",
+                label="Organization",
+                properties={
+                    "name": "Azerty Inc.",
+                    "created": 1999,
+                },
+            ),
+        ],
+        relationships=[
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="3",
+                type="KNOWS",
+            ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="10",
+                type="MANAGES",
+            ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="10",
+                type="WORKS_FOR",
+            ),
+        ],
+    )
+    await _test(extracted_graph, schema_dict, filtered_graph)
+
+
+@pytest.mark.asyncio
+async def test_graph_pruning_existence_constraint_relationship_property(
+    extracted_graph: Neo4jGraph,
+) -> None:
+    """Relationship-scoped EXISTENCE: KNOWS without ``firstMetIn`` is flagged when that property is mandatory."""
+    schema_dict = {
+        "node_types": [
+            {
+                "label": "Person",
+                "properties": [
+                    {"name": "name", "type": "STRING"},
+                    {"name": "height", "type": "INTEGER"},
+                ],
+                "additional_properties": True,
+            }
+        ],
+        "relationship_types": [
+            {
+                "label": "KNOWS",
+                "properties": [
+                    {"name": "firstMetIn", "type": "INTEGER"},
+                ],
+                "additional_properties": True,
+            }
+        ],
+        "patterns": [
+            ("Person", "KNOWS", "Person"),
+        ],
+        "constraints": [
+            {
+                "type": "EXISTENCE",
+                "node_type": "",
+                "property_names": ["firstMetIn"],
+                "relationship_type": "KNOWS",
+            }
+        ],
+        "additional_node_types": True,
+        "additional_relationship_types": True,
+        "additional_patterns": True,
+    }
+    schema = GraphSchema.model_validate(schema_dict)
+    assert schema.existence_property_names_for_relationship("KNOWS") == {"firstMetIn"}
+
+    pruner = GraphPruning()
+    res = await pruner.run(extracted_graph, schema)
+    assert any(
+        x.pruned_reason == PruningReason.MISSING_REQUIRED_PROPERTY
+        for x in res.pruning_stats.pruned_relationships
+    )
+
+
+@pytest.mark.asyncio
+async def test_graph_pruning_strict_properties_and_node_types(
+    extracted_graph: Neo4jGraph,
+) -> None:
+    """Additional properties on Person nodes are not allowed.
+    Additional node types are not allowed.
+
+    => we prune "Organization" nodes (not in schema)
+    and the "weight" property that was extracted for some persons.
+    """
+    schema_dict = {
+        "node_types": [
+            {
+                "label": "Person",
+                "properties": [
+                    {
+                        "name": "name",
+                        "type": "STRING",
+                    },
+                    {"name": "height", "type": "INTEGER"},
+                ],
+                # "additional_properties": False,  # default value
+            }
+        ],
+        "relationship_types": [
+            {
+                "label": "KNOWS",
+            }
+        ],
+        "patterns": [
+            ("Person", "KNOWS", "Person"),
+        ],
+        # "additional_node_types": False,  # default value
+        "additional_relationship_types": True,
+        "additional_patterns": True,
+    }
+    filtered_graph = Neo4jGraph(
+        nodes=[
+            Neo4jNode(
+                id="1",
+                label="Person",
+                properties={
+                    "name": "John Doe",
+                },
+            ),
+            Neo4jNode(
+                id="2",
+                label="Person",
+                properties={
+                    "height": 180,
+                },
+            ),
+            Neo4jNode(
+                id="3",
+                label="Person",
+                properties={
+                    "name": "Jane Doe",
+                    # weight not in listed properties
+                    # "weight": 90,
+                },
+            ),
+            # label "Organization" not in schema
+            # Neo4jNode(
+            #     id="10",
+            #     label="Organization",
+            #     properties={
+            #         "name": "Azerty Inc.",
+            #         "created": 1999,
+            #     }
+            # ),
+        ],
+        relationships=[
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="2",
+                type="KNOWS",
+                properties={"firstMetIn": 2025},
+            ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="3",
+                type="KNOWS",
+            ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="2",
+                type="MANAGES",
+            ),
+            # node "10" was pruned (label not allowed)
+            # Neo4jRelationship(
+            #     start_node_id="1",
+            #     end_node_id="10",
+            #     type="MANAGES",
+            # ),
+            # Neo4jRelationship(
+            #     start_node_id="1",
+            #     end_node_id="10",
+            #     type="WORKS_FOR",
+            # )
+        ],
+    )
+    await _test(extracted_graph, schema_dict, filtered_graph)
+
+
+@pytest.mark.asyncio
+async def test_graph_pruning_strict_patterns(extracted_graph: Neo4jGraph) -> None:
+    """Additional patterns not allowed:
+
+    - MANAGES: it's a known relationship type but without any pattern, it's pruned
+    - WORKS_FOR: it's not a known relationship type, and additional_relationship_types is allowed
+        so we keep it.
+    """
+    # - no additional patterns allowed
+    schema_dict = {
+        "node_types": [
+            {
+                "label": "Person",
+                "properties": [
+                    {
+                        "name": "name",
+                        "type": "STRING",
+                    },
+                    {"name": "height", "type": "INTEGER"},
+                ],
+                "additional_properties": True,
+            },
+            {
+                "label": "Organization",
+            },
+        ],
+        "relationship_types": [
+            {
+                "label": "KNOWS",
+            },
+            {
+                "label": "MANAGES",
+            },
+        ],
+        "patterns": (
+            ("Person", "KNOWS", "Person"),
+            ("Person", "KNOWS", "Organization"),
+        ),
+        "additional_node_types": True,
+        "additional_relationship_types": False,
+        "additional_patterns": False,
+    }
+    filtered_graph = Neo4jGraph(
+        nodes=[
+            Neo4jNode(
+                id="1",
+                label="Person",
+                properties={
+                    "name": "John Doe",
+                },
+            ),
+            Neo4jNode(
+                id="2",
+                label="Person",
+                properties={
+                    "height": 180,
+                },
+            ),
+            Neo4jNode(
+                id="3",
+                label="Person",
+                properties={
+                    "name": "Jane Doe",
+                    "weight": 90,
+                },
+            ),
+            Neo4jNode(
+                id="10",
+                label="Organization",
+                properties={
+                    "name": "Azerty Inc.",
+                    "created": 1999,
+                },
+            ),
+        ],
+        relationships=[
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="2",
+                type="KNOWS",
+                properties={"firstMetIn": 2025},
+            ),
+            Neo4jRelationship(
+                start_node_id="1",
+                end_node_id="3",
+                type="KNOWS",
+            ),
+            # invalid pattern (person, manages, person)
+            # Neo4jRelationship(
+            #     start_node_id="1",
+            #     end_node_id="2",
+            #     type="MANAGES",
+            # ),
+            # invalid pattern (person, works for, person)
+            # Neo4jRelationship(
+            #     start_node_id="1",
+            #     end_node_id="10",
+            #     type="WORKS_FOR",
+            # ),
+        ],
+    )
+    await _test(extracted_graph, schema_dict, filtered_graph)

참고/neo4j-graphrag-python-main/tests/e2e/experimental/test_kg_builder_pipeline_e2e.py

@@ -0,0 +1,600 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+import os
+from collections import Counter
+from unittest.mock import MagicMock
+
+import neo4j
+import pytest
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.exceptions import LLMGenerationError
+from neo4j_graphrag.experimental.components.embedder import TextChunkEmbedder
+from neo4j_graphrag.experimental.components.entity_relation_extractor import (
+    LLMEntityRelationExtractor,
+    OnError,
+)
+from neo4j_graphrag.experimental.components.kg_writer import Neo4jWriter
+from neo4j_graphrag.experimental.components.resolver import (
+    SinglePropertyExactMatchResolver,
+)
+from neo4j_graphrag.experimental.components.schema import (
+    SchemaBuilder,
+    NodeType,
+    PropertyType,
+    RelationshipType,
+)
+from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import (
+    FixedSizeSplitter,
+)
+from neo4j_graphrag.experimental.pipeline import Pipeline
+from neo4j_graphrag.experimental.pipeline.pipeline import PipelineResult
+from neo4j_graphrag.llm import LLMInterface, LLMResponse
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+@pytest.fixture
+def llm() -> LLMInterface:
+    llm = MagicMock(spec=LLMInterface)
+    return llm
+
+
+@pytest.fixture
+def embedder() -> Embedder:
+    embedder = MagicMock(spec=Embedder)
+    return embedder
+
+
+@pytest.fixture
+def schema_builder() -> SchemaBuilder:
+    return SchemaBuilder()
+
+
+@pytest.fixture
+def text_splitter() -> FixedSizeSplitter:
+    return FixedSizeSplitter(chunk_size=500, chunk_overlap=10)
+
+
+@pytest.fixture
+def chunk_embedder(embedder: Embedder) -> TextChunkEmbedder:
+    return TextChunkEmbedder(embedder=embedder)
+
+
+@pytest.fixture
+def entity_relation_extractor(llm: LLMInterface) -> LLMEntityRelationExtractor:
+    return LLMEntityRelationExtractor(
+        llm=llm,
+        on_error=OnError.RAISE,
+    )
+
+
+@pytest.fixture
+def kg_writer(driver: neo4j.Driver) -> Neo4jWriter:
+    return Neo4jWriter(driver)
+
+
+@pytest.fixture
+def entity_resolver(driver: neo4j.Driver) -> SinglePropertyExactMatchResolver:
+    return SinglePropertyExactMatchResolver(driver)
+
+
+@pytest.fixture
+def kg_builder_pipeline(
+    text_splitter: FixedSizeSplitter,
+    chunk_embedder: TextChunkEmbedder,
+    schema_builder: SchemaBuilder,
+    entity_relation_extractor: LLMEntityRelationExtractor,
+    kg_writer: Neo4jWriter,
+    entity_resolver: SinglePropertyExactMatchResolver,
+) -> Pipeline:
+    pipe = Pipeline()
+    # define the components
+    pipe.add_component(text_splitter, "splitter")
+    pipe.add_component(chunk_embedder, "embedder")
+    pipe.add_component(schema_builder, "schema")
+    pipe.add_component(entity_relation_extractor, "extractor")
+    pipe.add_component(kg_writer, "writer")
+    pipe.add_component(entity_resolver, "resolver")
+    # define the execution order of component
+    # and how the output of previous components must be used
+    pipe.connect("splitter", "embedder", input_config={"text_chunks": "splitter"})
+    # pipe.connect("splitter", "extractor", input_config={"chunks": "splitter"})
+    pipe.connect("schema", "extractor", input_config={"schema": "schema"})
+    pipe.connect("embedder", "extractor", input_config={"chunks": "embedder"})
+    pipe.connect(
+        "extractor",
+        "writer",
+        input_config={"graph": "extractor"},
+    )
+    pipe.connect("writer", "resolver", {})
+    return pipe
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_pipeline_builder_happy_path(
+    harry_potter_text: str,
+    llm: MagicMock,
+    embedder: MagicMock,
+    driver: neo4j.Driver,
+    kg_builder_pipeline: Pipeline,
+) -> None:
+    """When everything works as expected, extracted entities, relations and text
+    chunks must be in the DB
+    """
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    embedder.async_embed_query.return_value = [1, 2, 3]
+    llm.ainvoke.side_effect = [
+        LLMResponse(
+            content="""{
+                        "nodes": [
+                            {
+                                "id": "0",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Harry Potter"
+                                }
+                            },
+                            {
+                                "id": "1",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Alastor Mad-Eye Moody"
+                                }
+                            },
+                            {
+                                "id": "2",
+                                "label": "Organization",
+                                "properties": {
+                                    "name": "The Order of the Phoenix"
+                                }
+                            }
+                        ],
+                        "relationships": [
+                            {
+                                "type": "KNOWS",
+                                "start_node_id": "0",
+                                "end_node_id": "1"
+                            },
+                            {
+                                "type": "LED_BY",
+                                "start_node_id": "2",
+                                "end_node_id": "1"
+                            }
+                        ]
+                    }"""
+        ),
+        LLMResponse(content='{"nodes": [], "relationships": []}'),
+    ]
+
+    # user input:
+    # the initial text
+    # and the list of entities and relations we are looking for
+    pipe_inputs = {
+        "splitter": {"text": harry_potter_text},
+        "schema": {
+            "node_types": [
+                NodeType(
+                    label="Person",
+                    properties=[
+                        PropertyType(name="name", type="STRING"),
+                        PropertyType(name="place_of_birth", type="STRING"),
+                        PropertyType(name="date_of_birth", type="DATE"),
+                    ],
+                ),
+                NodeType(
+                    label="Organization",
+                    properties=[
+                        PropertyType(name="name", type="STRING"),
+                    ],
+                ),
+                NodeType(
+                    label="Potion",
+                    properties=[
+                        PropertyType(name="name", type="STRING"),
+                    ],
+                ),
+                NodeType(
+                    label="Location",
+                    properties=[
+                        PropertyType(name="address", type="STRING"),
+                    ],
+                ),
+            ],
+            "relationship_types": [
+                RelationshipType(
+                    label="KNOWS",
+                ),
+                RelationshipType(
+                    label="PART_OF",
+                ),
+                RelationshipType(
+                    label="LED_BY",
+                ),
+                RelationshipType(
+                    label="DRINKS",
+                ),
+            ],
+            "patterns": [
+                ("Person", "KNOWS", "Person"),
+                ("Person", "DRINKS", "Potion"),
+                ("Person", "PART_OF", "Organization"),
+                ("Organization", "LED_BY", "Person"),
+            ],
+        },
+        "extractor": {"document_info": {"path": "my document path"}},
+    }
+    res = await kg_builder_pipeline.run(pipe_inputs)
+    # llm must have been called for each chunk
+    assert llm.ainvoke.call_count == 2
+    # result must be success
+    assert isinstance(res, PipelineResult)
+    assert res.run_id is not None
+    assert "resolver" in res.result
+    # check component's results
+    chunks = await kg_builder_pipeline.store.get_result_for_component(
+        res.run_id, "splitter"
+    )
+    assert len(chunks["chunks"]) == 2
+    graph = await kg_builder_pipeline.store.get_result_for_component(
+        res.run_id, "extractor"
+    )
+    # 3 entities + 2 chunks + 1 document
+    nodes = graph["nodes"]
+    assert len(nodes) == 6
+    label_counts = dict(Counter([n["label"] for n in nodes]))
+    assert label_counts == {
+        "Chunk": 2,
+        "Document": 1,
+        "Person": 2,
+        "Organization": 1,
+    }
+    # 2 relationships between entities
+    # + 3 rels between entities and their chunk
+    # + 2 "NEXT_CHUNK" rels
+    relationships = graph["relationships"]
+    assert len(relationships) == 8
+    type_counts = dict(Counter([r["type"] for r in relationships]))
+    assert type_counts == {
+        "FROM_CHUNK": 3,
+        "FROM_DOCUMENT": 2,
+        "KNOWS": 1,
+        "LED_BY": 1,
+        "NEXT_CHUNK": 1,
+    }
+    # then check content of neo4j db
+    created_nodes = driver.execute_query("MATCH (n) RETURN n")
+    assert len(created_nodes.records) == 6
+    created_rels = driver.execute_query("MATCH ()-[r]->() RETURN r")
+    assert len(created_rels.records) == 8
+
+    created_chunks = driver.execute_query("MATCH (n:Chunk) RETURN n").records
+    assert len(created_chunks) == 2
+    for c in created_chunks:
+        node = c.get("n")
+        assert node.get("embedding") == [1, 2, 3]
+        assert node.get("text") is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_pipeline_builder_failing_chunk_raise(
+    harry_potter_text: str,
+    embedder: MagicMock,
+    llm: MagicMock,
+    driver: neo4j.Driver,
+    kg_builder_pipeline: Pipeline,
+) -> None:
+    """If on_error is set to "RAISE", any issue with the entity/relation
+    extractor should stop the process with an exception. Nothing should be
+    added to the DB
+    """
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    embedder.async_embed_query.return_value = [1, 2, 3]
+    llm.ainvoke.side_effect = [
+        LLMResponse(
+            content="""{
+                        "nodes": [
+                            {
+                                "id": "0",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Harry Potter"
+                                }
+                            },
+                            {
+                                "id": "1",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Alastor Mad-Eye Moody"
+                                }
+                            },
+                            {
+                                "id": "2",
+                                "label": "Organization",
+                                "properties": {
+                                    "name": "The Order of the Phoenix"
+                                }
+                            }
+                        ],
+                        "relationships": [
+                            {
+                                "type": "KNOWS",
+                                "start_node_id": "0",
+                                "end_node_id": "1"
+                            },
+                            {
+                                "type": "LED_BY",
+                                "start_node_id": "2",
+                                "end_node_id": "1"
+                            }
+                        ]
+                    }"""
+        ),
+        LLMResponse(content="invalid json"),
+    ]
+
+    # user input:
+    # the initial text
+    # and the list of entities and relations we are looking for
+    pipe_inputs = {
+        "splitter": {"text": harry_potter_text},
+        # note: schema not used in this test because
+        # we are mocking the LLM
+        "schema": {
+            "node_types": (),
+            "relationship_types": (),
+            "patterns": (),
+        },
+    }
+    with pytest.raises(LLMGenerationError):
+        await kg_builder_pipeline.run(pipe_inputs)
+    created_nodes = driver.execute_query("MATCH (n) RETURN n")
+    assert len(created_nodes.records) == 0
+    created_rels = driver.execute_query("MATCH ()-[r]->() RETURN r")
+    assert len(created_rels.records) == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_pipeline_builder_failing_chunk_do_not_raise(
+    harry_potter_text: str,
+    llm: MagicMock,
+    embedder: MagicMock,
+    driver: neo4j.Driver,
+    kg_builder_pipeline: Pipeline,
+) -> None:
+    """If on_error is set to "IGNORE", process must continue
+    and nodes/relationships created for the chunks that succeeded
+    """
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    embedder.async_embed_query.return_value = [1, 2, 3]
+    llm.ainvoke.side_effect = [
+        LLMResponse(content="invalid json"),
+        LLMResponse(
+            content="""{
+                        "nodes": [
+                            {
+                                "id": "0",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Harry Potter"
+                                }
+                            },
+                            {
+                                "id": "1",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Alastor Mad-Eye Moody"
+                                }
+                            },
+                            {
+                                "id": "2",
+                                "label": "Organization",
+                                "properties": {
+                                    "name": "The Order of the Phoenix"
+                                }
+                            }
+                        ],
+                        "relationships": [
+                            {
+                                "type": "KNOWS",
+                                "start_node_id": "0",
+                                "end_node_id": "1"
+                            },
+                            {
+                                "type": "LED_BY",
+                                "start_node_id": "2",
+                                "end_node_id": "1"
+                            }
+                        ]
+                    }"""
+        ),
+    ]
+
+    # user input:
+    # the initial text
+    # and the list of entities and relations we are looking for
+    pipe_inputs = {
+        "splitter": {"text": harry_potter_text},
+        # note: schema not used in this test because
+        # we are mocking the LLM
+        "schema": {
+            "node_types": (),
+            "relationship_types": (),
+            "patterns": (),
+        },
+    }
+    kg_builder_pipeline.get_node_by_name(
+        "extractor"
+    ).component.on_error = OnError.IGNORE  # type: ignore[attr-defined, unused-ignore]
+    res = await kg_builder_pipeline.run(pipe_inputs)
+    # llm must have been called for each chunk
+    assert llm.ainvoke.call_count == 2
+    # result must be success
+    assert isinstance(res, PipelineResult)
+    assert res.run_id is not None
+    assert res.result == {
+        "resolver": {"number_of_created_nodes": 3, "number_of_nodes_to_resolve": 3}
+    }
+    # check component's results
+    chunks = await kg_builder_pipeline.store.get_result_for_component(
+        res.run_id, "splitter"
+    )
+    assert len(chunks["chunks"]) == 2
+    graph = await kg_builder_pipeline.store.get_result_for_component(
+        res.run_id, "extractor"
+    )
+    # 3 entities + 2 chunks
+    nodes = graph["nodes"]
+    assert len(nodes) == 5
+    label_counts = dict(Counter([n["label"] for n in nodes]))
+    assert label_counts == {
+        "Chunk": 2,
+        "Person": 2,
+        "Organization": 1,
+    }
+    # 2 relationships between entities
+    # + 3 rels between entities and their chunk
+    # + 1 "NEXT_CHUNK" rels
+    relationships = graph["relationships"]
+    assert len(relationships) == 6
+    type_counts = dict(Counter([r["type"] for r in relationships]))
+    assert type_counts == {"FROM_CHUNK": 3, "KNOWS": 1, "LED_BY": 1, "NEXT_CHUNK": 1}
+    # then check content of neo4j db
+    created_nodes = driver.execute_query("MATCH (n) RETURN n")
+    assert len(created_nodes.records) == 5
+    created_rels = driver.execute_query("MATCH ()-[r]->() RETURN r")
+    assert len(created_rels.records) == 6
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_pipeline_builder_two_documents(
+    harry_potter_text_part1: str,
+    harry_potter_text_part2: str,
+    embedder: MagicMock,
+    llm: MagicMock,
+    driver: neo4j.Driver,
+    kg_builder_pipeline: Pipeline,
+) -> None:
+    """Run same pipeline on two documents. Check entity resolution.
+
+    First document:
+    2 chunks, entities Harry and The Order of the Phoenix, 1 relationship
+    Second document:
+    1 chunk, entities Harry and Alastor Mad-Eye Moody, 1 relationship
+
+    Should create:
+    1 document node
+    3 chunk nodes
+    3 entities (1 Harry + the other two)
+    ==> 7 nodes
+    3 relationships for lexical graph + 3 relationships for the entity graph
+    ==> 6 relationships
+    """
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    embedder.async_embed_query.return_value = [1, 2, 3]
+    llm.ainvoke.side_effect = [
+        LLMResponse(
+            content="""{
+                        "nodes": [
+                            {
+                                "id": "0",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Harry"
+                                }
+                            },
+                            {
+                                "id": "1",
+                                "label": "Organization",
+                                "properties": {
+                                    "name": "The Order of the Phoenix"
+                                }
+                            }
+                        ],
+                        "relationships": [
+                            {
+                                "type": "MEMBER_OF",
+                                "start_node_id": "0",
+                                "end_node_id": "1"
+                            }
+                        ]
+                    }"""
+        ),
+        LLMResponse(
+            content="""{
+                        "nodes": [
+                            {
+                                "id": "10",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Harry"
+                                }
+                            },
+                            {
+                                "id": "11",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Alastor Mad-Eye Moody"
+                                }
+                            }
+                        ],
+                        "relationships": [
+                            {
+                                "type": "KNOWS",
+                                "start_node_id": "10",
+                                "end_node_id": "11"
+                            }
+                        ]
+                    }"""
+        ),
+        LLMResponse(content='{"nodes": [], "relationships": []}'),
+    ]
+
+    # user input:
+    # the initial text
+    # and the list of entities and relations we are looking for
+    pipe_inputs_1 = {
+        "splitter": {"text": harry_potter_text_part1},
+        # note: schema not used in this test because
+        # we are mocking the LLM
+        "schema": {
+            "node_types": (),
+            "relationship_types": (),
+            "patterns": (),
+        },
+    }
+    pipe_inputs_2 = {
+        "splitter": {"text": harry_potter_text_part2},
+        # note: schema not used in this test because
+        # we are mocking the LLM
+        "schema": {
+            "node_types": (),
+            "relationship_types": (),
+            "patterns": (),
+        },
+    }
+    await kg_builder_pipeline.run(pipe_inputs_1)
+    await kg_builder_pipeline.run(pipe_inputs_2)
+    created_nodes = driver.execute_query("MATCH (n:__Entity__) RETURN n")
+    assert len(created_nodes.records) == 3
+    created_rels = driver.execute_query(
+        "MATCH (:__Entity__)-[r]->(:__Entity__) RETURN r"
+    )
+    assert len(created_rels.records) == 2

참고/neo4j-graphrag-python-main/tests/e2e/experimental/test_kg_writer_component_e2e.py

@@ -0,0 +1,549 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import logging
+import tempfile
+from pathlib import Path
+
+import neo4j
+import pytest
+
+from neo4j_graphrag.experimental.components.filename_collision_handler import (
+    FilenameCollisionHandler,
+)
+from neo4j_graphrag.experimental.components.kg_writer import Neo4jWriter, ParquetWriter
+from neo4j_graphrag.experimental.components.parquet_formatter import (
+    INTERNAL_ID_PROPERTY,
+)
+from neo4j_graphrag.experimental.components.types import (
+    LexicalGraphConfig,
+    Neo4jGraph,
+    Neo4jNode,
+    Neo4jRelationship,
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_kg_writer(driver: neo4j.Driver) -> None:
+    start_node = Neo4jNode(
+        id="1",
+        label="MyLabel",
+        properties={"id": "abc"},
+        embedding_properties={"vectorProperty": [1.0, 2.0, 3.0]},
+    )
+    end_node = Neo4jNode(
+        id="2",
+        label="MyLabel",
+        properties={"id": "def"},
+    )
+    relationship = Neo4jRelationship(
+        start_node_id="1", end_node_id="2", type="MY_RELATIONSHIP"
+    )
+    node_with_two_embeddings = Neo4jNode(
+        id="3",
+        label="MyLabel",
+        properties={"id": "ghi"},
+        embedding_properties={
+            "vectorProperty": [1.0, 2.0, 3.0],
+            "otherVectorProperty": [10.0, 20.0, 30.0],
+        },
+    )
+    graph = Neo4jGraph(
+        nodes=[start_node, end_node, node_with_two_embeddings],
+        relationships=[relationship],
+    )
+
+    neo4j_writer = Neo4jWriter(driver=driver)
+    res = await neo4j_writer.run(graph=graph)
+    assert res.status == "SUCCESS"
+
+    query = """
+    MATCH (a:MyLabel {id: 'abc'})-[r:MY_RELATIONSHIP]->(b:MyLabel {id: 'def'})
+    RETURN a, r, b
+    """
+    record = driver.execute_query(query).records[0]
+    assert "a" and "b" and "r" in record.keys()
+
+    node_a = record["a"]
+    assert start_node.label in list(node_a.labels)
+    assert start_node.properties.get("id") == str(node_a.get("id"))
+    for key, val in start_node.properties.items():
+        assert key in node_a.keys()
+        assert val == node_a.get(key)
+    if start_node.embedding_properties:  # for mypy
+        for emb_key, emb_val in start_node.embedding_properties.items():
+            assert emb_key in node_a.keys()
+            assert emb_val == node_a.get(emb_key)
+
+    node_b = record["b"]
+    assert end_node.label in list(node_b.labels)
+    assert end_node.properties.get("id") == str(node_b.get("id"))
+    for key, val in end_node.properties.items():
+        assert key in node_b.keys()
+        assert val == node_b.get(key)
+
+    rel = record["r"]
+    assert rel.type == relationship.type
+    assert rel.start_node.get("id") == start_node.properties.get("id")
+    assert rel.end_node.get("id") == end_node.properties.get("id")
+
+    query = """
+    MATCH (c:MyLabel {id: 'ghi'})
+    RETURN c
+    """
+    records = driver.execute_query(query).records
+    assert len(records) == 1
+    node_c = records[0]["c"]
+    if node_with_two_embeddings.embedding_properties:  # for mypy
+        for emb_key, emb_val in node_with_two_embeddings.embedding_properties.items():
+            assert emb_key in node_c.keys()
+            assert emb_val == node_c.get(emb_key)
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_kg_writer_no_neo4j_deprecation_warning(
+    driver: neo4j.Driver, caplog: pytest.LogCaptureFixture
+) -> None:
+    start_node = Neo4jNode(
+        id="1",
+        label="MyLabel",
+        properties={"chunk": 1},
+        embedding_properties={"vectorProperty": [1.0, 2.0, 3.0]},
+    )
+    end_node = Neo4jNode(
+        id="2",
+        label="MyLabel",
+        properties={},
+    )
+    relationship = Neo4jRelationship(
+        start_node_id="1", end_node_id="2", type="MY_RELATIONSHIP"
+    )
+    graph = Neo4jGraph(
+        nodes=[start_node, end_node],
+        relationships=[relationship],
+    )
+
+    neo4j_writer = Neo4jWriter(driver=driver)
+    with caplog.at_level(logging.WARNING):
+        res = await neo4j_writer.run(graph=graph)
+
+    for record in caplog.records:
+        if (
+            "Neo.ClientNotification.Statement.FeatureDeprecationWarning"
+            in record.message
+        ):
+            assert False, f"Deprecation warning found in logs: {record.message}"
+
+    assert res.status == "SUCCESS"
+
+
+class _LocalParquetDestination:
+    """E2E test-only implementation of ParquetOutputDestination for a local directory."""
+
+    def __init__(self, path: Path) -> None:
+        self._path = Path(path)
+        self._path.mkdir(parents=True, exist_ok=True)
+
+    @property
+    def output_path(self) -> str:
+        return str(self._path.resolve())
+
+    async def write(self, data: bytes, filename: str) -> None:
+        (self._path / filename).write_bytes(data)
+
+
+@pytest.mark.asyncio
+async def test_parquet_writer_e2e() -> None:
+    """E2E test for ParquetWriter: write graph to Parquet files and verify content."""
+    pyarrow = pytest.importorskip("pyarrow")
+
+    start_node = Neo4jNode(
+        id="p1",
+        label="Person",
+        properties={"name": "Alice", "age": 30},
+    )
+    end_node = Neo4jNode(
+        id="p2",
+        label="Person",
+        properties={"name": "Bob", "age": 25},
+    )
+    relationship = Neo4jRelationship(
+        start_node_id="p1",
+        end_node_id="p2",
+        type="KNOWS",
+    )
+    graph = Neo4jGraph(
+        nodes=[start_node, end_node],
+        relationships=[relationship],
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        output_path = Path(tmpdir)
+        dest = _LocalParquetDestination(output_path)
+        collision_handler = FilenameCollisionHandler()
+        writer = ParquetWriter(
+            nodes_dest=dest,
+            relationships_dest=dest,
+            collision_handler=collision_handler,
+            prefix="e2e_",
+        )
+        result = await writer.run(
+            graph=graph,
+            lexical_graph_config=LexicalGraphConfig(),
+        )
+
+        assert result.status == "SUCCESS"
+        assert result.metadata is not None
+        stats = result.metadata.get("statistics") or {}
+        assert stats["node_count"] == 2
+        assert stats["relationship_count"] == 1
+        assert stats["nodes_per_label"]["Person"] == 2
+        assert stats["rel_per_type"]["KNOWS"] == 1
+        assert "input_files_count" in stats
+        assert "input_files_total_size_bytes" in stats
+
+        files_meta = result.metadata.get("files") or []
+        assert len(files_meta) == 2, f"Expected 2 files, got {files_meta}"
+        node_file = rel_file = None
+        for f in files_meta:
+            path = Path(f["file_path"])
+            assert path.exists(), f"Expected file {path}"
+            table = pyarrow.parquet.read_table(path)
+            if "from" in table.column_names and "to" in table.column_names:
+                rel_file = path
+            elif "labels" in table.column_names:
+                node_file = path
+        assert node_file is not None, "No node Parquet file found"
+        assert rel_file is not None, "No relationship Parquet file found"
+
+        node_table = pyarrow.parquet.read_table(node_file)
+        assert node_table.num_rows == 2
+        assert "name" in node_table.column_names
+        assert "age" in node_table.column_names
+        name_values = node_table.column("name").to_pylist()
+        assert "Alice" in name_values and "Bob" in name_values
+
+        rel_table = pyarrow.parquet.read_table(rel_file)
+        assert rel_table.num_rows == 1
+        assert "from" in rel_table.column_names
+        assert "to" in rel_table.column_names
+        assert rel_table.column("type")[0].as_py() == "KNOWS"
+        assert rel_table.column("from")[0].as_py() == "p1"
+        assert rel_table.column("to")[0].as_py() == "p2"
+
+
+@pytest.mark.asyncio
+async def test_parquet_writer_preserves_lexical_graph_nodes_and_rels() -> None:
+    """ParquetWriter must preserve lexical graph nodes and lexical relationship types in Parquet output."""
+    pyarrow = pytest.importorskip("pyarrow")
+
+    config = LexicalGraphConfig(
+        document_node_label="__Document__",
+        chunk_node_label="__Chunk__",
+        chunk_to_document_relationship_type="__CHUNK_TO_DOCUMENT__",
+        next_chunk_relationship_type="__NEXT_CHUNK__",
+        node_to_chunk_relationship_type="__NODE_TO_CHUNK__",
+    )
+
+    # Lexical graph nodes (custom labels from config)
+    doc_node = Neo4jNode(
+        id="doc-1",
+        label=config.document_node_label,
+        properties={"id": "doc-1", "name": "MyDoc"},
+    )
+    chunk_a = Neo4jNode(
+        id="chunk-1",
+        label=config.chunk_node_label,
+        properties={"id": "chunk-1", "index": 0, "text": "First chunk."},
+    )
+    chunk_b = Neo4jNode(
+        id="chunk-2",
+        label=config.chunk_node_label,
+        properties={"id": "chunk-2", "index": 1, "text": "Second chunk."},
+    )
+    # Entity node (non-lexical)
+    person_node = Neo4jNode(
+        id="p1",
+        label="Person",
+        properties={"name": "Alice"},
+    )
+
+    # Lexical relationships: Chunk -> Document, Chunk -> Chunk, Person -> Chunk (node-to-chunk)
+    rel_chunk_to_doc = Neo4jRelationship(
+        start_node_id="chunk-1",
+        end_node_id="doc-1",
+        type=config.chunk_to_document_relationship_type,
+    )
+    rel_next_chunk = Neo4jRelationship(
+        start_node_id="chunk-1",
+        end_node_id="chunk-2",
+        type=config.next_chunk_relationship_type,
+    )
+    rel_node_to_chunk = Neo4jRelationship(
+        start_node_id="p1",
+        end_node_id="chunk-1",
+        type=config.node_to_chunk_relationship_type,
+    )
+    # Entity relationship
+    rel_knows = Neo4jRelationship(
+        start_node_id="p1",
+        end_node_id="p1",
+        type="KNOWS",
+    )
+
+    graph = Neo4jGraph(
+        nodes=[doc_node, chunk_a, chunk_b, person_node],
+        relationships=[rel_chunk_to_doc, rel_next_chunk, rel_node_to_chunk, rel_knows],
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        output_path = Path(tmpdir)
+        dest = _LocalParquetDestination(output_path)
+        collision_handler = FilenameCollisionHandler()
+        writer = ParquetWriter(
+            nodes_dest=dest,
+            relationships_dest=dest,
+            collision_handler=collision_handler,
+            prefix="lex_",
+        )
+        result = await writer.run(
+            graph=graph,
+            lexical_graph_config=config,
+        )
+
+        assert result.status == "SUCCESS"
+        assert result.metadata is not None
+        statistics = result.metadata.get("statistics") or {}
+        assert statistics["node_count"] == 4
+        assert statistics["relationship_count"] == 4
+        assert "input_files_count" in statistics
+        assert "input_files_total_size_bytes" in statistics
+
+        # Returned metadata must include lexical graph node and relationship entries in statistics
+        nodes_per_label = statistics["nodes_per_label"]
+        rel_per_type = statistics["rel_per_type"]
+
+        # Custom lexical node labels must appear in nodes_per_label
+        assert config.document_node_label in nodes_per_label, (
+            f"FAIL: Custom '{config.document_node_label}' label not found in nodes_per_label. "
+            f"Found: {list(nodes_per_label.keys())}. "
+            f"Lexical graph config was not applied correctly."
+        )
+        assert config.chunk_node_label in nodes_per_label, (
+            f"FAIL: Custom '{config.chunk_node_label}' label not found in nodes_per_label. "
+            f"Found: {list(nodes_per_label.keys())}. "
+            f"Lexical graph config was not applied correctly."
+        )
+
+        # Non-zero counts for lexical nodes
+        assert nodes_per_label[config.document_node_label] > 0, (
+            f"Expected at least 1 {config.document_node_label} node, "
+            f"got {nodes_per_label[config.document_node_label]}"
+        )
+        assert nodes_per_label[config.chunk_node_label] > 0, (
+            f"Expected at least 1 {config.chunk_node_label} node, "
+            f"got {nodes_per_label[config.chunk_node_label]}"
+        )
+
+        # Default labels must NOT be present
+        assert "Document" not in nodes_per_label, (
+            f"FAIL: Found default 'Document' label (should be '{config.document_node_label}'). "
+            f"Labels: {list(nodes_per_label.keys())}"
+        )
+        assert "Chunk" not in nodes_per_label, (
+            f"FAIL: Found default 'Chunk' label (should be '{config.chunk_node_label}'). "
+            f"Labels: {list(nodes_per_label.keys())}"
+        )
+
+        # Custom lexical relationship types must appear in rel_per_type (keyed by type name)
+        assert config.chunk_to_document_relationship_type in rel_per_type, (
+            f"FAIL: Custom '{config.chunk_to_document_relationship_type}' relationship not found. "
+            f"Found: {list(rel_per_type.keys())}"
+        )
+        assert config.next_chunk_relationship_type in rel_per_type, (
+            f"FAIL: Custom '{config.next_chunk_relationship_type}' relationship not found. "
+            f"Found: {list(rel_per_type.keys())}"
+        )
+        assert config.node_to_chunk_relationship_type in rel_per_type, (
+            f"FAIL: Custom '{config.node_to_chunk_relationship_type}' relationship not found. "
+            f"Found: {list(rel_per_type.keys())}"
+        )
+
+        # Non-zero counts for lexical relationships
+        assert (
+            rel_per_type.get(config.chunk_to_document_relationship_type, 0) > 0
+        ), f"Expected at least 1 {config.chunk_to_document_relationship_type} relationship"
+        assert (
+            rel_per_type.get(config.next_chunk_relationship_type, 0) > 0
+        ), f"Expected at least 1 {config.next_chunk_relationship_type} relationship"
+        assert (
+            rel_per_type.get(config.node_to_chunk_relationship_type, 0) > 0
+        ), f"Expected at least 1 {config.node_to_chunk_relationship_type} relationship"
+
+        # Default relationship type names must NOT be present
+        assert (
+            "FROM_DOCUMENT" not in rel_per_type
+        ), f"FAIL: Found default 'FROM_DOCUMENT' (should be '{config.chunk_to_document_relationship_type}')"
+        assert (
+            "FROM_CHUNK" not in rel_per_type
+        ), f"FAIL: Found default 'FROM_CHUNK' (should be '{config.node_to_chunk_relationship_type}')"
+
+        # metadata["files"] must include lexical graph node and relationship file entries
+        files_meta = result.metadata.get("files") or []
+        node_files_meta = [f for f in files_meta if f.get("is_node") is True]
+        rel_files_meta = [f for f in files_meta if f.get("is_node") is False]
+        assert any(
+            f.get("name") == config.document_node_label
+            or config.document_node_label in (f.get("labels") or [])
+            for f in node_files_meta
+        ), "metadata files should include lexical document node file"
+        assert any(
+            f.get("name") == config.chunk_node_label
+            or config.chunk_node_label in (f.get("labels") or [])
+            for f in node_files_meta
+        ), "metadata files should include lexical chunk node file"
+        assert any(
+            f.get("relationship_type") == config.chunk_to_document_relationship_type
+            for f in rel_files_meta
+        ), "metadata files should include chunk-to-document relationship file"
+        assert any(
+            f.get("relationship_type") == config.next_chunk_relationship_type
+            for f in rel_files_meta
+        ), "metadata files should include next-chunk relationship file"
+        assert any(
+            f.get("relationship_type") == config.node_to_chunk_relationship_type
+            for f in rel_files_meta
+        ), "metadata files should include node-to-chunk relationship file"
+
+        files_meta_paths = [
+            Path(f["file_path"]) for f in (result.metadata.get("files") or [])
+        ]
+        node_files = {}
+        rel_files = {}
+        for path in files_meta_paths:
+            assert path.exists(), f"Expected file {path}"
+            table = pyarrow.parquet.read_table(path)
+            if "from" in table.column_names and "to" in table.column_names:
+                rel_type = table.column("type")[0].as_py() if table.num_rows else None
+                key = (
+                    table.column("from_label")[0].as_py(),
+                    rel_type,
+                    table.column("to_label")[0].as_py(),
+                )
+                rel_files[key] = path
+            elif "labels" in table.column_names:
+                labels_col = table.column("labels")
+                first_labels = labels_col.slice(0, 1)
+                label_set = set(first_labels[0].as_py()) if first_labels else set()
+                if config.document_node_label in label_set:
+                    node_files[config.document_node_label] = path
+                elif (
+                    config.chunk_node_label in label_set
+                    and "__Entity__" not in label_set
+                ):
+                    node_files[config.chunk_node_label] = path
+                elif "Person" in label_set:
+                    node_files["Person"] = path
+
+        assert (
+            config.document_node_label in node_files
+        ), "Document node Parquet file should exist"
+        assert (
+            config.chunk_node_label in node_files
+        ), "Chunk node Parquet file should exist"
+        assert "Person" in node_files, "Person node Parquet file should exist"
+
+        # Lexical nodes: labels column must NOT contain __Entity__
+        doc_table = pyarrow.parquet.read_table(node_files[config.document_node_label])
+        assert doc_table.num_rows == 1
+        doc_labels = doc_table.column("labels")[0].as_py()
+        assert config.document_node_label in doc_labels
+        assert (
+            "__Entity__" not in doc_labels
+        ), "Lexical document node must not have __Entity__ label"
+        assert doc_table.column(INTERNAL_ID_PROPERTY)[0].as_py() == "doc-1"
+        assert doc_table.column("name")[0].as_py() == "MyDoc"
+
+        chunk_table = pyarrow.parquet.read_table(node_files[config.chunk_node_label])
+        assert chunk_table.num_rows == 2
+        for i in range(2):
+            chunk_labels = chunk_table.column("labels")[i].as_py()
+            assert config.chunk_node_label in chunk_labels
+            assert (
+                "__Entity__" not in chunk_labels
+            ), "Lexical chunk node must not have __Entity__ label"
+        ids = chunk_table.column(INTERNAL_ID_PROPERTY).to_pylist()
+        assert "chunk-1" in ids and "chunk-2" in ids
+        texts = chunk_table.column("text").to_pylist()
+        assert "First chunk." in texts and "Second chunk." in texts
+
+        # Entity node: must have __Entity__ in labels
+        person_table = pyarrow.parquet.read_table(node_files["Person"])
+        assert person_table.num_rows == 1
+        person_labels = person_table.column("labels")[0].as_py()
+        assert "Person" in person_labels
+        assert "__Entity__" in person_labels
+
+        # Lexical relationships preserved in Parquet files
+        chunk_doc_key = (
+            config.chunk_node_label,
+            config.chunk_to_document_relationship_type,
+            config.document_node_label,
+        )
+        chunk_chunk_key = (
+            config.chunk_node_label,
+            config.next_chunk_relationship_type,
+            config.chunk_node_label,
+        )
+        node_to_chunk_key = (
+            "Person",
+            config.node_to_chunk_relationship_type,
+            config.chunk_node_label,
+        )
+        assert (
+            chunk_doc_key in rel_files
+        ), "Chunk-to-document relationship file should exist"
+        assert chunk_chunk_key in rel_files, "Next-chunk relationship file should exist"
+        assert (
+            node_to_chunk_key in rel_files
+        ), "Node-to-chunk relationship file should exist"
+
+        from_doc_table = pyarrow.parquet.read_table(rel_files[chunk_doc_key])
+        assert from_doc_table.num_rows == 1
+        assert (
+            from_doc_table.column("type")[0].as_py()
+            == config.chunk_to_document_relationship_type
+        )
+        assert from_doc_table.column("from")[0].as_py() == "chunk-1"
+        assert from_doc_table.column("to")[0].as_py() == "doc-1"
+
+        next_chunk_table = pyarrow.parquet.read_table(rel_files[chunk_chunk_key])
+        assert next_chunk_table.num_rows == 1
+        assert (
+            next_chunk_table.column("type")[0].as_py()
+            == config.next_chunk_relationship_type
+        )
+        assert next_chunk_table.column("from")[0].as_py() == "chunk-1"
+        assert next_chunk_table.column("to")[0].as_py() == "chunk-2"
+
+        node_to_chunk_table = pyarrow.parquet.read_table(rel_files[node_to_chunk_key])
+        assert node_to_chunk_table.num_rows == 1
+        assert (
+            node_to_chunk_table.column("type")[0].as_py()
+            == config.node_to_chunk_relationship_type
+        )
+        assert node_to_chunk_table.column("from")[0].as_py() == "p1"
+        assert node_to_chunk_table.column("to")[0].as_py() == "chunk-1"

참고/neo4j-graphrag-python-main/tests/e2e/experimental/test_lexical_graph_builder_component_e2e.py

@@ -0,0 +1,59 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import neo4j
+import pytest
+from neo4j_graphrag.experimental.components.kg_writer import Neo4jWriter
+from neo4j_graphrag.experimental.components.lexical_graph import LexicalGraphBuilder
+from neo4j_graphrag.experimental.components.types import (
+    LexicalGraphConfig,
+    TextChunk,
+    TextChunks,
+)
+from neo4j_graphrag.experimental.pipeline import Pipeline
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_lexical_graph_component_alone_default_config(
+    driver: neo4j.Driver,
+) -> None:
+    pipe = Pipeline()
+    pipe.add_component(LexicalGraphBuilder(), "lexical_graph")
+    pipe.add_component(Neo4jWriter(driver), "writer")
+    pipe.connect("lexical_graph", "writer", {"graph": "lexical_graph.graph"})
+
+    result = await pipe.run(
+        {
+            "lexical_graph": {
+                "text_chunks": TextChunks(chunks=[TextChunk(text="my text", index=0)])
+            }
+        }
+    )
+    assert result.result["writer"]["status"] == "SUCCESS"
+    meta = result.result["writer"]["metadata"]
+    assert "statistics" in meta
+    assert meta["statistics"]["node_count"] == 1
+    assert meta["statistics"]["relationship_count"] == 0
+    default_config = LexicalGraphConfig()
+    created_chunks = driver.execute_query(
+        f"MATCH (n:{default_config.chunk_node_label}) RETURN n"
+    )
+    assert len(created_chunks.records) == 1
+    created_chunks = driver.execute_query(
+        f"MATCH (n:{default_config.document_node_label}) RETURN n"
+    )
+    assert len(created_chunks.records) == 0
+    created_rels = driver.execute_query("MATCH ()-[r]->() RETURN r")
+    assert len(created_rels.records) == 0

참고/neo4j-graphrag-python-main/tests/e2e/experimental/test_neo4j_reader_component_e2e.py

@@ -0,0 +1,124 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from unittest.mock import MagicMock
+
+import neo4j
+import pytest
+from neo4j_graphrag.experimental.components.entity_relation_extractor import (
+    LLMEntityRelationExtractor,
+)
+from neo4j_graphrag.experimental.components.neo4j_reader import Neo4jChunkReader
+from neo4j_graphrag.experimental.components.types import LexicalGraphConfig, TextChunk
+from neo4j_graphrag.experimental.pipeline import Pipeline
+from neo4j_graphrag.llm import LLMResponse
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction_with_chunks")
+async def test_neo4j_reader(driver: neo4j.Driver) -> None:
+    reader = Neo4jChunkReader(driver)
+    res = await reader.run()
+    assert len(res.chunks) == 2
+    assert res.chunks[0] == TextChunk(
+        index=0,
+        text="some text",
+        metadata={"embedding": None},
+        uid="0",
+    )
+    assert res.chunks[1] == TextChunk(
+        index=1,
+        text="some longer text",
+        metadata={"embedding": None},
+        uid="1",
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction_with_chunks")
+async def test_neo4j_reader_in_pipeline(driver: neo4j.Driver, llm: MagicMock) -> None:
+    llm.ainvoke.side_effect = [
+        LLMResponse(
+            content="""{
+                "nodes": [
+                    {
+                        "id": "0",
+                        "label": "Person",
+                        "properties": {
+                            "name": "Harry Potter"
+                        }
+                    }
+                ],
+                "relationships": []
+            }"""
+        ),
+        LLMResponse(content='{"nodes": [], "relationships": []}'),
+    ]
+    pipeline = Pipeline()
+    pipeline.add_component(Neo4jChunkReader(driver), "reader")
+    pipeline.add_component(
+        LLMEntityRelationExtractor(llm, create_lexical_graph=False), "extractor"
+    )
+    pipeline.connect("reader", "extractor", {"chunks": "reader"})
+    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
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction_with_chunks")
+async def test_neo4j_reader_in_pipeline_with_lexical_graph_config(
+    driver: neo4j.Driver, llm: MagicMock
+) -> None:
+    llm.ainvoke.side_effect = [
+        LLMResponse(
+            content="""{
+                "nodes": [
+                    {
+                        "id": "0",
+                        "label": "Person",
+                        "properties": {
+                            "name": "Harry Potter"
+                        }
+                    }
+                ],
+                "relationships": []
+            }"""
+        ),
+        LLMResponse(content='{"nodes": [], "relationships": []}'),
+    ]
+    pipeline = Pipeline()
+    pipeline.add_component(Neo4jChunkReader(driver), "reader")
+    pipeline.add_component(
+        LLMEntityRelationExtractor(llm, create_lexical_graph=False), "extractor"
+    )
+    pipeline.connect("reader", "extractor", {"chunks": "reader"})
+    lg_config = LexicalGraphConfig(node_to_chunk_relationship_type="COMES_FROM")
+    pipeline_output = await pipeline.run(
+        {
+            "extractor": {
+                "lexical_graph_config": lg_config,
+            }
+        }
+    )
+    created_graph = pipeline_output.result["extractor"]
+    assert len(created_graph["nodes"]) == 1
+    assert len(created_graph["relationships"]) == 1  # entity to chunk relationship
+    assert (
+        created_graph["relationships"][0]["type"]
+        == lg_config.node_to_chunk_relationship_type
+        == "COMES_FROM"
+    )

참고/neo4j-graphrag-python-main/tests/e2e/experimental/test_simplekgpipeline_e2e.py

@@ -0,0 +1,481 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock
+
+import neo4j
+import pytest
+from neo4j import Driver
+from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import (
+    FixedSizeSplitter,
+)
+from neo4j_graphrag.experimental.components.types import LexicalGraphConfig
+from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
+from neo4j_graphrag.llm import LLMResponse
+
+
+@pytest.fixture(scope="function", autouse=True)
+def clear_db(driver: Driver) -> Any:
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    yield
+
+
+@pytest.fixture(scope="module")
+def llm_json_response_3_nodes_2_relationships() -> str:
+    return """{
+                        "nodes": [
+                            {
+                                "id": "0",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Harry Potter"
+                                }
+                            },
+                            {
+                                "id": "1",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Alastor Mad-Eye Moody"
+                                }
+                            },
+                            {
+                                "id": "2",
+                                "label": "Organization",
+                                "properties": {
+                                    "name": "The Order of the Phoenix"
+                                }
+                            }
+                        ],
+                        "relationships": [
+                            {
+                                "type": "KNOWS",
+                                "start_node_id": "0",
+                                "end_node_id": "1"
+                            },
+                            {
+                                "type": "LED_BY",
+                                "start_node_id": "2",
+                                "end_node_id": "1"
+                            }
+                        ]
+                    }"""
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_pipeline_builder_happy_path_legacy_schema(
+    harry_potter_text: str,
+    llm: MagicMock,
+    embedder: MagicMock,
+    driver: neo4j.Driver,
+    llm_json_response_3_nodes_2_relationships: str,
+) -> None:
+    """When everything works as expected, extracted entities, relations and text
+    chunks must be in the DB
+    """
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    embedder.async_embed_query.return_value = [1, 2, 3]
+    llm.ainvoke.side_effect = [
+        LLMResponse(
+            content=llm_json_response_3_nodes_2_relationships,
+        )
+    ]
+
+    # Instantiate Entity and Relation objects
+    entities = ["Person", "Organization", "Horcrux", "Location"]
+    relations = ["SITUATED_AT", "INTERACTS", "OWNS", "LED_BY"]
+    potential_schema = [
+        ("Person", "SITUATED_AT", "Location"),
+        ("Person", "INTERACTS", "Person"),
+        ("Person", "OWNS", "Horcrux"),
+        ("Organization", "LED_BY", "Person"),
+    ]
+
+    # Additional arguments
+    lexical_graph_config = LexicalGraphConfig(chunk_node_label="chunkNodeLabel")
+    from_file = False
+    on_error = "RAISE"
+
+    # Create an instance of the SimpleKGPipeline
+    kg_builder_text = SimpleKGPipeline(
+        llm=llm,
+        driver=driver,
+        embedder=embedder,
+        entities=entities,
+        relations=relations,
+        potential_schema=potential_schema,
+        from_file=from_file,
+        on_error=on_error,
+        lexical_graph_config=lexical_graph_config,
+    )
+
+    # Run the knowledge graph building process with text input
+    await kg_builder_text.run_async(text=harry_potter_text)
+
+    # check the content of the graph:
+    # check lexical graph content
+    records, _, _ = driver.execute_query("MATCH (start:chunkNodeLabel) RETURN start")
+    assert len(records) == 1
+
+    # check entity -> chunk relationships
+    records, _, _ = driver.execute_query(
+        "MATCH (chunk:chunkNodeLabel)<-[rel:FROM_CHUNK]-(entity:__Entity__) RETURN chunk, rel, entity"
+    )
+    assert len(records) == 3  # three entities according to mocked LLMResponse
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_pipeline_builder_happy_path(
+    harry_potter_text: str,
+    llm: MagicMock,
+    embedder: MagicMock,
+    driver: neo4j.Driver,
+    llm_json_response_3_nodes_2_relationships: str,
+) -> None:
+    """When everything works as expected, extracted entities, relations and text
+    chunks must be in the DB
+    """
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    embedder.async_embed_query.return_value = [1, 2, 3]
+    llm.ainvoke.side_effect = [
+        LLMResponse(
+            content=llm_json_response_3_nodes_2_relationships,
+        )
+    ]
+
+    # Instantiate schema
+    entities = ["Person"]
+    relations: list[str] = []
+    potential_schema: list[tuple[str, str, str]] = []
+    schema = {
+        "node_types": entities,
+        "relationship_types": relations,
+        "patterns": potential_schema,
+        "additional_node_types": False,
+    }
+
+    # Additional arguments
+    lexical_graph_config = LexicalGraphConfig(chunk_node_label="chunkNodeLabel")
+    from_file = False
+    on_error = "RAISE"
+
+    # Create an instance of the SimpleKGPipeline
+    kg_builder_text = SimpleKGPipeline(
+        llm=llm,
+        driver=driver,
+        embedder=embedder,
+        schema=schema,  # type: ignore[arg-type]
+        from_file=from_file,
+        on_error=on_error,
+        lexical_graph_config=lexical_graph_config,
+    )
+
+    # Run the knowledge graph building process with text input
+    await kg_builder_text.run_async(text=harry_potter_text)
+
+    # check the content of the graph:
+    # check lexical graph content
+    records, _, _ = driver.execute_query("MATCH (start:chunkNodeLabel) RETURN start")
+    assert len(records) == 1
+
+    # check entity -> chunk relationships
+    records, _, _ = driver.execute_query(
+        "MATCH (chunk:chunkNodeLabel)<-[rel:FROM_CHUNK]-(entity:__Entity__) RETURN chunk, rel, entity"
+    )
+    assert len(records) == 2  # only two persons
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_pipeline_builder_two_documents(
+    harry_potter_text_part1: str,
+    harry_potter_text_part2: str,
+    llm: MagicMock,
+    embedder: MagicMock,
+    driver: neo4j.Driver,
+) -> None:
+    """When everything works as expected, extracted entities, relations and text
+    chunks must be in the DB
+    """
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    embedder.async_embed_query.return_value = [1, 2, 3]
+    llm.ainvoke.side_effect = [
+        # first document
+        # first chunk
+        LLMResponse(
+            content="""{
+                        "nodes": [
+                            {
+                                "id": "0",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Harry Potter"
+                                }
+                            },
+                        ],
+                        "relationships": []
+                    }"""
+        ),
+        # second chunk
+        LLMResponse(content='{"nodes": [], "relationships": []}'),
+        # second document
+        # first chunk
+        LLMResponse(
+            content="""{
+                        "nodes": [
+                            {
+                                "id": "0",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Hermione Granger"
+                                }
+                            },
+                        ],
+                        "relationships": []
+                    }"""
+        ),
+        # second chunk
+        LLMResponse(content='{"nodes": [], "relationships": []}'),
+    ]
+
+    # Create an instance of the SimpleKGPipeline
+    kg_builder_text = SimpleKGPipeline(
+        llm=llm,
+        driver=driver,
+        embedder=embedder,
+        from_file=False,
+        # provide minimal schema to bypass automatic schema extraction
+        entities=["Person"],
+        # in order to have 2 chunks:
+        text_splitter=FixedSizeSplitter(chunk_size=400, chunk_overlap=5),
+    )
+
+    # Run the knowledge graph building process with text input
+    await kg_builder_text.run_async(text=harry_potter_text_part1)
+    await kg_builder_text.run_async(text=harry_potter_text_part2)
+
+    # check graph content
+    # check lexical graph content
+    records, _, _ = driver.execute_query(
+        "MATCH (start:Chunk)-[rel:NEXT_CHUNK]->(end:Chunk) RETURN start, rel, end"
+    )
+    assert len(records) == 2  # one for each run
+
+    # check entity -> chunk relationships
+    records, _, _ = driver.execute_query(
+        "MATCH (chunk:Chunk)<-[rel:FROM_CHUNK]-(entity:__Entity__) RETURN chunk, rel, entity"
+    )
+    assert len(records) == 2  # two entities according to mocked LLMResponse
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_pipeline_builder_same_document_two_runs(
+    harry_potter_text_part1: str,
+    llm: MagicMock,
+    embedder: MagicMock,
+    driver: neo4j.Driver,
+) -> None:
+    """When everything works as expected, extracted entities, relations and text
+    chunks must be in the DB
+    """
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    embedder.async_embed_query.return_value = [1, 2, 3]
+    llm.ainvoke.side_effect = [
+        # first run
+        # first chunk
+        LLMResponse(
+            content="""{
+                        "nodes": [
+                            {
+                                "id": "0",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Harry Potter"
+                                }
+                            },
+                        ],
+                        "relationships": []
+                    }"""
+        ),
+        # second chunk
+        LLMResponse(content='{"nodes": [], "relationships": []}'),
+        # second run
+        # first chunk
+        LLMResponse(
+            content="""{
+                        "nodes": [
+                            {
+                                "id": "0",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Harry Potter"
+                                }
+                            },
+                        ],
+                        "relationships": []
+                    }"""
+        ),
+        # second chunk
+        LLMResponse(content='{"nodes": [], "relationships": []}'),
+    ]
+
+    # Create an instance of the SimpleKGPipeline
+    kg_builder_text = SimpleKGPipeline(
+        llm=llm,
+        driver=driver,
+        embedder=embedder,
+        from_file=False,
+        # provide minimal schema to bypass automatic schema extraction
+        entities=["Person"],
+        # in order to have 2 chunks:
+        text_splitter=FixedSizeSplitter(chunk_size=400, chunk_overlap=5),
+    )
+
+    # Run the knowledge graph building process with text input
+    await kg_builder_text.run_async(text=harry_potter_text_part1)
+    await kg_builder_text.run_async(text=harry_potter_text_part1)
+
+    # check lexical graph content
+    records, _, _ = driver.execute_query(
+        "MATCH (start:Chunk)-[rel:NEXT_CHUNK]->(end:Chunk) RETURN start, rel, end"
+    )
+    assert len(records) == 2  # one for each run
+
+    # check entity -> chunk relationships
+    records, _, _ = driver.execute_query(
+        "MATCH (chunk:Chunk)<-[rel:FROM_CHUNK]-(entity:__Entity__) RETURN chunk, rel, entity"
+    )
+    assert len(records) == 2  # two entities according to mocked LLMResponse
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_pipeline_builder_with_automatic_schema_extraction(
+    harry_potter_text_part1: str,
+    llm: MagicMock,
+    embedder: MagicMock,
+    driver: neo4j.Driver,
+) -> None:
+    """Test pipeline with automatic schema extraction (no schema provided).
+    This test verifies that the pipeline correctly handles automatic schema extraction.
+    """
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    embedder.async_embed_query.return_value = [1, 2, 3]
+
+    # set up mock LLM responses for both schema extraction and entity extraction
+    llm.ainvoke.side_effect = [
+        # first call - schema extraction response
+        LLMResponse(
+            content="""{
+                "node_types": [
+                    {
+                        "label": "Person",
+                        "description": "A character in the story",
+                        "properties": [
+                            {"name": "name", "type": "STRING"},
+                            {"name": "age", "type": "INTEGER"}
+                        ]
+                    },
+                    {
+                        "label": "Location",
+                        "description": "A place in the story",
+                        "properties": [
+                            {"name": "name", "type": "STRING"}
+                        ]
+                    }
+                ],
+                "relationship_types": [
+                    {
+                        "label": "LOCATED_AT",
+                        "description": "Indicates where a person is located",
+                        "properties": []
+                    }
+                ],
+                "patterns": [
+                    ["Person", "LOCATED_AT", "Location"]
+                ]
+            }"""
+        ),
+        # second call - entity extraction for first chunk
+        LLMResponse(
+            content="""{
+                "nodes": [
+                    {
+                        "id": "0",
+                        "label": "Person",
+                        "properties": {
+                            "name": "Harry Potter"
+                        }
+                    },
+                    {
+                        "id": "1",
+                        "label": "Location",
+                        "properties": {
+                            "name": "Hogwarts"
+                        }
+                    }
+                ],
+                "relationships": [
+                    {
+                        "type": "LOCATED_AT",
+                        "start_node_id": "0",
+                        "end_node_id": "1"
+                    }
+                ]
+            }"""
+        ),
+        # third call - entity extraction for second chunk (if text is split)
+        LLMResponse(content='{"nodes": [], "relationships": []}'),
+    ]
+
+    # create an instance of the SimpleKGPipeline with NO schema provided
+    kg_builder_text = SimpleKGPipeline(
+        llm=llm,
+        driver=driver,
+        embedder=embedder,
+        from_file=False,
+        # use smaller chunk size to ensure we have at least 2 chunks
+        text_splitter=FixedSizeSplitter(chunk_size=400, chunk_overlap=5),
+    )
+
+    # run the knowledge graph building process with text input
+    await kg_builder_text.run_async(text=harry_potter_text_part1)
+
+    # verify LLM was called for schema extraction
+    assert llm.ainvoke.call_count >= 2
+
+    # verify entities were created
+    records, _, _ = driver.execute_query("MATCH (n:Person) RETURN n")
+    assert len(records) == 1
+
+    # verify locations were created
+    records, _, _ = driver.execute_query("MATCH (n:Location) RETURN n")
+    assert len(records) == 1
+
+    # verify relationships were created
+    records, _, _ = driver.execute_query(
+        "MATCH (p:Person)-[r:LOCATED_AT]->(l:Location) RETURN p, r, l"
+    )
+    assert len(records) == 1
+
+    # verify chunks and relationships to entities
+    records, _, _ = driver.execute_query(
+        "MATCH (c:Chunk)<-[:FROM_CHUNK]-(e) RETURN c, e"
+    )
+    assert len(records) >= 1