796 lines
23 KiB
Python
796 lines
23 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,
|
||
)
|
||
)
|
||
|
||
# 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,
|
||
" ",
|
||
)
|