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

794 lines
18 KiB
Python

from __future__ import annotations
from tools.graphrag.graphrag import (
build_graph_payload,
extract_explicit_proposal_title,
extract_works,
person_kind_from_profile_path,
person_profile_path,
)
JAKUB_PATH = "pages/students/2016/" "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,
"title": title,
"categories": (
categories
if categories is not None
else [
"dp2021",
"bp2019",
]
),
"tags": (
tags
if tags is not None
else [
"spacy",
"ner",
]
),
"author": author,
"published": True,
}
def make_chunk(
index: int,
headings: list[str],
text: str,
*,
path: str = JAKUB_PATH,
title: str = "Jakub Maruniak",
) -> dict:
return {
"chunk_id": (f"{path}::chunk-{index}"),
"document_path": path,
"title": title,
"categories": [
"dp2021",
"bp2019",
],
"tags": [
"spacy",
"ner",
],
"author": "Daniel Hladek",
"published": True,
"chunk_index": index,
"heading_paths": [headings],
"text": text,
"text_length": len(text),
"token_count": 1,
"content_hash": f"hash-{index}",
}
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(
0,
["Jakub Maruniak"],
(
"Dokument: Jakub Maruniak\n\n"
"*Návrh na názov DP*:\n\n"
"Anotácia a rozpoznávanie "
"pomenovaných entít v "
"slovenskom jazyku."
),
)
]
title, source_chunk_ids = extract_explicit_proposal_title(
"diploma_thesis",
chunks,
)
assert title == ("Anotácia a rozpoznávanie " "pomenovaných entít v " "slovenskom jazyku.")
assert source_chunk_ids == [f"{JAKUB_PATH}::chunk-0"]
def test_extract_works_uses_dp_proposal_fallback() -> None:
document = make_document()
chunks = [
make_chunk(
0,
["Jakub Maruniak"],
(
"Dokument: Jakub Maruniak\n\n"
"*Rok začiatku štúdia*: 2016\n\n"
"*Návrh na názov DP*:\n\n"
"Anotácia a rozpoznávanie "
"pomenovaných entít v "
"slovenskom jazyku."
),
),
make_chunk(
1,
[
"Jakub Maruniak",
"Diplomová práca 2021",
],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Diplomová práca 2021\n\n"
"Stretnutie 12.3.\n\n"
"Stav:\n\n"
"- Anotovanie dát."
),
),
]
works = extract_works(
document,
chunks,
)
diploma_works = [work for work in works if work["work_type"] == "diploma_thesis"]
assert len(diploma_works) == 1
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"] == [
f"{JAKUB_PATH}::chunk-0",
f"{JAKUB_PATH}::chunk-1",
]
def test_existing_work_title_has_priority() -> None:
document = make_document()
chunks = [
make_chunk(
0,
["Jakub Maruniak"],
("Dokument: Jakub Maruniak\n\n" "Návrh na názov DP: " "Starší pracovný názov"),
),
make_chunk(
1,
[
"Jakub Maruniak",
"Diplomová práca 2021",
],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Diplomová práca 2021\n\n"
"Názov: Finálny názov práce"
),
),
]
works = extract_works(
document,
chunks,
)
work = next(work for work in works if work["work_type"] == "diploma_thesis")
assert work["title"] == ("Finálny názov práce")
assert work["source_chunk_ids"] == [f"{JAKUB_PATH}::chunk-1"]
def test_ambiguous_dp_proposals_are_not_used() -> None:
document = make_document()
chunks = [
make_chunk(
0,
["Jakub Maruniak"],
("Návrh na názov DP: " "Prvý návrh"),
),
make_chunk(
1,
["Jakub Maruniak"],
("Návrh na názov DP: " "Druhý návrh"),
),
make_chunk(
2,
[
"Jakub Maruniak",
"Diplomová práca 2021",
],
("Dokument: Jakub Maruniak\n" "Sekcia: Diplomová práca 2021\n\n" "Stretnutie 12.3."),
),
]
works = extract_works(
document,
chunks,
)
work = next(work for work in works if work["work_type"] == "diploma_thesis")
assert work["title"] is None
def test_proposal_is_not_used_for_diploma_project() -> None:
document = make_document()
chunks = [
make_chunk(
0,
["Jakub Maruniak"],
("Návrh na názov DP: " "Názov diplomovej práce"),
),
make_chunk(
1,
[
"Jakub Maruniak",
"Diplomový projekt 2020",
],
(
"Dokument: Jakub Maruniak\n"
"Sekcia: Diplomový projekt 2020\n\n"
"Stretnutie 9.3.2020\n\n"
"Úlohy:\n\n"
"- Pripraviť dáta."
),
),
]
works = extract_works(
document,
chunks,
)
work = next(work for work in works if work["work_type"] == "diploma_project")
assert work["title"] is None
def test_proposal_is_not_shared_between_multiple_dp_works() -> None:
document = make_document(
categories=[
"dp2021",
"dp2022",
]
)
chunks = [
make_chunk(
0,
["Jakub Maruniak"],
("Návrh na názov DP: " "Jeden návrh"),
),
make_chunk(
1,
[
"Jakub Maruniak",
"Diplomová práca 2021",
],
("Dokument: Jakub Maruniak\n" "Sekcia: Diplomová práca 2021\n\n" "Stretnutie 1.1.2021"),
),
make_chunk(
2,
[
"Jakub Maruniak",
"Diplomová práca 2022",
],
("Dokument: Jakub Maruniak\n" "Sekcia: Diplomová práca 2022\n\n" "Stretnutie 1.1.2022"),
),
]
works = extract_works(
document,
chunks,
)
diploma_works = [work for work in works if work["work_type"] == "diploma_thesis"]
assert len(diploma_works) == 2
assert all(work["title"] is None for work in diploma_works)
def test_graph_payload_contains_thesis_with_fallback_title() -> None:
document = make_document()
chunks = [
make_chunk(
0,
["Jakub Maruniak"],
(
"Dokument: Jakub Maruniak\n\n"
"Rok začiatku štúdia: 2016\n\n"
"Návrh na názov DP:\n\n"
"Anotácia a rozpoznávanie "
"pomenovaných entít v "
"slovenskom jazyku."
),
),
make_chunk(
1,
[
"Jakub Maruniak",
"Diplomová práca 2021",
],
("Dokument: Jakub Maruniak\n" "Sekcia: Diplomová práca 2021\n\n" "Stretnutie 12.3."),
),
]
payload = build_graph_payload(
[document],
chunks,
)
diploma_thesis = next(
thesis for thesis in payload["theses"] if thesis["thesis_type"] == "diploma_thesis"
)
assert diploma_thesis["title"] == (
"Anotácia a rozpoznávanie " "pomenovaných entít v " "slovenskom jazyku."
)
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"),
"source_document": path,
"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),
},
]