1329 lines
28 KiB
Python
1329 lines
28 KiB
Python
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,
|
||
)
|
||
)
|
||
|
||
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,
|
||
" ",
|
||
)
|