from __future__ import annotations from typing import Any import pytest import scripts.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, ) ) 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, "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, "work_type": None, "person_kind": "student", "source_document": ( "pages/students/2016/" "jakub_maruniak/README.md" ), }, { "element_id": "t1", "kind": "Topic", "id": "spacy", "name": "Spacy", "title": None, "heading": None, "start_year": None, "year": None, "work_type": None, "person_kind": None, "source_document": None, }, { "element_id": "t2", "kind": "Topic", "id": "rag", "name": "RAG", "title": None, "heading": None, "start_year": None, "year": None, "work_type": None, "person_kind": None, "source_document": None, }, { "element_id": "w1", "kind": "Work", "id": "work: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, "work_type": "diploma_thesis", "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[ "work_types" ] == [ "diploma_thesis" ] assert result[ "mentions_title" ] is True assert "Work" in result[ "target_kinds" ] 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_resolve_entities_finds_two_people_and_topic() -> None: result = graph_search.resolve_entities( "Akú spoločnú tému Spacy majú " "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 ( "Topic", "Spacy", ) in resolved def test_person_is_preferred_over_author_with_same_name() -> None: catalog = sample_catalog() + [ { "element_id": "a1", "kind": "Author", "id": "maros-harahus", "name": "Maroš Harahus", "title": None, "heading": None, "start_year": None, "year": 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_topic_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_work_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"] == "Work" and item["element_id"] == "w1" for item in result ) def test_more_complete_inflected_work_title_ranks_above_shorter_phrase() -> None: catalog = sample_catalog() + [ { "element_id": "w2", "kind": "Work", "id": "work:shorter", "name": ( "Rozpoznávanie pomenovaných " "entít v slovenskom jazyku" ), "title": ( "Rozpoznávanie pomenovaných " "entít v slovenskom jazyku" ), "heading": "Diplomová práca 2020", "start_year": None, "year": 2020, "work_type": "diploma_thesis", "person_kind": None, "source_document": "pages/students/other/README.md", } ] result = graph_search.resolve_entities( "Kto riešil Anotáciu a " "rozpoznávanie pomenovaných entít " "v slovenskom jazyku?", catalog, ) work_results = [ item for item in result if item["kind"] == "Work" ] assert work_results[0]["element_id"] == "w1" def test_search_graph_with_session_uses_topic_path_for_two_people() -> None: topic_path = { "nodes": [ { "element_id": "p1", "kind": "Person", "name": "Maroš Harahus", "path": None, }, { "element_id": "d1", "kind": "Document", "name": "Maroš Harahus", "path": ( "pages/students/2016/" "maros_harahus/README.md" ), "source_url": ( "https://zp.kemt.fei.tuke.sk/" "students/2016/maros_harahus" ), }, { "element_id": "t1", "kind": "Topic", "name": "Spacy", "id": "spacy", }, { "element_id": "d2", "kind": "Document", "name": "Jakub Maruniak", "path": ( "pages/students/2016/" "jakub_maruniak/README.md" ), "source_url": ( "https://zp.kemt.fei.tuke.sk/" "students/2016/jakub_maruniak" ), }, { "element_id": "p2", "kind": "Person", "name": "Jakub Maruniak", "path": None, }, ], "relationships": [ "HAS_DOCUMENT", "HAS_TOPIC", "HAS_TOPIC", "HAS_DOCUMENT", ], "hops": 4, } session = FakeSession( { graph_search.CATALOG_QUERY: ( sample_catalog() ), graph_search.NEIGHBORHOOD_QUERY: [], graph_search.PAIR_PATH_QUERY: [ topic_path ], } ) result = ( graph_search.search_graph_with_session( session, "Akú spoločnú tému majú " "Maroš Harahus a Jakub Maruniak?", ) ) assert result[ "constraints" ]["mentions_topic"] is True assert len( result["paths"] ) == 1 assert { item["document_path"] for item in result[ "documents" ] } == { ( "pages/students/2016/" "maros_harahus/README.md" ), ( "pages/students/2016/" "jakub_maruniak/README.md" ), } pair_calls = [ parameters for query, parameters in session.calls if query == graph_search.PAIR_PATH_QUERY ] assert len(pair_calls) == 1 assert pair_calls[0][ "require_topic" ] is True 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": "w1", "work_id": "work:jakub:dp2021", "title": ( "Anotácia a rozpoznávanie " "pomenovaných entít v " "slovenskom jazyku." ), "heading": "Diplomová práca 2021", "year": 2021, "work_type": "diploma_thesis", "work_type_label": "Diplomová práca", "document_path": ( "pages/students/2016/" "jakub_maruniak/README.md" ), "source_url": None, "source_chunk_ids": [ ( "pages/students/2016/" "jakub_maruniak/" "README.md::chunk-0" ) ], } session = FakeSession( { graph_search.CATALOG_QUERY: ( sample_catalog() ), graph_search.NEIGHBORHOOD_QUERY: [], graph_search.WORK_FACTS_QUERY: [ work_fact ], } ) result = ( graph_search.search_graph_with_session( session, "Aký bol názov diplomovej práce " "Jakuba Maruniaka v roku 2021?", ) ) assert any( fact["fact_type"] == "person_work" and fact["year"] == 2021 and fact["work_type"] == "diploma_thesis" for fact in result["facts"] ) work_calls = [ parameters for query, parameters in session.calls if query == graph_search.WORK_FACTS_QUERY ] assert len(work_calls) == 1 assert work_calls[0][ "years" ] == [2021] assert work_calls[0][ "work_types" ] == [ "diploma_thesis" ] assert work_calls[0][ "person_ids" ] == ["p2"] def test_search_graph_with_session_uses_start_year_facts() -> None: session = FakeSession( { graph_search.CATALOG_QUERY: ( sample_catalog() ), graph_search.NEIGHBORHOOD_QUERY: [], graph_search.PAIR_PATH_QUERY: [], graph_search.START_YEAR_FACTS_QUERY: [ { "person_element_id": "p1", "person": "Maroš Harahus", "start_year": 2016, "document_path": ( "pages/students/2016/" "maros_harahus/README.md" ), "source_url": None, }, { "person_element_id": "p2", "person": "Jakub Maruniak", "start_year": 2016, "document_path": ( "pages/students/2016/" "jakub_maruniak/README.md" ), "source_url": None, }, ], } ) result = ( graph_search.search_graph_with_session( session, "Mali Maroš Harahus a Jakub " "Maruniak rovnaký rok začiatku " "štúdia?", ) ) values = { fact["person"]: fact[ "start_year" ] for fact in result["facts"] if fact["fact_type"] == "person_start_year" } assert values == { "Maroš Harahus": 2016, "Jakub Maruniak": 2016, } def test_search_graph_rejects_empty_query() -> None: session = FakeSession({}) with pytest.raises( ValueError, match="prázdny", ): graph_search.search_graph_with_session( session, " ", ) def test_shared_topic_query_returns_structured_fact_and_no_generic_neighborhood() -> None: shared_topic_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", "topic_element_id": "t1", "topic_id": "spacy", "topic": "Spacy", } session = FakeSession( { graph_search.CATALOG_QUERY: sample_catalog(), graph_search.PAIR_PATH_QUERY: [], graph_search.SHARED_TOPIC_FACTS_QUERY: [ shared_topic_fact ], } ) result = graph_search.search_graph_with_session( session, "Akú spoločnú tému majú Maroš Harahus " "a Jakub Maruniak?", ) shared_facts = [ fact for fact in result["facts"] if fact["fact_type"] == "shared_topic" ] assert len(shared_facts) == 1 assert shared_facts[0]["topic_id"] == "spacy" assert { item["document_path"] for item in result["documents"] } == { "pages/students/2016/maros_harahus/README.md", "pages/students/2016/jakub_maruniak/README.md", } called_queries = { query for query, _ in session.calls } assert graph_search.NEIGHBORHOOD_QUERY not in called_queries assert graph_search.PAIR_PATH_QUERY in called_queries assert graph_search.SHARED_TOPIC_FACTS_QUERY in called_queries def test_work_query_uses_precise_work_paths_instead_of_generic_neighborhood() -> None: work_path = { "nodes": [ { "element_id": "p2", "kind": "Person", "name": "Jakub Maruniak", }, { "element_id": "w1", "kind": "Work", "name": ( "Anotácia a rozpoznávanie " "pomenovaných entít v " "slovenskom jazyku." ), "year": 2021, "work_type": "diploma_thesis", }, { "element_id": "d2", "kind": "Document", "path": ( "pages/students/2016/" "jakub_maruniak/README.md" ), }, ], "relationships": [ "HAS_WORK", "EVIDENCED_BY", ], "hops": 2, } session = FakeSession( { graph_search.CATALOG_QUERY: sample_catalog(), graph_search.WORK_PATHS_QUERY: [work_path], graph_search.WORK_FACTS_QUERY: [], } ) result = graph_search.search_graph_with_session( session, "Aký bol názov diplomovej práce " "Jakuba Maruniaka v roku 2021?", ) assert result["paths"] == [work_path] called_queries = { query for query, _ in session.calls } assert graph_search.WORK_PATHS_QUERY in called_queries assert graph_search.NEIGHBORHOOD_QUERY not in called_queries work_path_calls = [ parameters for query, parameters in session.calls if query == graph_search.WORK_PATHS_QUERY ] assert work_path_calls[0]["years"] == [2021] assert work_path_calls[0]["work_types"] == [ "diploma_thesis" ] assert work_path_calls[0]["person_ids"] == ["p2"] def test_start_year_comparison_is_structured_and_skips_irrelevant_paths() -> 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, }, ], } ) result = graph_search.search_graph_with_session( session, "Mali Maroš Harahus a Jakub Maruniak " "rovnaký rok začiatku štúdia?", ) assert result["paths"] == [] comparison = next( fact for fact in result["facts"] if fact["fact_type"] == "start_year_comparison" ) assert comparison["same_start_year"] is True assert comparison["shared_start_year"] == 2016 assert comparison["start_years"] == { "Jakub Maruniak": 2016, "Maroš Harahus": 2016, } called_queries = { query for query, _ in session.calls } assert graph_search.NEIGHBORHOOD_QUERY not in called_queries assert graph_search.PAIR_PATH_QUERY not in called_queries def test_resolved_work_entity_triggers_exact_work_lookup() -> None: session = FakeSession( { graph_search.CATALOG_QUERY: sample_catalog(), graph_search.WORK_PATHS_QUERY: [], graph_search.WORK_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.WORK_FACTS_QUERY ] assert len(calls) == 1 assert calls[0]["work_ids"] == ["w1"]