참고소스 수정본

참고/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