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

Finalizing GraphRAG model and search capabilities

Browse Source
This commit is contained in:
Ján Pták 2026-10-17 22:14:44 +02:00
parent 921642a16a
commit 9f3483d290
5 changed files with 4291 additions and 1033 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

160
tools/graphrag/build_graphrag.py
View File

@ -44,8 +44,8 @@ CONSTRAINTS = (
"id",
),
(
"zpwiki_author_id_unique",
"Author",
"zpwiki_document_author_id_unique",
"DocumentAuthor",
"id",
),
(
@ -58,6 +58,16 @@ CONSTRAINTS = (
"Topic",
"id",
),
(
"zpwiki_keyword_id_unique",
"Keyword",
"id",
),
(
"zpwiki_thesis_id_unique",
"Thesis",
"id",
),
(
"zpwiki_work_id_unique",
"Work",
@ -66,15 +76,20 @@ CONSTRAINTS = (
)
LEGACY_CONSTRAINTS = (
"zpwiki_author_id_unique",
)
NODE_QUERIES = {
"people": """
UNWIND $rows AS row
MERGE (n:Person {id: row.id})
SET n += row
""",
"authors": """
"document_authors": """
UNWIND $rows AS row
MERGE (n:Author {id: row.id})
MERGE (n:DocumentAuthor {id: row.id})
SET n += row
""",
"documents": """
@ -92,6 +107,16 @@ NODE_QUERIES = {
MERGE (n:Topic {id: row.id})
SET n += row
""",
"keywords": """
UNWIND $rows AS row
MERGE (n:Keyword {id: row.id})
SET n += row
""",
"theses": """
UNWIND $rows AS row
MERGE (n:Thesis {id: row.id})
SET n += row
""",
"works": """
UNWIND $rows AS row
MERGE (n:Work {id: row.id})
@ -103,19 +128,23 @@ NODE_QUERIES = {
RELATIONSHIP_QUERIES = {
"person_documents": """
UNWIND $rows AS row
MATCH (a:Person {id: row.person_id})
MATCH (a:Person {
id: row.person_id
})
MATCH (b:Document {
path: row.document_path
})
MERGE (a)-[:HAS_DOCUMENT]->(b)
""",
"author_documents": """
"document_author_documents": """
UNWIND $rows AS row
MATCH (a:Author {id: row.author_id})
MATCH (a:DocumentAuthor {
id: row.document_author_id
})
MATCH (b:Document {
path: row.document_path
})
MERGE (a)-[:AUTHORED]->(b)
MERGE (a)-[:AUTHORED_DOCUMENT]->(b)
""",
"document_categories": """
UNWIND $rows AS row
@ -127,15 +156,15 @@ RELATIONSHIP_QUERIES = {
})
MERGE (a)-[:IN_CATEGORY]->(b)
""",
"document_topics": """
"document_keywords": """
UNWIND $rows AS row
MATCH (a:Document {
path: row.document_path
})
MATCH (b:Topic {
id: row.topic_id
MATCH (b:Keyword {
id: row.keyword_id
})
MERGE (a)-[:HAS_TOPIC]->(b)
MERGE (a)-[:HAS_KEYWORD]->(b)
""",
"document_describes_topics": """
UNWIND $rows AS row
@ -147,6 +176,16 @@ RELATIONSHIP_QUERIES = {
})
MERGE (a)-[:DESCRIBES]->(b)
""",
"person_theses": """
UNWIND $rows AS row
MATCH (a:Person {
id: row.person_id
})
MATCH (b:Thesis {
id: row.thesis_id
})
MERGE (a)-[:AUTHORED]->(b)
""",
"person_works": """
UNWIND $rows AS row
MATCH (a:Person {
@ -157,6 +196,16 @@ RELATIONSHIP_QUERIES = {
})
MERGE (a)-[:HAS_WORK]->(b)
""",
"thesis_documents": """
UNWIND $rows AS row
MATCH (a:Thesis {
id: row.thesis_id
})
MATCH (b:Document {
path: row.document_path
})
MERGE (a)-[:EVIDENCED_BY]->(b)
""",
"work_documents": """
UNWIND $rows AS row
MATCH (a:Work {
@ -214,6 +263,18 @@ def create_constraints(
).consume()
def drop_legacy_constraints(
session,
) -> None:
for name in LEGACY_CONSTRAINTS:
session.run(
(
f"DROP CONSTRAINT "
f"{name} IF EXISTS"
)
).consume()
def clear_managed_graph(
session,
) -> None:
@ -252,7 +313,9 @@ def graph_counts(
MATCH (n)
WHERE n.graph_scope = $scope
UNWIND labels(n) AS label
RETURN label, count(*) AS count
RETURN
label,
count(*) AS count
ORDER BY label
""",
scope=GRAPH_SCOPE,
@ -295,16 +358,25 @@ def graph_counts(
return nodes, relationships
def work_quality_counts(
def titled_entity_counts(
session,
label: str,
) -> tuple[int, int]:
if label not in {
"Thesis",
"Work",
}:
raise ValueError(
f"Unsupported label: {label}"
)
record = session.run(
"""
MATCH (w:Work)
WHERE w.graph_scope = $scope
f"""
MATCH (n:{label})
WHERE n.graph_scope = $scope
RETURN
count(w) AS total,
count(w.title) AS with_title
count(n) AS total,
count(n.title) AS with_title
""",
scope=GRAPH_SCOPE,
).single()
@ -330,10 +402,12 @@ def print_payload_summary(
for key in (
"people",
"authors",
"document_authors",
"documents",
"categories",
"topics",
"keywords",
"theses",
"works",
):
print(
@ -341,6 +415,25 @@ def print_payload_summary(
f"{len(payload[key]):>6}"
)
print()
print("Relationships")
for key in (
"person_documents",
"document_author_documents",
"document_categories",
"document_keywords",
"document_describes_topics",
"person_theses",
"person_works",
"thesis_documents",
"work_documents",
):
print(
f"{key:<30}"
f"{len(payload[key]):>6}"
)
print("=" * 60)
@ -374,16 +467,31 @@ def print_database_summary(
f"{count:>6}"
)
total, with_title = (
work_quality_counts(
session
thesis_total, thesis_with_title = (
titled_entity_counts(
session,
"Thesis",
)
)
work_total, work_with_title = (
titled_entity_counts(
session,
"Work",
)
)
print()
print(
"Theses with extracted title: "
f"{thesis_with_title}/"
f"{thesis_total}"
)
print(
"Works with extracted title: "
f"{with_title}/{total}"
f"{work_with_title}/"
f"{work_total}"
)
print("=" * 60)
@ -427,6 +535,10 @@ def main() -> None:
with driver.session(
database=settings.database
) as session:
drop_legacy_constraints(
session
)
create_constraints(
session
)

2223
tools/graphrag/graph_search.py
View File

File diff suppressed because it is too large Load Diff

828
tools/graphrag/graphrag.py
View File

File diff suppressed because it is too large Load Diff

1456
tools/graphrag/tests/test_graph_search.py
View File

File diff suppressed because it is too large Load Diff

657
tools/graphrag/tests/test_graphrag.py
View File

@ -4,6 +4,8 @@ from tools.graphrag.graphrag import (
build_graph_payload,
extract_explicit_proposal_title,
extract_works,
person_kind_from_profile_path,
person_profile_path,
)
@ -12,12 +14,19 @@ JAKUB_PATH = (
"jakub_maruniak/README.md"
)
JAKUB_NESTED_PATH = (
"pages/students/2016/"
"jakub_maruniak/dp2021/README.md"
)
def make_document(
*,
path: str = JAKUB_PATH,
title: str = "Jakub Maruniak",
categories: list[str] | None = None,
tags: list[str] | None = None,
author: str = "Daniel Hladek",
) -> dict:
return {
"path": path,
@ -25,13 +34,20 @@ def make_document(
"categories": (
categories
if categories is not None
else ["dp2021", "bp2019"]
else [
"dp2021",
"bp2019",
]
),
"tags": [
"spacy",
"ner",
],
"author": "Daniel Hladek",
"tags": (
tags
if tags is not None
else [
"spacy",
"ner",
]
),
"author": author,
"published": True,
}
@ -42,13 +58,14 @@ def make_chunk(
text: str,
*,
path: str = JAKUB_PATH,
title: str = "Jakub Maruniak",
) -> dict:
return {
"chunk_id": (
f"{path}::chunk-{index}"
),
"document_path": path,
"title": "Jakub Maruniak",
"title": title,
"categories": [
"dp2021",
"bp2019",
@ -70,6 +87,70 @@ def make_chunk(
}
def test_person_profile_path_for_student_profile() -> None:
assert (
person_profile_path(
JAKUB_PATH
)
== JAKUB_PATH
)
def test_person_profile_path_for_nested_student_document() -> None:
assert (
person_profile_path(
JAKUB_NESTED_PATH
)
== JAKUB_PATH
)
def test_person_profile_path_for_non_person_document() -> None:
assert (
person_profile_path(
"pages/topics/nlp/README.md"
)
is None
)
def test_person_kind_accepts_only_profile_root() -> None:
assert (
person_kind_from_profile_path(
JAKUB_PATH
)
== "student"
)
assert (
person_kind_from_profile_path(
JAKUB_NESTED_PATH
)
is None
)
def test_person_kind_for_intern_profile() -> None:
path = (
"pages/interns/"
"john_doe/README.md"
)
assert (
person_profile_path(
path
)
== path
)
assert (
person_kind_from_profile_path(
path
)
== "intern"
)
def test_explicit_dp_proposal_is_found() -> None:
chunks = [
make_chunk(
@ -152,13 +233,16 @@ def test_extract_works_uses_dp_proposal_fallback() -> None:
work = diploma_works[0]
assert work["year"] == 2021
assert work["title"] == (
"Anotácia a rozpoznávanie "
"pomenovaných entít v "
"slovenskom jazyku."
)
assert work["source_chunk_ids"] == [
assert work[
"source_chunk_ids"
] == [
f"{JAKUB_PATH}::chunk-0",
f"{JAKUB_PATH}::chunk-1",
]
@ -207,7 +291,9 @@ def test_existing_work_title_has_priority() -> None:
"Finálny názov práce"
)
assert work["source_chunk_ids"] == [
assert work[
"source_chunk_ids"
] == [
f"{JAKUB_PATH}::chunk-1"
]
@ -359,14 +445,17 @@ def test_proposal_is_not_shared_between_multiple_dp_works() -> None:
== "diploma_thesis"
]
assert len(diploma_works) == 2
assert len(
diploma_works
) == 2
assert all(
work["title"] is None
for work in diploma_works
)
def test_graph_payload_contains_fallback_title() -> None:
def test_graph_payload_contains_thesis_with_fallback_title() -> None:
document = make_document()
chunks = [
@ -401,22 +490,556 @@ def test_graph_payload_contains_fallback_title() -> None:
chunks,
)
diploma_work = next(
work
for work in payload["works"]
if work["work_type"]
diploma_thesis = next(
thesis
for thesis
in payload["theses"]
if thesis["thesis_type"]
== "diploma_thesis"
)
assert diploma_work["title"] == (
assert diploma_thesis[
"title"
] == (
"Anotácia a rozpoznávanie "
"pomenovaných entít v "
"slovenskom jazyku."
)
assert diploma_work[
assert diploma_thesis[
"source_chunk_ids"
] == [
f"{JAKUB_PATH}::chunk-0",
f"{JAKUB_PATH}::chunk-1",
]
assert payload[
"works"
] == []
def test_tags_become_keywords_not_topics() -> None:
document = make_document(
tags=[
"spacy",
"ner",
"rag",
]
)
payload = build_graph_payload(
[document],
[],
)
keyword_ids = {
keyword["id"]
for keyword
in payload["keywords"]
}
assert keyword_ids == {
"spacy",
"ner",
"rag",
}
assert payload[
"topics"
] == []
assert len(
payload[
"document_keywords"
]
) == 3
def test_topic_is_created_from_topic_document() -> None:
path = (
"pages/topics/nlp/README.md"
)
document = make_document(
path=path,
title=(
"Nástroje a zdroje pre "
"spracovanie prirodzeného jazyka"
),
categories=[
"info",
],
tags=[],
)
payload = build_graph_payload(
[document],
[],
)
assert payload["topics"] == [
{
"id": "nlp",
"name": (
"Nástroje a zdroje pre "
"spracovanie prirodzeného jazyka"
),
"graph_scope": "zpwiki",
}
]
assert payload[
"document_describes_topics"
] == [
{
"document_path": path,
"topic_id": "nlp",
}
]
def test_metadata_author_becomes_document_author() -> None:
document = make_document(
author="Daniel Hladek"
)
payload = build_graph_payload(
[document],
[],
)
assert payload[
"document_authors"
] == [
{
"id": "daniel-hladek",
"name": "Daniel Hladek",
"graph_scope": "zpwiki",
}
]
assert payload[
"document_author_documents"
] == [
{
"document_author_id": (
"daniel-hladek"
),
"document_path": (
JAKUB_PATH
),
}
]
def test_nested_document_does_not_create_second_person() -> None:
profile = make_document()
nested = make_document(
path=JAKUB_NESTED_PATH,
title="Diplomová práca 2021",
categories=[
"dp2021",
],
tags=[
"spacy",
],
)
payload = build_graph_payload(
[
profile,
nested,
],
[],
)
assert len(
payload["people"]
) == 1
person = payload[
"people"
][0]
assert person["name"] == (
"Jakub Maruniak"
)
assert person[
"source_document"
] == JAKUB_PATH
assert payload[
"person_documents"
] == [
{
"person_id": (
f"person:{JAKUB_PATH}"
),
"document_path": (
JAKUB_PATH
),
},
{
"person_id": (
f"person:{JAKUB_PATH}"
),
"document_path": (
JAKUB_NESTED_PATH
),
},
]
def test_thesis_and_other_work_are_separated() -> None:
document = make_document(
categories=[
"dp2021",
"tp2020",
]
)
chunks = [
make_chunk(
0,
[
"Jakub Maruniak",
"Diplomová práca 2021",
],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Diplomová práca 2021\n\n"
"Názov: Diplomová práca"
),
),
make_chunk(
1,
[
"Jakub Maruniak",
"Tímový projekt 2020",
],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Tímový projekt 2020\n\n"
"Názov: Tímový projekt"
),
),
]
payload = build_graph_payload(
[document],
chunks,
)
assert len(
payload["theses"]
) == 1
assert len(
payload["works"]
) == 1
thesis = payload[
"theses"
][0]
work = payload[
"works"
][0]
assert thesis[
"thesis_type"
] == "diploma_thesis"
assert work[
"work_type"
] == "team_project"
assert payload[
"person_theses"
] == [
{
"person_id": (
f"person:{JAKUB_PATH}"
),
"thesis_id": thesis["id"],
}
]
assert payload[
"person_works"
] == [
{
"person_id": (
f"person:{JAKUB_PATH}"
),
"work_id": work["id"],
}
]
def test_bulleted_title_label_is_used() -> None:
document = make_document(
categories=["bp2020"]
)
chunks = [
make_chunk(
0,
["Bakalárska práca 2020"],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Bakalárska práca 2020\n\n"
"- Názov: Jednoduchý chatbot v jazyku Python,\n\n"
"Zadanie:\n\n"
"1. Prvá úloha."
),
)
]
work = extract_works(
document,
chunks,
)[0]
assert work["title"] == (
"Jednoduchý chatbot v jazyku Python"
)
def test_preliminary_title_is_used() -> None:
document = make_document(
categories=["bp2025"]
)
chunks = [
make_chunk(
0,
["Bakalárska práca 2025"],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Bakalárska práca 2025\n\n"
"Vedúci: doc.\n\n"
"Predbežný názov:\n\n"
"Testovanie hybridného klaudu"
),
)
]
work = extract_works(
document,
chunks,
)[0]
assert work["title"] == (
"Testovanie hybridného klaudu"
)
def test_proposed_topic_is_used() -> None:
document = make_document(
categories=["dp2027"]
)
chunks = [
make_chunk(
0,
["Diplomová práca 2027"],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Diplomová práca 2027\n\n"
"Návrh na tému:\n\n"
"Prepis reči pre zdravotný záznam\n\n"
"Ciele:\n\n"
"- Vytvoriť systém."
),
)
]
work = extract_works(
document,
chunks,
)[0]
assert work["title"] == (
"Prepis reči pre zdravotný záznam"
)
def test_leading_title_skips_status_lines() -> None:
document = make_document(
categories=["bp2020"]
)
chunks = [
make_chunk(
0,
["Bakalárska práca 2020"],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Bakalárska práca 2020\n\n"
"Práca úspešne obhájená.\n\n"
"Odovzdaný klaster mal zlé káble.\n\n"
"Pokusný Kubernetes klaster\n\n"
"1. Prvá úloha."
),
)
]
work = extract_works(
document,
chunks,
)[0]
assert work["title"] == (
"Pokusný Kubernetes klaster"
)
def test_goal_is_not_promoted_to_title() -> None:
document = make_document(
categories=["dp2025"]
)
chunks = [
make_chunk(
0,
["Diplomová práca 2025"],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Diplomová práca 2025\n\n"
"Cieľ:\n\n"
"Overiť možnosti generovania testov\n\n"
"Zadanie:\n\n"
"1. Prvá úloha."
),
)
]
work = extract_works(
document,
chunks,
)[0]
assert work["title"] is None
def test_deepest_work_heading_owns_nested_chunk() -> None:
document = make_document(
categories=[
"dp2027",
"bp2025",
]
)
chunks = [
make_chunk(
0,
["Diplomová práca 2027"],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Diplomová práca 2027\n\n"
"Téma:\n\n"
"Diplomová téma"
),
),
make_chunk(
1,
[
"Diplomová práca 2027",
"Bakalárska práca 2025",
],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Diplomová práca 2027 > "
"Bakalárska práca 2025\n\n"
"Téma:\n\n"
"Bakalárska téma"
),
),
]
works = extract_works(
document,
chunks,
)
diploma = next(
work
for work in works
if work["work_type"]
== "diploma_thesis"
)
bachelor = next(
work
for work in works
if work["work_type"]
== "bachelor_thesis"
)
assert diploma["title"] == (
"Diplomová téma"
)
assert bachelor["title"] == (
"Bakalárska téma"
)
def test_nested_document_does_not_create_academic_work() -> None:
profile = make_document(
categories=[]
)
nested = make_document(
path=JAKUB_NESTED_PATH,
title="Diplomový projekt",
categories=[],
tags=[],
)
nested_chunk = make_chunk(
0,
["Diplomový projekt 2020"],
(
"Dokument: Diplomový projekt\n"
"Sekcia: Diplomový projekt 2020\n\n"
"Názov: Falošná samostatná práca"
),
path=JAKUB_NESTED_PATH,
title="Diplomový projekt",
)
payload = build_graph_payload(
[
profile,
nested,
],
[nested_chunk],
)
assert payload["works"] == []
assert payload["theses"] == []
assert payload[
"person_documents"
] == [
{
"person_id": (
f"person:{JAKUB_PATH}"
),
"document_path": (
JAKUB_PATH
),
},
{
"person_id": (
f"person:{JAKUB_PATH}"
),
"document_path": (
JAKUB_NESTED_PATH
),
},
]