172 lines
3.1 KiB
Python
172 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
from tools.graphrag.graphrag import (
|
|
build_graph_payload,
|
|
extract_works,
|
|
)
|
|
|
|
|
|
DARIUS_PATH = (
|
|
"pages/students/2016/"
|
|
"darius_lindvai/README.md"
|
|
)
|
|
|
|
|
|
def make_darius_document() -> dict:
|
|
return {
|
|
"path": DARIUS_PATH,
|
|
"title": "Dárius Lindvai",
|
|
"categories": [
|
|
"dp2021",
|
|
"bp2019",
|
|
],
|
|
"tags": [
|
|
"nn",
|
|
"interpunction",
|
|
"nlp",
|
|
],
|
|
"author": "Daniel Hladek",
|
|
"published": True,
|
|
}
|
|
|
|
|
|
def make_root_chunk() -> dict:
|
|
text = (
|
|
"Dokument: Dárius Lindvai\n\n"
|
|
"Rok začiatku štúdia: 2016\n\n"
|
|
"Repozitár so zdrojovými kódmi\n\n"
|
|
"Názov: Obnovenie interpunkcie pomocou "
|
|
"hlbokých neurónových sietí\n\n"
|
|
"1. Vypracujte prehľad metód.\n"
|
|
"2. Vyberte vhodnú metódu."
|
|
)
|
|
|
|
return {
|
|
"chunk_id": (
|
|
f"{DARIUS_PATH}::chunk-0"
|
|
),
|
|
"document_path": DARIUS_PATH,
|
|
"title": "Dárius Lindvai",
|
|
"categories": [
|
|
"dp2021",
|
|
"bp2019",
|
|
],
|
|
"tags": [
|
|
"nn",
|
|
"interpunction",
|
|
"nlp",
|
|
],
|
|
"author": "Daniel Hladek",
|
|
"published": True,
|
|
"chunk_index": 0,
|
|
"heading_paths": [
|
|
[
|
|
"Dárius Lindvai",
|
|
]
|
|
],
|
|
"text": text,
|
|
"text_length": len(text),
|
|
"token_count": 1,
|
|
"content_hash": "darius-root",
|
|
}
|
|
|
|
|
|
def test_root_title_creates_missing_latest_thesis() -> None:
|
|
works = extract_works(
|
|
make_darius_document(),
|
|
[
|
|
make_root_chunk(),
|
|
],
|
|
)
|
|
|
|
diploma = [
|
|
work
|
|
for work in works
|
|
if (
|
|
work["work_type"]
|
|
== "diploma_thesis"
|
|
)
|
|
]
|
|
|
|
bachelor = [
|
|
work
|
|
for work in works
|
|
if (
|
|
work["work_type"]
|
|
== "bachelor_thesis"
|
|
)
|
|
]
|
|
|
|
assert len(
|
|
diploma
|
|
) == 1
|
|
|
|
assert bachelor == []
|
|
|
|
thesis = diploma[0]
|
|
|
|
assert thesis[
|
|
"year"
|
|
] == 2021
|
|
|
|
assert thesis[
|
|
"heading"
|
|
] == "Diplomová práca 2021"
|
|
|
|
assert thesis[
|
|
"title"
|
|
] == (
|
|
"Obnovenie interpunkcie pomocou "
|
|
"hlbokých neurónových sietí"
|
|
)
|
|
|
|
assert thesis[
|
|
"source_chunk_ids"
|
|
] == [
|
|
f"{DARIUS_PATH}::chunk-0"
|
|
]
|
|
|
|
|
|
def test_graph_payload_contains_darius_thesis() -> None:
|
|
payload = build_graph_payload(
|
|
[
|
|
make_darius_document(),
|
|
],
|
|
[
|
|
make_root_chunk(),
|
|
],
|
|
)
|
|
|
|
assert len(
|
|
payload[
|
|
"theses"
|
|
]
|
|
) == 1
|
|
|
|
thesis = payload[
|
|
"theses"
|
|
][0]
|
|
|
|
assert thesis[
|
|
"thesis_type"
|
|
] == "diploma_thesis"
|
|
|
|
assert thesis[
|
|
"year"
|
|
] == 2021
|
|
|
|
assert thesis[
|
|
"title"
|
|
] == (
|
|
"Obnovenie interpunkcie pomocou "
|
|
"hlbokých neurónových sietí"
|
|
)
|
|
|
|
assert thesis[
|
|
"source_document"
|
|
] == DARIUS_PATH
|
|
|
|
assert payload[
|
|
"works"
|
|
] == []
|