from __future__ import annotations from typing import Any import pytest import tools.graphrag.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, ) ) # Existing unit fixtures omit the separate stored-source lookup. Model # that lookup as present documents; publication enforcement itself is # exercised by real Neo4j fixtures with explicit publication flags. if query == graph_search.DOCUMENT_SOURCES_QUERY and query not in self.responses: urls = {} for rows in self.responses.values(): for row in rows: for prefix in ["", "left_", "right_"]: if row.get(prefix + "document_path"): urls[row[prefix + "document_path"]] = row.get(prefix + "source_url") for node in row.get("nodes", []): if node.get("path"): urls[node["path"]] = node.get("source_url") return [ {"document_path": path, "source_url": urls.get(path)} for path in parameters.get("paths", []) ] 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, "thesis_type": 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, "thesis_type": None, "work_type": None, "person_kind": "student", "source_document": ("pages/students/2016/" "jakub_maruniak/README.md"), }, { "element_id": "k1", "kind": "Keyword", "id": "spacy", "name": "Spacy", "title": None, "heading": None, "start_year": None, "year": None, "thesis_type": None, "work_type": None, "person_kind": None, "source_document": None, }, { "element_id": "k2", "kind": "Keyword", "id": "rag", "name": "RAG", "title": None, "heading": None, "start_year": None, "year": None, "thesis_type": None, "work_type": None, "person_kind": None, "source_document": None, }, { "element_id": "t1", "kind": "Topic", "id": "nlp", "name": ("Spracovanie prirodzeného " "jazyka"), "title": None, "heading": None, "start_year": None, "year": None, "thesis_type": None, "work_type": None, "person_kind": None, "source_document": None, }, { "element_id": "c1", "kind": "Category", "id": None, "name": "dp2021", "title": None, "heading": None, "start_year": None, "year": 2021, "thesis_type": None, "work_type": None, "person_kind": None, "source_document": None, }, { "element_id": "th1", "kind": "Thesis", "id": "thesis: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, "thesis_type": ("diploma_thesis"), "work_type": None, "person_kind": None, "source_document": ("pages/students/2016/" "jakub_maruniak/README.md"), }, { "element_id": "th2", "kind": "Thesis", "id": "thesis:artem:bp2022", "name": ("Rozpoznávanie pomenovaných " "entít v slovenskom jazyku"), "title": ("Rozpoznávanie pomenovaných " "entít v slovenskom jazyku"), "heading": ("Bakalárska práca 2022"), "start_year": None, "year": 2022, "thesis_type": ("bachelor_thesis"), "work_type": None, "person_kind": None, "source_document": ("pages/students/2019/" "artem_yatsenko/README.md"), }, { "element_id": "w1", "kind": "Work", "id": "work:jakub:tp2019", "name": ("Tímový projekt 2019"), "title": None, "heading": ("Tímový projekt 2019"), "start_year": None, "year": 2019, "thesis_type": None, "work_type": ("team_project"), "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["thesis_types"] == ["diploma_thesis"] assert result["work_types"] == [] assert result["mentions_title"] is True assert result["target_kinds"] == ["Thesis"] def test_extract_query_constraints_for_diploma_project() -> None: result = graph_search.extract_query_constraints( "Aký diplomový projekt mal " "Jakub Maruniak v roku 2020?" ) assert result["thesis_types"] == [] assert result["work_types"] == ["diploma_project"] assert result["target_kinds"] == ["Work"] def test_extract_query_constraints_for_shared_keywords() -> None: result = graph_search.extract_query_constraints( "Aké spoločné kľúčové slová " "majú Maroš Harahus a " "Jakub Maruniak?" ) assert result["mentions_keyword"] is True assert result["mentions_shared"] is True assert result["target_kinds"] == ["Keyword"] def test_extract_query_constraints_for_shared_topic() -> None: result = graph_search.extract_query_constraints( "Akú spoločnú tému majú " "Maroš Harahus a " "Jakub Maruniak?" ) assert result["mentions_topic"] is True assert result["mentions_shared"] is True assert result["target_kinds"] == [ "Keyword", "Topic", ] 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_who_query_targets_person() -> None: result = graph_search.extract_query_constraints("Kto súvisí s RAG?") assert result["mentions_person_answer"] is True assert result["target_kinds"] == ["Person"] def test_who_category_query_targets_person() -> None: result = graph_search.extract_query_constraints("Kto súvisí s kategóriou dp2021?") assert result["mentions_category"] is True assert result["target_kinds"] == ["Person"] def test_resolve_entities_finds_people_keyword_and_topic() -> None: result = graph_search.resolve_entities( ("Spacy, Spracovanie " "prirodzeného jazyka, " "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 ( "Keyword", "Spacy", ) in resolved assert ( "Topic", "Spracovanie prirodzeného jazyka", ) in resolved def test_person_is_preferred_over_document_author_with_same_name() -> None: catalog = sample_catalog() + [ { "element_id": "a1", "kind": ("DocumentAuthor"), "id": ("maros-harahus"), "name": ("Maroš Harahus"), "title": None, "heading": None, "start_year": None, "year": None, "thesis_type": 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_keyword_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_thesis_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"] == "Thesis" and item["element_id"] == "th1" for item in result) def test_longer_exact_thesis_suppresses_nested_shorter_title() -> None: result = graph_search.resolve_entities( ("Kto riešil Anotáciu a " "rozpoznávanie pomenovaných " "entít v slovenskom jazyku?"), sample_catalog(), ) thesis_ids = [item["element_id"] for item in result if item["kind"] == "Thesis"] assert thesis_ids == ["th1"] def test_search_graph_with_session_uses_thesis_constraints() -> None: thesis_fact = { "person_element_id": "p2", "person": "Jakub Maruniak", "start_year": 2016, "thesis_element_id": "th1", "thesis_id": ("thesis:jakub:dp2021"), "title": ("Anotácia a rozpoznávanie " "pomenovaných entít v " "slovenskom jazyku."), "heading": ("Diplomová práca 2021"), "year": 2021, "thesis_type": ("diploma_thesis"), "thesis_type_label": ("Diplomová práca"), "document_path": ("pages/students/2016/" "jakub_maruniak/README.md"), "source_url": ("https://zp.kemt.fei.tuke.sk/" "students/2016/jakub_maruniak"), "source_chunk_ids": [], } session = FakeSession( { graph_search.CATALOG_QUERY: (sample_catalog()), graph_search.THESIS_PATHS_QUERY: [], graph_search.THESIS_FACTS_QUERY: [thesis_fact], } ) result = graph_search.search_graph_with_session( session, ("Aký bol názov diplomovej " "práce Jakuba Maruniaka " "v roku 2021?"), ) thesis_facts = [fact for fact in result["facts"] if fact["fact_type"] == "person_thesis"] assert len(thesis_facts) == 1 assert thesis_facts[0]["year"] == 2021 assert thesis_facts[0]["thesis_type"] == "diploma_thesis" calls = [ parameters for query, parameters in session.calls if query == graph_search.THESIS_FACTS_QUERY ] assert len(calls) == 1 assert calls[0]["years"] == [2021] assert calls[0]["thesis_types"] == ["diploma_thesis"] assert calls[0]["person_ids"] == ["p2"] def test_diploma_thesis_query_does_not_call_work_queries() -> None: session = FakeSession( { graph_search.CATALOG_QUERY: (sample_catalog()), graph_search.THESIS_PATHS_QUERY: [], graph_search.THESIS_FACTS_QUERY: [], } ) graph_search.search_graph_with_session( session, ("Aký bol názov diplomovej " "práce Jakuba Maruniaka " "v roku 2021?"), ) called = {query for query, _ in session.calls} assert graph_search.THESIS_PATHS_QUERY in called assert graph_search.THESIS_FACTS_QUERY in called assert graph_search.WORK_PATHS_QUERY not in called assert graph_search.WORK_FACTS_QUERY not in called assert graph_search.NEIGHBORHOOD_QUERY not in called 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": "w2", "work_id": ("work:jakub:dp2020"), "title": None, "heading": ("Diplomový projekt 1 2020"), "year": 2020, "work_type": ("diploma_project"), "work_type_label": ("Diplomový projekt"), "document_path": ("pages/students/2016/" "jakub_maruniak/README.md"), "source_url": None, "source_chunk_ids": [], } session = FakeSession( { graph_search.CATALOG_QUERY: (sample_catalog()), graph_search.WORK_PATHS_QUERY: [], graph_search.WORK_FACTS_QUERY: [work_fact], } ) result = graph_search.search_graph_with_session( session, ("Aký diplomový projekt mal " "Jakub Maruniak " "v roku 2020?"), ) assert any( fact["fact_type"] == "person_work" and fact["year"] == 2020 and fact["work_type"] == "diploma_project" for fact in result["facts"] ) def test_person_keyword_query_returns_structured_facts() -> None: keyword_rows = [ { "person_element_id": "p2", "person": "Jakub Maruniak", "keyword_element_id": "k1", "keyword_id": "spacy", "keyword": "Spacy", "document_path": ("jakub.md"), "source_url": None, }, { "person_element_id": "p2", "person": "Jakub Maruniak", "keyword_element_id": "k1", "keyword_id": "spacy", "keyword": "Spacy", "document_path": ("duplicate.md"), "source_url": None, }, ] session = FakeSession( { graph_search.CATALOG_QUERY: (sample_catalog()), graph_search.PERSON_KEYWORD_FACTS_QUERY: (keyword_rows), } ) result = graph_search.search_graph_with_session( session, "Aké kľúčové slová má Jakub Maruniak?", ) facts = [fact for fact in result["facts"] if fact["fact_type"] == "person_keyword"] assert len(facts) == 1 assert facts[0]["keyword_id"] == "spacy" called = {query for query, _ in session.calls} assert graph_search.PERSON_KEYWORD_FACTS_QUERY in called assert graph_search.NEIGHBORHOOD_QUERY not in called def test_direct_keyword_query_returns_people() -> None: keyword_rows = [ { "person_element_id": "p2", "person": "Jakub Maruniak", "keyword_element_id": "k2", "keyword_id": "rag", "keyword": "RAG", "document_path": ("jakub.md"), "source_url": None, } ] session = FakeSession( { graph_search.CATALOG_QUERY: (sample_catalog()), graph_search.PERSON_KEYWORD_FACTS_QUERY: (keyword_rows), } ) result = graph_search.search_graph_with_session( session, "Kto súvisí s RAG?", ) facts = [fact for fact in result["facts"] if fact["fact_type"] == "person_keyword"] assert len(facts) == 1 assert facts[0]["person"] == "Jakub Maruniak" calls = [ parameters for query, parameters in session.calls if query == graph_search.PERSON_KEYWORD_FACTS_QUERY ] assert len(calls) == 1 assert calls[0]["keyword_ids"] == ["k2"] def test_direct_category_query_returns_people() -> None: category_rows = [ { "person_element_id": "p2", "person": "Jakub Maruniak", "category_element_id": "c1", "category": "dp2021", "category_type": ("diploma_thesis"), "category_year": 2021, "document_path": ("jakub.md"), "source_url": None, } ] session = FakeSession( { graph_search.CATALOG_QUERY: (sample_catalog()), graph_search.PERSON_CATEGORY_FACTS_QUERY: (category_rows), } ) result = graph_search.search_graph_with_session( session, "Kto súvisí s kategóriou dp2021?", ) facts = [fact for fact in result["facts"] if fact["fact_type"] == "person_category"] assert len(facts) == 1 assert facts[0]["person"] == "Jakub Maruniak" calls = [ parameters for query, parameters in session.calls if query == graph_search.PERSON_CATEGORY_FACTS_QUERY ] assert len(calls) == 1 assert calls[0]["category_ids"] == ["c1"] def test_shared_keyword_query_returns_structured_fact_without_generic_traversal() -> None: shared_keyword_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"), "keyword_element_id": "k1", "keyword_id": "spacy", "keyword": "Spacy", } session = FakeSession( { graph_search.CATALOG_QUERY: (sample_catalog()), graph_search.SHARED_KEYWORD_FACTS_QUERY: [shared_keyword_fact], } ) result = graph_search.search_graph_with_session( session, ("Aké spoločné kľúčové slová " "majú Maroš Harahus a " "Jakub Maruniak?"), ) shared_facts = [fact for fact in result["facts"] if fact["fact_type"] == "shared_keyword"] assert len(shared_facts) == 1 assert shared_facts[0]["keyword_id"] == "spacy" called = {query for query, _ in session.calls} assert graph_search.SHARED_KEYWORD_FACTS_QUERY in called assert graph_search.NEIGHBORHOOD_QUERY not in called assert graph_search.PAIR_PATH_QUERY not in called def test_shared_topic_word_uses_shared_keyword_semantics_for_people() -> None: session = FakeSession( { graph_search.CATALOG_QUERY: (sample_catalog()), graph_search.SHARED_KEYWORD_FACTS_QUERY: [], } ) graph_search.search_graph_with_session( session, ("Akú spoločnú tému majú " "Maroš Harahus a " "Jakub Maruniak?"), ) called = {query for query, _ in session.calls} assert graph_search.SHARED_KEYWORD_FACTS_QUERY in called def test_start_year_query_returns_one_fact_per_person() -> 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, }, { "person_element_id": "p2", "person": ("Jakub Maruniak"), "start_year": 2016, "document_path": ("duplicate.md"), "source_url": None, }, ], } ) result = graph_search.search_graph_with_session( session, ("Mali Maroš Harahus a " "Jakub Maruniak rovnaký " "rok začiatku štúdia?"), ) person_facts = [fact for fact in result["facts"] if fact["fact_type"] == "person_start_year"] assert len(person_facts) == 2 comparison = next( fact for fact in result["facts"] if fact["fact_type"] == "start_year_comparison" ) assert comparison["people"] == [ "Jakub Maruniak", "Maroš Harahus", ] assert comparison["start_years"] == { "Jakub Maruniak": 2016, "Maroš Harahus": 2016, } assert comparison["same_start_year"] is True assert comparison["shared_start_year"] == 2016 def test_start_year_query_skips_graph_traversal() -> None: session = FakeSession( { graph_search.CATALOG_QUERY: (sample_catalog()), graph_search.START_YEAR_FACTS_QUERY: [], } ) graph_search.search_graph_with_session( session, ("Mali Maroš Harahus a " "Jakub Maruniak rovnaký " "rok začiatku štúdia?"), ) called = {query for query, _ in session.calls} assert graph_search.NEIGHBORHOOD_QUERY not in called assert graph_search.PAIR_PATH_QUERY not in called def test_direct_thesis_entity_triggers_exact_thesis_lookup() -> None: session = FakeSession( { graph_search.CATALOG_QUERY: (sample_catalog()), graph_search.THESIS_PATHS_QUERY: [], graph_search.THESIS_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.THESIS_FACTS_QUERY ] assert len(calls) == 1 assert calls[0]["thesis_ids"] == ["th1"] def test_search_graph_rejects_empty_query() -> None: session = FakeSession({}) with pytest.raises( ValueError, match="prázdny", ): graph_search.search_graph_with_session( session, " ", )