Add embedder for TextChunkEmbedder to SimpleKGPipeline (#166)

* Add embedder for TextChunkEmbedder to SimpleKGPipeline

* E2E test

* llm back to MagicMock

* from_pdf = False in e2e test

* Remove async client close in E2E test

Author: willtai

README.md

@@ -64,6 +64,7 @@ llm = OpenAILLM(
 kg_builder = SimpleKGPipeline(
     llm=llm,
     driver=driver,
+    embedder=OpenAIEmbeddings(),
     file_path=file_path,
     entities=entities,
     relations=relations,

examples/pipeline/kg_builder_example.py

@@ -18,6 +18,7 @@
 import logging
 
 import neo4j
+from neo4j_graphrag.embeddings import OpenAIEmbeddings
 from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
 from neo4j_graphrag.llm.openai_llm import OpenAILLM
 
@@ -44,10 +45,14 @@ async def main(neo4j_driver: neo4j.Driver) -> None:
         },
     )
 
+    # Use OpenAIEmbeddings as embedder
+    embedder = OpenAIEmbeddings()
+
     # Create an instance of the SimpleKGPipeline
     kg_builder_pdf = SimpleKGPipeline(
         llm=llm,
         driver=neo4j_driver,
+        embedder=embedder,
         entities=entities,
         relations=relations,
         potential_schema=potential_schema,
@@ -64,6 +69,7 @@ async def main(neo4j_driver: neo4j.Driver) -> None:
     kg_builder_text = SimpleKGPipeline(
         llm=llm,
         driver=neo4j_driver,
+        embedder=embedder,
         entities=entities,
         relations=relations,
         potential_schema=potential_schema,

src/neo4j_graphrag/experimental/pipeline/kg_builder.py

@@ -20,6 +20,8 @@
 import neo4j
 from pydantic import BaseModel, ConfigDict, Field
 
+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,
@@ -47,6 +49,7 @@ class SimpleKGPipelineConfig(BaseModel):
     llm: LLMInterface
     driver: neo4j.Driver
     from_pdf: bool
+    embedder: Embedder
     entities: list[SchemaEntity] = Field(default_factory=list)
     relations: list[SchemaRelation] = Field(default_factory=list)
     potential_schema: list[tuple[str, str, str]] = Field(default_factory=list)
@@ -68,6 +71,7 @@ class SimpleKGPipeline:
     Args:
         llm (LLMInterface): An instance of an LLM to use for entity and relation extraction.
         driver (neo4j.Driver): A Neo4j driver instance for database connection.
+        embedder (Embedder): An instance of an embedder used to generate chunk embeddings from text chunks.
         entities (Optional[List[str]]): A list of entity labels as strings.
         relations (Optional[List[str]]): A list of relation labels as strings.
         potential_schema (Optional[List[tuple]]): A list of potential schema relationships.
@@ -79,12 +83,15 @@ class SimpleKGPipeline:
         kg_writer (Optional[Any]): A knowledge graph writer component. Defaults to Neo4jWriter().
         on_error (str): Error handling strategy. Defaults to "RAISE". Possible values: "RAISE" or "IGNORE".
         perform_entity_resolution (bool): Merge entities with same label and name. Default: True
+        text_splitter (Optional[Any]): A text splitter component. Defaults to FixedSizeSplitter().
+        prompt_template (str): A custom prompt template to use for extraction.
     """
 
     def __init__(
         self,
         llm: LLMInterface,
         driver: neo4j.Driver,
+        embedder: Embedder,
         entities: Optional[List[str]] = None,
         relations: Optional[List[str]] = None,
         potential_schema: Optional[List[tuple[str, str, str]]] = None,
@@ -119,12 +126,14 @@ def __init__(
             text_splitter=text_splitter,
             on_error=on_error_enum,
             prompt_template=prompt_template,
+            embedder=embedder,
             perform_entity_resolution=perform_entity_resolution,
         )
 
         self.from_pdf = config.from_pdf
         self.llm = config.llm
         self.driver = config.driver
+        self.embedder = config.embedder
         self.text_splitter = config.text_splitter or FixedSizeSplitter()
         self.on_error = config.on_error
         self.pdf_loader = config.pdf_loader if pdf_loader is not None else PdfLoader()
@@ -149,6 +158,7 @@ def _build_pipeline(self) -> Pipeline:
             ),
             "extractor",
         )
+        pipe.add_component(TextChunkEmbedder(embedder=self.embedder), "chunk_embedder")
         pipe.add_component(self.kg_writer, "writer")
 
         if self.from_pdf:
@@ -178,9 +188,11 @@ def _build_pipeline(self) -> Pipeline:
             )
 
         pipe.connect(
-            "splitter",
-            "extractor",
-            input_config={"chunks": "splitter"},
+            "splitter", "chunk_embedder", input_config={"text_chunks": "splitter"}
+        )
+
+        pipe.connect(
+            "chunk_embedder", "extractor", input_config={"chunks": "chunk_embedder"}
         )
 
         # Connect extractor to writer

tests/e2e/test_simplekgpipeline_e2e.py

@@ -0,0 +1,128 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+import os
+from unittest.mock import MagicMock
+
+import neo4j
+import pytest
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
+from neo4j_graphrag.llm import LLMInterface, LLMResponse
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+@pytest.fixture
+def llm() -> LLMInterface:
+    llm = MagicMock(spec=LLMInterface)
+    return llm
+
+
+@pytest.fixture
+def embedder() -> Embedder:
+    embedder = MagicMock(spec=Embedder)
+    return embedder
+
+
+@pytest.fixture
+def harry_potter_text() -> str:
+    with open(os.path.join(BASE_DIR, "data/harry_potter.txt"), "r") as f:
+        text = f.read()
+    return text
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("setup_neo4j_for_kg_construction")
+async def test_pipeline_builder_happy_path(
+    harry_potter_text: str,
+    llm: MagicMock,
+    embedder: MagicMock,
+    driver: neo4j.Driver,
+) -> None:
+    """When everything works as expected, extracted entities, relations and text
+    chunks must be in the DB
+    """
+    driver.execute_query("MATCH (n) DETACH DELETE n")
+    embedder.embed_query.return_value = [1, 2, 3]
+    llm.ainvoke.side_effect = [
+        LLMResponse(
+            content="""{
+                        "nodes": [
+                            {
+                                "id": "0",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Harry Potter"
+                                }
+                            },
+                            {
+                                "id": "1",
+                                "label": "Person",
+                                "properties": {
+                                    "name": "Alastor Mad-Eye Moody"
+                                }
+                            },
+                            {
+                                "id": "2",
+                                "label": "Organization",
+                                "properties": {
+                                    "name": "The Order of the Phoenix"
+                                }
+                            }
+                        ],
+                        "relationships": [
+                            {
+                                "type": "KNOWS",
+                                "start_node_id": "0",
+                                "end_node_id": "1"
+                            },
+                            {
+                                "type": "LED_BY",
+                                "start_node_id": "2",
+                                "end_node_id": "1"
+                            }
+                        ]
+                    }"""
+        ),
+        LLMResponse(content='{"nodes": [], "relationships": []}'),
+    ]
+
+    # Instantiate Entity and Relation objects
+    entities = ["PERSON", "ORGANIZATION", "HORCRUX", "LOCATION"]
+    relations = ["SITUATED_AT", "INTERACTS", "OWNS", "LED_BY"]
+    potential_schema = [
+        ("PERSON", "SITUATED_AT", "LOCATION"),
+        ("PERSON", "INTERACTS", "PERSON"),
+        ("PERSON", "OWNS", "HORCRUX"),
+        ("ORGANIZATION", "LED_BY", "PERSON"),
+    ]
+
+    # Create an instance of the SimpleKGPipeline
+    kg_builder_text = SimpleKGPipeline(
+        llm=llm,
+        driver=driver,
+        embedder=embedder,
+        entities=entities,
+        relations=relations,
+        potential_schema=potential_schema,
+        from_pdf=False,
+        on_error="RAISE",
+    )
+
+    # Run the knowledge graph building process with text input
+    text_input = "John Doe lives in New York City."
+    await kg_builder_text.run_async(text=text_input)

tests/unit/experimental/pipeline/test_kg_builder.py

@@ -16,6 +16,7 @@
 
 import neo4j
 import pytest
+from neo4j_graphrag.embeddings import Embedder
 from neo4j_graphrag.experimental.components.entity_relation_extractor import OnError
 from neo4j_graphrag.experimental.components.schema import SchemaEntity, SchemaRelation
 from neo4j_graphrag.experimental.pipeline.exceptions import PipelineDefinitionError
@@ -28,15 +29,18 @@
 async def test_knowledge_graph_builder_init_with_text() -> None:
     llm = MagicMock(spec=LLMInterface)
     driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
 
     kg_builder = SimpleKGPipeline(
         llm=llm,
         driver=driver,
+        embedder=embedder,
         from_pdf=False,
     )
 
     assert kg_builder.llm == llm
     assert kg_builder.driver == driver
+    assert kg_builder.embedder == embedder
     assert kg_builder.from_pdf is False
     assert kg_builder.entities == []
     assert kg_builder.relations == []
@@ -60,10 +64,12 @@ async def test_knowledge_graph_builder_init_with_text() -> None:
 async def test_knowledge_graph_builder_init_with_file_path() -> None:
     llm = MagicMock(spec=LLMInterface)
     driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
 
     kg_builder = SimpleKGPipeline(
         llm=llm,
         driver=driver,
+        embedder=embedder,
         from_pdf=True,
     )
 
@@ -92,10 +98,12 @@ async def test_knowledge_graph_builder_init_with_file_path() -> None:
 async def test_knowledge_graph_builder_run_with_both_inputs() -> None:
     llm = MagicMock(spec=LLMInterface)
     driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
 
     kg_builder = SimpleKGPipeline(
         llm=llm,
         driver=driver,
+        embedder=embedder,
         from_pdf=True,
     )
 
@@ -114,11 +122,13 @@ async def test_knowledge_graph_builder_run_with_both_inputs() -> None:
 async def test_knowledge_graph_builder_run_with_no_inputs() -> None:
     llm = MagicMock(spec=LLMInterface)
     driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
 
     kg_builder = SimpleKGPipeline(
         llm=llm,
         driver=driver,
-        from_pdf=True,  # or False
+        embedder=embedder,
+        from_pdf=True,
     )
 
     with pytest.raises(PipelineDefinitionError) as exc_info:
@@ -133,10 +143,12 @@ async def test_knowledge_graph_builder_run_with_no_inputs() -> None:
 async def test_knowledge_graph_builder_document_info_with_file() -> None:
     llm = MagicMock(spec=LLMInterface)
     driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
 
     kg_builder = SimpleKGPipeline(
         llm=llm,
         driver=driver,
+        embedder=embedder,
         from_pdf=True,
     )
 
@@ -159,10 +171,12 @@ async def test_knowledge_graph_builder_document_info_with_file() -> None:
 async def test_knowledge_graph_builder_document_info_with_text() -> None:
     llm = MagicMock(spec=LLMInterface)
     driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
 
     kg_builder = SimpleKGPipeline(
         llm=llm,
         driver=driver,
+        embedder=embedder,
         from_pdf=False,
     )
 
@@ -184,6 +198,7 @@ async def test_knowledge_graph_builder_document_info_with_text() -> None:
 async def test_knowledge_graph_builder_with_entities_and_file() -> None:
     llm = MagicMock(spec=LLMInterface)
     driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
 
     entities = ["Document", "Section"]
     relations = ["CONTAINS"]
@@ -192,6 +207,7 @@ async def test_knowledge_graph_builder_with_entities_and_file() -> None:
     kg_builder = SimpleKGPipeline(
         llm=llm,
         driver=driver,
+        embedder=embedder,
         entities=entities,
         relations=relations,
         potential_schema=potential_schema,
@@ -221,10 +237,12 @@ async def test_knowledge_graph_builder_with_entities_and_file() -> None:
 def test_simple_kg_pipeline_on_error_conversion() -> None:
     llm = MagicMock(spec=LLMInterface)
     driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
 
     kg_builder = SimpleKGPipeline(
         llm=llm,
         driver=driver,
+        embedder=embedder,
         on_error="RAISE",
     )
 
@@ -234,11 +252,13 @@ def test_simple_kg_pipeline_on_error_conversion() -> None:
 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) as exc_info:
         SimpleKGPipeline(
             llm=llm,
             driver=driver,
+            embedder=embedder,
             on_error="IGNORE",
         )
 
@@ -248,9 +268,14 @@ def test_simple_kg_pipeline_on_error_invalid_value() -> None:
 def test_simple_kg_pipeline_no_entity_resolution() -> None:
     llm = MagicMock(spec=LLMInterface)
     driver = MagicMock(spec=neo4j.Driver)
+    embedder = MagicMock(spec=Embedder)
 
     kg_builder = SimpleKGPipeline(
-        llm=llm, driver=driver, on_error="CONTINUE", perform_entity_resolution=False
+        llm=llm,
+        driver=driver,
+        embedder=embedder,
+        on_error="CONTINUE",
+        perform_entity_resolution=False,
     )
 
     assert "resolver" not in kg_builder.pipeline