diff --git a/scripts/graph_search.py b/scripts/graph_search.py new file mode 100644 index 0000000..a6bd3a1 --- /dev/null +++ b/scripts/graph_search.py @@ -0,0 +1,1346 @@ +from __future__ import annotations + +import argparse +import json +import re +import unicodedata +from typing import Any, Iterable + + +GRAPH_SCOPE = "zpwiki" +DEFAULT_LIMIT = 40 +DEFAULT_MAX_ENTITIES = 8 + +RELATIONSHIP_TYPES = ( + "HAS_DOCUMENT|AUTHORED|IN_CATEGORY|HAS_TOPIC|" + "DESCRIBES|HAS_WORK|EVIDENCED_BY" +) + +CATALOG_QUERY = """ +MATCH (n) +WHERE n.graph_scope = $scope + AND ( + n:Person + OR n:Topic + OR n:Category + OR n:Author + OR n:Work + ) +RETURN + elementId(n) AS element_id, + CASE + WHEN n:Person THEN 'Person' + WHEN n:Topic THEN 'Topic' + WHEN n:Category THEN 'Category' + WHEN n:Author THEN 'Author' + WHEN n:Work THEN 'Work' + ELSE 'Other' + END AS kind, + n.id AS id, + n.name AS name, + n.title AS title, + n.heading AS heading, + n.start_year AS start_year, + n.year AS year, + n.work_type AS work_type, + n.person_kind AS person_kind, + n.source_document AS source_document +ORDER BY kind, name, title, heading +""" + +NEIGHBORHOOD_QUERY = f""" +MATCH (start) +WHERE elementId(start) IN $start_ids +MATCH p=(start)-[:{RELATIONSHIP_TYPES}*1..2]-(target) +WHERE all(node IN nodes(p) WHERE node.graph_scope = $scope) + AND ( + size($target_kinds) = 0 + OR any( + node IN nodes(p) + WHERE + ('Topic' IN $target_kinds AND node:Topic) + OR ('Category' IN $target_kinds AND node:Category) + OR ('Author' IN $target_kinds AND node:Author) + OR ('Work' IN $target_kinds AND node:Work) + OR ('Document' IN $target_kinds AND node:Document) + OR ('Person' IN $target_kinds AND node:Person) + ) + ) +RETURN + [node IN nodes(p) | {{ + element_id: elementId(node), + kind: CASE + WHEN node:Person THEN 'Person' + WHEN node:Document THEN 'Document' + WHEN node:Work THEN 'Work' + WHEN node:Topic THEN 'Topic' + WHEN node:Category THEN 'Category' + WHEN node:Author THEN 'Author' + ELSE 'Other' + END, + id: node.id, + name: node.name, + title: node.title, + heading: node.heading, + path: node.path, + start_year: node.start_year, + year: node.year, + work_type: node.work_type, + source_document: node.source_document, + source_url: node.source_url + }}] AS nodes, + [relationship IN relationships(p) | type(relationship)] AS relationships, + length(p) AS hops +ORDER BY hops +LIMIT $limit +""" + +PAIR_PATH_QUERY = f""" +MATCH (left), (right) +WHERE elementId(left) IN $start_ids + AND elementId(right) IN $start_ids + AND elementId(left) < elementId(right) +MATCH p=allShortestPaths( + (left)-[:{RELATIONSHIP_TYPES}*1..4]-(right) +) +WHERE all(node IN nodes(p) WHERE node.graph_scope = $scope) + AND (NOT $require_topic OR any(node IN nodes(p) WHERE node:Topic)) + AND (NOT $require_category OR any(node IN nodes(p) WHERE node:Category)) + AND (NOT $require_author OR any(node IN nodes(p) WHERE node:Author)) +RETURN + [node IN nodes(p) | {{ + element_id: elementId(node), + kind: CASE + WHEN node:Person THEN 'Person' + WHEN node:Document THEN 'Document' + WHEN node:Work THEN 'Work' + WHEN node:Topic THEN 'Topic' + WHEN node:Category THEN 'Category' + WHEN node:Author THEN 'Author' + ELSE 'Other' + END, + id: node.id, + name: node.name, + title: node.title, + heading: node.heading, + path: node.path, + start_year: node.start_year, + year: node.year, + work_type: node.work_type, + source_document: node.source_document, + source_url: node.source_url + }}] AS nodes, + [relationship IN relationships(p) | type(relationship)] AS relationships, + length(p) AS hops +ORDER BY hops +LIMIT $limit +""" + +WORK_FACTS_QUERY = """ +MATCH (person:Person)-[:HAS_WORK]->(work:Work)-[:EVIDENCED_BY]->(document:Document) +WHERE person.graph_scope = $scope + AND work.graph_scope = $scope + AND document.graph_scope = $scope + AND ( + size($person_ids) = 0 + OR elementId(person) IN $person_ids + ) + AND ( + size($work_ids) = 0 + OR elementId(work) IN $work_ids + ) + AND ( + size($years) = 0 + OR work.year IN $years + ) + AND ( + size($work_types) = 0 + OR work.work_type IN $work_types + ) +RETURN + elementId(person) AS person_element_id, + person.name AS person, + person.start_year AS start_year, + elementId(work) AS work_element_id, + work.id AS work_id, + work.title AS title, + work.heading AS heading, + work.year AS year, + work.work_type AS work_type, + work.work_type_label AS work_type_label, + document.path AS document_path, + document.source_url AS source_url, + work.source_chunk_ids AS source_chunk_ids +ORDER BY person, year, work_type, title +LIMIT $limit +""" + +START_YEAR_FACTS_QUERY = """ +MATCH (person:Person)-[:HAS_DOCUMENT]->(document:Document) +WHERE person.graph_scope = $scope + AND document.graph_scope = $scope + AND ( + size($person_ids) = 0 + OR elementId(person) IN $person_ids + ) + AND ( + size($years) = 0 + OR person.start_year IN $years + ) +RETURN + elementId(person) AS person_element_id, + person.name AS person, + person.start_year AS start_year, + document.path AS document_path, + document.source_url AS source_url +ORDER BY person +LIMIT $limit +""" + +WORK_PATHS_QUERY = """ +MATCH p=(person:Person)-[:HAS_WORK]->(work:Work)-[:EVIDENCED_BY]->(document:Document) +WHERE person.graph_scope = $scope + AND work.graph_scope = $scope + AND document.graph_scope = $scope + AND ( + size($person_ids) = 0 + OR elementId(person) IN $person_ids + ) + AND ( + size($work_ids) = 0 + OR elementId(work) IN $work_ids + ) + AND ( + size($years) = 0 + OR work.year IN $years + ) + AND ( + size($work_types) = 0 + OR work.work_type IN $work_types + ) +RETURN + [node IN nodes(p) | { + element_id: elementId(node), + kind: CASE + WHEN node:Person THEN 'Person' + WHEN node:Document THEN 'Document' + WHEN node:Work THEN 'Work' + ELSE 'Other' + END, + id: node.id, + name: node.name, + title: node.title, + heading: node.heading, + path: node.path, + start_year: node.start_year, + year: node.year, + work_type: node.work_type, + source_document: node.source_document, + source_url: node.source_url + }] AS nodes, + [relationship IN relationships(p) | type(relationship)] AS relationships, + length(p) AS hops +ORDER BY person.name, work.year, work.work_type, work.title +LIMIT $limit +""" + +SHARED_TOPIC_FACTS_QUERY = """ +MATCH (left:Person)-[:HAS_DOCUMENT]->(left_document:Document)-[:HAS_TOPIC]->(topic:Topic) +MATCH (right:Person)-[:HAS_DOCUMENT]->(right_document:Document)-[:HAS_TOPIC]->(topic) +WHERE left.graph_scope = $scope + AND right.graph_scope = $scope + AND left_document.graph_scope = $scope + AND right_document.graph_scope = $scope + AND topic.graph_scope = $scope + AND elementId(left) IN $person_ids + AND elementId(right) IN $person_ids + AND elementId(left) < elementId(right) +RETURN + elementId(left) AS left_person_element_id, + left.name AS left_person, + left_document.path AS left_document_path, + left_document.source_url AS left_source_url, + elementId(right) AS right_person_element_id, + right.name AS right_person, + right_document.path AS right_document_path, + right_document.source_url AS right_source_url, + elementId(topic) AS topic_element_id, + topic.id AS topic_id, + topic.name AS topic +ORDER BY left_person, right_person, topic +LIMIT $limit +""" + + +def normalize_for_match(value: str) -> str: + decomposed = unicodedata.normalize( + "NFKD", + str(value), + ) + + without_marks = "".join( + character + for character in decomposed + if not unicodedata.combining(character) + ) + + normalized = re.sub( + r"[^a-zA-Z0-9]+", + " ", + without_marks.casefold(), + ) + + return re.sub( + r"\s+", + " ", + normalized, + ).strip() + + +def normalized_tokens(value: str) -> list[str]: + normalized = normalize_for_match(value) + + if not normalized: + return [] + + return normalized.split() + + +def has_stem( + tokens: Iterable[str], + stem: str, +) -> bool: + return any( + token.startswith(stem) + for token in tokens + ) + + +def extract_query_constraints( + query: str, +) -> dict[str, Any]: + normalized = normalize_for_match(query) + tokens = normalized.split() + token_set = set(tokens) + + years = sorted( + { + int(value) + for value in re.findall( + r"\b(20\d{2})\b", + normalized, + ) + } + ) + + work_types: list[str] = [] + + has_work_word = has_stem( + tokens, + "prac", + ) + + has_project_word = has_stem( + tokens, + "projekt", + ) + + if ( + ( + has_stem(tokens, "bakalar") + and has_work_word + ) + or "bp" in token_set + ): + work_types.append( + "bachelor_thesis" + ) + + if ( + ( + has_stem(tokens, "diplom") + and has_work_word + ) + or "dp" in token_set + ): + work_types.append( + "diploma_thesis" + ) + + if ( + has_stem(tokens, "diplom") + and has_project_word + ): + work_types.append( + "diploma_project" + ) + + if ( + has_stem(tokens, "timov") + and has_project_word + ): + work_types.append( + "team_project" + ) + + if ( + has_stem(tokens, "vyskum") + and has_project_word + ): + work_types.append( + "research_project" + ) + + if ( + ( + has_stem(tokens, "dizert") + and has_work_word + ) + or "phd" in token_set + ): + work_types.append( + "doctoral_thesis" + ) + + if "tp" in token_set: + work_types.append( + "team_project" + ) + + if "vp" in token_set: + work_types.append( + "research_project" + ) + + work_types = list( + dict.fromkeys(work_types) + ) + + mentions_start_year = ( + has_stem(tokens, "zaciat") + and has_stem(tokens, "stud") + ) + + mentions_topic = ( + has_stem(tokens, "tem") + or "tag" in token_set + or "tagy" in token_set + ) + + mentions_category = has_stem( + tokens, + "kategori", + ) + + mentions_author = ( + has_stem(tokens, "autor") + or has_stem(tokens, "veduc") + ) + + mentions_title = has_stem( + tokens, + "nazov", + ) + + mentions_shared = ( + has_stem(tokens, "spoloc") + or has_stem(tokens, "rovnak") + ) + + target_kinds: list[str] = [] + + if mentions_topic: + target_kinds.append("Topic") + + if mentions_category: + target_kinds.append("Category") + + if mentions_author: + target_kinds.append("Author") + + if ( + mentions_title + or work_types + or ( + years + and not mentions_start_year + ) + ): + target_kinds.append("Work") + + return { + "years": years, + "work_types": work_types, + "mentions_start_year": ( + mentions_start_year + ), + "mentions_topic": mentions_topic, + "mentions_category": ( + mentions_category + ), + "mentions_author": mentions_author, + "mentions_title": mentions_title, + "mentions_shared": mentions_shared, + "target_kinds": list( + dict.fromkeys(target_kinds) + ), + } + + +def entity_aliases( + entity: dict[str, Any], +) -> list[str]: + kind = str( + entity.get("kind") + or "" + ) + + values: list[Any] = [] + + if kind in { + "Person", + "Author", + }: + values.append( + entity.get("name") + ) + + elif kind == "Topic": + values.extend( + [ + entity.get("name"), + entity.get("id"), + ] + ) + + elif kind == "Category": + values.append( + entity.get("name") + ) + + elif kind == "Work": + values.append( + entity.get("title") + ) + + aliases: list[str] = [] + + for value in values: + if value is None: + continue + + normalized = normalize_for_match( + str(value) + ) + + if normalized: + aliases.append(normalized) + + return list( + dict.fromkeys(aliases) + ) + + +def tokens_match( + left: str, + right: str, +) -> bool: + if left == right: + return True + + shorter = min( + len(left), + len(right), + ) + + if shorter <= 2: + return False + + required = ( + 3 + if shorter <= 4 + else 4 + if shorter <= 6 + else 5 + ) + + return ( + left[:required] + == right[:required] + ) + + +def alias_match_score( + normalized_query: str, + query_tokens: set[str], + alias: str, +) -> float: + alias_tokens = alias.split() + + if not alias_tokens: + return 0.0 + + if len(alias_tokens) == 1: + if alias not in query_tokens: + return 0.0 + + return 50.0 + min( + len(alias), + 20, + ) / 100.0 + + padded_query = ( + " " + + normalized_query + + " " + ) + + padded_alias = ( + " " + + alias + + " " + ) + + if padded_alias in padded_query: + score = ( + 100.0 + + len(alias_tokens) * 5.0 + + min(len(alias), 100) / 100.0 + ) + + if normalized_query == alias: + score += 50.0 + + return score + + if all( + any( + tokens_match( + alias_token, + query_token, + ) + for query_token in query_tokens + ) + for alias_token in alias_tokens + ): + return ( + 95.0 + + len(alias_tokens) * 5.0 + + min(len(alias), 100) / 100.0 + ) + + return 0.0 + + +def resolve_entities( + query: str, + catalog: list[dict[str, Any]], + *, + max_entities: int = DEFAULT_MAX_ENTITIES, +) -> list[dict[str, Any]]: + normalized_query = normalize_for_match( + query + ) + + query_tokens = set( + normalized_query.split() + ) + + matches: list[dict[str, Any]] = [] + + for entity in catalog: + best_score = 0.0 + best_alias = "" + + for alias in entity_aliases( + entity + ): + score = alias_match_score( + normalized_query, + query_tokens, + alias, + ) + + if score > best_score: + best_score = score + best_alias = alias + + if best_score <= 0.0: + continue + + item = dict(entity) + item["matched_alias"] = best_alias + item["match_score"] = round( + best_score, + 3, + ) + + matches.append(item) + + person_aliases = { + str(item.get("matched_alias") or "") + for item in matches + if item.get("kind") == "Person" + } + + matches = [ + item + for item in matches + if not ( + item.get("kind") == "Author" + and str( + item.get("matched_alias") + or "" + ) + in person_aliases + ) + ] + + matches.sort( + key=lambda item: ( + -float(item["match_score"]), + str(item.get("kind") or ""), + str( + item.get("name") + or item.get("title") + or item.get("id") + or "" + ), + ) + ) + + selected: list[dict[str, Any]] = [] + seen_ids: set[str] = set() + + for item in matches: + element_id = str( + item.get("element_id") + or "" + ) + + if not element_id: + continue + + if element_id in seen_ids: + continue + + seen_ids.add(element_id) + selected.append(item) + + if len(selected) >= max_entities: + break + + return selected + + +def rows_from_result( + result: Any, +) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + + for record in result: + if isinstance(record, dict): + rows.append( + dict(record) + ) + continue + + data = getattr( + record, + "data", + None, + ) + + if callable(data): + rows.append( + dict(data()) + ) + continue + + rows.append( + dict(record) + ) + + return rows + + +def run_rows( + session: Any, + query: str, + **parameters: Any, +) -> list[dict[str, Any]]: + return rows_from_result( + session.run( + query, + **parameters, + ) + ) + + +def deduplicate_paths( + paths: list[dict[str, Any]], +) -> list[dict[str, Any]]: + result: list[dict[str, Any]] = [] + seen: set[str] = set() + + for path in paths: + key = json.dumps( + { + "nodes": path.get( + "nodes", + [], + ), + "relationships": path.get( + "relationships", + [], + ), + }, + ensure_ascii=False, + sort_keys=True, + default=str, + ) + + if key in seen: + continue + + seen.add(key) + result.append(path) + + return result + + +def collect_documents( + resolved_entities: list[dict[str, Any]], + paths: list[dict[str, Any]], + facts: list[dict[str, Any]], +) -> list[dict[str, Any]]: + documents: dict[ + str, + dict[str, Any], + ] = {} + + def add_document( + path: Any, + source_url: Any = None, + ) -> None: + value = str( + path + or "" + ).strip() + + if not value: + return + + current = documents.setdefault( + value, + { + "document_path": value, + "source_url": None, + }, + ) + + if ( + source_url + and not current.get( + "source_url" + ) + ): + current["source_url"] = str( + source_url + ) + + for entity in resolved_entities: + add_document( + entity.get( + "source_document" + ) + ) + + for path in paths: + for node in path.get( + "nodes", + [], + ): + if not isinstance( + node, + dict, + ): + continue + + if node.get("kind") == "Document": + add_document( + node.get("path"), + node.get("source_url"), + ) + + add_document( + node.get( + "source_document" + ) + ) + + for fact in facts: + add_document( + fact.get( + "document_path" + ), + fact.get("source_url"), + ) + + add_document( + fact.get( + "left_document_path" + ), + fact.get( + "left_source_url" + ), + ) + + add_document( + fact.get( + "right_document_path" + ), + fact.get( + "right_source_url" + ), + ) + + return sorted( + documents.values(), + key=lambda item: item[ + "document_path" + ], + ) + + +def build_facts( + work_rows: list[dict[str, Any]], + start_year_rows: list[dict[str, Any]], + shared_topic_rows: list[dict[str, Any]], + *, + mentions_shared: bool, +) -> list[dict[str, Any]]: + facts: list[dict[str, Any]] = [] + + for row in work_rows: + fact = dict(row) + fact["fact_type"] = ( + "person_work" + ) + facts.append(fact) + + for row in start_year_rows: + fact = dict(row) + fact["fact_type"] = ( + "person_start_year" + ) + facts.append(fact) + + for row in shared_topic_rows: + fact = dict(row) + fact["fact_type"] = ( + "shared_topic" + ) + facts.append(fact) + + if ( + mentions_shared + and len(start_year_rows) >= 2 + ): + ordered = sorted( + start_year_rows, + key=lambda row: str( + row.get("person") + or "" + ), + ) + + years = [ + row.get("start_year") + for row in ordered + ] + + known_years = [ + year + for year in years + if year is not None + ] + + same = ( + len(known_years) + == len(years) + and len(set(known_years)) + == 1 + ) + + facts.append( + { + "fact_type": ( + "start_year_comparison" + ), + "people": [ + str( + row.get("person") + or "" + ) + for row in ordered + ], + "start_years": { + str( + row.get("person") + or "" + ): row.get( + "start_year" + ) + for row in ordered + }, + "same_start_year": same, + "shared_start_year": ( + known_years[0] + if same + else None + ), + } + ) + + return facts + + +def search_graph_with_session( + session: Any, + query: str, + *, + limit: int = DEFAULT_LIMIT, + max_entities: int = DEFAULT_MAX_ENTITIES, +) -> dict[str, Any]: + clean_query = str( + query + or "" + ).strip() + + if not clean_query: + raise ValueError( + "Query nesmie byť prázdny" + ) + + if limit < 1: + raise ValueError( + "Limit musí byť aspoň 1" + ) + + if max_entities < 1: + raise ValueError( + "max_entities musí byť aspoň 1" + ) + + constraints = extract_query_constraints( + clean_query + ) + + catalog = run_rows( + session, + CATALOG_QUERY, + scope=GRAPH_SCOPE, + ) + + resolved_entities = resolve_entities( + clean_query, + catalog, + max_entities=max_entities, + ) + + start_ids = [ + str(entity["element_id"]) + for entity in resolved_entities + if entity.get("element_id") + ] + + person_ids = [ + str(entity["element_id"]) + for entity in resolved_entities + if ( + entity.get("element_id") + and entity.get("kind") + == "Person" + ) + ] + + work_ids = [ + str(entity["element_id"]) + for entity in resolved_entities + if ( + entity.get("element_id") + and entity.get("kind") + == "Work" + ) + ] + + work_specific = bool( + constraints["work_types"] + or constraints["mentions_title"] + or work_ids + or ( + constraints["years"] + and not constraints[ + "mentions_start_year" + ] + ) + ) + + shared_relation_query = bool( + constraints["mentions_shared"] + and len(person_ids) >= 2 + and ( + constraints["mentions_topic"] + or constraints[ + "mentions_category" + ] + or constraints[ + "mentions_author" + ] + ) + ) + + paths: list[dict[str, Any]] = [] + + if work_specific: + paths.extend( + run_rows( + session, + WORK_PATHS_QUERY, + scope=GRAPH_SCOPE, + person_ids=person_ids, + work_ids=work_ids, + years=constraints["years"], + work_types=( + constraints[ + "work_types" + ] + ), + limit=limit, + ) + ) + + elif shared_relation_query: + paths.extend( + run_rows( + session, + PAIR_PATH_QUERY, + scope=GRAPH_SCOPE, + start_ids=person_ids, + require_topic=( + constraints[ + "mentions_topic" + ] + ), + require_category=( + constraints[ + "mentions_category" + ] + ), + require_author=( + constraints[ + "mentions_author" + ] + ), + limit=limit, + ) + ) + + elif not constraints[ + "mentions_start_year" + ] and start_ids: + paths.extend( + run_rows( + session, + NEIGHBORHOOD_QUERY, + scope=GRAPH_SCOPE, + start_ids=start_ids, + target_kinds=( + constraints[ + "target_kinds" + ] + ), + limit=limit, + ) + ) + + work_rows: list[dict[str, Any]] = [] + + if work_specific: + work_rows = run_rows( + session, + WORK_FACTS_QUERY, + scope=GRAPH_SCOPE, + person_ids=person_ids, + work_ids=work_ids, + years=constraints["years"], + work_types=( + constraints[ + "work_types" + ] + ), + limit=limit, + ) + + start_year_rows: list[ + dict[str, Any] + ] = [] + + if constraints[ + "mentions_start_year" + ]: + start_year_rows = run_rows( + session, + START_YEAR_FACTS_QUERY, + scope=GRAPH_SCOPE, + person_ids=person_ids, + years=constraints["years"], + limit=limit, + ) + + shared_topic_rows: list[ + dict[str, Any] + ] = [] + + if ( + constraints["mentions_shared"] + and constraints["mentions_topic"] + and len(person_ids) >= 2 + ): + shared_topic_rows = run_rows( + session, + SHARED_TOPIC_FACTS_QUERY, + scope=GRAPH_SCOPE, + person_ids=person_ids, + limit=limit, + ) + + paths = deduplicate_paths( + paths + ) + + facts = build_facts( + work_rows, + start_year_rows, + shared_topic_rows, + mentions_shared=( + constraints[ + "mentions_shared" + ] + ), + ) + + documents = collect_documents( + resolved_entities, + paths, + facts, + ) + + strategies: list[str] = [] + + if resolved_entities: + strategies.append( + "entity_resolution" + ) + + if paths: + strategies.append( + "graph_traversal" + ) + + if work_rows: + strategies.append( + "work_constraints" + ) + + if start_year_rows: + strategies.append( + "start_year_constraints" + ) + + if shared_topic_rows: + strategies.append( + "shared_topic_constraints" + ) + + return { + "query": clean_query, + "strategies": strategies, + "constraints": constraints, + "resolved_entities": ( + resolved_entities + ), + "paths": paths[:limit], + "facts": facts[:limit], + "documents": documents, + } + + +def search_graph( + query: str, + *, + limit: int = DEFAULT_LIMIT, + max_entities: int = DEFAULT_MAX_ENTITIES, +) -> dict[str, Any]: + from scripts.graph_db import ( + create_neo4j_driver, + get_neo4j_settings, + ) + + settings = get_neo4j_settings() + + with create_neo4j_driver( + settings + ) as driver: + driver.verify_connectivity() + + with driver.session( + database=settings.database + ) as session: + return search_graph_with_session( + session, + query, + limit=limit, + max_entities=max_entities, + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + + parser.add_argument( + "query", + ) + + parser.add_argument( + "--limit", + type=int, + default=DEFAULT_LIMIT, + ) + + parser.add_argument( + "--max-entities", + type=int, + default=DEFAULT_MAX_ENTITIES, + ) + + args = parser.parse_args() + + result = search_graph( + args.query, + limit=args.limit, + max_entities=args.max_entities, + ) + + print( + json.dumps( + result, + ensure_ascii=False, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/test/test_graph_search.py b/test/test_graph_search.py new file mode 100644 index 0000000..fe241c6 --- /dev/null +++ b/test/test_graph_search.py @@ -0,0 +1,784 @@ +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"]