참고소스 수정본

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/__init__.py

@@ -0,0 +1,14 @@
+#  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.

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/sample_data/hello.md

@@ -0,0 +1,3 @@
+# Hello
+
+Markdown **content** for tests.

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_data_loader.py

@@ -0,0 +1,194 @@
+#  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 importlib
+import sys
+from pathlib import Path
+from typing import Optional, Union
+from unittest.mock import patch
+
+import pytest
+from fsspec import AbstractFileSystem
+from fsspec.implementations.local import LocalFileSystem
+from neo4j_graphrag.exceptions import MarkdownLoadError, PdfLoaderError
+from neo4j_graphrag.experimental.components.data_loader import (
+    MarkdownLoader,
+    PdfLoader,
+)
+from neo4j_graphrag.experimental.components.types import DocumentType, LoadedDocument
+
+BASE_DIR = Path(__file__).parent
+
+
+@pytest.fixture
+def pdf_loader() -> PdfLoader:
+    return PdfLoader()
+
+
+@pytest.fixture
+def dummy_pdf_path() -> str:
+    return str(BASE_DIR / "sample_data/lorem_ipsum.pdf")
+
+
+@pytest.fixture
+def dummy_md_path() -> str:
+    return str(BASE_DIR / "sample_data/hello.md")
+
+
+def test_pdf_loading(pdf_loader: PdfLoader, dummy_pdf_path: str) -> None:
+    expected_content = "Lorem ipsum dolor sit amet."
+    actual_content = pdf_loader.load_file(dummy_pdf_path, fs=LocalFileSystem())
+    assert actual_content == expected_content
+
+
+def test_pdf_processing_error(pdf_loader: PdfLoader, dummy_pdf_path: str) -> None:
+    with patch(
+        "fsspec.implementations.local.LocalFileSystem.open",
+        side_effect=Exception("Failed to open"),
+    ):
+        with pytest.raises(PdfLoaderError):
+            pdf_loader.load_file(dummy_pdf_path, fs=LocalFileSystem())
+
+
+def test_markdown_processing_error(dummy_md_path: str) -> None:
+    with patch(
+        "fsspec.implementations.local.LocalFileSystem.open",
+        side_effect=Exception("Failed to open"),
+    ):
+        with pytest.raises(MarkdownLoadError):
+            MarkdownLoader.load_file(dummy_md_path, fs=LocalFileSystem())
+
+
+def test_markdown_loading() -> None:
+    md_path = str(BASE_DIR / "sample_data/hello.md")
+    text = MarkdownLoader.load_file(md_path, fs=LocalFileSystem())
+    assert "# Hello" in text
+    assert "Markdown **content**" in text
+
+
+@pytest.mark.asyncio
+async def test_markdown_loader_run() -> None:
+    md_path = BASE_DIR / "sample_data/hello.md"
+    loader = MarkdownLoader()
+    doc = await loader.run(filepath=md_path)
+    assert doc.document_info.document_type == DocumentType.MARKDOWN
+    assert "# Hello" in doc.text
+
+
+@pytest.mark.asyncio
+async def test_pdf_loader_run() -> None:
+    """``PdfLoader.run`` wraps ``load_file`` with :class:`DocumentInfo` (default ``fs``)."""
+    pdf_path = BASE_DIR / "sample_data/lorem_ipsum.pdf"
+    loader = PdfLoader()
+    doc = await loader.run(filepath=pdf_path)
+    assert doc.document_info.document_type == DocumentType.PDF
+    assert doc.document_info.path == str(pdf_path)
+    assert doc.text == "Lorem ipsum dolor sit amet."
+
+
+@pytest.mark.asyncio
+async def test_pdf_loader_run_fs_string_resolves_with_fsspec(
+    dummy_pdf_path: str,
+) -> None:
+    """``fs`` may be a protocol name passed to ``fsspec.filesystem`` (e.g. ``\"file\"``)."""
+    loader = PdfLoader()
+    doc = await loader.run(filepath=dummy_pdf_path, fs="file")
+    assert "Lorem ipsum" in doc.text
+
+
+@pytest.mark.asyncio
+async def test_markdown_loader_run_fs_string() -> None:
+    md_path = str(BASE_DIR / "sample_data/hello.md")
+    loader = MarkdownLoader()
+    doc = await loader.run(filepath=md_path, fs="file")
+    assert doc.document_info.document_type == DocumentType.MARKDOWN
+    assert "# Hello" in doc.text
+
+
+@pytest.mark.asyncio
+async def test_run_passes_metadata_to_document_info(dummy_pdf_path: str) -> None:
+    loader = PdfLoader()
+    meta = {"source": "unit-test", "lang": "en"}
+    doc = await loader.run(filepath=dummy_pdf_path, metadata=meta)
+    assert doc.document_info.metadata == meta
+
+
+class _PdfLoaderWithDerivedMetadata(PdfLoader):
+    """Exercise :meth:`DataLoader.get_document_metadata` override."""
+
+    async def run(
+        self,
+        filepath: Union[str, Path],
+        metadata: Optional[dict[str, str]] = None,
+        fs: Optional[Union[AbstractFileSystem, str]] = None,
+    ) -> LoadedDocument:
+        return await super().run(filepath=filepath, metadata=metadata, fs=fs)
+
+    def get_document_metadata(
+        self, text: str, metadata: dict[str, str] | None = None
+    ) -> dict[str, str] | None:
+        base = dict(metadata or {})
+        base["text_length"] = str(len(text))
+        return base
+
+
+@pytest.mark.asyncio
+async def test_get_document_metadata_override_merges_into_document_info(
+    dummy_pdf_path: str,
+) -> None:
+    loader = _PdfLoaderWithDerivedMetadata()
+    doc = await loader.run(
+        filepath=dummy_pdf_path,
+        metadata={"source": "derived-test"},
+    )
+    assert doc.document_info.metadata is not None
+    assert doc.document_info.metadata["source"] == "derived-test"
+    assert doc.document_info.metadata["text_length"] == str(len(doc.text))
+
+
+def test_pdf_loader_non_local_filesystem_branch_uses_bytesio(
+    dummy_pdf_path: str,
+) -> None:
+    """Non-\"default\" local FS (``auto_mkdir=True``) reads into BytesIO for pypdf."""
+    from neo4j_graphrag.experimental.components.data_loader import is_default_fs
+
+    fs = LocalFileSystem(auto_mkdir=True)
+    assert is_default_fs(fs) is False
+    text = PdfLoader.load_file(dummy_pdf_path, fs=fs)
+    assert text == "Lorem ipsum dolor sit amet."
+
+
+def test_pdf_loader_backward_compat_reexport_module() -> None:
+    """``pdf_loader`` submodule re-exports the same classes as ``data_loader``."""
+    from neo4j_graphrag.experimental.components.data_loader import (
+        DataLoader as DataLoaderDirect,
+        PdfLoader as PdfLoaderDirect,
+    )
+
+    with pytest.warns(DeprecationWarning, match="pdf_loader"):
+        from neo4j_graphrag.experimental.components.pdf_loader import (
+            DataLoader as DataLoaderReexport,
+            PdfLoader as PdfLoaderReexport,
+        )
+
+    assert PdfLoaderDirect is PdfLoaderReexport
+    assert DataLoaderDirect is DataLoaderReexport
+
+
+def test_pdf_loader_module_emits_import_time_deprecation_warning() -> None:
+    module_name = "neo4j_graphrag.experimental.components.pdf_loader"
+    sys.modules.pop(module_name, None)
+    with pytest.warns(DeprecationWarning, match="Importing from .*pdf_loader"):
+        importlib.import_module(module_name)

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_embedder.py

@@ -0,0 +1,41 @@
+#  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 pytest
+from neo4j_graphrag.experimental.components.embedder import TextChunkEmbedder
+from neo4j_graphrag.experimental.components.types import (
+    TextChunk,
+    TextChunks,
+)
+
+
+@pytest.mark.asyncio
+async def test_text_chunk_embedder_run(embedder: MagicMock) -> None:
+    embedder.async_embed_query.return_value = [1.0, 2.0, 3.0]
+    text_chunk_embedder = TextChunkEmbedder(embedder=embedder)
+    text_chunks = TextChunks(
+        chunks=[TextChunk(text="may thy knife chip and shatter", index=0)]
+    )
+    embedded_chunks = await text_chunk_embedder.run(text_chunks)
+    embedder.async_embed_query.assert_called_once_with("may thy knife chip and shatter")
+    assert isinstance(embedded_chunks, TextChunks)
+    for chunk in embedded_chunks.chunks:
+        assert isinstance(chunk, TextChunk)
+        assert chunk.metadata is not None
+        assert "embedding" in chunk.metadata.keys()
+        assert isinstance(chunk.metadata["embedding"], list)
+        for i in chunk.metadata["embedding"]:
+            assert isinstance(i, float)

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_entity_relation_extractor.py

@@ -0,0 +1,492 @@
+#  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 json
+from unittest.mock import MagicMock
+
+import pytest
+from neo4j_graphrag.exceptions import LLMGenerationError
+from neo4j_graphrag.experimental.components.entity_relation_extractor import (
+    LLMEntityRelationExtractor,
+    OnError,
+    balance_curly_braces,
+    fix_invalid_json,
+)
+from neo4j_graphrag.experimental.components.types import (
+    DocumentInfo,
+    Neo4jGraph,
+    TextChunk,
+    TextChunks,
+)
+from neo4j_graphrag.experimental.pipeline.exceptions import InvalidJSONError
+from neo4j_graphrag.llm import LLMInterface, LLMResponse
+from neo4j_graphrag.llm import AnthropicLLM, OpenAILLM, VertexAILLM
+from unittest.mock import patch
+
+
+@pytest.mark.asyncio
+async def test_extractor_happy_path_no_entities_no_document() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(content='{"nodes": [], "relationships": []}')
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    result = await extractor.run(chunks=chunks)
+    assert isinstance(result, Neo4jGraph)
+    # only one Chunk node (no document info provided)
+    assert len(result.nodes) == 1
+    assert result.nodes[0].label == "Chunk"
+    assert result.relationships == []
+
+
+@pytest.mark.asyncio
+async def test_extractor_happy_path_no_entities() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(content='{"nodes": [], "relationships": []}')
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    document_info = DocumentInfo(path="path")
+    result = await extractor.run(chunks=chunks, document_info=document_info)
+    assert isinstance(result, Neo4jGraph)
+    # one Chunk node and one Document node
+    assert len(result.nodes) == 2
+    assert set(n.label for n in result.nodes) == {"Chunk", "Document"}
+    assert len(result.relationships) == 1
+    assert result.relationships[0].type == "FROM_DOCUMENT"
+
+
+@pytest.mark.asyncio
+async def test_extractor_happy_path_no_entities_no_lexical_graph() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(content='{"nodes": [], "relationships": []}')
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+        create_lexical_graph=False,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    document_info = DocumentInfo(path="path")
+    graph = await extractor.run(chunks=chunks, document_info=document_info)
+    assert graph.nodes == []
+    assert graph.relationships == []
+
+
+@pytest.mark.asyncio
+async def test_extractor_happy_path_non_empty_result() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(
+        content='{"nodes": [{"id": "0", "label": "Person", "properties": {}}], "relationships": []}'
+    )
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    document_info = DocumentInfo(path="path")
+    result = await extractor.run(chunks=chunks, document_info=document_info)
+    assert isinstance(result, Neo4jGraph)
+    assert len(result.nodes) == 3
+    doc = result.nodes[0]
+    assert doc.label == "Document"
+    chunk_entity = result.nodes[1]
+    assert chunk_entity.label == "Chunk"
+    entity = result.nodes[2]
+    assert entity.id == f"{chunk_entity.id}:0"
+    assert entity.label == "Person"
+    assert len(result.relationships) == 2
+    assert result.relationships[0].type == "FROM_DOCUMENT"
+    assert result.relationships[0].start_node_id == f"{chunk_entity.id}"
+    assert result.relationships[0].end_node_id == f"{doc.id}"
+    assert result.relationships[1].type == "FROM_CHUNK"
+
+
+@pytest.mark.asyncio
+async def test_extractor_missing_entity_id() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(
+        content='{"nodes": [{"label": "Person", "properties": {}}], "relationships": []}'
+    )
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    with pytest.raises(LLMGenerationError):
+        await extractor.run(chunks=chunks)
+
+
+@pytest.mark.asyncio
+async def test_extractor_llm_ainvoke_failed() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.side_effect = LLMGenerationError()
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    with pytest.raises(LLMGenerationError):
+        await extractor.run(chunks=chunks)
+
+
+@pytest.mark.asyncio
+async def test_extractor_llm_unfixable_json() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(
+        content='{"nodes": [{"id": "0", "label": "Person", "properties": {}}], "relationships": }'
+    )
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+
+    with pytest.raises(LLMGenerationError):
+        await extractor.run(chunks=chunks)
+
+
+@pytest.mark.asyncio
+async def test_extractor_llm_invalid_json() -> None:
+    """Test what happens when the returned JSON is valid JSON but
+    does not match the expected Pydantic model"""
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(
+        # missing "label" for entity
+        content='{"nodes": [{"id": 0, "entity_type": "Person", "properties": {}}], "relationships": []}'
+    )
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    with pytest.raises(LLMGenerationError):
+        await extractor.run(chunks=chunks)
+
+
+@pytest.mark.asyncio
+async def test_extractor_llm_invalid_json_is_a_list() -> None:
+    """Test what happens when the returned JSON is a valid JSON list,
+    but it does not match the expected Pydantic model"""
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(
+        # missing "label" for entity
+        content='[{"nodes": [{"id": 0, "entity_type": "Person", "properties": {}}], "relationships": []}]'
+    )
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    with pytest.raises(LLMGenerationError):
+        await extractor.run(chunks=chunks)
+
+
+@pytest.mark.asyncio
+async def test_extractor_llm_badly_formatted_json_gets_fixed() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(
+        content='{"nodes": [{"id": "0", "label": "Person", "properties": {}}], "relationships": [}'
+    )
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+        on_error=OnError.IGNORE,
+        create_lexical_graph=False,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    res = await extractor.run(chunks=chunks)
+
+    assert len(res.nodes) == 1
+    assert res.nodes[0].label == "Person"
+    assert res.nodes[0].embedding_properties == {}
+    assert res.relationships == []
+
+
+@pytest.mark.asyncio
+async def test_extractor_custom_prompt() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(content='{"nodes": [], "relationships": []}')
+
+    extractor = LLMEntityRelationExtractor(llm=llm, prompt_template="this is my prompt")
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    await extractor.run(chunks=chunks)
+    llm.ainvoke.assert_called_once_with("this is my prompt")
+
+
+def test_fix_invalid_json_empty_result() -> None:
+    json_string = "invalid json"
+
+    with patch("json_repair.repair_json", return_value=""):
+        with pytest.raises(InvalidJSONError):
+            fix_invalid_json(json_string)
+
+
+def test_fix_unquoted_keys() -> None:
+    json_string = '{name: "John", age: "30"}'
+    expected_result = '{"name": "John", "age": "30"}'
+
+    fixed_json = fix_invalid_json(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_fix_unquoted_string_values() -> None:
+    json_string = '{"name": John, "age": 30}'
+    expected_result = '{"name": "John", "age": 30}'
+
+    fixed_json = fix_invalid_json(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_remove_trailing_commas() -> None:
+    json_string = '{"name": "John", "age": 30,}'
+    expected_result = '{"name": "John", "age": 30}'
+
+    fixed_json = fix_invalid_json(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_fix_excessive_braces() -> None:
+    json_string = '{{"name": "John"}}'
+    expected_result = '{"name": "John"}'
+
+    fixed_json = fix_invalid_json(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_fix_multiple_issues() -> None:
+    json_string = '{name: John, "hobbies": ["reading", "swimming",], "age": 30}'
+    expected_result = '{"name": "John", "hobbies": ["reading", "swimming"], "age": 30}'
+
+    fixed_json = fix_invalid_json(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_fix_null_values() -> None:
+    json_string = '{"name": John, "nickname": null}'
+    expected_result = '{"name": "John", "nickname": null}'
+
+    fixed_json = fix_invalid_json(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_fix_numeric_values() -> None:
+    json_string = '{"age": 30, "score": 95.5}'
+    expected_result = '{"age": 30, "score": 95.5}'
+
+    fixed_json = fix_invalid_json(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_balance_curly_braces_missing_closing() -> None:
+    json_string = '{"name": "John", "hobbies": {"reading": "yes"'
+    expected_result = '{"name": "John", "hobbies": {"reading": "yes"}}'
+
+    fixed_json = balance_curly_braces(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_balance_curly_braces_extra_closing() -> None:
+    json_string = '{"name": "John", "hobbies": {"reading": "yes"}}}'
+    expected_result = '{"name": "John", "hobbies": {"reading": "yes"}}'
+
+    fixed_json = balance_curly_braces(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_balance_curly_braces_balanced_input() -> None:
+    json_string = '{"name": "John", "hobbies": {"reading": "yes"}, "age": 30}'
+    expected_result = json_string
+
+    fixed_json = balance_curly_braces(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_balance_curly_braces_nested_structure() -> None:
+    json_string = '{"person": {"name": "John", "hobbies": {"reading": "yes"}}}'
+    expected_result = json_string
+
+    fixed_json = balance_curly_braces(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_balance_curly_braces_unbalanced_nested() -> None:
+    json_string = '{"person": {"name": "John", "hobbies": {"reading": "yes"}}'
+    expected_result = '{"person": {"name": "John", "hobbies": {"reading": "yes"}}}'
+
+    fixed_json = balance_curly_braces(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_balance_curly_braces_unmatched_openings() -> None:
+    json_string = '{"name": "John", "hobbies": {"reading": "yes"'
+    expected_result = '{"name": "John", "hobbies": {"reading": "yes"}}'
+
+    fixed_json = balance_curly_braces(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_balance_curly_braces_unmatched_closings() -> None:
+    json_string = '{"name": "John", "hobbies": {"reading": "yes"}}}'
+    expected_result = '{"name": "John", "hobbies": {"reading": "yes"}}'
+
+    fixed_json = balance_curly_braces(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_balance_curly_braces_complex_structure() -> None:
+    json_string = (
+        '{"name": "John", "details": {"age": 30, "hobbies": {"reading": "yes"}}}'
+    )
+    expected_result = json_string
+
+    fixed_json = balance_curly_braces(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_balance_curly_braces_incorrect_nested_closings() -> None:
+    json_string = '{"key1": {"key2": {"reading": "yes"}}, "key3": {"age": 30}}}'
+    expected_result = '{"key1": {"key2": {"reading": "yes"}}, "key3": {"age": 30}}'
+
+    fixed_json = balance_curly_braces(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_balance_curly_braces_braces_inside_string() -> None:
+    json_string = '{"name": "John", "example": "a{b}c", "age": 30}'
+    expected_result = json_string
+
+    fixed_json = balance_curly_braces(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+def test_balance_curly_braces_unbalanced_with_string() -> None:
+    json_string = '{"name": "John", "example": "a{b}c", "hobbies": {"reading": "yes"'
+    expected_result = (
+        '{"name": "John", "example": "a{b}c", "hobbies": {"reading": "yes"}}'
+    )
+
+    fixed_json = balance_curly_braces(json_string)
+
+    assert json.loads(fixed_json)
+    assert fixed_json == expected_result
+
+
+@pytest.mark.asyncio
+async def test_extractor_structured_output_with_openai() -> None:
+    """Test that use_structured_output=True works with OpenAILLM."""
+    llm = MagicMock(spec=OpenAILLM)
+    llm.ainvoke.return_value = LLMResponse(content='{"nodes": [], "relationships": []}')
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+        use_structured_output=True,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    await extractor.run(chunks=chunks)
+
+    # Verify ainvoke was called with response_format=Neo4jGraph
+    llm.ainvoke.assert_called_once()
+    call_args = llm.ainvoke.call_args
+    assert call_args[1]["response_format"] == Neo4jGraph
+
+
+@pytest.mark.asyncio
+async def test_extractor_structured_output_with_vertexai() -> None:
+    """Test that use_structured_output=True works with VertexAILLM."""
+    llm = MagicMock(spec=VertexAILLM)
+    llm.ainvoke.return_value = LLMResponse(content='{"nodes": [], "relationships": []}')
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+        use_structured_output=True,
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    await extractor.run(chunks=chunks)
+
+    # Verify ainvoke was called with response_format=Neo4jGraph
+    llm.ainvoke.assert_called_once()
+    call_args = llm.ainvoke.call_args
+    assert call_args[1]["response_format"] == Neo4jGraph
+
+
+def test_extractor_structured_output_unsupported_llm() -> None:
+    """Test that use_structured_output=True raises error with unsupported LLMs."""
+    llm = AnthropicLLM(api_key="test", model_name="claude-3-opus")
+
+    with pytest.raises(ValueError) as exc_info:
+        LLMEntityRelationExtractor(
+            llm=llm,
+            use_structured_output=True,
+        )
+
+    assert "Structured output is not supported" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_extractor_structured_output_false_uses_v1() -> None:
+    """Test that use_structured_output=False uses V1 interface (prompt-based)."""
+    llm = MagicMock(spec=LLMInterface)
+    llm.ainvoke.return_value = LLMResponse(content='{"nodes": [], "relationships": []}')
+
+    extractor = LLMEntityRelationExtractor(
+        llm=llm,
+        use_structured_output=False,  # Default behavior
+    )
+    chunks = TextChunks(chunks=[TextChunk(text="some text", index=0)])
+    await extractor.run(chunks=chunks)
+
+    # Verify ainvoke was called with just a string prompt (V1), not response_format
+    llm.ainvoke.assert_called_once()
+    call_args = llm.ainvoke.call_args
+    assert isinstance(call_args[0][0], str)  # First arg is string prompt
+    assert "response_format" not in call_args[1]  # No response_format kwarg

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_graph_pruning.py

@@ -0,0 +1,758 @@
+#  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 datetime
+from typing import Any, Optional
+from unittest.mock import ANY, Mock, patch
+
+import pytest
+from neo4j_graphrag.experimental.components.graph_pruning import (
+    GraphPruning,
+    GraphPruningResult,
+    PruningReason,
+    PruningStats,
+)
+from neo4j_graphrag.experimental.components.schema import (
+    GraphConstraintType,
+    GraphSchema,
+    NodeType,
+    Pattern,
+    PropertyType,
+    RelationshipType,
+)
+from neo4j_graphrag.experimental.components.types import (
+    LexicalGraphConfig,
+    Neo4jGraph,
+    Neo4jNode,
+    Neo4jRelationship,
+)
+
+
+@pytest.fixture(scope="module")
+def lexical_graph_config() -> LexicalGraphConfig:
+    return LexicalGraphConfig(
+        chunk_node_label="Paragraph",
+    )
+
+
+@pytest.mark.parametrize(
+    "properties, valid_properties, additional_properties, expected_filtered_properties",
+    [
+        (
+            # no required, additional allowed
+            {
+                "name": "John Does",
+                "age": 25,
+            },
+            [
+                PropertyType(
+                    name="name",
+                    type="STRING",
+                )
+            ],
+            True,
+            {
+                "name": "John Does",
+                "age": 25,
+            },
+        ),
+        (
+            # no required, additional not allowed
+            {
+                "name": "John Does",
+                "age": 25,
+            },
+            [
+                PropertyType(
+                    name="name",
+                    type="STRING",
+                )
+            ],
+            False,
+            {
+                "name": "John Does",
+            },
+        ),
+    ],
+)
+def test_graph_pruning_filter_properties(
+    properties: dict[str, Any],
+    valid_properties: list[PropertyType],
+    additional_properties: bool,
+    expected_filtered_properties: dict[str, Any],
+) -> None:
+    pruner = GraphPruning()
+    filtered_properties = pruner._filter_properties(
+        properties,
+        valid_properties,
+        additional_properties=additional_properties,
+        node_label="Label",
+        pruning_stats=PruningStats(),
+    )
+    assert filtered_properties == expected_filtered_properties
+
+
+@pytest.mark.parametrize(
+    "properties, expected_filtered_properties",
+    [
+        (
+            # all good, no bad types
+            {
+                "name": "John Does",
+                "age": 25,
+                "is_active": True,
+            },
+            {
+                "name": "John Does",
+                "age": 25,
+                "is_active": True,
+            },
+        ),
+        (
+            # map must be serialized
+            {
+                "age": {"dob": datetime.date(2000, 1, 1), "age_in_2025": 25},
+            },
+            {
+                "age": '{"dob": "2000-01-01", "age_in_2025": 25}',
+            },
+        ),
+    ],
+)
+def test_graph_pruning_ensure_property_type(
+    properties: dict[str, Any],
+    expected_filtered_properties: dict[str, Any],
+) -> None:
+    pruner = GraphPruning()
+    type_safe_properties = pruner._ensure_property_types(
+        properties,
+    )
+    assert type_safe_properties == expected_filtered_properties
+
+
+@pytest.fixture(scope="module")
+def node_type_no_properties() -> NodeType:
+    return NodeType(label="Person")
+
+
+@pytest.fixture(scope="module")
+def node_type_required_name() -> NodeType:
+    return NodeType(
+        label="Person",
+        properties=[
+            PropertyType(name="name", type="STRING", required=True),
+            PropertyType(name="age", type="INTEGER"),
+        ],
+    )
+
+
+def _graph_schema_for_node_entity(entity: NodeType | None) -> GraphSchema:
+    """Build a GraphSchema from a node type fixture (applies required→EXISTENCE migration)."""
+    if entity is None:
+        return GraphSchema(node_types=tuple())
+    return GraphSchema.model_validate({"node_types": [entity.model_dump()]})
+
+
+def _schema_for_relationship_validation(
+    patterns: tuple[Pattern, ...],
+) -> GraphSchema:
+    """Minimal valid GraphSchema for _validate_relationship tests (REL + Person/Location)."""
+    return GraphSchema.model_validate(
+        {
+            "node_types": [
+                {
+                    "label": "Person",
+                    "properties": [{"name": "name", "type": "STRING"}],
+                },
+                {
+                    "label": "Location",
+                    "properties": [{"name": "name", "type": "STRING"}],
+                },
+            ],
+            "relationship_types": [{"label": "REL"}],
+            "patterns": [tuple(p) for p in patterns],
+        }
+    )
+
+
+@pytest.mark.parametrize(
+    "node, entity, additional_node_types, expected_node",
+    [
+        # all good, with default values
+        (
+            Neo4jNode(id="1", label="Person", properties={"name": "John Doe"}),
+            "node_type_no_properties",
+            True,
+            Neo4jNode(id="1", label="Person", properties={"name": "John Doe"}),
+        ),
+        # properties empty (missing default)
+        (
+            Neo4jNode(id="1", label="Person", properties={"age": 45}),
+            "node_type_required_name",
+            True,
+            None,
+        ),
+        # node label not is schema, additional not allowed
+        (
+            Neo4jNode(id="1", label="Location", properties={"name": "New York"}),
+            None,
+            False,
+            None,
+        ),
+        # node label not is schema, additional allowed
+        (
+            Neo4jNode(id="1", label="Location", properties={"name": "New York"}),
+            None,
+            True,
+            Neo4jNode(id="1", label="Location", properties={"name": "New York"}),
+        ),
+        # node label not valid
+        (
+            Neo4jNode(id="1", label="", properties={"name": "New York"}),
+            "node_type_required_name",
+            True,
+            None,
+        ),
+        # node ID not valid
+        (
+            Neo4jNode(id="", label="Location", properties={"name": "New York"}),
+            "node_type_required_name",
+            True,
+            None,
+        ),
+    ],
+)
+def test_graph_pruning_validate_node(
+    node: Neo4jNode,
+    entity: str,
+    additional_node_types: bool,
+    expected_node: Neo4jNode,
+    request: pytest.FixtureRequest,
+) -> None:
+    e_fixture = request.getfixturevalue(entity) if entity else None
+    schema = _graph_schema_for_node_entity(e_fixture)
+    e = schema.node_type_from_label(node.label) if node.label else None
+
+    pruner = GraphPruning()
+    result = pruner._validate_node(
+        node, PruningStats(), e, schema, additional_node_types
+    )
+    if expected_node is not None:
+        assert result == expected_node
+    else:
+        assert result is None
+
+
+def test_graph_pruning_enforce_nodes_lexical_graph(
+    lexical_graph_config: LexicalGraphConfig,
+) -> None:
+    pruner = GraphPruning()
+    result = pruner._enforce_nodes(
+        nodes=[
+            Neo4jNode(id="1", label="Paragraph"),
+        ],
+        schema=GraphSchema(node_types=tuple(), additional_node_types=False),
+        lexical_graph_config=lexical_graph_config,
+        pruning_stats=PruningStats(),
+    )
+    assert len(result) == 1
+    assert result[0].label == "Paragraph"
+
+
+def test_graph_pruning_enforce_relationships_lexical_graph_with_pruned_nodes(
+    lexical_graph_config: LexicalGraphConfig,
+) -> None:
+    """Test that lexical relationships are pruned when their nodes are pruned."""
+    pruner = GraphPruning()
+
+    # Create nodes: Chunk (lexical) and Person (extracted, will be pruned)
+    nodes = [
+        Neo4jNode(id="chunk-1", label="Paragraph", properties={"text": "test"}),
+        Neo4jNode(
+            id="person-1", label="Person", properties={"age": 30}
+        ),  # Missing required 'name'
+    ]
+
+    # Person node type: name must exist (EXISTENCE constraint)
+    schema = GraphSchema.model_validate(
+        {
+            "node_types": [
+                {
+                    "label": "Person",
+                    "properties": [
+                        {"name": "name", "type": "STRING"},
+                        {"name": "age", "type": "INTEGER"},
+                    ],
+                }
+            ],
+            "constraints": [
+                {
+                    "type": GraphConstraintType.EXISTENCE.value,
+                    "node_type": "Person",
+                    "property_name": "name",
+                    "relationship_type": None,
+                }
+            ],
+        }
+    )
+
+    # Filter nodes - Person should be pruned due to missing required property
+    pruning_stats = PruningStats()
+    filtered_nodes = pruner._enforce_nodes(
+        nodes=nodes,
+        schema=schema,
+        lexical_graph_config=lexical_graph_config,
+        pruning_stats=pruning_stats,
+    )
+
+    # Only Chunk should remain
+    assert len(filtered_nodes) == 1
+    assert filtered_nodes[0].id == "chunk-1"
+    assert pruning_stats.number_of_pruned_nodes == 1
+
+    # Create relationships including FROM_CHUNK to the pruned Person node
+    relationships = [
+        Neo4jRelationship(
+            start_node_id="person-1",
+            end_node_id="chunk-1",
+            type="FROM_CHUNK",  # Lexical relationship
+        ),
+        Neo4jRelationship(
+            start_node_id="chunk-1",
+            end_node_id="person-1",
+            type="NEXT_CHUNK",  # Another lexical relationship
+        ),
+    ]
+
+    # Filter relationships
+    pruning_stats_rels = PruningStats()
+    filtered_rels = pruner._enforce_relationships(
+        relationships=relationships,
+        filtered_nodes=filtered_nodes,
+        schema=schema,
+        lexical_graph_config=lexical_graph_config,
+        pruning_stats=pruning_stats_rels,
+    )
+
+    # Both lexical relationships should be pruned because person-1 node is missing
+    assert len(filtered_rels) == 0
+    assert pruning_stats_rels.number_of_pruned_relationships == 2
+
+
+def test_graph_pruning_key_constraint_prunes_missing_property(
+    lexical_graph_config: LexicalGraphConfig,
+) -> None:
+    """KEY constraints require presence like EXISTENCE (Neo4j key mandatory properties)."""
+    pruner = GraphPruning()
+    nodes = [
+        Neo4jNode(id="chunk-1", label="Paragraph", properties={"text": "test"}),
+        Neo4jNode(
+            id="person-1", label="Person", properties={"age": 30}
+        ),  # Missing KEY 'email'
+    ]
+    schema = GraphSchema.model_validate(
+        {
+            "node_types": [
+                {
+                    "label": "Person",
+                    "properties": [
+                        {"name": "email", "type": "STRING"},
+                        {"name": "age", "type": "INTEGER"},
+                    ],
+                }
+            ],
+            "constraints": [
+                {
+                    "type": GraphConstraintType.KEY.value,
+                    "node_type": "Person",
+                    "property_name": "email",
+                    "relationship_type": None,
+                }
+            ],
+        }
+    )
+    pruning_stats = PruningStats()
+    filtered_nodes = pruner._enforce_nodes(
+        nodes=nodes,
+        schema=schema,
+        lexical_graph_config=lexical_graph_config,
+        pruning_stats=pruning_stats,
+    )
+    assert len(filtered_nodes) == 1
+    assert filtered_nodes[0].id == "chunk-1"
+    assert pruning_stats.number_of_pruned_nodes == 1
+
+
+def _schema_with_relationship_key_constraint() -> GraphSchema:
+    return GraphSchema.model_validate(
+        {
+            "node_types": [
+                {"label": "Person", "properties": [{"name": "name", "type": "STRING"}]},
+                {
+                    "label": "Company",
+                    "properties": [{"name": "name", "type": "STRING"}],
+                },
+            ],
+            "relationship_types": [
+                {
+                    "label": "WORKS_FOR",
+                    "properties": [
+                        {"name": "since", "type": "STRING"},
+                        {"name": "role", "type": "STRING"},
+                    ],
+                }
+            ],
+            "patterns": [("Person", "WORKS_FOR", "Company")],
+            "constraints": [
+                {
+                    "type": GraphConstraintType.KEY.value,
+                    "node_type": "",
+                    "property_name": "since",
+                    "relationship_type": "WORKS_FOR",
+                }
+            ],
+        }
+    )
+
+
+def test_graph_pruning_key_constraint_on_relationship_mandatory_enforced() -> None:
+    """KEY on a relationship contributes to mandatory props (like EXISTENCE)."""
+    schema = _schema_with_relationship_key_constraint()
+    assert schema.mandatory_property_names_for_relationship("WORKS_FOR") == {"since"}
+
+    pruner = GraphPruning()
+    rel = Neo4jRelationship(
+        start_node_id="p1",
+        end_node_id="c1",
+        type="WORKS_FOR",
+        properties={"role": "engineer"},
+    )
+    valid_nodes = {"p1": "Person", "c1": "Company"}
+    pruning_stats = PruningStats()
+    rel_type = schema.relationship_type_from_label("WORKS_FOR")
+    assert rel_type is not None
+    out = pruner._validate_relationship(
+        rel,
+        valid_nodes,
+        pruning_stats,
+        rel_type,
+        schema.additional_relationship_types,
+        schema.patterns,
+        schema.additional_patterns,
+        schema,
+    )
+    assert out is not None
+    assert out.properties == {}
+    assert len(pruning_stats.pruned_relationships) == 1
+    assert (
+        pruning_stats.pruned_relationships[0].pruned_reason
+        == PruningReason.MISSING_REQUIRED_PROPERTY
+    )
+    assert pruning_stats.pruned_relationships[0].metadata.get(
+        "missing_required_properties"
+    ) == ["since"]
+
+
+def test_graph_pruning_key_constraint_on_relationship_kept_when_present() -> None:
+    schema = _schema_with_relationship_key_constraint()
+    pruner = GraphPruning()
+    rel = Neo4jRelationship(
+        start_node_id="p1",
+        end_node_id="c1",
+        type="WORKS_FOR",
+        properties={"since": "2020-01-01", "role": "engineer"},
+    )
+    pruning_stats = PruningStats()
+    rel_type = schema.relationship_type_from_label("WORKS_FOR")
+    assert rel_type is not None
+    out = pruner._validate_relationship(
+        rel,
+        {"p1": "Person", "c1": "Company"},
+        pruning_stats,
+        rel_type,
+        schema.additional_relationship_types,
+        schema.patterns,
+        schema.additional_patterns,
+        schema,
+    )
+    assert out is not None
+    assert out.properties == {"since": "2020-01-01", "role": "engineer"}
+    assert pruning_stats.number_of_pruned_relationships == 0
+
+
+@pytest.fixture
+def neo4j_relationship() -> Neo4jRelationship:
+    return Neo4jRelationship(
+        start_node_id="1",
+        end_node_id="2",
+        type="REL",
+        properties={},
+    )
+
+
+@pytest.fixture
+def neo4j_relationship_invalid_type() -> Neo4jRelationship:
+    return Neo4jRelationship(
+        start_node_id="1",
+        end_node_id="2",
+        type="",
+        properties={},
+    )
+
+
+@pytest.fixture
+def neo4j_reversed_relationship(
+    neo4j_relationship: Neo4jRelationship,
+) -> Neo4jRelationship:
+    return Neo4jRelationship(
+        start_node_id=neo4j_relationship.end_node_id,
+        end_node_id=neo4j_relationship.start_node_id,
+        type=neo4j_relationship.type,
+        properties=neo4j_relationship.properties,
+    )
+
+
+@pytest.mark.parametrize(
+    "relationship, valid_nodes, relationship_type, additional_relationship_types, patterns, additional_patterns, expected_relationship",
+    [
+        # all good
+        (
+            "neo4j_relationship",  # relationship,
+            {  # valid_nodes
+                "1": "Person",
+                "2": "Location",
+            },
+            RelationshipType(  # relationship_type
+                label="REL",
+            ),
+            True,  # additional_relationship_types
+            (
+                Pattern(source="Person", relationship="REL", target="Location"),
+            ),  # patterns
+            True,  # additional_patterns
+            "neo4j_relationship",  # expected_relationship
+        ),
+        # reverse relationship
+        (
+            "neo4j_reversed_relationship",
+            {
+                "1": "Person",
+                "2": "Location",
+            },
+            RelationshipType(
+                label="REL",
+            ),
+            True,  # additional_relationship_types
+            (Pattern(source="Person", relationship="REL", target="Location"),),
+            True,  # additional_patterns
+            "neo4j_relationship",
+        ),
+        # invalid start node ID
+        (
+            "neo4j_reversed_relationship",
+            {
+                "10": "Person",
+                "2": "Location",
+            },
+            RelationshipType(
+                label="REL",
+            ),
+            True,  # additional_relationship_types
+            (Pattern(source="Person", relationship="REL", target="Location"),),
+            True,  # additional_patterns
+            None,
+        ),
+        # invalid type, addition allowed
+        (
+            "neo4j_relationship",
+            {
+                "1": "Person",
+                "2": "Location",
+            },
+            None,  # relationship_type
+            True,  # additional_relationship_types
+            (Pattern(source="Person", relationship="REL", target="Location"),),
+            True,  # additional_patterns
+            "neo4j_relationship",
+        ),
+        # invalid type, addition allowed but invalid node ID
+        (
+            "neo4j_relationship",
+            {
+                "1": "Person",
+            },
+            None,  # relationship_type
+            True,  # additional_relationship_types
+            (Pattern(source="Person", relationship="REL", target="Location"),),
+            True,  # additional_patterns
+            None,
+        ),
+        # invalid type, addition not allowed
+        (
+            "neo4j_relationship",
+            {
+                "1": "Person",
+                "2": "Location",
+            },
+            None,  # relationship_type
+            False,  # additional_relationship_types
+            (Pattern(source="Person", relationship="REL", target="Location"),),
+            True,  # additional_patterns
+            None,
+        ),
+        # invalid pattern, addition allowed
+        (
+            "neo4j_relationship",
+            {
+                "1": "Person",
+                "2": "Location",
+            },
+            RelationshipType(
+                label="REL",
+            ),
+            True,  # additional_relationship_types
+            (Pattern(source="Person", relationship="REL", target="Person"),),
+            True,  # additional_patterns
+            "neo4j_relationship",
+        ),
+        # invalid pattern, addition not allowed
+        (
+            "neo4j_relationship",
+            {
+                "1": "Person",
+                "2": "Location",
+            },
+            RelationshipType(
+                label="REL",
+            ),
+            True,  # additional_relationship_types
+            (Pattern(source="Person", relationship="REL", target="Person"),),
+            False,  # additional_patterns
+            None,
+        ),
+        # invalid extracted type
+        (
+            "neo4j_relationship_invalid_type",  # relationship,
+            {  # valid_nodes
+                "1": "Person",
+                "2": "Location",
+            },
+            RelationshipType(  # relationship_type
+                label="REL",
+            ),
+            True,  # additional_relationship_types
+            (
+                Pattern(source="Person", relationship="REL", target="Location"),
+            ),  # patterns
+            True,  # additional_patterns
+            None,  # expected_relationship
+        ),
+    ],
+)
+def test_graph_pruning_validate_relationship(
+    relationship: str,
+    valid_nodes: dict[str, str],
+    relationship_type: RelationshipType,
+    additional_relationship_types: bool,
+    patterns: tuple[Pattern, ...],
+    additional_patterns: bool,
+    expected_relationship: Optional[str],
+    request: pytest.FixtureRequest,
+) -> None:
+    relationship_obj = request.getfixturevalue(relationship)
+    expected_relationship_obj = (
+        request.getfixturevalue(expected_relationship)
+        if expected_relationship
+        else None
+    )
+
+    pruner = GraphPruning()
+    schema = _schema_for_relationship_validation(patterns)
+    assert (
+        pruner._validate_relationship(
+            relationship_obj,
+            valid_nodes,
+            PruningStats(),
+            relationship_type,
+            additional_relationship_types,
+            patterns,
+            additional_patterns,
+            schema,
+        )
+        == expected_relationship_obj
+    )
+
+
+@patch("neo4j_graphrag.experimental.components.graph_pruning.GraphPruning._clean_graph")
+@pytest.mark.asyncio
+async def test_graph_pruning_run_happy_path(
+    mock_clean_graph: Mock,
+    node_type_required_name: NodeType,
+    lexical_graph_config: LexicalGraphConfig,
+) -> None:
+    initial_graph = Neo4jGraph(
+        nodes=[Neo4jNode(id="1", label="Person"), Neo4jNode(id="2", label="Location")],
+    )
+    schema = _graph_schema_for_node_entity(node_type_required_name)
+    cleaned_graph = Neo4jGraph(nodes=[Neo4jNode(id="1", label="Person")])
+    mock_clean_graph.return_value = (cleaned_graph, PruningStats())
+    pruner = GraphPruning()
+    pruner_result = await pruner.run(
+        graph=initial_graph,
+        schema=schema,
+        lexical_graph_config=lexical_graph_config,
+    )
+    assert isinstance(pruner_result, GraphPruningResult)
+    assert pruner_result.graph == cleaned_graph
+    mock_clean_graph.assert_called_once_with(
+        initial_graph, schema, lexical_graph_config
+    )
+
+
+@pytest.mark.asyncio
+async def test_graph_pruning_run_no_schema() -> None:
+    initial_graph = Neo4jGraph(nodes=[Neo4jNode(id="1", label="Person")])
+    pruner = GraphPruning()
+    pruner_result = await pruner.run(
+        graph=initial_graph,
+        schema=None,
+    )
+    assert isinstance(pruner_result, GraphPruningResult)
+    assert pruner_result.graph == initial_graph
+
+
+@patch(
+    "neo4j_graphrag.experimental.components.graph_pruning.GraphPruning._enforce_nodes"
+)
+def test_graph_pruning_clean_graph(
+    mock_enforce_nodes: Mock,
+    lexical_graph_config: LexicalGraphConfig,
+) -> None:
+    mock_enforce_nodes.return_value = []
+    initial_graph = Neo4jGraph(nodes=[Neo4jNode(id="1", label="Person")])
+    schema = GraphSchema(node_types=())
+    pruner = GraphPruning()
+    cleaned_graph, pruning_stats = pruner._clean_graph(
+        initial_graph, schema, lexical_graph_config
+    )
+    assert cleaned_graph == Neo4jGraph()
+    assert isinstance(pruning_stats, PruningStats)
+    mock_enforce_nodes.assert_called_once_with(
+        [Neo4jNode(id="1", label="Person")],
+        schema,
+        lexical_graph_config,
+        ANY,
+    )

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_graph_schema_extraction.py

@@ -0,0 +1,325 @@
+#  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.
+
+"""Behavioral tests for schema-from-text extraction wire types and conversion."""
+
+from __future__ import annotations
+
+import pytest
+
+from neo4j_graphrag.exceptions import SchemaExtractionError
+from neo4j_graphrag.experimental.components.graph_schema_extraction import (
+    ExtractedConstraintType,
+    ExtractedNodeType,
+    ExtractedPropertyType,
+    ExtractedRelationshipType,
+    GraphSchemaExtractionOutput,
+    wire_extraction_constraints_for_graph_schema,
+)
+from neo4j_graphrag.experimental.components.schema import GraphSchema, Pattern
+
+
+def test_wire_extraction_constraints_empty_and_null_become_none() -> None:
+    """Maps ``\"\"`` and legacy ``null`` to ``None`` for :class:`ConstraintType`."""
+    out = wire_extraction_constraints_for_graph_schema(
+        [
+            {
+                "type": "UNIQUENESS",
+                "node_type": "Person",
+                "property_names": ["id"],
+                "relationship_type": "",
+            },
+            {
+                "type": "UNIQUENESS",
+                "node_type": "Org",
+                "property_names": ["id"],
+                "relationship_type": None,
+            },
+        ]
+    )
+    assert out[0]["relationship_type"] is None
+    assert out[1]["relationship_type"] is None
+
+
+def test_wire_extraction_constraints_preserves_relationship_scoped_existence() -> None:
+    out = wire_extraction_constraints_for_graph_schema(
+        [
+            {
+                "type": "EXISTENCE",
+                "node_type": "",
+                "property_names": ["since"],
+                "relationship_type": "KNOWS",
+            }
+        ]
+    )
+    assert out[0]["relationship_type"] == "KNOWS"
+
+
+def test_uniqueness_with_relationship_type_fails_at_graph_schema_validation() -> None:
+    """If a bad constraint survives extraction filters, :class:`ConstraintType` rejects it."""
+    dto = GraphSchemaExtractionOutput(
+        node_types=[
+            ExtractedNodeType(
+                label="Person",
+                properties=[ExtractedPropertyType(name="name", type="STRING")],
+            )
+        ],
+        relationship_types=[],
+        patterns=[],
+        constraints=[
+            ExtractedConstraintType(
+                type="UNIQUENESS",
+                node_type="Person",
+                property_names=["name"],
+                relationship_type="KNOWS",
+            ),
+        ],
+    )
+    with pytest.raises(SchemaExtractionError):
+        GraphSchema.from_extraction_output(dto)
+
+
+def test_invalid_constraints_dropped_by_extraction_filters_without_error() -> None:
+    """Cross-reference filters drop semantically invalid constraints (see ``_extraction_filter_invalid_constraints``)."""
+    dto = GraphSchemaExtractionOutput(
+        node_types=[
+            ExtractedNodeType(
+                label="Person",
+                properties=[ExtractedPropertyType(name="name", type="STRING")],
+            )
+        ],
+        relationship_types=[],
+        patterns=[],
+        constraints=[
+            ExtractedConstraintType(
+                type="UNIQUENESS",
+                node_type="",
+                property_names=["name"],
+                relationship_type="",
+            ),
+            ExtractedConstraintType(
+                type="EXISTENCE",
+                node_type="Person",
+                property_names=["name"],
+                relationship_type="KNOWS",
+            ),
+            ExtractedConstraintType(
+                type="EXISTENCE",
+                node_type="",
+                property_names=["name"],
+                relationship_type="",
+            ),
+        ],
+    )
+    gs = GraphSchema.from_extraction_output(dto)
+    assert len(gs.constraints) == 0
+
+
+def test_from_extraction_output_relationship_existence_constraint() -> None:
+    """Relationship-scoped EXISTENCE uses empty ``node_type`` and non-empty ``relationship_type``."""
+    dto = GraphSchemaExtractionOutput(
+        node_types=[
+            ExtractedNodeType(
+                label="Person",
+                properties=[ExtractedPropertyType(name="name", type="STRING")],
+            )
+        ],
+        relationship_types=[
+            ExtractedRelationshipType(
+                label="KNOWS",
+                properties=[ExtractedPropertyType(name="since", type="LOCAL_DATETIME")],
+            )
+        ],
+        patterns=[
+            Pattern(
+                source="Person",
+                relationship="KNOWS",
+                target="Person",
+            ),
+        ],
+        constraints=[
+            ExtractedConstraintType(
+                type="EXISTENCE",
+                node_type="",
+                property_names=["since"],
+                relationship_type="KNOWS",
+            ),
+        ],
+    )
+    gs = GraphSchema.from_extraction_output(dto)
+    assert gs.existence_property_names_for_relationship("KNOWS") == {"since"}
+    assert gs.existence_property_names_for_node("Person") == set()
+
+
+def test_from_extraction_output_two_constraints_same_property_distinct_kinds() -> None:
+    """Sanity: UNIQUENESS + EXISTENCE on the same property name yields two runtime constraints."""
+    dto = GraphSchemaExtractionOutput(
+        node_types=[
+            ExtractedNodeType(
+                label="Person",
+                properties=[ExtractedPropertyType(name="name", type="STRING")],
+            )
+        ],
+        relationship_types=[],
+        patterns=[],
+        constraints=[
+            ExtractedConstraintType(
+                type="UNIQUENESS",
+                node_type="Person",
+                property_names=["name"],
+            ),
+            ExtractedConstraintType(
+                type="EXISTENCE",
+                node_type="Person",
+                property_names=["name"],
+                relationship_type="",
+            ),
+        ],
+    )
+    gs = GraphSchema.from_extraction_output(dto)
+    assert len(gs.constraints) == 2
+
+
+def test_from_extraction_output_key_constraint_node() -> None:
+    dto = GraphSchemaExtractionOutput(
+        node_types=[
+            ExtractedNodeType(
+                label="Person",
+                properties=[ExtractedPropertyType(name="email", type="STRING")],
+            )
+        ],
+        relationship_types=[],
+        patterns=[],
+        constraints=[
+            ExtractedConstraintType(
+                type="KEY",
+                node_type="Person",
+                property_names=["email"],
+                relationship_type="",
+            ),
+        ],
+    )
+    gs = GraphSchema.from_extraction_output(dto)
+    assert gs.key_property_names_for_node("Person") == {"email"}
+    assert gs.mandatory_property_names_for_node("Person") == {"email"}
+
+
+def test_from_extraction_output_composite_key_constraint() -> None:
+    """Composite KEY constraint from LLM extraction preserves multi-property grouping."""
+    dto = GraphSchemaExtractionOutput(
+        node_types=[
+            ExtractedNodeType(
+                label="Actor",
+                properties=[
+                    ExtractedPropertyType(name="firstname", type="STRING"),
+                    ExtractedPropertyType(name="surname", type="STRING"),
+                ],
+            )
+        ],
+        relationship_types=[],
+        patterns=[],
+        constraints=[
+            ExtractedConstraintType(
+                type="KEY",
+                node_type="Actor",
+                property_names=["firstname", "surname"],
+            ),
+        ],
+    )
+    gs = GraphSchema.from_extraction_output(dto)
+    assert gs.key_property_names_for_node("Actor") == {"firstname", "surname"}
+    assert gs.mandatory_property_names_for_node("Actor") == {"firstname", "surname"}
+    assert len(gs.constraints) == 1
+    assert gs.constraints[0].property_names == ("firstname", "surname")
+
+
+def test_from_extraction_output_composite_uniqueness_constraint() -> None:
+    """Composite UNIQUENESS constraint from LLM extraction."""
+    dto = GraphSchemaExtractionOutput(
+        node_types=[
+            ExtractedNodeType(
+                label="Book",
+                properties=[
+                    ExtractedPropertyType(name="title", type="STRING"),
+                    ExtractedPropertyType(name="year", type="INTEGER"),
+                ],
+            )
+        ],
+        relationship_types=[],
+        patterns=[],
+        constraints=[
+            ExtractedConstraintType(
+                type="UNIQUENESS",
+                node_type="Book",
+                property_names=["title", "year"],
+            ),
+        ],
+    )
+    gs = GraphSchema.from_extraction_output(dto)
+    assert gs.uniqueness_property_names_for_node("Book") == {"title", "year"}
+    assert len(gs.constraints) == 1
+    assert gs.constraints[0].property_names == ("title", "year")
+
+
+def test_extraction_filter_rejects_composite_existence() -> None:
+    """Composite EXISTENCE from LLM output is filtered out (Neo4j only supports single-property EXISTENCE)."""
+    dto = GraphSchemaExtractionOutput(
+        node_types=[
+            ExtractedNodeType(
+                label="Person",
+                properties=[
+                    ExtractedPropertyType(name="name", type="STRING"),
+                    ExtractedPropertyType(name="email", type="STRING"),
+                ],
+            )
+        ],
+        relationship_types=[],
+        patterns=[],
+        constraints=[
+            ExtractedConstraintType(
+                type="EXISTENCE",
+                node_type="Person",
+                property_names=["name", "email"],
+            ),
+        ],
+    )
+    gs = GraphSchema.from_extraction_output(dto)
+    assert len(gs.constraints) == 0
+
+
+def test_extraction_filter_drops_composite_with_nonexistent_property() -> None:
+    """Composite constraint where one property doesn't exist on the node type is filtered out."""
+    dto = GraphSchemaExtractionOutput(
+        node_types=[
+            ExtractedNodeType(
+                label="Actor",
+                properties=[
+                    ExtractedPropertyType(name="firstname", type="STRING"),
+                ],
+            )
+        ],
+        relationship_types=[],
+        patterns=[],
+        constraints=[
+            ExtractedConstraintType(
+                type="KEY",
+                node_type="Actor",
+                property_names=["firstname", "surname"],
+            ),
+        ],
+    )
+    gs = GraphSchema.from_extraction_output(dto)
+    # "surname" doesn't exist on Actor, so the entire composite constraint is dropped
+    assert len(gs.constraints) == 0

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_graph_schema_extraction_contract.py

@@ -0,0 +1,101 @@
+#  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.
+
+"""Contract tests: keep :class:`GraphSchemaExtractionOutput` aligned with :class:`GraphSchema` / runtime models.
+
+If these fail after a refactor, update the extraction models **and** conversion together.
+"""
+
+from __future__ import annotations
+
+from neo4j_graphrag.experimental.components.graph_schema_extraction import (
+    ExtractedConstraintType,
+    ExtractedNodeType,
+    ExtractedPropertyType,
+    ExtractedRelationshipType,
+    GraphSchemaExtractionOutput,
+)
+from neo4j_graphrag.experimental.components.schema import (
+    ConstraintType,
+    GraphSchema,
+    NodeType,
+    PropertyType,
+    RelationshipType,
+)
+
+
+def test_extracted_property_type_field_names_match_property_type() -> None:
+    """Extraction wire format matches :class:`PropertyType` except deprecated ``required`` (use EXISTENCE)."""
+    assert set(ExtractedPropertyType.model_fields) == set(PropertyType.model_fields) - {
+        "required"
+    }
+
+
+def test_extracted_property_type_uses_same_type_annotation_as_property_type() -> None:
+    """Both models must share the exact ``type`` annotation (``Neo4jPropertyTypeName`` in schema)."""
+    ext_ann = ExtractedPropertyType.model_fields["type"].annotation
+    prop_ann = PropertyType.model_fields["type"].annotation
+    assert ext_ann is prop_ann
+
+
+def test_extracted_node_type_has_core_node_type_fields_only() -> None:
+    """Lean model: same core keys as :class:`NodeType` except ``additional_properties`` (set at conversion)."""
+    assert set(ExtractedNodeType.model_fields) == {"label", "description", "properties"}
+    assert set(ExtractedNodeType.model_fields).issubset(set(NodeType.model_fields))
+
+
+def test_extracted_relationship_type_has_core_relationship_type_fields_only() -> None:
+    assert set(ExtractedRelationshipType.model_fields) == {
+        "label",
+        "description",
+        "properties",
+    }
+    assert set(ExtractedRelationshipType.model_fields).issubset(
+        set(RelationshipType.model_fields)
+    )
+
+
+def test_graph_schema_extraction_output_root_keys_match_validate_payload() -> None:
+    """Keys must match the dict passed to :meth:`GraphSchema.model_validate` from extraction (no ``additional_*``)."""
+    assert set(GraphSchemaExtractionOutput.model_fields) == {
+        "node_types",
+        "relationship_types",
+        "patterns",
+        "constraints",
+    }
+
+
+def test_extracted_constraint_type_aligns_with_runtime_constraint_type() -> None:
+    """Wire model avoids nullable ``relationship_type`` and the deprecated ``property_name``;
+    maps to :class:`ConstraintType` after conversion."""
+    wire_keys = set(ExtractedConstraintType.model_fields)
+    runtime_keys = set(ConstraintType.model_fields)
+    assert wire_keys == runtime_keys - {"property_name"}
+    assert (
+        ExtractedConstraintType.model_fields["type"].annotation
+        != ConstraintType.model_fields["type"].annotation
+    )
+
+
+def test_graph_schema_model_fields_contain_extraction_superset() -> None:
+    """Runtime schema adds pipeline-only fields; extraction is a strict subset at the root."""
+    extraction_keys = set(GraphSchemaExtractionOutput.model_fields)
+    graph_keys = set(GraphSchema.model_fields)
+    assert extraction_keys.issubset(graph_keys)
+    assert {
+        "additional_node_types",
+        "additional_relationship_types",
+        "additional_patterns",
+    }.issubset(graph_keys)

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_lexical_graph_builder.py

@@ -0,0 +1,166 @@
+#  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 uuid
+
+import pytest
+from neo4j_graphrag.experimental.components.lexical_graph import LexicalGraphBuilder
+from neo4j_graphrag.experimental.components.types import (
+    DEFAULT_CHUNK_NODE_LABEL,
+    DEFAULT_CHUNK_TO_DOCUMENT_RELATIONSHIP_TYPE,
+    DEFAULT_DOCUMENT_NODE_LABEL,
+    DEFAULT_NEXT_CHUNK_RELATIONSHIP_TYPE,
+    DocumentInfo,
+    DocumentType,
+    GraphResult,
+    LexicalGraphConfig,
+    Neo4jNode,
+    TextChunk,
+    TextChunks,
+)
+
+
+def test_lexical_graph_builder_create_chunk_node_no_metadata() -> None:
+    builder = LexicalGraphBuilder()
+    node = builder.create_chunk_node(chunk=TextChunk(text="text chunk", index=0))
+    assert isinstance(node, Neo4jNode)
+    assert node.id is not None
+    assert node.properties == {"index": 0, "text": "text chunk"}
+    assert node.embedding_properties == {}
+
+
+def test_lexical_graph_builder_create_chunk_node_metadata_no_embedding() -> None:
+    builder = LexicalGraphBuilder()
+    node = builder.create_chunk_node(
+        chunk=TextChunk(text="text chunk", index=0, metadata={"status": "ok"})
+    )
+    assert isinstance(node, Neo4jNode)
+    assert node.id is not None
+    assert node.properties == {"index": 0, "text": "text chunk", "status": "ok"}
+    assert node.embedding_properties == {}
+
+
+def test_lexical_graph_builder_create_chunk_node_metadata_embedding() -> None:
+    builder = LexicalGraphBuilder()
+    node = builder.create_chunk_node(
+        chunk=TextChunk(
+            text="text chunk",
+            index=0,
+            metadata={"status": "ok", "embedding": [1, 2, 3]},
+        ),
+    )
+    assert isinstance(node, Neo4jNode)
+    assert node.id is not None
+    assert node.properties == {"index": 0, "text": "text chunk", "status": "ok"}
+    assert node.embedding_properties == {"embedding": [1, 2, 3]}
+
+
+@pytest.mark.asyncio
+async def test_lexical_graph_builder_run_with_document() -> None:
+    lexical_graph_builder = LexicalGraphBuilder()
+    doc_uid = str(uuid.uuid4())
+    result = await lexical_graph_builder.run(
+        text_chunks=TextChunks(
+            chunks=[
+                TextChunk(text="text chunk 1", index=0),
+                TextChunk(text="text chunk 1", index=1),
+            ]
+        ),
+        document_info=DocumentInfo(
+            path="test_lexical_graph",
+            uid=doc_uid,
+            document_type=DocumentType.PDF,
+        ),
+    )
+    assert isinstance(result, GraphResult)
+    graph = result.graph
+    nodes = graph.nodes
+    assert len(nodes) == 3
+    document = nodes[0]
+    assert document.id == doc_uid
+    assert document.label == DEFAULT_DOCUMENT_NODE_LABEL
+    assert document.properties["path"] == "test_lexical_graph"
+    assert document.properties["createdAt"] is not None
+    assert document.properties["document_type"] == "pdf"
+    chunk1 = nodes[1]
+    assert chunk1.label == DEFAULT_CHUNK_NODE_LABEL
+    chunk2 = nodes[2]
+    assert chunk2.label == DEFAULT_CHUNK_NODE_LABEL
+    assert len(graph.relationships) == 3
+    assert graph.relationships[0].type == DEFAULT_CHUNK_TO_DOCUMENT_RELATIONSHIP_TYPE
+    assert graph.relationships[1].type == DEFAULT_NEXT_CHUNK_RELATIONSHIP_TYPE
+    assert graph.relationships[2].type == DEFAULT_CHUNK_TO_DOCUMENT_RELATIONSHIP_TYPE
+
+
+@pytest.mark.asyncio
+async def test_lexical_graph_builder_run_no_document() -> None:
+    lexical_graph_builder = LexicalGraphBuilder()
+    result = await lexical_graph_builder.run(
+        text_chunks=TextChunks(
+            chunks=[
+                TextChunk(text="text chunk 1", index=0),
+                TextChunk(text="text chunk 1", index=1),
+            ]
+        ),
+    )
+    assert isinstance(result, GraphResult)
+    graph = result.graph
+    nodes = graph.nodes
+    assert len(nodes) == 2
+    chunk1 = nodes[0]
+    assert chunk1.label == DEFAULT_CHUNK_NODE_LABEL
+    chunk2 = nodes[1]
+    assert chunk2.label == DEFAULT_CHUNK_NODE_LABEL
+    assert len(graph.relationships) == 1
+    assert graph.relationships[0].type == DEFAULT_NEXT_CHUNK_RELATIONSHIP_TYPE
+
+
+@pytest.mark.asyncio
+async def test_lexical_graph_builder_run_custom_labels() -> None:
+    lexical_graph_builder = LexicalGraphBuilder(
+        config=LexicalGraphConfig(
+            document_node_label="Report",
+            chunk_node_label="Page",
+            chunk_to_document_relationship_type="IN_REPORT",
+            next_chunk_relationship_type="NEXT_PAGE",
+        ),
+    )
+    doc_uid = str(uuid.uuid4())
+    result = await lexical_graph_builder.run(
+        text_chunks=TextChunks(
+            chunks=[
+                TextChunk(text="text chunk 1", index=0),
+                TextChunk(text="text chunk 1", index=1),
+            ]
+        ),
+        document_info=DocumentInfo(path="test_lexical_graph", uid=doc_uid),
+    )
+    assert isinstance(result, GraphResult)
+    graph = result.graph
+    nodes = graph.nodes
+    assert len(nodes) == 3
+    document = nodes[0]
+    assert document.id == doc_uid
+    assert document.label == "Report"
+    assert document.properties["path"] == "test_lexical_graph"
+    chunk1 = nodes[1]
+    assert chunk1.label == "Page"
+    chunk2 = nodes[2]
+    assert chunk2.label == "Page"
+    assert len(graph.relationships) == 3
+    assert graph.relationships[0].type == "IN_REPORT"
+    assert graph.relationships[1].type == "NEXT_PAGE"
+    assert graph.relationships[2].type == "IN_REPORT"

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_neo4j_reader.py

@@ -0,0 +1,129 @@
+#  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 Mock
+
+import neo4j
+import pytest
+from neo4j_graphrag.experimental.components.neo4j_reader import (
+    Neo4jChunkReader,
+)
+from neo4j_graphrag.experimental.components.types import LexicalGraphConfig, TextChunks
+
+
+@pytest.mark.asyncio
+async def test_neo4j_chunk_reader(driver: Mock) -> None:
+    driver.execute_query.return_value = (
+        [neo4j.Record({"chunk": {"index": 0, "text": "some text", "id": "azerty"}})],
+        None,
+        None,
+    )
+    chunk_reader = Neo4jChunkReader(driver, neo4j_database="mydb")
+    res = await chunk_reader.run()
+
+    driver.execute_query.assert_called_once_with(
+        "MATCH (c:`Chunk`) RETURN c { .*, embedding: null } as chunk ORDER BY c.index",
+        database_="mydb",
+        routing_=neo4j.RoutingControl.READ,
+    )
+
+    assert isinstance(res, TextChunks)
+    assert len(res.chunks) == 1
+    chunk = res.chunks[0]
+    assert chunk.uid == "azerty"
+    assert chunk.text == "some text"
+    assert chunk.index == 0
+    assert chunk.metadata == {}
+
+
+@pytest.mark.asyncio
+async def test_neo4j_chunk_reader_custom_lg_config(driver: Mock) -> None:
+    driver.execute_query.return_value = (
+        [
+            neo4j.Record(
+                {
+                    "chunk": {
+                        "k": 0,
+                        "content": "some text",
+                        "id": "azerty",
+                        "other": "property",
+                    }
+                }
+            )
+        ],
+        None,
+        None,
+    )
+    chunk_reader = Neo4jChunkReader(driver)
+    res = await chunk_reader.run(
+        lexical_graph_config=LexicalGraphConfig(
+            chunk_node_label="Page",
+            chunk_text_property="content",
+            chunk_index_property="k",
+        )
+    )
+
+    driver.execute_query.assert_called_once_with(
+        "MATCH (c:`Page`) RETURN c { .*, embedding: null } as chunk ORDER BY c.k",
+        database_=None,
+        routing_=neo4j.RoutingControl.READ,
+    )
+
+    assert isinstance(res, TextChunks)
+    assert len(res.chunks) == 1
+    chunk = res.chunks[0]
+    assert chunk.uid == "azerty"
+    assert chunk.text == "some text"
+    assert chunk.index == 0
+    assert chunk.metadata == {"other": "property"}
+
+
+@pytest.mark.asyncio
+async def test_neo4j_chunk_reader_fetch_embedding(driver: Mock) -> None:
+    driver.execute_query.return_value = (
+        [
+            neo4j.Record(
+                {
+                    "chunk": {
+                        "index": 0,
+                        "text": "some text",
+                        "other": "property",
+                        "embedding": [1.0, 2.0, 3.0],
+                        "id": "azerty",
+                    }
+                }
+            )
+        ],
+        None,
+        None,
+    )
+    chunk_reader = Neo4jChunkReader(driver, fetch_embeddings=True)
+    res = await chunk_reader.run()
+
+    driver.execute_query.assert_called_once_with(
+        "MATCH (c:`Chunk`) RETURN c { .* } as chunk ORDER BY c.index",
+        database_=None,
+        routing_=neo4j.RoutingControl.READ,
+    )
+
+    assert isinstance(res, TextChunks)
+    assert len(res.chunks) == 1
+    chunk = res.chunks[0]
+    assert chunk.uid == "azerty"
+    assert chunk.text == "some text"
+    assert chunk.index == 0
+    assert chunk.metadata == {
+        "other": "property",
+        "embedding": [1.0, 2.0, 3.0],
+    }

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_property_type_deprecation.py

@@ -0,0 +1,115 @@
+#  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.
+
+"""Explicit coverage for :class:`PropertyType` ``required`` deprecation (see EXISTENCE constraints)."""
+
+from __future__ import annotations
+
+import pytest
+
+from neo4j_graphrag.experimental.components.schema import (
+    ConstraintType,
+    GraphConstraintType,
+    GraphSchema,
+    NodeType,
+    PropertyType,
+    RelationshipType,
+)
+
+
+def test_property_type_required_field_emits_deprecation_on_access() -> None:
+    """Reading ``.required`` on a model instance should warn (Pydantic deprecated field)."""
+    prop = PropertyType(name="x", type="STRING", required=True)
+    with pytest.warns(DeprecationWarning, match="EXISTENCE"):
+        assert prop.required is True
+
+
+def test_legacy_required_true_in_json_becomes_existence_constraint() -> None:
+    """Loading JSON with ``required: true`` migrates to EXISTENCE (avoid asserting ``.required``; any read warns)."""
+    schema = GraphSchema.model_validate(
+        {
+            "node_types": [
+                {
+                    "label": "Person",
+                    "properties": [
+                        {"name": "name", "type": "STRING", "required": True},
+                    ],
+                }
+            ],
+        }
+    )
+    assert schema.existence_property_names_for_node("Person") == {"name"}
+    assert any(
+        c.type == "EXISTENCE"
+        and c.node_type == "Person"
+        and c.property_names == ("name",)
+        for c in schema.constraints
+    )
+
+
+def test_programmatic_node_type_required_migrates_to_existence() -> None:
+    """``GraphSchema(node_types=(NodeType(..., required=True),))`` migrates like dict input."""
+    nt = NodeType(
+        label="Person",
+        properties=[PropertyType(name="name", type="STRING", required=True)],
+    )
+    schema = GraphSchema(node_types=(nt,))
+    assert schema.existence_property_names_for_node("Person") == {"name"}
+    assert len(schema.constraints) == 1
+    assert schema.constraints[0].type == GraphConstraintType.EXISTENCE
+    assert schema.constraints[0].node_type == "Person"
+    assert schema.constraints[0].property_names == ("name",)
+
+
+def test_programmatic_relationship_type_required_migrates_to_existence() -> None:
+    """``RelationshipType`` with ``PropertyType(required=True)`` migrates to relationship-scoped EXISTENCE."""
+    person = NodeType(
+        label="Person",
+        properties=[PropertyType(name="id", type="STRING")],
+    )
+    knows = RelationshipType(
+        label="KNOWS",
+        properties=[
+            PropertyType(name="since", type="LOCAL_DATETIME", required=True),
+        ],
+    )
+    schema = GraphSchema(node_types=(person,), relationship_types=(knows,))
+    assert schema.existence_property_names_for_relationship("KNOWS") == {"since"}
+    assert any(
+        c.type == GraphConstraintType.EXISTENCE
+        and c.relationship_type == "KNOWS"
+        and c.property_names == ("since",)
+        for c in schema.constraints
+    )
+
+
+def test_programmatic_required_deduped_when_existence_constraint_already_present() -> (
+    None
+):
+    """Pre-existing EXISTENCE constraint does not duplicate when ``required=True`` on instances."""
+    nt = NodeType(
+        label="Person",
+        properties=[PropertyType(name="name", type="STRING", required=True)],
+    )
+    existing = ConstraintType(
+        type=GraphConstraintType.EXISTENCE,
+        node_type="Person",
+        property_names=("name",),
+        property_name="name",
+        relationship_type=None,
+    )
+    schema = GraphSchema(node_types=(nt,), constraints=(existing,))
+    assert len(schema.constraints) == 1
+    assert schema.existence_property_names_for_node("Person") == {"name"}

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_resolver.py

@@ -0,0 +1,303 @@
+#  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 types import SimpleNamespace
+from unittest.mock import MagicMock, call, patch
+
+import neo4j
+import numpy as np
+import pytest
+from neo4j_graphrag.experimental.components.resolver import (
+    FuzzyMatchResolver,
+    SinglePropertyExactMatchResolver,
+    SpaCySemanticMatchResolver,
+)
+from neo4j_graphrag.experimental.components.types import ResolutionStats
+
+
+class FakeNLPModel:
+    """
+    Stand-in for a spaCy NLP model for unit tests.
+    It returns an object with a `.vector` attribute so the resolver can compute cosine similarity.
+    """
+
+    def __call__(self, text: str) -> SimpleNamespace:
+        if "23-45-6789" in text or text == "Alice":
+            return SimpleNamespace(vector=np.array([1.0, 0.0], dtype=np.float64))
+        if text == "Bob":
+            return SimpleNamespace(vector=np.array([0.0, 1.0], dtype=np.float64))
+        return SimpleNamespace(vector=np.array([0.0, 0.0], dtype=np.float64))
+
+
+@pytest.mark.asyncio
+async def test_simple_resolver(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        ([neo4j.Record({"c": 2})], None, None),
+        ([neo4j.Record({"c": 1})], None, None),
+    ]
+    resolver = SinglePropertyExactMatchResolver(driver=driver)
+    res = await resolver.run()
+    assert isinstance(res, ResolutionStats)
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 1
+    assert driver.execute_query.call_count == 2
+    driver.execute_query.assert_has_calls(
+        [call("MATCH (entity:__Entity__)  RETURN count(entity) as c", database_=None)]
+    )
+
+
+@pytest.mark.asyncio
+async def test_simple_resolver_custom_filter(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        ([neo4j.Record({"c": 2})], None, None),
+        ([neo4j.Record({"c": 1})], None, None),
+    ]
+    resolver = SinglePropertyExactMatchResolver(
+        driver=driver, filter_query="WHERE not entity:Resolved"
+    )
+    await resolver.run()
+    driver.execute_query.assert_has_calls(
+        [
+            call(
+                "MATCH (entity:__Entity__) WHERE not entity:Resolved RETURN count(entity) as c",
+                database_=None,
+            )
+        ]
+    )
+
+
+@pytest.mark.asyncio
+async def test_spacy_resolver_match_on_name_property(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 1, "name": "Alice"},
+                            {"id": 2, "name": "Alice"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 1})],
+            None,
+            None,
+        ),
+    ]
+
+    resolver = SpaCySemanticMatchResolver(driver=driver, nlp=FakeNLPModel())
+
+    res = await resolver.run()
+    assert isinstance(res, ResolutionStats)
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 1
+
+    assert driver.execute_query.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_spacy_resolver_no_merge(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 1, "name": "Alice"},
+                            {"id": 2, "name": "Bob"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        ),
+    ]
+
+    resolver = SpaCySemanticMatchResolver(driver=driver, nlp=FakeNLPModel())
+
+    res = await resolver.run()
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 0
+
+    assert driver.execute_query.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_spacy_resolver_match_on_multiple_text_properties(
+    driver: MagicMock,
+) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 10, "name": "John Smith", "ssn": "23-45-6789"},
+                            {"id": 11, "name": "Jonathan Smith", "ssn": "23-45-6789"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 10})],
+            None,
+            None,
+        ),
+    ]
+
+    resolver = SpaCySemanticMatchResolver(
+        driver=driver, resolve_properties=["name", "ssn"], nlp=FakeNLPModel()
+    )
+    res = await resolver.run()
+    assert isinstance(res, ResolutionStats)
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 1
+
+    assert driver.execute_query.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_fuzzy_match_resolver_no_merge(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 1, "name": "Alice"},
+                            {"id": 2, "name": "Bob"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        )
+    ]
+
+    resolver = FuzzyMatchResolver(driver=driver)
+
+    res = await resolver.run()
+    assert isinstance(res, ResolutionStats)
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 0
+
+    assert driver.execute_query.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_fuzzy_match_resolver_multiple_properties(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 10, "name": "John Smith", "ssn": "123-45-6789"},
+                            {"id": 11, "name": "Jon Smith", "ssn": "123-45-6789"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 10})],
+            None,
+            None,
+        ),
+    ]
+
+    resolver = FuzzyMatchResolver(driver=driver, resolve_properties=["name", "ssn"])
+
+    res = await resolver.run()
+    assert isinstance(res, ResolutionStats)
+
+    assert res.number_of_nodes_to_resolve == 2
+    assert res.number_of_created_nodes == 1
+
+    assert driver.execute_query.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_fuzzy_match_resolver_normalization(driver: MagicMock) -> None:
+    # instantiate with a dummy driver
+    resolver = FuzzyMatchResolver(driver=driver)
+
+    sim = resolver.compute_similarity("  ALICE  ", "alice!")
+    assert sim == 1
+
+
+@pytest.mark.asyncio
+async def test_spacy_resolver_caching(driver: MagicMock) -> None:
+    driver.execute_query.side_effect = [
+        (
+            [
+                neo4j.Record(
+                    {
+                        "lab": "Person",
+                        "labelCluster": [
+                            {"id": 1, "name": "Alice"},
+                            {"id": 2, "name": "Alice"},
+                            {"id": 3, "name": "Bob"},
+                        ],
+                    }
+                )
+            ],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 1})],
+            None,
+            None,
+        ),
+        (
+            [neo4j.Record({"id(node)": 3})],
+            None,
+            None,
+        ),
+    ]
+
+    resolver = SpaCySemanticMatchResolver(driver=driver, nlp=FakeNLPModel())
+
+    # patch spaCy NLP call to track how often embeddings are computed
+    with patch.object(resolver, "nlp", wraps=resolver.nlp) as mock_nlp:
+        await resolver.run()
+
+    # "Alice" should be embedded only once, despite being used twice.
+    # "Bob" should be embedded once.
+    assert mock_nlp.call_count == 2, (
+        f"Expected spaCy to embed each unique text once. Got {mock_nlp.call_count} "
+        f"calls."
+    )
+
+    # "Alice" and "Bob" are expected to be the only two distinct texts passed to spaCy.
+    called_texts = {call.args[0] for call in mock_nlp.call_args_list}
+    assert called_texts == {"Alice", "Bob"}

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/test_types.py

@@ -0,0 +1,55 @@
+#  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 pytest
+
+
+def test_pdf_document_alias_is_loaded_document_and_warns() -> None:
+    """Accessing PdfDocument emits DeprecationWarning and returns LoadedDocument."""
+    import neo4j_graphrag.experimental.components.types as types_mod
+
+    with pytest.warns(DeprecationWarning, match="PdfDocument is deprecated"):
+        PdfDocument = types_mod.PdfDocument
+
+    from neo4j_graphrag.experimental.components.types import LoadedDocument
+
+    assert PdfDocument is LoadedDocument
+
+
+def test_pdf_document_from_import_backward_compat_warns() -> None:
+    """``from ...types import PdfDocument`` still works and emits DeprecationWarning."""
+    with pytest.warns(DeprecationWarning, match="PdfDocument is deprecated"):
+        from neo4j_graphrag.experimental.components.types import (
+            LoadedDocument,
+            PdfDocument,
+        )
+
+    assert PdfDocument is LoadedDocument
+
+
+def test_types_dir_includes_pdf_document() -> None:
+    import neo4j_graphrag.experimental.components.types as types_mod
+
+    assert "PdfDocument" in dir(types_mod)
+    assert "LoadedDocument" in dir(types_mod)
+
+
+def test_getattr_unknown_raises() -> None:
+    import neo4j_graphrag.experimental.components.types as types_mod
+
+    with pytest.raises(AttributeError, match="no attribute 'not_a_real_export'"):
+        _ = types_mod.not_a_real_export

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/text_splitters/test_fixed_size_splitter.py

@@ -0,0 +1,214 @@
+#  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 itertools import zip_longest
+
+import pytest
+from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import (
+    FixedSizeSplitter,
+    _adjust_chunk_end,
+    _adjust_chunk_start,
+)
+from neo4j_graphrag.experimental.components.types import TextChunk
+
+
+@pytest.mark.asyncio
+async def test_split_text_no_overlap() -> None:
+    text = "may thy knife chip and shatter"
+    chunk_size = 5
+    chunk_overlap = 0
+    approximate = False
+    splitter = FixedSizeSplitter(chunk_size, chunk_overlap, approximate)
+    chunks = await splitter.run(text)
+    expected_chunks = [
+        TextChunk(text="may t", index=0),
+        TextChunk(text="hy kn", index=1),
+        TextChunk(text="ife c", index=2),
+        TextChunk(text="hip a", index=3),
+        TextChunk(text="nd sh", index=4),
+        TextChunk(text="atter", index=5),
+    ]
+    for actual, expected in zip_longest(chunks.chunks, expected_chunks):
+        assert actual.text == expected.text
+        assert actual.index == expected.index
+        assert expected.uid is not None
+
+
+@pytest.mark.asyncio
+async def test_split_text_with_overlap() -> None:
+    text = "may thy knife chip and shatter"
+    chunk_size = 10
+    chunk_overlap = 2
+    approximate = False
+    splitter = FixedSizeSplitter(chunk_size, chunk_overlap, approximate)
+    chunks = await splitter.run(text)
+    expected_chunks = [
+        TextChunk(text="may thy kn", index=0),
+        TextChunk(text="knife chip", index=1),
+        TextChunk(text="ip and sha", index=2),
+        TextChunk(text="hatter", index=3),
+    ]
+    for actual, expected in zip_longest(chunks.chunks, expected_chunks):
+        assert actual.text == expected.text
+        assert actual.index == expected.index
+        assert expected.uid is not None
+
+
+@pytest.mark.asyncio
+async def test_split_text_empty_string() -> None:
+    text = ""
+    chunk_size = 5
+    chunk_overlap = 1
+    approximate = False
+    splitter = FixedSizeSplitter(chunk_size, chunk_overlap, approximate)
+    chunks = await splitter.run(text)
+    assert chunks.chunks == []
+
+
+def test_invalid_chunk_overlap() -> None:
+    with pytest.raises(ValueError) as excinfo:
+        FixedSizeSplitter(5, 5)
+    assert "chunk_overlap must be strictly less than chunk_size" in str(excinfo)
+
+
+def test_invalid_chunk_size() -> None:
+    with pytest.raises(ValueError) as excinfo:
+        FixedSizeSplitter(0, 0)
+    assert "chunk_size must be strictly greater than 0" in str(excinfo)
+
+
+@pytest.mark.parametrize(
+    "text, approximate_start, expected_start",
+    [
+        # Case: approximate_start is at word boundary already
+        ("Hello World", 6, 6),
+        # Case: approximate_start is at a whitespace already
+        ("Hello World", 5, 5),
+        # Case: approximate_start is at the middle of word and no whitespace is found
+        ("Hello World", 2, 2),
+        # Case: approximate_start is at the middle of a word
+        ("Hello World", 8, 6),
+        # Case: approximate_start = 0
+        ("Hello World", 0, 0),
+    ],
+)
+def test_adjust_chunk_start(
+    text: str, approximate_start: int, expected_start: int
+) -> None:
+    """
+    Test that the _adjust_chunk_start function correctly shifts
+    the start index to avoid breaking words, unless no whitespace is found.
+    """
+    result = _adjust_chunk_start(text, approximate_start)
+    assert result == expected_start
+
+
+@pytest.mark.parametrize(
+    "text, start, approximate_end, expected_end",
+    [
+        # Case: approximate_end is at word boundary already
+        ("Hello World", 0, 5, 5),
+        # Case: approximate_end is at the middle of a word
+        ("Hello World", 0, 8, 6),
+        # Case: approximate_end is at the middle of word and no whitespace is found
+        ("Hello World", 0, 3, 3),
+        # Case: adjusted_end == start => fallback to approximate_end
+        ("Hello World", 6, 7, 7),
+        # Case: end>=len(text)
+        ("Hello World", 6, 15, 15),
+    ],
+)
+def test_adjust_chunk_end(
+    text: str, start: int, approximate_end: int, expected_end: int
+) -> None:
+    """
+    Test that the _adjust_chunk_end function correctly shifts
+    the end index to avoid breaking words, unless no whitespace is found.
+    """
+    result = _adjust_chunk_end(text, start, approximate_end)
+    assert result == expected_end
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "text, chunk_size, chunk_overlap, approximate, expected_chunks",
+    [
+        # Case: approximate fixed size splitting
+        (
+            "Hello World, this is a test message.",
+            10,
+            2,
+            True,
+            ["Hello ", "World, ", "this is a ", "a test ", "message."],
+        ),
+        # Case: fixed size splitting
+        (
+            "Hello World, this is a test message.",
+            10,
+            2,
+            False,
+            ["Hello Worl", "rld, this ", "s is a tes", "est messag", "age."],
+        ),
+        # Case: short text => only one chunk
+        (
+            "Short text",
+            20,
+            5,
+            True,
+            ["Short text"],
+        ),
+        # Case: short text => only one chunk
+        (
+            "Short text",
+            12,
+            4,
+            True,
+            ["Short text"],
+        ),
+        # Case: text with no spaces
+        (
+            "1234567890",
+            5,
+            1,
+            True,
+            ["12345", "56789", "90"],
+        ),
+    ],
+)
+async def test_fixed_size_splitter_run(
+    text: str,
+    chunk_size: int,
+    chunk_overlap: int,
+    approximate: bool,
+    expected_chunks: list[str],
+) -> None:
+    """
+    Test that 'FixedSizeSplitter.run' returns the expected chunks
+    for different configurations.
+    """
+    splitter = FixedSizeSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        approximate=approximate,
+    )
+    text_chunks = await splitter.run(text)
+
+    # Verify number of chunks
+    assert len(text_chunks.chunks) == len(expected_chunks)
+
+    # Verify content of each chunk
+    for i, expected_text in enumerate(expected_chunks):
+        assert text_chunks.chunks[i].text == expected_text
+        assert isinstance(text_chunks.chunks[i], TextChunk)
+        assert text_chunks.chunks[i].index == i

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/text_splitters/test_langchain.py

@@ -0,0 +1,40 @@
+#  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 pytest
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from neo4j_graphrag.experimental.components.text_splitters.langchain import (
+    LangChainTextSplitterAdapter,
+)
+from neo4j_graphrag.experimental.components.types import TextChunk, TextChunks
+
+text = """
+Lorem ipsum dolor sit amet, consectetur adipiscing elit.
+In cursus erat quis ornare condimentum. Ut sollicitudin libero nec quam vestibulum, non tristique augue tempor.
+Nulla fringilla, augue ac fermentum ultricies, mauris tellus tempor orci, at tincidunt purus arcu vitae nisl.
+Nunc suscipit neque vitae ipsum viverra, eu interdum tortor iaculis.
+Suspendisse sit amet quam non ipsum molestie euismod finibus eu nisi. Quisque sit amet aliquet leo, vel auctor dolor.
+Sed auctor enim at tempus eleifend. Suspendisse potenti. Suspendisse congue tellus id justo bibendum, at commodo sapien porta.
+Nam sagittis nisl vitae nibh pellentesque, et convallis turpis ultrices.
+"""
+
+
+@pytest.mark.asyncio
+async def test_langchain_adapter() -> None:
+    text_splitter = LangChainTextSplitterAdapter(RecursiveCharacterTextSplitter())
+    text_chunks = await text_splitter.run(text)
+    assert isinstance(text_chunks, TextChunks)
+    for text_chunk in text_chunks.chunks:
+        assert isinstance(text_chunk, TextChunk)
+        assert text_chunk.text in text

참고/neo4j-graphrag-python-main/tests/unit/experimental/components/text_splitters/test_llamaindex.py

@@ -0,0 +1,40 @@
+#  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 pytest
+from llama_index.core.node_parser.text.sentence import SentenceSplitter
+from neo4j_graphrag.experimental.components.text_splitters.llamaindex import (
+    LlamaIndexTextSplitterAdapter,
+)
+from neo4j_graphrag.experimental.components.types import TextChunk, TextChunks
+
+text = """
+Lorem ipsum dolor sit amet, consectetur adipiscing elit.
+In cursus erat quis ornare condimentum. Ut sollicitudin libero nec quam vestibulum, non tristique augue tempor.
+Nulla fringilla, augue ac fermentum ultricies, mauris tellus tempor orci, at tincidunt purus arcu vitae nisl.
+Nunc suscipit neque vitae ipsum viverra, eu interdum tortor iaculis.
+Suspendisse sit amet quam non ipsum molestie euismod finibus eu nisi. Quisque sit amet aliquet leo, vel auctor dolor.
+Sed auctor enim at tempus eleifend. Suspendisse potenti. Suspendisse congue tellus id justo bibendum, at commodo sapien porta.
+Nam sagittis nisl vitae nibh pellentesque, et convallis turpis ultrices.
+"""
+
+
+@pytest.mark.asyncio
+async def test_llamaindex_adapter() -> None:
+    text_splitter = LlamaIndexTextSplitterAdapter(SentenceSplitter())
+    text_chunks = await text_splitter.run(text)
+    assert isinstance(text_chunks, TextChunks)
+    for text_chunk in text_chunks.chunks:
+        assert isinstance(text_chunk, TextChunk)
+        assert text_chunk.text in text

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/components.py

@@ -0,0 +1,65 @@
+#  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 asyncio
+
+from neo4j_graphrag.experimental.pipeline import Component, DataModel
+from neo4j_graphrag.experimental.pipeline.types.context import RunContext
+
+
+class StringResultModel(DataModel):
+    result: str
+
+
+class IntResultModel(DataModel):
+    result: int
+
+
+class ComponentNoParam(Component):
+    async def run(self) -> StringResultModel:
+        return StringResultModel(result="")
+
+
+class ComponentPassThrough(Component):
+    async def run(self, value: str) -> StringResultModel:
+        return StringResultModel(result=f"value is: {value}")
+
+
+class ComponentAdd(Component):
+    async def run(self, number1: int, number2: int) -> IntResultModel:
+        return IntResultModel(result=number1 + number2)
+
+
+class ComponentMultiply(Component):
+    async def run(self, number1: int, number2: int = 2) -> IntResultModel:
+        return IntResultModel(result=number1 * number2)
+
+
+class ComponentMultiplyWithContext(Component):
+    async def run_with_context(
+        self, context_: RunContext, number1: int, number2: int = 2
+    ) -> IntResultModel:
+        await context_.notify(
+            message="my message", data={"number1": number1, "number2": number2}
+        )
+        return IntResultModel(result=number1 * number2)
+
+
+class SlowComponentMultiply(Component):
+    def __init__(self, sleep: float = 1.0) -> None:
+        self.sleep = sleep
+
+    async def run(self, number1: int, number2: int = 2) -> IntResultModel:
+        await asyncio.sleep(self.sleep)
+        return IntResultModel(result=number1 * number2)

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/config/template_pipeline/test_simple_kg_builder.py

@@ -0,0 +1,594 @@
+#  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 re
+from unittest.mock import Mock, patch
+
+import neo4j
+import pytest
+from neo4j_graphrag.embeddings import Embedder
+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.data_loader import MarkdownLoader, PdfLoader
+from neo4j_graphrag.experimental.components.schema import (
+    SchemaBuilder,
+    SchemaFromTextExtractor,
+    GraphSchema,
+)
+from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import (
+    FixedSizeSplitter,
+)
+from neo4j_graphrag.experimental.pipeline.config.object_config import (
+    ComponentConfig,
+    ComponentType,
+)
+from neo4j_graphrag.experimental.pipeline.config.template_pipeline import (
+    SimpleKGPipelineConfig,
+)
+from neo4j_graphrag.experimental.pipeline.exceptions import PipelineDefinitionError
+from neo4j_graphrag.experimental.pipeline.types.schema import (
+    EntityInputType,
+    RelationInputType,
+)
+from neo4j_graphrag.experimental.components.types import DocumentType
+from neo4j_graphrag.generation.prompts import ERExtractionTemplate
+from neo4j_graphrag.llm import LLMInterface
+
+
+def test_simple_kg_pipeline_config_file_loader_from_file_is_false() -> None:
+    config = SimpleKGPipelineConfig(from_file=False)
+    assert config._get_file_loader() is None
+
+
+def test_simple_kg_pipeline_config_file_loader_from_file_is_true() -> None:
+    config = SimpleKGPipelineConfig(from_file=True)
+    assert config._get_file_loader() is not None
+
+
+@pytest.mark.asyncio
+async def test_simple_kg_pipeline_config_default_file_loader_supports_markdown() -> (
+    None
+):
+    config = SimpleKGPipelineConfig(from_file=True)
+    loader = config._get_file_loader()
+    assert loader is not None
+    doc = await loader.run(
+        filepath="tests/unit/experimental/components/sample_data/hello.md"
+    )
+    assert doc.document_info.document_type == DocumentType.MARKDOWN
+
+
+def test_simple_kg_pipeline_config_from_pdf_deprecated_maps_to_from_file() -> None:
+    with pytest.warns(DeprecationWarning, match="from_pdf"):
+        config = SimpleKGPipelineConfig(from_pdf=True)
+    assert config.from_file is True
+    with pytest.warns(DeprecationWarning, match="from_pdf"):
+        config_false = SimpleKGPipelineConfig(from_pdf=False)
+    assert config_false.from_file is False
+
+
+def test_simple_kg_pipeline_config_pdf_loader_deprecated_maps_to_file_loader() -> None:
+    my_loader = PdfLoader()
+    with pytest.warns(DeprecationWarning, match="pdf_loader"):
+        config = SimpleKGPipelineConfig(
+            from_file=True,
+            pdf_loader=ComponentType(my_loader),
+        )
+    assert config._get_file_loader() == my_loader
+
+
+def test_simple_kg_pipeline_config_pdf_loader_and_file_loader_conflict() -> None:
+    with pytest.raises(ValueError, match="pdf_loader"):
+        SimpleKGPipelineConfig(
+            from_file=True,
+            file_loader=ComponentType(PdfLoader()),
+            pdf_loader=ComponentType(PdfLoader()),
+        )
+
+
+def test_simple_kg_pipeline_config_file_loader_from_file_is_true_class_overwrite() -> (
+    None
+):
+    my_file_loader = MarkdownLoader()
+    config = SimpleKGPipelineConfig(
+        from_file=True, file_loader=ComponentType(my_file_loader)
+    )
+    assert config._get_file_loader() == my_file_loader
+
+
+def test_simple_kg_pipeline_config_file_loader_class_overwrite_but_from_file_is_false() -> (
+    None
+):
+    my_file_loader = PdfLoader()
+    config = SimpleKGPipelineConfig(
+        from_file=False, file_loader=ComponentType(my_file_loader)
+    )
+    assert config._get_file_loader() is None
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.ComponentType.parse")
+def test_simple_kg_pipeline_config_file_loader_from_file_is_true_class_overwrite_from_config(
+    mock_component_parse: Mock,
+) -> None:
+    my_file_loader_config = ComponentConfig(
+        class_="",
+    )
+    my_file_loader = PdfLoader()
+    mock_component_parse.return_value = my_file_loader
+    config = SimpleKGPipelineConfig(
+        from_file=True,
+        file_loader=ComponentType(my_file_loader_config),
+    )
+    assert config._get_file_loader() == my_file_loader
+
+
+def test_simple_kg_pipeline_config_text_splitter() -> None:
+    config = SimpleKGPipelineConfig()
+    assert isinstance(config._get_splitter(), FixedSizeSplitter)
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.ComponentType.parse")
+def test_simple_kg_pipeline_config_text_splitter_overwrite(
+    mock_component_parse: Mock,
+) -> None:
+    my_text_splitter_config = ComponentConfig(
+        class_="",
+    )
+    my_text_splitter = FixedSizeSplitter()
+    mock_component_parse.return_value = my_text_splitter
+    config = SimpleKGPipelineConfig(
+        text_splitter=my_text_splitter_config,  # type: ignore
+    )
+    assert config._get_splitter() == my_text_splitter
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.template_pipeline.simple_kg_builder.SimpleKGPipelineConfig.get_default_embedder"
+)
+def test_simple_kg_pipeline_config_chunk_embedder(
+    mock_embedder: Mock, embedder: Embedder
+) -> None:
+    mock_embedder.return_value = embedder
+    config = SimpleKGPipelineConfig()
+    chunk_embedder = config._get_chunk_embedder()
+    assert isinstance(chunk_embedder, TextChunkEmbedder)
+    assert chunk_embedder._embedder == embedder
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.template_pipeline.simple_kg_builder.SimpleKGPipelineConfig.get_default_llm"
+)
+def test_simple_kg_pipeline_config_automatic_schema(
+    mock_llm: Mock, llm: LLMInterface
+) -> None:
+    mock_llm.return_value = llm
+    config = SimpleKGPipelineConfig()
+    schema = config._get_schema()
+    assert isinstance(schema, SchemaFromTextExtractor)
+    assert schema._llm == llm
+    assert schema.use_structured_output is False
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.template_pipeline.simple_kg_builder.SimpleKGPipelineConfig.get_default_llm"
+)
+def test_simple_kg_pipeline_config_automatic_schema_structured_output(
+    mock_llm: Mock, llm: LLMInterface
+) -> None:
+    llm.supports_structured_output = True
+    mock_llm.return_value = llm
+    config = SimpleKGPipelineConfig()
+    schema = config._get_schema()
+    assert isinstance(schema, SchemaFromTextExtractor)
+    assert schema.use_structured_output is True
+
+
+def test_simple_kg_pipeline_config_manual_schema() -> None:
+    config = SimpleKGPipelineConfig(entities=["Person"])
+    assert isinstance(config._get_schema(), SchemaBuilder)
+
+
+def test_simple_kg_pipeline_config_literal_schema_validation() -> None:
+    config = SimpleKGPipelineConfig(schema="FREE")  # type: ignore
+    assert config.schema_ == GraphSchema.create_empty()
+
+    config = SimpleKGPipelineConfig(schema="EXTRACTED")  # type: ignore
+    assert config.schema_ is None
+
+
+def test_simple_kg_pipeline_config_schema_run_params() -> None:
+    config = SimpleKGPipelineConfig(
+        entities=["Person"],
+        relations=["KNOWS"],
+        potential_schema=[("Person", "KNOWS", "Person")],
+    )
+    assert config._get_run_params_for_schema() == {
+        "node_types": ["Person"],
+        "relationship_types": ["KNOWS"],
+        "patterns": [
+            ("Person", "KNOWS", "Person"),
+        ],
+    }
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.template_pipeline.simple_kg_builder.SimpleKGPipelineConfig.get_default_llm"
+)
+def test_simple_kg_pipeline_config_extractor(mock_llm: Mock, llm: LLMInterface) -> None:
+    mock_llm.return_value = llm
+    config = SimpleKGPipelineConfig(
+        on_error="IGNORE",  # type: ignore
+        prompt_template=ERExtractionTemplate(template="my template {text}"),
+    )
+    extractor = config._get_extractor()
+    assert isinstance(extractor, LLMEntityRelationExtractor)
+    assert extractor.llm == llm
+    assert extractor.on_error == OnError.IGNORE
+    assert extractor.prompt_template.template == "my template {text}"
+    assert extractor.use_structured_output is False
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.template_pipeline.simple_kg_builder.SimpleKGPipelineConfig.get_default_llm"
+)
+def test_simple_kg_pipeline_config_extractor_structured_output(
+    mock_llm: Mock, llm: LLMInterface
+) -> None:
+    llm.supports_structured_output = True
+    mock_llm.return_value = llm
+    config = SimpleKGPipelineConfig()
+    extractor = config._get_extractor()
+    assert isinstance(extractor, LLMEntityRelationExtractor)
+    assert extractor.use_structured_output is True
+
+
+@patch(
+    "neo4j_graphrag.experimental.components.kg_writer.get_version",
+    return_value=((5, 23, 0), False, False),
+)
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.template_pipeline.simple_kg_builder.SimpleKGPipelineConfig.get_default_neo4j_driver"
+)
+def test_simple_kg_pipeline_config_writer(
+    mock_driver: Mock,
+    _: Mock,
+    driver: neo4j.Driver,
+) -> None:
+    mock_driver.return_value = driver
+    config = SimpleKGPipelineConfig(
+        neo4j_database="my_db",
+    )
+    writer = config._get_writer()
+    assert isinstance(writer, Neo4jWriter)
+    assert writer.driver == driver
+    assert writer.neo4j_database == "my_db"
+
+
+@patch(
+    "neo4j_graphrag.experimental.components.kg_writer.get_version",
+    return_value=((5, 23, 0), False, False),
+)
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.ComponentType.parse")
+def test_simple_kg_pipeline_config_writer_overwrite(
+    mock_component_parse: Mock,
+    _: Mock,
+    driver: neo4j.Driver,
+) -> None:
+    my_writer_config = ComponentConfig(
+        class_="",
+    )
+    my_writer = Neo4jWriter(driver, neo4j_database="my_db")
+    mock_component_parse.return_value = my_writer
+    config = SimpleKGPipelineConfig(
+        kg_writer=my_writer_config,  # type: ignore
+        neo4j_database="my_other_db",
+    )
+    writer: Neo4jWriter = config._get_writer()  # type: ignore
+    assert writer == my_writer
+    # database not changed:
+    assert writer.neo4j_database == "my_db"
+
+
+def test_simple_kg_pipeline_config_connections_from_file() -> None:
+    config = SimpleKGPipelineConfig(
+        from_file=True,
+        perform_entity_resolution=False,
+    )
+    connections = config._get_connections()
+    assert len(connections) == 7
+    expected_connections = [
+        ("file_loader", "splitter"),
+        ("file_loader", "schema"),
+        ("schema", "extractor"),
+        ("splitter", "chunk_embedder"),
+        ("chunk_embedder", "extractor"),
+        ("extractor", "pruner"),
+        ("pruner", "writer"),
+    ]
+    for actual, expected in zip(connections, expected_connections):
+        assert (actual.start, actual.end) == expected
+
+
+def test_simple_kg_pipeline_config_connections_from_text() -> None:
+    config = SimpleKGPipelineConfig(
+        from_file=False,
+        perform_entity_resolution=False,
+    )
+    connections = config._get_connections()
+    assert len(connections) == 5
+    expected_connections = [
+        ("schema", "extractor"),
+        ("splitter", "chunk_embedder"),
+        ("chunk_embedder", "extractor"),
+        ("extractor", "pruner"),
+        ("pruner", "writer"),
+    ]
+    for actual, expected in zip(connections, expected_connections):
+        assert (actual.start, actual.end) == expected
+
+
+def test_simple_kg_pipeline_config_connections_with_er() -> None:
+    config = SimpleKGPipelineConfig(
+        from_file=True,
+        perform_entity_resolution=True,
+    )
+    connections = config._get_connections()
+    assert len(connections) == 8
+    expected_connections = [
+        ("file_loader", "splitter"),
+        ("file_loader", "schema"),
+        ("schema", "extractor"),
+        ("splitter", "chunk_embedder"),
+        ("chunk_embedder", "extractor"),
+        ("extractor", "pruner"),
+        ("pruner", "writer"),
+        ("writer", "resolver"),
+    ]
+    for actual, expected in zip(connections, expected_connections):
+        assert (actual.start, actual.end) == expected
+
+
+def test_simple_kg_pipeline_config_run_params_from_file_file_path() -> None:
+    config = SimpleKGPipelineConfig(from_file=True)
+    assert config.get_run_params({"file_path": "my_file"}) == {
+        "file_loader": {"filepath": "my_file", "metadata": None}
+    }
+
+
+def test_simple_kg_pipeline_config_run_params_from_text_text() -> None:
+    config = SimpleKGPipelineConfig(from_file=False)
+    run_params = config.get_run_params({"text": "my text"})
+    assert run_params["splitter"] == {"text": "my text"}
+    assert run_params["schema"] == {"text": "my text"}
+    assert run_params["extractor"]["document_info"]["path"] == "document.txt"
+
+
+def test_simple_kg_pipeline_config_run_params_from_file_text() -> None:
+    config = SimpleKGPipelineConfig(from_file=True)
+    with pytest.raises(PipelineDefinitionError) as excinfo:
+        config.get_run_params({"text": "my text"})
+    assert (
+        "Expected 'file_path' to a PDF or Markdown file when 'from_file' is True"
+        in str(excinfo)
+    )
+
+
+def test_simple_kg_pipeline_config_run_params_from_text_file_path() -> None:
+    config = SimpleKGPipelineConfig(from_file=False)
+    with pytest.raises(PipelineDefinitionError) as excinfo:
+        config.get_run_params({"file_path": "my file"})
+    assert "Expected 'text' argument when 'from_file' is False" in str(excinfo)
+
+
+def test_simple_kg_pipeline_config_run_params_no_file_no_text() -> None:
+    config = SimpleKGPipelineConfig(from_file=False)
+    with pytest.raises(
+        PipelineDefinitionError,
+        match=re.escape(
+            "At least one of `text` (when from_file=False) or `file_path` (when from_file=True) argument must be provided."
+        ),
+    ):
+        config.get_run_params({})
+
+
+def test_simple_kg_pipeline_config_process_schema_with_precedence_legacy() -> None:
+    entities: list[EntityInputType] = [
+        "Person",
+        {
+            "label": "Organization",
+            "description": "A group of persons",
+            "properties": [
+                {
+                    "name": "name",
+                    "type": "STRING",
+                }
+            ],
+        },
+    ]
+    relations: list[RelationInputType] = [
+        "WORKS_FOR",
+        {
+            "label": "CREATED",
+            "description": "A person created an organization",
+            "properties": [
+                {
+                    "name": "date",
+                    "description": "The date the organization was created",
+                    "type": "DATE",
+                },
+                {"name": "isActive", "type": "BOOLEAN"},
+            ],
+        },
+    ]
+    potential_schema = [
+        ("Person", "WORKS_FOR", "Organization"),
+        ("Person", "CREATED", "Organization"),
+    ]
+    config = SimpleKGPipelineConfig(
+        entities=entities,
+        relations=relations,
+        potential_schema=potential_schema,
+    )
+    schema_dict = config._process_schema_with_precedence()
+    node_types = schema_dict["node_types"]
+    relationship_types = schema_dict["relationship_types"]
+    patterns = schema_dict["patterns"]
+    assert len(node_types) == 2
+    assert node_types[0] == "Person"
+    assert node_types[1]["label"] == "Organization"
+    assert len(node_types[1]["properties"]) == 1
+    assert relationship_types is not None
+    assert len(relationship_types) == 2
+    assert relationship_types[0] == "WORKS_FOR"
+    assert relationship_types[1]["label"] == "CREATED"
+    assert len(relationship_types[1]["properties"]) == 2
+    assert patterns is not None
+    assert len(patterns) == 2
+    assert "additional_node_types" not in schema_dict
+
+
+def test_simple_kg_pipeline_config_process_schema_with_precedence_schema_dict() -> None:
+    entities = [
+        "Person",
+        {
+            "label": "Organization",
+            "description": "A group of persons",
+            "properties": [
+                {
+                    "name": "name",
+                    "type": "STRING",
+                }
+            ],
+        },
+    ]
+    relations = [
+        "WORKS_FOR",
+        {
+            "label": "CREATED",
+            "description": "A person created an organization",
+            "properties": [
+                {
+                    "name": "date",
+                    "description": "The date the organization was created",
+                    "type": "DATE",
+                },
+                {"name": "isActive", "type": "BOOLEAN"},
+            ],
+        },
+    ]
+    potential_schema = [
+        ("Person", "WORKS_FOR", "Organization"),
+        ("Person", "CREATED", "Organization"),
+    ]
+    config = SimpleKGPipelineConfig(
+        schema={  # type: ignore
+            "node_types": entities,
+            "relationship_types": relations,
+            "patterns": potential_schema,
+            "additional_node_types": False,
+        }
+    )
+    schema_dict = config._process_schema_with_precedence()
+    node_types = schema_dict["node_types"]
+    relationship_types = schema_dict["relationship_types"]
+    patterns = schema_dict["patterns"]
+    assert len(node_types) == 2
+    assert node_types[0]["label"] == "Person"
+    # String input gets default "name" property and additional_properties=True
+    assert len(node_types[0]["properties"]) == 1
+    assert node_types[0]["properties"][0]["name"] == "name"
+    assert node_types[0]["additional_properties"] is True
+    assert node_types[1]["label"] == "Organization"
+    assert len(node_types[1]["properties"]) == 1
+    assert relationship_types is not None
+    assert len(relationship_types) == 2
+    assert relationship_types[0]["label"] == "WORKS_FOR"
+    assert len(relationship_types[0]["properties"]) == 0
+    assert relationship_types[1]["label"] == "CREATED"
+    assert len(relationship_types[1]["properties"]) == 2
+    assert patterns is not None
+    assert len(patterns) == 2
+    assert schema_dict["additional_node_types"] is False
+
+
+def test_simple_kg_pipeline_config_process_schema_with_precedence_schema_object() -> (
+    None
+):
+    entities = [
+        "Person",
+        {
+            "label": "Organization",
+            "description": "A group of persons",
+            "properties": [
+                {
+                    "name": "name",
+                    "type": "STRING",
+                }
+            ],
+        },
+    ]
+    relations = [
+        "WORKS_FOR",
+        {
+            "label": "CREATED",
+            "description": "A person created an organization",
+            "properties": [
+                {
+                    "name": "date",
+                    "description": "The date the organization was created",
+                    "type": "DATE",
+                },
+                {"name": "isActive", "type": "BOOLEAN"},
+            ],
+        },
+    ]
+    potential_schema = [
+        ("Person", "WORKS_FOR", "Organization"),
+        ("Person", "CREATED", "Organization"),
+    ]
+    config = SimpleKGPipelineConfig(
+        schema=GraphSchema.model_validate(
+            {
+                "node_types": entities,
+                "relationship_types": relations,
+                "patterns": potential_schema,
+                "additional_node_types": False,
+            }
+        )
+    )
+    schema_dict = config._process_schema_with_precedence()
+    node_types = schema_dict["node_types"]
+    relationship_types = schema_dict["relationship_types"]
+    patterns = schema_dict["patterns"]
+    assert len(node_types) == 2
+    assert node_types[0]["label"] == "Person"
+    # String input gets default "name" property and additional_properties=True
+    assert len(node_types[0]["properties"]) == 1
+    assert node_types[0]["properties"][0]["name"] == "name"
+    assert node_types[0]["additional_properties"] is True
+    assert node_types[1]["label"] == "Organization"
+    assert len(node_types[1]["properties"]) == 1
+    assert relationship_types is not None
+    assert len(relationship_types) == 2
+    assert relationship_types[0]["label"] == "WORKS_FOR"
+    assert len(relationship_types[0]["properties"]) == 0
+    assert relationship_types[1]["label"] == "CREATED"
+    assert len(relationship_types[1]["properties"]) == 2
+    assert patterns is not None
+    assert len(patterns) == 2
+    assert schema_dict["additional_node_types"] is False

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/config/test_base.py

@@ -0,0 +1,37 @@
+#  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 patch
+
+from neo4j_graphrag.experimental.pipeline.config.base import AbstractConfig
+from neo4j_graphrag.experimental.pipeline.config.param_resolver import (
+    ParamToResolveConfig,
+)
+
+
+def test_resolve_param_with_param_to_resolve_object() -> None:
+    c = AbstractConfig()
+    with patch(
+        "neo4j_graphrag.experimental.pipeline.config.param_resolver.ParamToResolveConfig",
+        spec=ParamToResolveConfig,
+    ) as mock_param_class:
+        mock_param = mock_param_class.return_value
+        mock_param.resolve.return_value = 1
+        assert c.resolve_param(mock_param) == 1
+        mock_param.resolve.assert_called_once_with({})
+
+
+def test_resolve_param_with_other_object() -> None:
+    c = AbstractConfig()
+    assert c.resolve_param("value") == "value"

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/config/test_object_config.py

@@ -0,0 +1,192 @@
+#  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 importlib
+import sys
+from abc import ABC
+from typing import ClassVar
+from unittest.mock import patch
+
+import neo4j
+import pytest
+from neo4j_graphrag.embeddings import Embedder, OpenAIEmbeddings
+from neo4j_graphrag.experimental.pipeline import Pipeline
+from neo4j_graphrag.experimental.pipeline.config.object_config import (
+    EmbedderConfig,
+    EmbedderType,
+    LLMConfig,
+    LLMType,
+    Neo4jDriverConfig,
+    Neo4jDriverType,
+    ObjectConfig,
+)
+from neo4j_graphrag.llm import LLMInterface, OpenAILLM
+
+
+def test_get_class_no_optional_module() -> None:
+    c: ObjectConfig[object] = ObjectConfig()
+    klass = c._get_class("neo4j_graphrag.experimental.pipeline.Pipeline")
+    assert klass == Pipeline
+
+
+def test_get_class_optional_module() -> None:
+    c: ObjectConfig[object] = ObjectConfig()
+    klass = c._get_class(
+        "Pipeline", optional_module="neo4j_graphrag.experimental.pipeline"
+    )
+    assert klass == Pipeline
+
+
+def test_get_class_path_and_optional_module() -> None:
+    c: ObjectConfig[object] = ObjectConfig()
+    klass = c._get_class(
+        "pipeline.Pipeline", optional_module="neo4j_graphrag.experimental"
+    )
+    assert klass == Pipeline
+
+
+def test_get_class_wrong_path() -> None:
+    c: ObjectConfig[object] = ObjectConfig()
+    with pytest.raises(ValueError):
+        c._get_class("MyClass")
+
+
+class _MyClass:
+    def __init__(self, param: str) -> None:
+        self.param = param
+
+
+class _MyInterface(ABC): ...
+
+
+def test_parse_after_module_reload() -> None:
+    class MyClassConfig(ObjectConfig[_MyClass]):
+        DEFAULT_MODULE: ClassVar[str] = __name__
+        INTERFACE: ClassVar[type] = _MyClass
+
+    param_value = "value"
+    config = MyClassConfig.model_validate(
+        {"class_": f"{__name__}.{_MyClass.__name__}", "params_": {"param": param_value}}
+    )
+    importlib.reload(sys.modules[__name__])
+
+    my_obj = config.parse()
+    assert isinstance(my_obj, _MyClass)
+    assert my_obj.param == param_value
+
+
+def test_neo4j_driver_config() -> None:
+    config = Neo4jDriverConfig.model_validate(
+        {
+            "params_": {
+                "uri": "bolt://",
+                "user": "a user",
+                "password": "a password",
+            }
+        }
+    )
+    assert config.class_ == "not used"
+    assert config.params_ == {
+        "uri": "bolt://",
+        "user": "a user",
+        "password": "a password",
+    }
+    with patch(
+        "neo4j_graphrag.experimental.pipeline.config.object_config.neo4j.GraphDatabase.driver"
+    ) as driver_mock:
+        driver_mock.return_value = "a driver"
+        d = config.parse()
+        driver_mock.assert_called_once_with("bolt://", auth=("a user", "a password"))
+        assert d == "a driver"  # type: ignore
+
+
+def test_neo4j_driver_type_with_driver(driver: neo4j.Driver) -> None:
+    driver_type = Neo4jDriverType(driver)
+    assert driver_type.parse() == driver
+
+
+def test_neo4j_driver_type_with_config() -> None:
+    driver_type = Neo4jDriverType(
+        Neo4jDriverConfig(
+            params_={
+                "uri": "bolt://",
+                "user": "",
+                "password": "",
+            }
+        )
+    )
+    driver = driver_type.parse()
+    assert isinstance(driver, neo4j.Driver)
+
+
+def test_llm_config() -> None:
+    config = LLMConfig.model_validate(
+        {
+            "class_": "OpenAILLM",
+            "params_": {"model_name": "gpt-5", "api_key": "my-api-key"},
+        }
+    )
+    assert config.class_ == "OpenAILLM"
+    assert config.get_module() == "neo4j_graphrag.llm"
+    assert config.get_interface() == LLMInterface
+    assert config.params_ == {"model_name": "gpt-5", "api_key": "my-api-key"}
+    d = config.parse()
+    assert isinstance(d, OpenAILLM)
+
+
+def test_llm_type_with_driver(llm: LLMInterface) -> None:
+    llm_type = LLMType(llm)
+    assert llm_type.parse() == llm
+
+
+def test_llm_type_with_config() -> None:
+    llm_type = LLMType(
+        LLMConfig(
+            class_="OpenAILLM",
+            params_={"model_name": "gpt-5", "api_key": "my-api-key"},
+        )
+    )
+    llm = llm_type.parse()
+    assert isinstance(llm, OpenAILLM)
+
+
+def test_embedder_config() -> None:
+    config = EmbedderConfig.model_validate(
+        {
+            "class_": "OpenAIEmbeddings",
+            "params_": {"api_key": "my-api-key"},
+        }
+    )
+    assert config.class_ == "OpenAIEmbeddings"
+    assert config.get_module() == "neo4j_graphrag.embeddings"
+    assert config.get_interface() == Embedder
+    assert config.params_ == {"api_key": "my-api-key"}
+    d = config.parse()
+    assert isinstance(d, OpenAIEmbeddings)
+
+
+def test_embedder_type_with_embedder(embedder: Embedder) -> None:
+    embedder_type = EmbedderType(embedder)
+    assert embedder_type.parse() == embedder
+
+
+def test_embedder_type_with_config() -> None:
+    embedder_type = EmbedderType(
+        EmbedderConfig(
+            class_="OpenAIEmbeddings",
+            params_={"api_key": "my-api-key"},
+        )
+    )
+    embedder = embedder_type.parse()
+    assert isinstance(embedder, OpenAIEmbeddings)

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/config/test_param_resolver.py

@@ -0,0 +1,56 @@
+#  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 unittest.mock import patch
+
+import pytest
+from neo4j_graphrag.experimental.pipeline.config.param_resolver import (
+    ParamFromEnvConfig,
+    ParamFromKeyConfig,
+)
+
+
+@patch.dict(os.environ, {"MY_KEY": "my_value"}, clear=True)
+def test_env_param_config_happy_path() -> None:
+    resolver = ParamFromEnvConfig(var_="MY_KEY")
+    assert resolver.resolve({}) == "my_value"
+
+
+@patch.dict(os.environ, {}, clear=True)
+def test_env_param_config_missing_env_var() -> None:
+    resolver = ParamFromEnvConfig(var_="MY_KEY")
+    assert resolver.resolve({}) is None
+
+
+def test_config_key_param_simple_key() -> None:
+    resolver = ParamFromKeyConfig(key_="my_key")
+    assert resolver.resolve({"my_key": "my_value"}) == "my_value"
+
+
+def test_config_key_param_missing_key() -> None:
+    resolver = ParamFromKeyConfig(key_="my_key")
+    with pytest.raises(KeyError):
+        resolver.resolve({})
+
+
+def test_config_complex_key_param() -> None:
+    resolver = ParamFromKeyConfig(key_="my_key.my_sub_key")
+    assert resolver.resolve({"my_key": {"my_sub_key": "value"}}) == "value"
+
+
+def test_config_complex_key_param_missing_subkey() -> None:
+    resolver = ParamFromKeyConfig(key_="my_key.my_sub_key")
+    with pytest.raises(KeyError):
+        resolver.resolve({"my_key": {}})

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/config/test_pipeline_config.py

@@ -0,0 +1,378 @@
+#  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 Mock, patch
+
+import neo4j
+from neo4j_graphrag.embeddings import Embedder
+from neo4j_graphrag.experimental.pipeline import Component
+from neo4j_graphrag.experimental.pipeline.config.object_config import (
+    ComponentConfig,
+    ComponentType,
+    Neo4jDriverConfig,
+    Neo4jDriverType,
+)
+from neo4j_graphrag.experimental.pipeline.config.param_resolver import (
+    ParamFromEnvConfig,
+    ParamFromKeyConfig,
+)
+from neo4j_graphrag.experimental.pipeline.config.pipeline_config import (
+    AbstractPipelineConfig,
+)
+from neo4j_graphrag.experimental.pipeline.types.definitions import ComponentDefinition
+from neo4j_graphrag.llm import LLMInterface
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.object_config.Neo4jDriverConfig.parse"
+)
+def test_abstract_pipeline_config_neo4j_config_is_a_dict_with_params_(
+    mock_neo4j_config: Mock,
+) -> None:
+    mock_neo4j_config.return_value = "text"
+    config = AbstractPipelineConfig.model_validate(
+        {
+            "neo4j_config": {
+                "params_": {
+                    "uri": "bolt://",
+                    "user": "",
+                    "password": "",
+                }
+            }
+        }
+    )
+    assert isinstance(config.neo4j_config, dict)
+    assert "default" in config.neo4j_config
+    config.parse()
+    mock_neo4j_config.assert_called_once()
+    assert config._global_data["neo4j_config"]["default"] == "text"
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.object_config.Neo4jDriverConfig.parse"
+)
+def test_abstract_pipeline_config_neo4j_config_is_a_dict_with_names(
+    mock_neo4j_config: Mock,
+) -> None:
+    mock_neo4j_config.return_value = "text"
+    config = AbstractPipelineConfig.model_validate(
+        {
+            "neo4j_config": {
+                "my_driver": {
+                    "params_": {
+                        "uri": "bolt://",
+                        "user": "",
+                        "password": "",
+                    }
+                }
+            }
+        }
+    )
+    assert isinstance(config.neo4j_config, dict)
+    assert "my_driver" in config.neo4j_config
+    config.parse()
+    mock_neo4j_config.assert_called_once()
+    assert config._global_data["neo4j_config"]["my_driver"] == "text"
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.object_config.Neo4jDriverConfig.parse"
+)
+def test_abstract_pipeline_config_neo4j_config_is_a_dict_with_driver(
+    mock_neo4j_config: Mock, driver: neo4j.Driver
+) -> None:
+    config = AbstractPipelineConfig.model_validate(
+        {
+            "neo4j_config": {
+                "my_driver": driver,
+            }
+        }
+    )
+    assert isinstance(config.neo4j_config, dict)
+    assert "my_driver" in config.neo4j_config
+    config.parse()
+    assert not mock_neo4j_config.called
+    assert config._global_data["neo4j_config"]["my_driver"] == driver
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.object_config.Neo4jDriverConfig.parse"
+)
+def test_abstract_pipeline_config_neo4j_config_is_a_driver(
+    mock_neo4j_config: Mock, driver: neo4j.Driver
+) -> None:
+    config = AbstractPipelineConfig.model_validate(
+        {
+            "neo4j_config": driver,
+        }
+    )
+    assert isinstance(config.neo4j_config, dict)
+    assert "default" in config.neo4j_config
+    config.parse()
+    assert not mock_neo4j_config.called
+    assert config._global_data["neo4j_config"]["default"] == driver
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.LLMConfig.parse")
+def test_abstract_pipeline_config_llm_config_is_a_dict_with_params_(
+    mock_llm_config: Mock,
+) -> None:
+    mock_llm_config.return_value = "text"
+    config = AbstractPipelineConfig.model_validate(
+        {"llm_config": {"class_": "OpenAILLM", "params_": {"model_name": "gpt-5"}}}
+    )
+    assert isinstance(config.llm_config, dict)
+    assert "default" in config.llm_config
+    config.parse()
+    mock_llm_config.assert_called_once()
+    assert config._global_data["llm_config"]["default"] == "text"
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.LLMConfig.parse")
+def test_abstract_pipeline_config_llm_config_is_a_dict_with_names(
+    mock_llm_config: Mock,
+) -> None:
+    mock_llm_config.return_value = "text"
+    config = AbstractPipelineConfig.model_validate(
+        {
+            "llm_config": {
+                "my_llm": {"class_": "OpenAILLM", "params_": {"model_name": "gpt-5"}}
+            }
+        }
+    )
+    assert isinstance(config.llm_config, dict)
+    assert "my_llm" in config.llm_config
+    config.parse()
+    mock_llm_config.assert_called_once()
+    assert config._global_data["llm_config"]["my_llm"] == "text"
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.LLMConfig.parse")
+def test_abstract_pipeline_config_llm_config_is_a_dict_with_llm(
+    mock_llm_config: Mock, llm: LLMInterface
+) -> None:
+    config = AbstractPipelineConfig.model_validate(
+        {
+            "llm_config": {
+                "my_llm": llm,
+            }
+        }
+    )
+    assert isinstance(config.llm_config, dict)
+    assert "my_llm" in config.llm_config
+    config.parse()
+    assert not mock_llm_config.called
+    assert config._global_data["llm_config"]["my_llm"] == llm
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.LLMConfig.parse")
+def test_abstract_pipeline_config_llm_config_is_a_llm(
+    mock_llm_config: Mock, llm: LLMInterface
+) -> None:
+    config = AbstractPipelineConfig.model_validate(
+        {
+            "llm_config": llm,
+        }
+    )
+    assert isinstance(config.llm_config, dict)
+    assert "default" in config.llm_config
+    config.parse()
+    assert not mock_llm_config.called
+    assert config._global_data["llm_config"]["default"] == llm
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.EmbedderConfig.parse")
+def test_abstract_pipeline_config_embedder_config_is_a_dict_with_params_(
+    mock_embedder_config: Mock,
+) -> None:
+    mock_embedder_config.return_value = "text"
+    config = AbstractPipelineConfig.model_validate(
+        {"embedder_config": {"class_": "OpenAIEmbeddings", "params_": {}}}
+    )
+    assert isinstance(config.embedder_config, dict)
+    assert "default" in config.embedder_config
+    config.parse()
+    mock_embedder_config.assert_called_once()
+    assert config._global_data["embedder_config"]["default"] == "text"
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.EmbedderConfig.parse")
+def test_abstract_pipeline_config_embedder_config_is_a_dict_with_names(
+    mock_embedder_config: Mock,
+) -> None:
+    mock_embedder_config.return_value = "text"
+    config = AbstractPipelineConfig.model_validate(
+        {
+            "embedder_config": {
+                "my_embedder": {"class_": "OpenAIEmbeddings", "params_": {}}
+            }
+        }
+    )
+    assert isinstance(config.embedder_config, dict)
+    assert "my_embedder" in config.embedder_config
+    config.parse()
+    mock_embedder_config.assert_called_once()
+    assert config._global_data["embedder_config"]["my_embedder"] == "text"
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.EmbedderConfig.parse")
+def test_abstract_pipeline_config_embedder_config_is_a_dict_with_llm(
+    mock_embedder_config: Mock, embedder: Embedder
+) -> None:
+    config = AbstractPipelineConfig.model_validate(
+        {
+            "embedder_config": {
+                "my_embedder": embedder,
+            }
+        }
+    )
+    assert isinstance(config.embedder_config, dict)
+    assert "my_embedder" in config.embedder_config
+    config.parse()
+    assert not mock_embedder_config.called
+    assert config._global_data["embedder_config"]["my_embedder"] == embedder
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.EmbedderConfig.parse")
+def test_abstract_pipeline_config_embedder_config_is_an_embedder(
+    mock_embedder_config: Mock, embedder: Embedder
+) -> None:
+    config = AbstractPipelineConfig.model_validate(
+        {
+            "embedder_config": embedder,
+        }
+    )
+    assert isinstance(config.embedder_config, dict)
+    assert "default" in config.embedder_config
+    config.parse()
+    assert not mock_embedder_config.called
+    assert config._global_data["embedder_config"]["default"] == embedder
+
+
+def test_abstract_pipeline_config_parse_global_data_no_extras(driver: Mock) -> None:
+    config = AbstractPipelineConfig(
+        neo4j_config={"my_driver": Neo4jDriverType(driver)},
+    )
+    gd = config._parse_global_data()
+    assert gd == {
+        "extras": {},
+        "neo4j_config": {
+            "my_driver": driver,
+        },
+        "llm_config": {},
+        "embedder_config": {},
+    }
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.param_resolver.ParamFromEnvConfig.resolve"
+)
+def test_abstract_pipeline_config_parse_global_data_extras(
+    mock_param_resolver: Mock,
+) -> None:
+    mock_param_resolver.return_value = "my value"
+    config = AbstractPipelineConfig(
+        extras={"my_extra_var": ParamFromEnvConfig(var_="some key")},
+    )
+    gd = config._parse_global_data()
+    assert gd == {
+        "extras": {"my_extra_var": "my value"},
+        "neo4j_config": {},
+        "llm_config": {},
+        "embedder_config": {},
+    }
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.param_resolver.ParamFromEnvConfig.resolve"
+)
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.object_config.Neo4jDriverType.parse"
+)
+def test_abstract_pipeline_config_parse_global_data_use_extras_in_other_config(
+    mock_neo4j_parser: Mock,
+    mock_param_resolver: Mock,
+) -> None:
+    """Parser is able to read variables in the 'extras' section of config
+    to instantiate another object (neo4j.Driver in this test case)
+    """
+    mock_param_resolver.side_effect = ["bolt://myhost", "myuser", "mypwd"]
+    mock_neo4j_parser.return_value = "my driver"
+    config = AbstractPipelineConfig(
+        extras={
+            "my_extra_uri": ParamFromEnvConfig(var_="some key"),
+            "my_extra_user": ParamFromEnvConfig(var_="some key"),
+            "my_extra_pwd": ParamFromEnvConfig(var_="some key"),
+        },
+        neo4j_config={
+            "my_driver": Neo4jDriverType(
+                Neo4jDriverConfig(
+                    params_=dict(
+                        uri=ParamFromKeyConfig(key_="extras.my_extra_uri"),
+                        user=ParamFromKeyConfig(key_="extras.my_extra_user"),
+                        password=ParamFromKeyConfig(key_="extras.my_extra_pwd"),
+                    )
+                )
+            )
+        },
+    )
+    gd = config._parse_global_data()
+    expected_extras = {
+        "my_extra_uri": "bolt://myhost",
+        "my_extra_user": "myuser",
+        "my_extra_pwd": "mypwd",
+    }
+    assert gd["extras"] == expected_extras
+    assert gd["neo4j_config"] == {"my_driver": "my driver"}
+    mock_neo4j_parser.assert_called_once_with({"extras": expected_extras})
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.ComponentType.parse")
+def test_abstract_pipeline_config_resolve_component_definition_no_run_params(
+    mock_component_parse: Mock,
+    component: Component,
+) -> None:
+    mock_component_parse.return_value = component
+    config = AbstractPipelineConfig()
+    component_type = ComponentType(component)
+    component_definition = config._resolve_component_definition("name", component_type)
+    assert isinstance(component_definition, ComponentDefinition)
+    mock_component_parse.assert_called_once_with({})
+    assert component_definition.name == "name"
+    assert component_definition.component == component
+    assert component_definition.run_params == {}
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.config.pipeline_config.AbstractPipelineConfig.resolve_params"
+)
+@patch("neo4j_graphrag.experimental.pipeline.config.object_config.ComponentType.parse")
+def test_abstract_pipeline_config_resolve_component_definition_with_run_params(
+    mock_component_parse: Mock,
+    mock_resolve_params: Mock,
+    component: Component,
+) -> None:
+    mock_component_parse.return_value = component
+    mock_resolve_params.return_value = {"param": "resolver param result"}
+    config = AbstractPipelineConfig()
+    component_type = ComponentType(
+        ComponentConfig(class_="", params_={}, run_params_={"param1": "value1"})
+    )
+    component_definition = config._resolve_component_definition("name", component_type)
+    assert isinstance(component_definition, ComponentDefinition)
+    mock_component_parse.assert_called_once_with({})
+    assert component_definition.name == "name"
+    assert component_definition.component == component
+    assert component_definition.run_params == {"param": "resolver param result"}
+    mock_resolve_params.assert_called_once_with({"param1": "value1"})

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/config/test_runner.py

@@ -0,0 +1,56 @@
+#  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 Mock, patch
+
+from neo4j_graphrag.experimental.pipeline import Pipeline
+from neo4j_graphrag.experimental.pipeline.config.pipeline_config import PipelineConfig
+from neo4j_graphrag.experimental.pipeline.config.runner import PipelineRunner
+from neo4j_graphrag.experimental.pipeline.types.definitions import PipelineDefinition
+
+
+@patch("neo4j_graphrag.experimental.pipeline.pipeline.Pipeline.from_definition")
+def test_pipeline_runner_from_def_empty(mock_from_definition: Mock) -> None:
+    mock_from_definition.return_value = Pipeline()
+    runner = PipelineRunner(
+        pipeline_definition=PipelineDefinition(components=[], connections=[])
+    )
+    assert runner.config is None
+    assert runner.pipeline is not None
+    assert runner.pipeline._nodes == {}
+    assert runner.pipeline._edges == []
+    assert runner.run_params == {}
+    mock_from_definition.assert_called_once()
+
+
+def test_pipeline_runner_from_config() -> None:
+    config = PipelineConfig(component_config={}, connection_config=[])
+    runner = PipelineRunner.from_config(config)
+    assert runner.config is not None
+    assert runner.pipeline is not None
+    assert runner.pipeline._nodes == {}
+    assert runner.pipeline._edges == []
+    assert runner.run_params == {}
+
+
+@patch("neo4j_graphrag.experimental.pipeline.config.runner.PipelineRunner.from_config")
+@patch("neo4j_graphrag.utils.file_handler.FileHandler.read")
+def test_pipeline_runner_from_config_file(
+    mock_read: Mock, mock_from_config: Mock
+) -> None:
+    mock_read.return_value = {"dict": "with data"}
+    PipelineRunner.from_config_file("file.yaml")
+
+    mock_read.assert_called_once_with("file.yaml")
+    mock_from_config.assert_called_once_with({"dict": "with data"}, do_cleaning=True)

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/test_component.py

@@ -0,0 +1,89 @@
+#  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 AsyncMock
+
+import pytest
+
+from neo4j_graphrag.experimental.pipeline import Component
+from neo4j_graphrag.experimental.pipeline.types.context import RunContext
+from .components import ComponentMultiply, ComponentMultiplyWithContext, IntResultModel
+
+
+def test_component_inputs() -> None:
+    inputs = ComponentMultiply.component_inputs
+    assert "number1" in inputs
+    assert inputs["number1"]["has_default"] is False
+    assert "number2" in inputs
+    assert inputs["number2"]["has_default"] is True
+
+
+def test_component_outputs() -> None:
+    outputs = ComponentMultiply.component_outputs
+    assert "result" in outputs
+    assert outputs["result"]["has_default"] is True
+    assert outputs["result"]["annotation"] == int
+
+
+@pytest.mark.asyncio
+async def test_component_run() -> None:
+    c = ComponentMultiply()
+    result = await c.run(number1=1, number2=2)
+    assert isinstance(result, IntResultModel)
+    assert isinstance(
+        result.result,
+        # we know this is a type and not a bool or str:
+        ComponentMultiply.component_outputs["result"]["annotation"],  # type: ignore
+    )
+
+
+@pytest.mark.asyncio
+async def test_component_run_with_context_default_implementation() -> None:
+    c = ComponentMultiply()
+    result = await c.run_with_context(
+        # context can not be null in the function signature,
+        # but it's ignored in this case
+        None,  # type: ignore
+        number1=1,
+        number2=2,
+    )
+    # the type checker doesn't know about the type
+    # because the method is not re-declared
+    assert result.result == 2  # type: ignore
+
+
+@pytest.mark.asyncio
+async def test_component_run_with_context() -> None:
+    c = ComponentMultiplyWithContext()
+    notifier_mock = AsyncMock()
+    result = await c.run_with_context(
+        RunContext(run_id="run_id", task_name="task_name", notifier=notifier_mock),
+        number1=1,
+        number2=2,
+    )
+    assert result.result == 2
+    notifier_mock.assert_awaited_once()
+
+
+def test_component_missing_method() -> None:
+    with pytest.raises(RuntimeError) as e:
+
+        class WrongComponent(Component):
+            # we must have either run or run_with_context
+            pass
+
+    assert (
+        "You must implement either `run` or `run_with_context` in Component 'WrongComponent'"
+        in str(e)
+    )

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/test_kg_builder.py

@@ -0,0 +1,272 @@
+#  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 import mock
+from unittest.mock import MagicMock, Mock, patch
+
+import neo4j
+import pytest
+from neo4j_graphrag.embeddings import Embedder
+from neo4j_graphrag.experimental.components.data_loader import PdfLoader
+from neo4j_graphrag.experimental.components.types import (
+    LexicalGraphConfig,
+)
+from neo4j_graphrag.experimental.pipeline.exceptions import PipelineDefinitionError
+from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
+from neo4j_graphrag.experimental.pipeline.pipeline import PipelineResult
+from neo4j_graphrag.llm.base import LLMInterface
+
+
+@mock.patch(
+    "neo4j_graphrag.experimental.components.kg_writer.get_version",
+    return_value=((5, 23, 0), False, False),
+)
+@pytest.mark.asyncio
+async def test_knowledge_graph_builder_from_pdf_deprecated_kwarg(_: Mock) -> None:
+    llm = MagicMock(spec=LLMInterface)
+    driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
+
+    with pytest.warns(DeprecationWarning, match="from_pdf"):
+        kg_builder = SimpleKGPipeline(
+            llm=llm,
+            driver=driver,
+            embedder=embedder,
+            from_pdf=True,
+        )
+
+    file_path = "path/to/test.pdf"
+    with patch.object(
+        kg_builder.runner.pipeline,
+        "run",
+        return_value=PipelineResult(run_id="test_run", result=None),
+    ) as mock_run:
+        await kg_builder.run_async(file_path=file_path)
+
+    pipe_inputs = mock_run.call_args[1]["data"]
+    assert "file_loader" in pipe_inputs
+
+
+@mock.patch(
+    "neo4j_graphrag.experimental.components.kg_writer.get_version",
+    return_value=((5, 23, 0), False, False),
+)
+@pytest.mark.asyncio
+async def test_knowledge_graph_builder_pdf_loader_deprecated_kwarg(_: Mock) -> None:
+    llm = MagicMock(spec=LLMInterface)
+    driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
+    loader = PdfLoader()
+
+    with pytest.warns(DeprecationWarning, match="pdf_loader"):
+        kg_builder = SimpleKGPipeline(
+            llm=llm,
+            driver=driver,
+            embedder=embedder,
+            pdf_loader=loader,
+        )
+
+    file_path = "path/to/test.pdf"
+    with patch.object(
+        kg_builder.runner.pipeline,
+        "run",
+        return_value=PipelineResult(run_id="test_run", result=None),
+    ) as mock_run:
+        await kg_builder.run_async(file_path=file_path)
+
+    pipe_inputs = mock_run.call_args[1]["data"]
+    assert "file_loader" in pipe_inputs
+
+
+@mock.patch(
+    "neo4j_graphrag.experimental.components.kg_writer.get_version",
+    return_value=((5, 23, 0), False, False),
+)
+@pytest.mark.asyncio
+async def test_knowledge_graph_builder_document_info_with_file(_: Mock) -> None:
+    llm = MagicMock(spec=LLMInterface)
+    driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
+
+    kg_builder = SimpleKGPipeline(
+        llm=llm,
+        driver=driver,
+        embedder=embedder,
+        from_file=True,
+    )
+
+    file_path = "path/to/test.pdf"
+
+    with patch.object(
+        kg_builder.runner.pipeline,
+        "run",
+        return_value=PipelineResult(run_id="test_run", result=None),
+    ) as mock_run:
+        await kg_builder.run_async(
+            file_path=file_path, document_metadata={"source": "google drive"}
+        )
+
+        pipe_inputs = mock_run.call_args[1]["data"]
+        assert "file_loader" in pipe_inputs
+        assert pipe_inputs["file_loader"] == {
+            "filepath": file_path,
+            "metadata": {"source": "google drive"},
+        }
+        assert "extractor" not in pipe_inputs
+
+
+@mock.patch(
+    "neo4j_graphrag.experimental.components.kg_writer.get_version",
+    return_value=((5, 23, 0), False, False),
+)
+@pytest.mark.asyncio
+async def test_knowledge_graph_builder_document_info_with_text(_: Mock) -> None:
+    llm = MagicMock(spec=LLMInterface)
+    driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
+
+    kg_builder = SimpleKGPipeline(
+        llm=llm,
+        driver=driver,
+        embedder=embedder,
+        from_file=False,
+    )
+
+    text_input = "May thy knife chip and shatter."
+
+    with patch.object(
+        kg_builder.runner.pipeline,
+        "run",
+        return_value=PipelineResult(run_id="test_run", result=None),
+    ) as mock_run:
+        await kg_builder.run_async(
+            text=text_input,
+            file_path="my_document.txt",
+            document_metadata={"source": "google drive"},
+        )
+
+        pipe_inputs = mock_run.call_args[1]["data"]
+        assert "splitter" in pipe_inputs
+        assert pipe_inputs["splitter"] == {"text": text_input}
+        assert pipe_inputs["extractor"]["document_info"]["path"] == "my_document.txt"
+        assert pipe_inputs["extractor"]["document_info"]["metadata"] == {
+            "source": "google drive"
+        }
+
+
+@mock.patch(
+    "neo4j_graphrag.experimental.components.kg_writer.get_version",
+    return_value=((5, 23, 0), False, False),
+)
+@pytest.mark.asyncio
+async def test_knowledge_graph_builder_with_entities_and_file(_: Mock) -> None:
+    llm = MagicMock(spec=LLMInterface)
+    driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
+
+    entities = ["Document", "Section"]
+    relations = ["CONTAINS"]
+    potential_schema = [("Document", "CONTAINS", "Section")]
+
+    kg_builder = SimpleKGPipeline(
+        llm=llm,
+        driver=driver,
+        embedder=embedder,
+        entities=entities,
+        relations=relations,
+        potential_schema=potential_schema,
+        from_file=True,
+    )
+
+    file_path = "path/to/test.pdf"
+
+    with patch.object(
+        kg_builder.runner.pipeline,
+        "run",
+        return_value=PipelineResult(run_id="test_run", result=None),
+    ) as mock_run:
+        await kg_builder.run_async(file_path=file_path)
+        pipe_inputs = mock_run.call_args[1]["data"]
+        assert pipe_inputs["schema"]["node_types"] == entities
+        assert pipe_inputs["schema"]["relationship_types"] == relations
+        assert pipe_inputs["schema"]["patterns"] == potential_schema
+
+
+def test_simple_kg_pipeline_on_error_invalid_value() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
+
+    with pytest.raises(PipelineDefinitionError):
+        SimpleKGPipeline(
+            llm=llm,
+            driver=driver,
+            embedder=embedder,
+            on_error="INVALID_VALUE",
+        )
+
+
+def test_knowledge_graph_builder_pdf_loader_and_file_loader_conflict() -> None:
+    llm = MagicMock(spec=LLMInterface)
+    driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
+
+    with pytest.raises(ValueError, match="pdf_loader"):
+        SimpleKGPipeline(
+            llm=llm,
+            driver=driver,
+            embedder=embedder,
+            file_loader=PdfLoader(),
+            pdf_loader=PdfLoader(),
+        )
+
+
+@mock.patch(
+    "neo4j_graphrag.experimental.components.kg_writer.get_version",
+    return_value=((5, 23, 0), False, False),
+)
+@pytest.mark.asyncio
+async def test_knowledge_graph_builder_with_lexical_graph_config(_: Mock) -> None:
+    llm = MagicMock(spec=LLMInterface)
+    driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
+
+    chunk_node_label = "TestChunk"
+    document_nodel_label = "TestDocument"
+    lexical_graph_config = LexicalGraphConfig(
+        chunk_node_label=chunk_node_label, document_node_label=document_nodel_label
+    )
+
+    kg_builder = SimpleKGPipeline(
+        llm=llm,
+        driver=driver,
+        embedder=embedder,
+        from_file=False,
+        lexical_graph_config=lexical_graph_config,
+    )
+
+    text_input = "May thy knife chip and shatter."
+
+    with patch.object(
+        kg_builder.runner.pipeline,
+        "run",
+        return_value=PipelineResult(run_id="test_run", result=None),
+    ) as mock_run:
+        await kg_builder.run_async(text=text_input)
+
+        pipe_inputs = mock_run.call_args[1]["data"]
+        assert "extractor" in pipe_inputs
+        assert pipe_inputs["extractor"]["lexical_graph_config"] == lexical_graph_config
+        assert pipe_inputs["extractor"]["document_info"] is not None
+        assert pipe_inputs["extractor"]["document_info"]["path"] == "document.txt"

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/test_orchestrator.py

@@ -0,0 +1,384 @@
+#  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 Mock, patch
+
+import pytest
+from neo4j_graphrag.experimental.pipeline import (
+    Component,
+    Pipeline,
+)
+from neo4j_graphrag.experimental.pipeline.exceptions import (
+    PipelineDefinitionError,
+    PipelineMissingDependencyError,
+    PipelineStatusUpdateError,
+)
+from neo4j_graphrag.experimental.pipeline.orchestrator import Orchestrator
+from neo4j_graphrag.experimental.pipeline.types.orchestration import RunStatus
+
+from tests.unit.experimental.pipeline.components import (
+    ComponentNoParam,
+    ComponentPassThrough,
+)
+
+
+def test_orchestrator_get_input_config_for_task_pipeline_not_validated() -> None:
+    pipe = Pipeline()
+    pipe.add_component(ComponentPassThrough(), "a")
+    pipe.add_component(ComponentPassThrough(), "b")
+    orchestrator = Orchestrator(pipe)
+    with pytest.raises(PipelineDefinitionError) as exc:
+        orchestrator.get_input_config_for_task(pipe.get_node_by_name("a"))
+    assert "You must validate the pipeline input config first" in str(exc.value)
+
+
+@pytest.mark.asyncio
+async def test_orchestrator_get_component_inputs_from_user_only() -> None:
+    """Components take all their inputs from user input."""
+    pipe = Pipeline()
+    pipe.add_component(ComponentPassThrough(), "a")
+    pipe.add_component(ComponentPassThrough(), "b")
+    orchestrator = Orchestrator(pipe)
+    input_data = {
+        "a": {"value": "user input for component a"},
+        "b": {"value": "user input for component b"},
+    }
+    data = await orchestrator.get_component_inputs("a", {}, input_data)
+    assert data == {"value": "user input for component a"}
+    data = await orchestrator.get_component_inputs("b", {}, input_data)
+    assert data == {"value": "user input for component b"}
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.get_results_for_component"
+)
+@pytest.mark.asyncio
+async def test_orchestrator_get_component_inputs_from_parent_specific(
+    mock_result: Mock,
+) -> None:
+    """Propagate one specific output field from parent to a child component."""
+    pipe = Pipeline()
+    pipe.add_component(ComponentPassThrough(), "a")
+    pipe.add_component(ComponentPassThrough(), "b")
+    pipe.connect("a", "b", input_config={"value": "a.result"})
+
+    # component "a" already run and results stored:
+    mock_result.return_value = {"result": "output from component a"}
+
+    orchestrator = Orchestrator(pipe)
+    data = await orchestrator.get_component_inputs(
+        "b", {"value": {"component": "a", "param": "result"}}, {}
+    )
+    assert data == {"value": "output from component a"}
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.get_results_for_component"
+)
+@pytest.mark.asyncio
+async def test_orchestrator_get_component_inputs_from_parent_all(
+    mock_result: Mock,
+) -> None:
+    """Use the component name to get the full output
+    (without extracting a specific field).
+    """
+    pipe = Pipeline()
+    pipe.add_component(ComponentNoParam(), "a")
+    pipe.add_component(ComponentPassThrough(), "b")
+    pipe.connect("a", "b", input_config={"value": "a"})
+
+    # component "a" already run and results stored:
+    mock_result.return_value = {"result": "output from component a"}
+
+    orchestrator = Orchestrator(pipe)
+    data = await orchestrator.get_component_inputs(
+        "b", {"value": {"component": "a"}}, {}
+    )
+    assert data == {"value": {"result": "output from component a"}}
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.get_results_for_component"
+)
+@pytest.mark.asyncio
+async def test_orchestrator_get_component_inputs_from_parent_and_input(
+    mock_result: Mock,
+) -> None:
+    """Some parameters from user input, some other parameter from previous component."""
+    pipe = Pipeline()
+    pipe.add_component(ComponentNoParam(), "a")
+    pipe.add_component(ComponentPassThrough(), "b")
+    pipe.connect("a", "b", input_config={"value": "a"})
+
+    # component "a" already run and results stored:
+    mock_result.return_value = {"result": "output from component a"}
+
+    orchestrator = Orchestrator(pipe)
+    data = await orchestrator.get_component_inputs(
+        "b",
+        {"value": {"component": "a"}},
+        {"b": {"other_value": "user input for component b 'other_value' param"}},
+    )
+    assert data == {
+        "value": {"result": "output from component a"},
+        "other_value": "user input for component b 'other_value' param",
+    }
+
+
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.get_results_for_component"
+)
+@pytest.mark.asyncio
+async def test_orchestrator_get_component_inputs_ignore_user_input_if_input_def_provided(
+    mock_result: Mock,
+) -> None:
+    """If a parameter is defined both in the user input and in an input definition
+    (ie propagated from a previous component), the user input is ignored and a
+    warning is raised.
+    """
+    pipe = Pipeline()
+    pipe.add_component(ComponentNoParam(), "a")
+    pipe.add_component(ComponentPassThrough(), "b")
+    pipe.connect("a", "b", input_config={"value": "a"})
+
+    # component "a" already run and results stored:
+    mock_result.return_value = {"result": "output from component a"}
+
+    orchestrator = Orchestrator(pipe)
+    with pytest.warns(Warning) as w:
+        data = await orchestrator.get_component_inputs(
+            "b",
+            {"value": {"component": "a"}},
+            {"b": {"value": "user input for component a"}},
+        )
+        assert data == {"value": {"result": "output from component a"}}
+        assert (
+            w[0].message.args[0]  # type: ignore[union-attr]
+            == "In component 'b', parameter 'value' from user input will be ignored and replaced by 'a.value'"
+        )
+
+
+@pytest.mark.asyncio
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.get_status_for_component"
+)
+@pytest.mark.parametrize(
+    "old_status, new_status, result",
+    [
+        # Normal path: from UNKNOWN to RUNNING to DONE
+        (RunStatus.UNKNOWN, RunStatus.RUNNING, "ok"),
+        (RunStatus.RUNNING, RunStatus.DONE, "ok"),
+        # Error: status is already set to this value
+        (RunStatus.RUNNING, RunStatus.RUNNING, "Status is already RunStatus.RUNNING"),
+        (RunStatus.DONE, RunStatus.DONE, "Status is already RunStatus.DONE"),
+        # Error: can't go back in time
+        (
+            RunStatus.DONE,
+            RunStatus.RUNNING,
+            "Can't go from RunStatus.DONE to RunStatus.RUNNING",
+        ),
+        (
+            RunStatus.RUNNING,
+            RunStatus.UNKNOWN,
+            "Can't go from RunStatus.RUNNING to RunStatus.UNKNOWN",
+        ),
+        (
+            RunStatus.DONE,
+            RunStatus.UNKNOWN,
+            "Can't go from RunStatus.DONE to RunStatus.UNKNOWN",
+        ),
+    ],
+)
+async def test_orchestrator_set_component_status(
+    mock_status: Mock,
+    old_status: RunStatus,
+    new_status: RunStatus,
+    result: str,
+) -> None:
+    pipe = Pipeline()
+    orchestrator = Orchestrator(pipeline=pipe)
+    mock_status.side_effect = [
+        old_status,
+    ]
+    if result == "ok":
+        await orchestrator.set_task_status("task_name", new_status)
+    else:
+        with pytest.raises(PipelineStatusUpdateError) as exc:
+            await orchestrator.set_task_status("task_name", new_status)
+        assert result in str(exc)
+
+
+@pytest.fixture(scope="function")
+def pipeline_branch() -> Pipeline:
+    pipe = Pipeline()
+    pipe.add_component(Component(), "a")  # type: ignore[abstract,unused-ignore]
+    pipe.add_component(Component(), "b")  # type: ignore[abstract,unused-ignore]
+    pipe.add_component(Component(), "c")  # type: ignore[abstract,unused-ignore]
+    pipe.connect("a", "b")
+    pipe.connect("a", "c")
+    return pipe
+
+
+@pytest.fixture(scope="function")
+def pipeline_aggregation() -> Pipeline:
+    pipe = Pipeline()
+    pipe.add_component(Component(), "a")  # type: ignore[abstract,unused-ignore]
+    pipe.add_component(Component(), "b")  # type: ignore[abstract,unused-ignore]
+    pipe.add_component(Component(), "c")  # type: ignore[abstract,unused-ignore]
+    pipe.connect("a", "c")
+    pipe.connect("b", "c")
+    return pipe
+
+
+@pytest.mark.asyncio
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.get_status_for_component"
+)
+async def test_orchestrator_check_dependency_complete(
+    mock_status: Mock, pipeline_branch: Pipeline
+) -> None:
+    """a -> b, c"""
+    orchestrator = Orchestrator(pipeline=pipeline_branch)
+    node_a = pipeline_branch.get_node_by_name("a")
+    await orchestrator.check_dependencies_complete(node_a)
+    node_b = pipeline_branch.get_node_by_name("b")
+    # dependency is DONE:
+    mock_status.side_effect = [RunStatus.DONE]
+    await orchestrator.check_dependencies_complete(node_b)
+    # dependency is not DONE:
+    mock_status.side_effect = [RunStatus.RUNNING]
+    with pytest.raises(PipelineMissingDependencyError):
+        await orchestrator.check_dependencies_complete(node_b)
+
+
+@pytest.mark.asyncio
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.get_status_for_component"
+)
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.check_dependencies_complete",
+)
+async def test_orchestrator_next_task_branch_no_missing_dependencies(
+    mock_dep: Mock, mock_status: Mock, pipeline_branch: Pipeline
+) -> None:
+    """a -> b, c"""
+    orchestrator = Orchestrator(pipeline=pipeline_branch)
+    node_a = pipeline_branch.get_node_by_name("a")
+    mock_status.side_effect = [
+        # next "b"
+        RunStatus.UNKNOWN,
+        # next "c"
+        RunStatus.UNKNOWN,
+    ]
+    mock_dep.side_effect = [
+        None,  # "b" has no missing dependencies
+        None,  # "c" has no missing dependencies
+    ]
+    next_tasks = [n async for n in orchestrator.next(node_a)]
+    next_task_names = [n.name for n in next_tasks]
+    assert next_task_names == ["b", "c"]
+
+
+@pytest.mark.asyncio
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.get_status_for_component"
+)
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.check_dependencies_complete",
+)
+async def test_orchestrator_next_task_branch_missing_dependencies(
+    mock_dep: Mock, mock_status: Mock, pipeline_branch: Pipeline
+) -> None:
+    """a -> b, c"""
+    orchestrator = Orchestrator(pipeline=pipeline_branch)
+    node_a = pipeline_branch.get_node_by_name("a")
+    mock_status.side_effect = [
+        # next "b"
+        RunStatus.UNKNOWN,
+        # next "c"
+        RunStatus.UNKNOWN,
+    ]
+    mock_dep.side_effect = [
+        PipelineMissingDependencyError,  # "b" has missing dependencies
+        None,  # "c" has no missing dependencies
+    ]
+    next_tasks = [n async for n in orchestrator.next(node_a)]
+    next_task_names = [n.name for n in next_tasks]
+    assert next_task_names == ["c"]
+
+
+@pytest.mark.asyncio
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.get_status_for_component"
+)
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.check_dependencies_complete",
+)
+async def test_orchestrator_next_task_aggregation_no_missing_dependencies(
+    mock_dep: Mock, mock_status: Mock, pipeline_aggregation: Pipeline
+) -> None:
+    """a, b -> c"""
+    orchestrator = Orchestrator(pipeline=pipeline_aggregation)
+    node_a = pipeline_aggregation.get_node_by_name("a")
+    mock_status.side_effect = [
+        RunStatus.UNKNOWN,  # status for "c", not started
+    ]
+    mock_dep.side_effect = [
+        None,  # no missing deps
+    ]
+    # then "c" can start
+    next_tasks = [n async for n in orchestrator.next(node_a)]
+    next_task_names = [n.name for n in next_tasks]
+    assert next_task_names == ["c"]
+
+
+@pytest.mark.asyncio
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.get_status_for_component"
+)
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.check_dependencies_complete",
+)
+async def test_orchestrator_next_task_aggregation_missing_dependency(
+    mock_dep: Mock, mock_status: Mock, pipeline_aggregation: Pipeline
+) -> None:
+    """a, b -> c"""
+    orchestrator = Orchestrator(pipeline=pipeline_aggregation)
+    node_a = pipeline_aggregation.get_node_by_name("a")
+    mock_status.side_effect = [
+        RunStatus.UNKNOWN,  # status for "c" is unknown, it's a possible next
+    ]
+    mock_dep.side_effect = [
+        PipelineMissingDependencyError,  # some dependencies are not done yet
+    ]
+    next_task_names = [n.name async for n in orchestrator.next(node_a)]
+    # "c" dependencies not ready yet
+    assert next_task_names == []
+
+
+@pytest.mark.asyncio
+@patch(
+    "neo4j_graphrag.experimental.pipeline.pipeline.Orchestrator.get_status_for_component"
+)
+async def test_orchestrator_next_task_aggregation_next_already_started(
+    mock_status: Mock, pipeline_aggregation: Pipeline
+) -> None:
+    """a, b -> c"""
+    orchestrator = Orchestrator(pipeline=pipeline_aggregation)
+    node_a = pipeline_aggregation.get_node_by_name("a")
+    mock_status.side_effect = [
+        RunStatus.RUNNING,  # status for "c" is already running, do not start it again
+    ]
+    next_task_names = [n.name async for n in orchestrator.next(node_a)]
+    assert next_task_names == []

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/test_pipeline.py

@@ -0,0 +1,607 @@
+#  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 asyncio
+import datetime
+import tempfile
+from typing import Sized
+from unittest import mock
+from unittest.mock import AsyncMock, call, patch
+
+import pytest
+from neo4j_graphrag.experimental.pipeline import Component, Pipeline
+from neo4j_graphrag.experimental.pipeline.exceptions import PipelineDefinitionError
+from neo4j_graphrag.experimental.pipeline.notification import (
+    EventCallbackProtocol,
+    EventType,
+    PipelineEvent,
+    TaskEvent,
+    Event,
+)
+from neo4j_graphrag.experimental.pipeline.types.orchestration import RunResult
+
+from .components import (
+    ComponentAdd,
+    ComponentMultiply,
+    ComponentNoParam,
+    ComponentPassThrough,
+    StringResultModel,
+    SlowComponentMultiply,
+)
+
+
+@pytest.mark.asyncio
+async def test_simple_pipeline_two_components() -> None:
+    pipe = Pipeline()
+    component_a = ComponentNoParam()
+    component_b = ComponentNoParam()
+    pipe.add_component(
+        component_a,
+        "a",
+    )
+    pipe.add_component(
+        component_b,
+        "b",
+    )
+    pipe.connect("a", "b", {})
+    with mock.patch(
+        "tests.unit.experimental.pipeline.test_pipeline.ComponentNoParam.run"
+    ) as mock_run:
+        mock_run.side_effect = [
+            StringResultModel(result="1"),
+            StringResultModel(result="2"),
+        ]
+        res = await pipe.run({})
+        mock_run.assert_awaited_with(**{})
+        mock_run.assert_awaited_with(**{})
+    assert "b" in res.result
+    assert res.result["b"] == {"result": "2"}
+
+
+@pytest.mark.asyncio
+async def test_pipeline_parameter_propagation() -> None:
+    pipe = Pipeline()
+    component_a = ComponentPassThrough()
+    component_b = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    # first component output product goes to second component input number1
+    pipe.connect("a", "b", {"value": "a.result"})
+    with mock.patch(
+        "tests.unit.experimental.pipeline.test_pipeline.ComponentPassThrough.run"
+    ) as mock_run:
+        mock_run.side_effect = [
+            StringResultModel(result="1"),
+            StringResultModel(result="2"),
+        ]
+        res = await pipe.run({"a": {"value": "text"}})
+        mock_run.assert_has_awaits([call(**{"value": "text"}), call(**{"value": "1"})])
+    assert res.result == {"b": {"result": "2"}}
+
+
+def test_pipeline_parameter_validation_no_expected_params() -> None:
+    pipe = Pipeline()
+    component_a = ComponentNoParam()
+    pipe.add_component(component_a, "a")
+    is_valid = pipe.validate_parameter_mapping_for_task(pipe.get_node_by_name("a"))
+    assert is_valid is True
+
+
+def test_pipeline_parameter_validation_one_component_all_good() -> None:
+    pipe = Pipeline()
+    component_a = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    is_valid = pipe.validate_parameter_mapping_for_task(pipe.get_node_by_name("a"))
+    assert is_valid is True
+
+
+def test_pipeline_invalidate() -> None:
+    pipe = Pipeline()
+    pipe.is_validated = True
+    pipe.param_mapping = {"a": {"key": {"component": "component", "param": "param"}}}
+    pipe.missing_inputs = {"a": ["other_key"]}
+    pipe.invalidate()
+    assert pipe.is_validated is False
+    assert len(pipe.param_mapping) == 0
+    assert len(pipe.missing_inputs) == 0
+
+
+def test_pipeline_parameter_validation_called_twice() -> None:
+    pipe = Pipeline()
+    component_a = ComponentPassThrough()
+    component_b = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {"value": "a.result"})
+    is_valid = pipe.validate_parameter_mapping_for_task(pipe.get_node_by_name("b"))
+    assert is_valid is True
+    with pytest.raises(PipelineDefinitionError):
+        pipe.validate_parameter_mapping_for_task(pipe.get_node_by_name("b"))
+    pipe.invalidate()
+    is_valid = pipe.validate_parameter_mapping_for_task(pipe.get_node_by_name("b"))
+    assert is_valid is True
+
+
+def test_pipeline_parameter_validation_one_component_input_param_missing() -> None:
+    pipe = Pipeline()
+    component_a = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    pipe.validate_parameter_mapping_for_task(pipe.get_node_by_name("a"))
+    assert pipe.missing_inputs["a"] == ["value"]
+
+
+def test_pipeline_parameter_validation_param_mapped_twice() -> None:
+    pipe = Pipeline()
+    component_a = ComponentPassThrough()
+    component_b = ComponentPassThrough()
+    component_c = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.add_component(component_c, "c")
+    pipe.connect("a", "c", {"value": "a.result"})
+    pipe.connect("b", "c", {"value": "b.result"})
+    with pytest.raises(PipelineDefinitionError) as excinfo:
+        pipe.validate_parameter_mapping_for_task(pipe.get_node_by_name("c"))
+        assert (
+            "Parameter 'value' already mapped to {'component': 'a', 'param': 'result'}"
+            in str(excinfo)
+        )
+
+
+def test_pipeline_parameter_validation_unexpected_input() -> None:
+    pipe = Pipeline()
+    component_a = ComponentPassThrough()
+    component_b = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {"unexpected_input_name": "a.result"})
+    with pytest.raises(PipelineDefinitionError) as excinfo:
+        pipe.validate_parameter_mapping_for_task(pipe.get_node_by_name("b"))
+        assert (
+            "Parameter 'unexpected_input_name' is not a valid input for component 'b' of type 'ComponentPassThrough'"
+            in str(excinfo)
+        )
+
+
+def test_pipeline_parameter_validation_connected_components_input() -> None:
+    """Parameter for component 'b' comes from the pipeline inputs"""
+    pipe = Pipeline()
+    component_a = ComponentNoParam()
+    component_b = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {})
+    is_valid = pipe.validate_parameter_mapping_for_task(pipe.get_node_by_name("b"))
+    assert is_valid is True
+    assert dict(pipe.missing_inputs) == {"b": ["value"]}
+
+
+def test_pipeline_parameter_validation_connected_components_result() -> None:
+    """Parameter for component 'b' comes from the result of component 'a'"""
+    pipe = Pipeline()
+    component_a = ComponentNoParam()
+    component_b = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {"value": "b.result"})
+    is_valid = pipe.validate_parameter_mapping_for_task(pipe.get_node_by_name("b"))
+    assert is_valid is True
+    assert pipe.missing_inputs == {"b": []}
+
+
+def test_pipeline_parameter_validation_connected_components_missing_input() -> None:
+    """Parameter for component 'b' is missing"""
+    pipe = Pipeline()
+    component_a = ComponentNoParam()
+    component_b = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {})
+    is_valid = pipe.validate_parameter_mapping_for_task(pipe.get_node_by_name("b"))
+    assert is_valid is True
+    assert pipe.missing_inputs["b"] == ["value"]
+
+
+def test_pipeline_parameter_validation_full_missing_inputs_in_user_data() -> None:
+    pipe = Pipeline()
+    component_a = ComponentNoParam()
+    component_b = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {})
+    is_valid = pipe.validate_input_data(data={"b": {"value": "input for b"}})
+    assert is_valid is True
+
+
+def test_pipeline_parameter_validation_full_missing_inputs_in_component_name() -> None:
+    pipe = Pipeline()
+    component_a = ComponentNoParam()
+    component_b = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {})
+    with pytest.raises(PipelineDefinitionError):
+        pipe.validate_input_data(data={"b": {}})
+
+
+def test_pipeline_parameter_validation_full_missing_inputs() -> None:
+    pipe = Pipeline()
+    component_a = ComponentNoParam()
+    component_b = ComponentPassThrough()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {})
+    with pytest.raises(PipelineDefinitionError):
+        pipe.validate_input_data(data={})
+
+
+@pytest.mark.asyncio
+async def test_pipeline_branches() -> None:
+    pipe = Pipeline()
+    component_a = AsyncMock(spec=Component)
+    component_a.run_with_context = AsyncMock(return_value={})
+    component_b = AsyncMock(spec=Component)
+    component_b.run_with_context = AsyncMock(return_value={})
+    component_c = AsyncMock(spec=Component)
+    component_c.run_with_context = AsyncMock(return_value={})
+
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.add_component(component_c, "c")
+    pipe.connect("a", "b")
+    pipe.connect("a", "c")
+    pipeline_result = await pipe.run({})
+    res = pipeline_result.result
+    assert "b" in res
+    assert "c" in res
+
+
+@pytest.mark.asyncio
+async def test_pipeline_aggregation() -> None:
+    pipe = Pipeline()
+    component_a = AsyncMock(spec=Component)
+    component_a.run_with_context = AsyncMock(return_value={})
+    component_b = AsyncMock(spec=Component)
+    component_b.run_with_context = AsyncMock(return_value={})
+    component_c = AsyncMock(spec=Component)
+    component_c.run_with_context = AsyncMock(return_value={})
+
+    pipe.add_component(
+        component_a,
+        "a",
+    )
+    pipe.add_component(
+        component_b,
+        "b",
+    )
+    pipe.add_component(component_c, "c")
+    pipe.connect("a", "c")
+    pipe.connect("b", "c")
+    pipeline_result = await pipe.run({})
+    res = pipeline_result.result
+    assert "c" in res
+
+
+@pytest.mark.asyncio
+async def test_pipeline_missing_param_on_init() -> None:
+    pipe = Pipeline()
+    component_a = ComponentAdd()
+    component_b = ComponentAdd()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {"number1": "a.result"})
+    with pytest.raises(PipelineDefinitionError) as excinfo:
+        await pipe.run({"a": {"number1": 1}})
+        assert (
+            "Missing input parameters for a: Expected parameters: ['number1', 'number2']. Got: ['number1']"
+            in str(excinfo.value)
+        )
+
+
+@pytest.mark.asyncio
+async def test_pipeline_missing_param_on_connect() -> None:
+    pipe = Pipeline()
+    component_a = ComponentAdd()
+    component_b = ComponentAdd()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {"number1": "a.result"})
+    with pytest.raises(PipelineDefinitionError) as excinfo:
+        await pipe.run({"a": {"number1": 1, "number2": 2}})
+        assert (
+            "Missing input parameters for b: Expected parameters: ['number1', 'number2']. Got: ['number1']"
+            in str(excinfo.value)
+        )
+
+
+@pytest.mark.asyncio
+async def test_pipeline_with_default_params() -> None:
+    pipe = Pipeline()
+    component_a = ComponentAdd()
+    component_b = ComponentMultiply()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {"number1": "a.result"})
+    pipeline_result = await pipe.run({"a": {"number1": 1, "number2": 2}})
+    res = pipeline_result.result
+    assert res == {"b": {"result": 6}}  # (1+2)*2
+
+
+@pytest.mark.asyncio
+async def test_pipeline_cycle() -> None:
+    pipe = Pipeline()
+    component_a = ComponentNoParam()
+    component_b = ComponentNoParam()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {})
+    with pytest.raises(PipelineDefinitionError) as excinfo:
+        pipe.connect("b", "a", {})
+        assert "Cycles are not allowed" in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_pipeline_wrong_component_name() -> None:
+    pipe = Pipeline()
+    component_a = ComponentNoParam()
+    component_b = ComponentNoParam()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    with pytest.raises(PipelineDefinitionError) as excinfo:
+        pipe.connect("a", "c", {})
+        assert "a or c not in the Pipeline" in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_pipeline_async() -> None:
+    pipe = Pipeline()
+    pipe.add_component(ComponentAdd(), "add")
+    run_params = [[1, 20], [10, 2]]
+    runs = []
+    for a, b in run_params:
+        runs.append(pipe.run({"add": {"number1": a, "number2": b}}))
+    pipeline_result = await asyncio.gather(*runs)
+    assert len(pipeline_result) == 2
+    assert pipeline_result[0].run_id != pipeline_result[1].run_id
+    assert pipeline_result[0].result == {"add": {"result": 21}}
+    assert pipeline_result[1].result == {"add": {"result": 12}}
+
+
+def test_pipeline_to_viz() -> None:
+    pipe = Pipeline()
+    component_a = ComponentAdd()
+    component_b = ComponentMultiply()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {"number1": "a.result"})
+    g = pipe._get_neo4j_viz_graph()
+    # 3 nodes:
+    #   - 2 components 'a' and 'b'
+    #   - 1 output 'a.result'
+    assert len(g.nodes) == 3
+    g = pipe._get_neo4j_viz_graph(hide_unused_outputs=False)
+    # 4 nodes:
+    #   - 2 components 'a' and 'b'
+    #   - 2 output 'a.result' and 'b.result'
+    assert len(g.nodes) == 4
+
+
+def test_pipeline_draw() -> None:
+    pipe = Pipeline()
+    pipe.add_component(ComponentAdd(), "add")
+    t = tempfile.NamedTemporaryFile(suffix=".html")
+    pipe.draw(t.name)
+    content = t.file.read()
+    assert len(content) > 0
+
+
+@patch("neo4j_graphrag.experimental.pipeline.pipeline.neo4j_viz_available", False)
+def test_pipeline_draw_missing_neo4j_viz_dep() -> None:
+    pipe = Pipeline()
+    pipe.add_component(ComponentAdd(), "add")
+    t = tempfile.NamedTemporaryFile(suffix=".html")
+    with pytest.raises(ImportError):
+        pipe.draw(t.name)
+
+
+def test_run_result_no_warning(recwarn: Sized) -> None:
+    RunResult()
+    assert len(recwarn) == 0
+
+
+@pytest.mark.asyncio
+async def test_pipeline_event_notification() -> None:
+    callback = AsyncMock(spec=EventCallbackProtocol)
+    pipe = Pipeline(callback=callback)
+    component_a = ComponentMultiply()
+    pipe.add_component(
+        component_a,
+        "a",
+    )
+    a_input_data = {"number1": 2, "number2": 3}
+    pipeline_result = await pipe.run({"a": a_input_data})
+
+    await_calls = callback.await_args_list
+
+    expected_event_list = [
+        PipelineEvent(
+            event_type=EventType.PIPELINE_STARTED,
+            run_id=pipeline_result.run_id,
+            timestamp=datetime.datetime.now(),
+            message=None,
+            payload={"a": a_input_data},
+        ),
+        TaskEvent(
+            event_type=EventType.TASK_STARTED,
+            run_id=pipeline_result.run_id,
+            task_name="a",
+            timestamp=datetime.datetime.now(),
+            message=None,
+            payload=a_input_data,
+        ),
+        TaskEvent(
+            event_type=EventType.TASK_FINISHED,
+            run_id=pipeline_result.run_id,
+            task_name="a",
+            timestamp=datetime.datetime.now(),
+            message=None,
+            payload={"result": 6},
+        ),
+        PipelineEvent(
+            event_type=EventType.PIPELINE_FINISHED,
+            run_id=pipeline_result.run_id,
+            timestamp=datetime.datetime.now(),
+            message=None,
+            payload={"a": {"result": 6}},
+        ),
+    ]
+    assert len(await_calls) == len(expected_event_list)
+
+    previous_ts = None
+    for await_call, expected_event in zip(await_calls, expected_event_list):
+        actual_event = await_call[0][0]
+        assert isinstance(actual_event, type(expected_event))
+        assert actual_event.event_type == expected_event.event_type
+        assert actual_event.run_id == expected_event.run_id
+        assert actual_event.message == expected_event.message
+        assert actual_event.payload == expected_event.payload
+        if previous_ts:
+            assert actual_event.timestamp > previous_ts
+        previous_ts = actual_event.timestamp
+
+
+@pytest.mark.asyncio
+async def test_pipeline_event_notification_error_in_pipeline_run() -> None:
+    callback = AsyncMock(spec=EventCallbackProtocol)
+    pipe = Pipeline(callback=callback)
+    component_a = ComponentAdd()
+    component_b = ComponentAdd()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {"number1": "a.result"})
+
+    with pytest.raises(PipelineDefinitionError):
+        await pipe.run({"a": {"number1": 1, "number2": 2}})
+    assert len(callback.await_args_list) == 2
+    assert callback.await_args_list[0][0][0].event_type == EventType.PIPELINE_STARTED
+    assert callback.await_args_list[1][0][0].event_type == EventType.PIPELINE_FAILED
+
+
+def test_event_model_no_warning(recwarn: Sized) -> None:
+    event = Event(
+        event_type=EventType.PIPELINE_STARTED,
+        run_id="run_id",
+        message=None,
+        payload=None,
+    )
+    assert event.timestamp is not None
+    assert len(recwarn) == 0
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_no_user_callback_happy_path() -> None:
+    pipe = Pipeline()
+    events = []
+    async for e in pipe.stream({}):
+        events.append(e)
+    assert len(events) == 2
+    assert events[0].event_type == EventType.PIPELINE_STARTED
+    assert events[1].event_type == EventType.PIPELINE_FINISHED
+    assert len(pipe.event_notifier.callbacks) == 0
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_with_user_callback_happy_path() -> None:
+    callback = AsyncMock()
+    pipe = Pipeline(callback=callback)
+    events = []
+    async for e in pipe.stream({}):
+        events.append(e)
+    assert len(events) == 2
+    assert len(callback.call_args_list) == 2
+    assert len(pipe.event_notifier.callbacks) == 1
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_very_long_running_user_callback() -> None:
+    async def callback(event: Event) -> None:
+        await asyncio.sleep(2)
+
+    pipe = Pipeline(callback=callback)
+    events = []
+    async for e in pipe.stream({}):
+        events.append(e)
+    assert len(events) == 2
+    assert len(pipe.event_notifier.callbacks) == 1
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_very_long_running_pipeline() -> None:
+    slow_component = SlowComponentMultiply()
+    pipe = Pipeline()
+    pipe.add_component(slow_component, "slow_component")
+    events = []
+    async for e in pipe.stream({"slow_component": {"number1": 1, "number2": 2}}):
+        events.append(e)
+    assert len(events) == 4
+    last_event = events[-1]
+    assert last_event.event_type == EventType.PIPELINE_FINISHED
+    assert last_event.payload == {"slow_component": {"result": 2}}
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_error_in_pipeline_definition() -> None:
+    pipe = Pipeline()
+    component_a = ComponentAdd()
+    component_b = ComponentAdd()
+    pipe.add_component(component_a, "a")
+    pipe.add_component(component_b, "b")
+    pipe.connect("a", "b", {"number1": "a.result"})
+    events = []
+    with pytest.raises(PipelineDefinitionError):
+        async for e in pipe.stream({"a": {"number1": 1, "number2": 2}}):
+            events.append(e)
+    assert len(events) == 2
+    assert events[0].event_type == EventType.PIPELINE_STARTED
+    assert events[1].event_type == EventType.PIPELINE_FAILED
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_error_in_component() -> None:
+    component = ComponentMultiply()
+    pipe = Pipeline()
+    pipe.add_component(component, "component")
+    events = []
+    with pytest.raises(TypeError):
+        async for e in pipe.stream({"component": {"number1": None, "number2": 2}}):
+            events.append(e)
+    assert len(events) == 3
+    assert events[0].event_type == EventType.PIPELINE_STARTED
+    assert events[1].event_type == EventType.TASK_STARTED
+    assert events[2].event_type == EventType.PIPELINE_FAILED
+
+
+@pytest.mark.asyncio
+async def test_pipeline_streaming_error_in_user_callback() -> None:
+    async def callback(event: Event) -> None:
+        raise Exception("error in callback")
+
+    pipe = Pipeline(callback=callback)
+    events = []
+    async for e in pipe.stream({}):
+        events.append(e)
+    assert len(events) == 2
+    assert len(pipe.event_notifier.callbacks) == 1

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/test_pipeline_graph.py

@@ -0,0 +1,137 @@
+#  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 pytest
+from neo4j_graphrag.experimental.pipeline.pipeline_graph import (
+    PipelineEdge,
+    PipelineGraph,
+    PipelineNode,
+)
+
+
+def test_node_alone() -> None:
+    n = PipelineNode(name="node", data={})
+    assert n.is_root() is True
+    assert n.is_leaf() is True
+
+
+def test_node_not_root() -> None:
+    n = PipelineNode(name="node", data={})
+    n.parents = ["other_node"]
+    assert n.is_root() is False
+    assert n.is_leaf() is True
+
+
+def test_node_not_leaf() -> None:
+    n = PipelineNode(name="node", data={})
+    n.children = ["other_node"]
+    assert n.is_root() is True
+    assert n.is_leaf() is False
+
+
+def test_graph_add_nodes() -> None:
+    g: PipelineGraph[PipelineNode, PipelineEdge] = PipelineGraph()
+    n1 = PipelineNode("n1", {})
+    n2 = PipelineNode("n2", {})
+    g.add_node(n1)
+    g.add_node(n2)
+    assert len(g._nodes) == 2
+    edge = PipelineEdge(n1.name, n2.name, {"key": "value"})
+    g.add_edge(edge)
+    assert len(g._edges) == 1
+
+    assert n1.children == [n2.name]
+    assert n2.parents == [n1.name]
+
+
+@pytest.fixture(scope="function")
+def graph() -> PipelineGraph[PipelineNode, PipelineEdge]:
+    g: PipelineGraph[PipelineNode, PipelineEdge] = PipelineGraph()
+    n1 = PipelineNode("n1", {})
+    n2 = PipelineNode("n2", {})
+    g.add_node(n1)
+    g.add_node(n2)
+    edge = PipelineEdge(n1.name, n2.name, {"key": "value"})
+    g.add_edge(edge)
+    return g
+
+
+def test_graph_roots(graph: PipelineGraph[PipelineNode, PipelineEdge]) -> None:
+    roots = graph.roots()
+    assert len(roots) == 1
+    assert roots[0].name == "n1"
+
+
+def test_graph_next_edge(graph: PipelineGraph[PipelineNode, PipelineEdge]) -> None:
+    start = graph._nodes["n1"]
+    next_edges = graph.next_edges(start.name)
+    assert len(next_edges) == 1
+    next_edge = next_edges[0]
+    assert isinstance(next_edge, PipelineEdge)
+    assert next_edge.start == "n1"
+    assert next_edge.end == "n2"
+
+
+def test_graph_prev_edge(graph: PipelineGraph[PipelineNode, PipelineEdge]) -> None:
+    start = graph._nodes["n2"]
+    next_edges = graph.previous_edges(start.name)
+    assert len(next_edges) == 1
+    next_edge = next_edges[0]
+    assert isinstance(next_edge, PipelineEdge)
+    assert next_edge.start == "n1"
+    assert next_edge.end == "n2"
+
+
+def test_graph_contains(graph: PipelineGraph[PipelineNode, PipelineEdge]) -> None:
+    start = graph._nodes["n2"]
+    assert start in graph
+
+
+def test_graph_is_cyclic(graph: PipelineGraph[PipelineNode, PipelineEdge]) -> None:
+    g: PipelineGraph[PipelineNode, PipelineEdge] = PipelineGraph()
+    n1 = PipelineNode("n1", {})
+    n2 = PipelineNode("n2", {})
+    g.add_node(n1)
+    g.add_node(n2)
+    edge = PipelineEdge(n1.name, n2.name, {})
+    g.add_edge(edge)
+    assert g.is_cyclic() is False
+
+    edge = PipelineEdge(n2.name, n1.name, {})
+    g.add_edge(edge)
+    assert g.is_cyclic() is True
+
+
+def test_graph_set_node(graph: PipelineGraph[PipelineNode, PipelineEdge]) -> None:
+    new_node = PipelineNode("n1", {})
+    graph.set_node(new_node)
+    new_node_from_graph = graph.get_node_by_name("n1")
+    assert new_node_from_graph.parents == []
+    assert new_node_from_graph.children == ["n2"]
+
+
+def test_graph_validate_edge_bad_node_name(
+    graph: PipelineGraph[PipelineNode, PipelineEdge],
+) -> None:
+    with pytest.raises(KeyError):
+        graph.add_edge(PipelineEdge("n0", "n1", {}))
+    with pytest.raises(KeyError):
+        graph.add_edge(PipelineEdge("n1", "n12", {}))
+
+
+def test_graph_validate_edge_no_parallel_edges(
+    graph: PipelineGraph[PipelineNode, PipelineEdge],
+) -> None:
+    with pytest.raises(ValueError):
+        graph.add_edge(PipelineEdge("n1", "n2", {}))

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/test_store.py

@@ -0,0 +1,15 @@
+import pytest
+from neo4j_graphrag.experimental.pipeline.stores import InMemoryStore
+
+
+@pytest.mark.asyncio
+async def test_memory_store() -> None:
+    store = InMemoryStore()
+    await store.add("key", "value")
+    res = await store.get("key")
+    assert res == "value"
+
+    with pytest.raises(KeyError):
+        await store.add("key", "value", overwrite=False)
+
+    assert store.all() == {"key": "value"}

참고/neo4j-graphrag-python-main/tests/unit/experimental/utils/__init__.py

@@ -0,0 +1,14 @@
+#  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.

참고/neo4j-graphrag-python-main/tests/unit/experimental/utils/test_schema.py

@@ -0,0 +1,107 @@
+#  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
+from unittest.mock import patch
+
+import pytest
+from pydantic import ValidationError
+
+from neo4j_viz import VisualizationGraph
+from neo4j_graphrag.experimental.components.schema import GraphSchema
+from neo4j_graphrag.experimental.utils.schema import schema_visualization
+
+
+@pytest.fixture(scope="module")
+def valid_schema_dict() -> dict[str, Any]:
+    return {
+        "node_types": [
+            "Location",
+            {
+                "label": "Person",
+                "properties": [
+                    {"name": "name", "type": "STRING", "required": True},
+                    {"name": "birthYear", "type": "INTEGER"},
+                ],
+            },
+        ],
+        "relationship_types": [
+            "BORN_IN",
+            {
+                "label": "KNOWS",
+                "properties": [
+                    {"name": "since", "type": "LOCAL_DATETIME"},
+                ],
+            },
+        ],
+        "patterns": [
+            ("Person", "BORN_IN", "Location"),
+            ("Person", "KNOWS", "Person"),
+        ],
+    }
+
+
+@pytest.fixture(scope="module")
+def invalid_schema_dict() -> dict[str, Any]:
+    return {
+        "node_types": [
+            {
+                "label": "Person",
+                "properties": [
+                    {"name": "name", "type": "STRING", "required": True},
+                    {"name": "birthYear", "type": "INTEGER"},
+                ],
+            },
+        ],
+        "relationship_types": [
+            "BORN_IN",
+        ],
+        "patterns": [
+            (
+                "Person",
+                "BORN_IN",
+                "Location",
+            ),  # invalid pattern, "Location" node type not defined
+        ],
+    }
+
+
+@patch("neo4j_graphrag.experimental.utils.schema.VisualizationGraph", None)
+def test_schema_visualization_import_error() -> None:
+    with pytest.raises(ImportError):
+        schema_visualization({})
+
+
+def test_schema_visualization_invalid_schema_dict(
+    invalid_schema_dict: dict[str, Any],
+) -> None:
+    with pytest.raises(ValidationError):
+        schema_visualization(invalid_schema_dict)
+
+
+def test_schema_visualization_valid_schema_dict(
+    valid_schema_dict: dict[str, Any],
+) -> None:
+    g = schema_visualization(valid_schema_dict)
+    assert isinstance(g, VisualizationGraph)
+    assert len(g.nodes) == 2
+    assert len(g.relationships) == 2
+
+
+def test_schema_visualization_schema_object(valid_schema_dict: dict[str, Any]) -> None:
+    schema = GraphSchema.model_validate(valid_schema_dict)
+    g = schema_visualization(schema)
+    assert isinstance(g, VisualizationGraph)
+    assert len(g.nodes) == 2
+    assert len(g.relationships) == 2