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