dp-zp-agent/tools/graphrag/tests/test_graph_search_neo4j.py

185 lines
6.8 KiB
Python

"""Optional real Cypher regressions; all fixtures are rolled back."""
from __future__ import annotations
import os
import uuid
import pytest
from tools.graphrag import graph_search as gs
from tools.graphrag.graph_db import create_neo4j_driver, get_neo4j_settings
pytestmark = pytest.mark.skipif(
os.getenv("RUN_NEO4J_TESTS") != "1",
reason="Set RUN_NEO4J_TESTS=1 for isolated Neo4j regressions.",
)
@pytest.fixture
def graph(monkeypatch):
scope = "regression-" + uuid.uuid4().hex
monkeypatch.setattr(gs, "GRAPH_SCOPE", scope)
category = "regcat-" + uuid.uuid4().hex
settings = get_neo4j_settings()
with (
create_neo4j_driver(settings) as driver,
driver.session(database=settings.database) as session,
):
tx = session.begin_transaction()
try:
for suffix, name, year, tag in [
("a", "Alena Nová", 2024, "alpha"),
("b", "Boris Iný", 2024, "beta"),
("c", "Cyril Malý", 2025, "alpha"),
]:
tx.run(
"CREATE (p:Person {id:$id, name:$name, source_document:$path, graph_scope:$scope}) "
"CREATE (d:Document {path:$path, source_url:$url, graph_scope:$scope}) "
'CREATE (t:Thesis {id:$thesis, title:$title, year:$year, thesis_type:"bachelor_thesis", graph_scope:$scope}) '
"MERGE (k:Keyword {id:$keyword, name:$tag, graph_scope:$scope}) "
"CREATE (p)-[:HAS_DOCUMENT]->(d)-[:HAS_KEYWORD]->(k) "
"CREATE (p)-[:AUTHORED]->(t)-[:EVIDENCED_BY]->(d)",
id=scope + suffix,
name=name,
path=scope + "/" + suffix + ".md",
url="https://example.test/" + suffix,
scope=scope,
thesis=scope + "th" + suffix,
title="Výskum " + suffix,
year=year,
keyword=scope + tag,
tag=tag,
).consume()
tx.run(
"MATCH (p:Person {id:$person})-[:HAS_DOCUMENT]->(d) "
"CREATE (c:Category {name:$category, graph_scope:$scope}) "
'CREATE (a:DocumentAuthor {id:$author, name:"Autor Správny", graph_scope:$scope}) '
"CREATE (d)-[:IN_CATEGORY]->(c) CREATE (a)-[:AUTHORED_DOCUMENT]->(d) "
"CREATE (child:Document {path:$child, graph_scope:$scope}) "
'CREATE (wrong:DocumentAuthor {id:$wrong, name:"Autor Iný", graph_scope:$scope}) '
"CREATE (p)-[:HAS_DOCUMENT]->(child)<-[:AUTHORED_DOCUMENT]-(wrong)",
person=scope + "a",
category=category,
scope=scope,
author=scope + "author",
child=scope + "/child.md",
wrong=scope + "wrong",
).consume()
yield tx, category
finally:
tx.rollback()
def test_real_keyword_year_conjunction(graph):
tx, _ = graph
result = gs.search_graph_with_session(
tx, "Ktorá z osôb Alena Nová a Boris Iný má kľúčové slovo alpha a jej práca je v roku 2024?"
)
assert {r["person"] for r in result["facts"] if r["fact_type"] == "person_thesis"} == {
"Alena Nová"
}
assert {r["person"] for r in result["facts"] if r["fact_type"] == "person_keyword"} == {
"Alena Nová"
}
def test_real_count_and_rejected_candidate_provenance(graph):
tx, _ = graph
result = gs.search_graph_with_session(
tx,
"Z kandidátov Alena Nová, Boris Iný, Cyril Malý vyber osoby s prácou v roku 2024. Koľko ich je?",
)
assert result["aggregation"] == {
"count": 2,
"people": ["Alena Nová", "Boris Iný"],
"candidate_count": 3,
"complete": True,
}
assert "https://example.test/c" in {r["source_url"] for r in result["documents"]}
def test_real_author_and_category_must_share_document(graph):
tx, category = graph
result = gs.search_graph_with_session(
tx, f"Ktorý autor dokumentu je pri osobe Alena Nová spojený s kategóriou {category}?"
)
authors = [
r["document_author"] for r in result["facts"] if r["fact_type"] == "person_document_author"
]
assert authors == ["Autor Správny"]
def test_real_shared_year_does_not_invent_intersection(graph):
tx, _ = graph
result = gs.search_graph_with_session(
tx, "V ktorom rovnakom roku je práca pri osobách Alena Nová a Cyril Malý?"
)
assert not any(r["fact_type"] in {"person_thesis", "person_work"} for r in result["facts"])
assert (
next(r for r in result["facts"] if r["fact_type"] == "academic_year_comparison")[
"shared_years"
]
== []
)
def test_real_shared_keyword_uses_requested_keyword(graph):
tx, _ = graph
result = gs.search_graph_with_session(
tx,
"Sú osoby Alena Nová a Cyril Malý prepojené cez kľúčové slovo alpha? Uveď spoločné kľúčové slovo.",
)
assert [r["keyword"] for r in result["facts"] if r["fact_type"] == "shared_keyword"] == [
"alpha"
]
def test_real_unknown_keyword_is_not_a_global_listing(graph):
tx, _ = graph
result = gs.search_graph_with_session(
tx, "Aký názov práce má Alena Nová, ak má kľúčové slovo missingtag?"
)
assert result["facts"] == []
@pytest.mark.parametrize(
"question",
[
"V ktorom rovnakom roku je práca pri osobách Alena Nová a Boris Iný?",
"Aký rok práce spája osoby Alena Nová a Boris Iný?",
],
)
def test_real_shared_year_paraphrases_have_same_answer(graph, question):
tx, _ = graph
result = gs.search_graph_with_session(tx, question)
assert next(r for r in result["facts"] if r["fact_type"] == "academic_year_comparison")[
"shared_years"
] == [2024]
assert {r["year"] for r in result["facts"] if r["fact_type"] == "person_thesis"} == {2024}
def test_real_title_and_keyword_are_both_returned(graph):
tx, _ = graph
result = gs.search_graph_with_session(
tx, "Aký je názov práce osoby Alena Nová, ak jej dokument obsahuje kľúčové slovo alpha?"
)
assert [
(r["person"], r["title"]) for r in result["facts"] if r["fact_type"] == "person_thesis"
] == [("Alena Nová", "Výskum a")]
assert [r["keyword"] for r in result["facts"] if r["fact_type"] == "person_keyword"] == [
"alpha"
]
def test_real_count_is_not_cut_off_by_presentation_limit(graph):
tx, _ = graph
result = gs.search_graph_with_session(
tx,
"Z kandidátov Alena Nová, Boris Iný, Cyril Malý vyber osoby s prácou v roku 2024. Koľko ich je?",
limit=1,
)
assert result["aggregation"]["count"] == 2
assert len(result["facts"]) == 1
assert result["truncated"]