Finalizing GraphRAG model and search capabilities

tools/graphrag/build_graphrag.py

@@ -44,8 +44,8 @@ CONSTRAINTS = (
         "id",
     ),
     (
-        "zpwiki_author_id_unique",
+        "zpwiki_document_author_id_unique",
-        "Author",
+        "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 {
+        MATCH (b:Keyword {
-            id: row.topic_id
+            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(
-        """
+        f"""
-        MATCH (w:Work)
+        MATCH (n:{label})
-        WHERE w.graph_scope = $scope
+        WHERE n.graph_scope = $scope
         RETURN
-            count(w) AS total,
+            count(n) AS total,
-            count(w.title) AS with_title
+            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 = (
+    thesis_total, thesis_with_title = (
-        work_quality_counts(
+        titled_entity_counts(
-            session
+            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
             )

tools/graphrag/tests/test_graphrag.py

@@ -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": [
+        "tags": (
+            tags
+            if tags is not None
+            else [
                 "spacy",
                 "ner",
-        ],
+            ]
-        "author": "Daniel Hladek",
+        ),
+        "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(
+    diploma_thesis = next(
-        work
+        thesis
-        for work in payload["works"]
+        for thesis
-        if work["work_type"]
+        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
+            ),
+        },
+    ]