785 lines
20 KiB
Python
785 lines
20 KiB
Python
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"]
|