jp170na/dp-zp-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add GraphRAG search

Browse Source
This commit is contained in:
Ján Pták 2026-10-11 00:04:06 +02:00
parent 9b59994d2c
commit f38ebf65f7
2 changed files with 2130 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1346
scripts/graph_search.py Normal file
View File

File diff suppressed because it is too large Load Diff

784
test/test_graph_search.py Normal file
View File

@ -0,0 +1,784 @@
from __future__ import annotations
from typing import Any
import pytest
import scripts.graph_search as graph_search
class FakeSession:
def __init__(
self,
responses: dict[
str,
list[dict[str, Any]],
],
) -> None:
self.responses = responses
self.calls: list[
tuple[str, dict[str, Any]]
] = []
def run(
self,
query: str,
**parameters: Any,
) -> list[dict[str, Any]]:
self.calls.append(
(
query,
parameters,
)
)
return self.responses.get(
query,
[],
)
def sample_catalog() -> list[dict[str, Any]]:
return [
{
"element_id": "p1",
"kind": "Person",
"id": "person:maros",
"name": "Maroš Harahus",
"title": None,
"heading": None,
"start_year": 2016,
"year": None,
"work_type": None,
"person_kind": "student",
"source_document": (
"pages/students/2016/"
"maros_harahus/README.md"
),
},
{
"element_id": "p2",
"kind": "Person",
"id": "person:jakub",
"name": "Jakub Maruniak",
"title": None,
"heading": None,
"start_year": 2016,
"year": None,
"work_type": None,
"person_kind": "student",
"source_document": (
"pages/students/2016/"
"jakub_maruniak/README.md"
),
},
{
"element_id": "t1",
"kind": "Topic",
"id": "spacy",
"name": "Spacy",
"title": None,
"heading": None,
"start_year": None,
"year": None,
"work_type": None,
"person_kind": None,
"source_document": None,
},
{
"element_id": "t2",
"kind": "Topic",
"id": "rag",
"name": "RAG",
"title": None,
"heading": None,
"start_year": None,
"year": None,
"work_type": None,
"person_kind": None,
"source_document": None,
},
{
"element_id": "w1",
"kind": "Work",
"id": "work:jakub:dp2021",
"name": (
"Anotácia a rozpoznávanie "
"pomenovaných entít v "
"slovenskom jazyku."
),
"title": (
"Anotácia a rozpoznávanie "
"pomenovaných entít v "
"slovenskom jazyku."
),
"heading": "Diplomová práca 2021",
"start_year": None,
"year": 2021,
"work_type": "diploma_thesis",
"person_kind": None,
"source_document": (
"pages/students/2016/"
"jakub_maruniak/README.md"
),
},
]
def test_normalize_for_match_removes_diacritics_and_punctuation() -> None:
assert (
graph_search.normalize_for_match(
" Ján Pták – Diplomová práca! "
)
== "jan ptak diplomova praca"
)
def test_extract_query_constraints_for_diploma_thesis_year() -> None:
result = (
graph_search.extract_query_constraints(
"Aký bol názov diplomovej práce "
"Jakuba Maruniaka v roku 2021?"
)
)
assert result["years"] == [
2021
]
assert result[
"work_types"
] == [
"diploma_thesis"
]
assert result[
"mentions_title"
] is True
assert "Work" in result[
"target_kinds"
]
def test_extract_query_constraints_for_start_year() -> None:
result = (
graph_search.extract_query_constraints(
"Mali Maroš Harahus a Jakub "
"Maruniak rovnaký rok začiatku "
"štúdia?"
)
)
assert result[
"mentions_start_year"
] is True
assert result[
"mentions_shared"
] is True
assert result[
"target_kinds"
] == []
def test_resolve_entities_finds_two_people_and_topic() -> None:
result = graph_search.resolve_entities(
"Akú spoločnú tému Spacy majú "
"Maroš Harahus a Jakub Maruniak?",
sample_catalog(),
)
resolved = {
(
item["kind"],
item.get("name"),
)
for item in result
}
assert (
"Person",
"Maroš Harahus",
) in resolved
assert (
"Person",
"Jakub Maruniak",
) in resolved
assert (
"Topic",
"Spacy",
) in resolved
def test_person_is_preferred_over_author_with_same_name() -> None:
catalog = sample_catalog() + [
{
"element_id": "a1",
"kind": "Author",
"id": "maros-harahus",
"name": "Maroš Harahus",
"title": None,
"heading": None,
"start_year": None,
"year": None,
"work_type": None,
"person_kind": None,
"source_document": None,
}
]
result = graph_search.resolve_entities(
"Čomu sa venuje Maroš Harahus?",
catalog,
)
matching = [
item
for item in result
if item.get("matched_alias")
== "maros harahus"
]
assert len(matching) == 1
assert matching[0]["kind"] == "Person"
def test_single_word_topic_does_not_match_inside_another_word() -> None:
result = graph_search.resolve_entities(
"Toto je pragmatický prístup.",
sample_catalog(),
)
assert all(
item.get("id") != "rag"
for item in result
)
def test_work_title_can_be_resolved_directly() -> None:
result = graph_search.resolve_entities(
"Kto riešil Anotáciu a "
"rozpoznávanie pomenovaných entít "
"v slovenskom jazyku?",
sample_catalog(),
)
assert any(
item["kind"] == "Work"
and item["element_id"] == "w1"
for item in result
)
def test_more_complete_inflected_work_title_ranks_above_shorter_phrase() -> None:
catalog = sample_catalog() + [
{
"element_id": "w2",
"kind": "Work",
"id": "work:shorter",
"name": (
"Rozpoznávanie pomenovaných "
"entít v slovenskom jazyku"
),
"title": (
"Rozpoznávanie pomenovaných "
"entít v slovenskom jazyku"
),
"heading": "Diplomová práca 2020",
"start_year": None,
"year": 2020,
"work_type": "diploma_thesis",
"person_kind": None,
"source_document": "pages/students/other/README.md",
}
]
result = graph_search.resolve_entities(
"Kto riešil Anotáciu a "
"rozpoznávanie pomenovaných entít "
"v slovenskom jazyku?",
catalog,
)
work_results = [
item
for item in result
if item["kind"] == "Work"
]
assert work_results[0]["element_id"] == "w1"
def test_search_graph_with_session_uses_topic_path_for_two_people() -> None:
topic_path = {
"nodes": [
{
"element_id": "p1",
"kind": "Person",
"name": "Maroš Harahus",
"path": None,
},
{
"element_id": "d1",
"kind": "Document",
"name": "Maroš Harahus",
"path": (
"pages/students/2016/"
"maros_harahus/README.md"
),
"source_url": (
"https://zp.kemt.fei.tuke.sk/"
"students/2016/maros_harahus"
),
},
{
"element_id": "t1",
"kind": "Topic",
"name": "Spacy",
"id": "spacy",
},
{
"element_id": "d2",
"kind": "Document",
"name": "Jakub Maruniak",
"path": (
"pages/students/2016/"
"jakub_maruniak/README.md"
),
"source_url": (
"https://zp.kemt.fei.tuke.sk/"
"students/2016/jakub_maruniak"
),
},
{
"element_id": "p2",
"kind": "Person",
"name": "Jakub Maruniak",
"path": None,
},
],
"relationships": [
"HAS_DOCUMENT",
"HAS_TOPIC",
"HAS_TOPIC",
"HAS_DOCUMENT",
],
"hops": 4,
}
session = FakeSession(
{
graph_search.CATALOG_QUERY: (
sample_catalog()
),
graph_search.NEIGHBORHOOD_QUERY: [],
graph_search.PAIR_PATH_QUERY: [
topic_path
],
}
)
result = (
graph_search.search_graph_with_session(
session,
"Akú spoločnú tému majú "
"Maroš Harahus a Jakub Maruniak?",
)
)
assert result[
"constraints"
]["mentions_topic"] is True
assert len(
result["paths"]
) == 1
assert {
item["document_path"]
for item in result[
"documents"
]
} == {
(
"pages/students/2016/"
"maros_harahus/README.md"
),
(
"pages/students/2016/"
"jakub_maruniak/README.md"
),
}
pair_calls = [
parameters
for query, parameters
in session.calls
if query
== graph_search.PAIR_PATH_QUERY
]
assert len(pair_calls) == 1
assert pair_calls[0][
"require_topic"
] is True
def test_search_graph_with_session_uses_work_constraints() -> None:
work_fact = {
"person_element_id": "p2",
"person": "Jakub Maruniak",
"start_year": 2016,
"work_element_id": "w1",
"work_id": "work:jakub:dp2021",
"title": (
"Anotácia a rozpoznávanie "
"pomenovaných entít v "
"slovenskom jazyku."
),
"heading": "Diplomová práca 2021",
"year": 2021,
"work_type": "diploma_thesis",
"work_type_label": "Diplomová práca",
"document_path": (
"pages/students/2016/"
"jakub_maruniak/README.md"
),
"source_url": None,
"source_chunk_ids": [
(
"pages/students/2016/"
"jakub_maruniak/"
"README.md::chunk-0"
)
],
}
session = FakeSession(
{
graph_search.CATALOG_QUERY: (
sample_catalog()
),
graph_search.NEIGHBORHOOD_QUERY: [],
graph_search.WORK_FACTS_QUERY: [
work_fact
],
}
)
result = (
graph_search.search_graph_with_session(
session,
"Aký bol názov diplomovej práce "
"Jakuba Maruniaka v roku 2021?",
)
)
assert any(
fact["fact_type"]
== "person_work"
and fact["year"] == 2021
and fact["work_type"]
== "diploma_thesis"
for fact in result["facts"]
)
work_calls = [
parameters
for query, parameters
in session.calls
if query
== graph_search.WORK_FACTS_QUERY
]
assert len(work_calls) == 1
assert work_calls[0][
"years"
] == [2021]
assert work_calls[0][
"work_types"
] == [
"diploma_thesis"
]
assert work_calls[0][
"person_ids"
] == ["p2"]
def test_search_graph_with_session_uses_start_year_facts() -> None:
session = FakeSession(
{
graph_search.CATALOG_QUERY: (
sample_catalog()
),
graph_search.NEIGHBORHOOD_QUERY: [],
graph_search.PAIR_PATH_QUERY: [],
graph_search.START_YEAR_FACTS_QUERY: [
{
"person_element_id": "p1",
"person": "Maroš Harahus",
"start_year": 2016,
"document_path": (
"pages/students/2016/"
"maros_harahus/README.md"
),
"source_url": None,
},
{
"person_element_id": "p2",
"person": "Jakub Maruniak",
"start_year": 2016,
"document_path": (
"pages/students/2016/"
"jakub_maruniak/README.md"
),
"source_url": None,
},
],
}
)
result = (
graph_search.search_graph_with_session(
session,
"Mali Maroš Harahus a Jakub "
"Maruniak rovnaký rok začiatku "
"štúdia?",
)
)
values = {
fact["person"]: fact[
"start_year"
]
for fact in result["facts"]
if fact["fact_type"]
== "person_start_year"
}
assert values == {
"Maroš Harahus": 2016,
"Jakub Maruniak": 2016,
}
def test_search_graph_rejects_empty_query() -> None:
session = FakeSession({})
with pytest.raises(
ValueError,
match="prázdny",
):
graph_search.search_graph_with_session(
session,
" ",
)
def test_shared_topic_query_returns_structured_fact_and_no_generic_neighborhood() -> None:
shared_topic_fact = {
"left_person_element_id": "p1",
"left_person": "Maroš Harahus",
"left_document_path": (
"pages/students/2016/"
"maros_harahus/README.md"
),
"left_source_url": "https://example.test/maros",
"right_person_element_id": "p2",
"right_person": "Jakub Maruniak",
"right_document_path": (
"pages/students/2016/"
"jakub_maruniak/README.md"
),
"right_source_url": "https://example.test/jakub",
"topic_element_id": "t1",
"topic_id": "spacy",
"topic": "Spacy",
}
session = FakeSession(
{
graph_search.CATALOG_QUERY: sample_catalog(),
graph_search.PAIR_PATH_QUERY: [],
graph_search.SHARED_TOPIC_FACTS_QUERY: [
shared_topic_fact
],
}
)
result = graph_search.search_graph_with_session(
session,
"Akú spoločnú tému majú Maroš Harahus "
"a Jakub Maruniak?",
)
shared_facts = [
fact
for fact in result["facts"]
if fact["fact_type"] == "shared_topic"
]
assert len(shared_facts) == 1
assert shared_facts[0]["topic_id"] == "spacy"
assert {
item["document_path"]
for item in result["documents"]
} == {
"pages/students/2016/maros_harahus/README.md",
"pages/students/2016/jakub_maruniak/README.md",
}
called_queries = {
query
for query, _ in session.calls
}
assert graph_search.NEIGHBORHOOD_QUERY not in called_queries
assert graph_search.PAIR_PATH_QUERY in called_queries
assert graph_search.SHARED_TOPIC_FACTS_QUERY in called_queries
def test_work_query_uses_precise_work_paths_instead_of_generic_neighborhood() -> None:
work_path = {
"nodes": [
{
"element_id": "p2",
"kind": "Person",
"name": "Jakub Maruniak",
},
{
"element_id": "w1",
"kind": "Work",
"name": (
"Anotácia a rozpoznávanie "
"pomenovaných entít v "
"slovenskom jazyku."
),
"year": 2021,
"work_type": "diploma_thesis",
},
{
"element_id": "d2",
"kind": "Document",
"path": (
"pages/students/2016/"
"jakub_maruniak/README.md"
),
},
],
"relationships": [
"HAS_WORK",
"EVIDENCED_BY",
],
"hops": 2,
}
session = FakeSession(
{
graph_search.CATALOG_QUERY: sample_catalog(),
graph_search.WORK_PATHS_QUERY: [work_path],
graph_search.WORK_FACTS_QUERY: [],
}
)
result = graph_search.search_graph_with_session(
session,
"Aký bol názov diplomovej práce "
"Jakuba Maruniaka v roku 2021?",
)
assert result["paths"] == [work_path]
called_queries = {
query
for query, _ in session.calls
}
assert graph_search.WORK_PATHS_QUERY in called_queries
assert graph_search.NEIGHBORHOOD_QUERY not in called_queries
work_path_calls = [
parameters
for query, parameters in session.calls
if query == graph_search.WORK_PATHS_QUERY
]
assert work_path_calls[0]["years"] == [2021]
assert work_path_calls[0]["work_types"] == [
"diploma_thesis"
]
assert work_path_calls[0]["person_ids"] == ["p2"]
def test_start_year_comparison_is_structured_and_skips_irrelevant_paths() -> None:
session = FakeSession(
{
graph_search.CATALOG_QUERY: sample_catalog(),
graph_search.START_YEAR_FACTS_QUERY: [
{
"person_element_id": "p1",
"person": "Maroš Harahus",
"start_year": 2016,
"document_path": "maros.md",
"source_url": None,
},
{
"person_element_id": "p2",
"person": "Jakub Maruniak",
"start_year": 2016,
"document_path": "jakub.md",
"source_url": None,
},
],
}
)
result = graph_search.search_graph_with_session(
session,
"Mali Maroš Harahus a Jakub Maruniak "
"rovnaký rok začiatku štúdia?",
)
assert result["paths"] == []
comparison = next(
fact
for fact in result["facts"]
if fact["fact_type"]
== "start_year_comparison"
)
assert comparison["same_start_year"] is True
assert comparison["shared_start_year"] == 2016
assert comparison["start_years"] == {
"Jakub Maruniak": 2016,
"Maroš Harahus": 2016,
}
called_queries = {
query
for query, _ in session.calls
}
assert graph_search.NEIGHBORHOOD_QUERY not in called_queries
assert graph_search.PAIR_PATH_QUERY not in called_queries
def test_resolved_work_entity_triggers_exact_work_lookup() -> None:
session = FakeSession(
{
graph_search.CATALOG_QUERY: sample_catalog(),
graph_search.WORK_PATHS_QUERY: [],
graph_search.WORK_FACTS_QUERY: [],
}
)
graph_search.search_graph_with_session(
session,
"Kto riešil Anotáciu a rozpoznávanie "
"pomenovaných entít v slovenskom jazyku?",
)
calls = [
parameters
for query, parameters in session.calls
if query == graph_search.WORK_FACTS_QUERY
]
assert len(calls) == 1
assert calls[0]["work_ids"] == ["w1"]