참고소스 수정본

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

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