diff --git a/tools/graphrag/build_graphrag.py b/tools/graphrag/build_graphrag.py index 29896f3..b8a8fba 100644 --- a/tools/graphrag/build_graphrag.py +++ b/tools/graphrag/build_graphrag.py @@ -44,8 +44,8 @@ CONSTRAINTS = ( "id", ), ( - "zpwiki_author_id_unique", - "Author", + "zpwiki_document_author_id_unique", + "DocumentAuthor", "id", ), ( @@ -58,6 +58,16 @@ CONSTRAINTS = ( "Topic", "id", ), + ( + "zpwiki_keyword_id_unique", + "Keyword", + "id", + ), + ( + "zpwiki_thesis_id_unique", + "Thesis", + "id", + ), ( "zpwiki_work_id_unique", "Work", @@ -66,15 +76,20 @@ CONSTRAINTS = ( ) +LEGACY_CONSTRAINTS = ( + "zpwiki_author_id_unique", +) + + NODE_QUERIES = { "people": """ UNWIND $rows AS row MERGE (n:Person {id: row.id}) SET n += row """, - "authors": """ + "document_authors": """ UNWIND $rows AS row - MERGE (n:Author {id: row.id}) + MERGE (n:DocumentAuthor {id: row.id}) SET n += row """, "documents": """ @@ -92,6 +107,16 @@ NODE_QUERIES = { MERGE (n:Topic {id: row.id}) SET n += row """, + "keywords": """ + UNWIND $rows AS row + MERGE (n:Keyword {id: row.id}) + SET n += row + """, + "theses": """ + UNWIND $rows AS row + MERGE (n:Thesis {id: row.id}) + SET n += row + """, "works": """ UNWIND $rows AS row MERGE (n:Work {id: row.id}) @@ -103,19 +128,23 @@ NODE_QUERIES = { RELATIONSHIP_QUERIES = { "person_documents": """ UNWIND $rows AS row - MATCH (a:Person {id: row.person_id}) + MATCH (a:Person { + id: row.person_id + }) MATCH (b:Document { path: row.document_path }) MERGE (a)-[:HAS_DOCUMENT]->(b) """, - "author_documents": """ + "document_author_documents": """ UNWIND $rows AS row - MATCH (a:Author {id: row.author_id}) + MATCH (a:DocumentAuthor { + id: row.document_author_id + }) MATCH (b:Document { path: row.document_path }) - MERGE (a)-[:AUTHORED]->(b) + MERGE (a)-[:AUTHORED_DOCUMENT]->(b) """, "document_categories": """ UNWIND $rows AS row @@ -127,15 +156,15 @@ RELATIONSHIP_QUERIES = { }) MERGE (a)-[:IN_CATEGORY]->(b) """, - "document_topics": """ + "document_keywords": """ UNWIND $rows AS row MATCH (a:Document { path: row.document_path }) - MATCH (b:Topic { - id: row.topic_id + MATCH (b:Keyword { + id: row.keyword_id }) - MERGE (a)-[:HAS_TOPIC]->(b) + MERGE (a)-[:HAS_KEYWORD]->(b) """, "document_describes_topics": """ UNWIND $rows AS row @@ -147,6 +176,16 @@ RELATIONSHIP_QUERIES = { }) MERGE (a)-[:DESCRIBES]->(b) """, + "person_theses": """ + UNWIND $rows AS row + MATCH (a:Person { + id: row.person_id + }) + MATCH (b:Thesis { + id: row.thesis_id + }) + MERGE (a)-[:AUTHORED]->(b) + """, "person_works": """ UNWIND $rows AS row MATCH (a:Person { @@ -157,6 +196,16 @@ RELATIONSHIP_QUERIES = { }) MERGE (a)-[:HAS_WORK]->(b) """, + "thesis_documents": """ + UNWIND $rows AS row + MATCH (a:Thesis { + id: row.thesis_id + }) + MATCH (b:Document { + path: row.document_path + }) + MERGE (a)-[:EVIDENCED_BY]->(b) + """, "work_documents": """ UNWIND $rows AS row MATCH (a:Work { @@ -214,6 +263,18 @@ def create_constraints( ).consume() +def drop_legacy_constraints( + session, +) -> None: + for name in LEGACY_CONSTRAINTS: + session.run( + ( + f"DROP CONSTRAINT " + f"{name} IF EXISTS" + ) + ).consume() + + def clear_managed_graph( session, ) -> None: @@ -252,7 +313,9 @@ def graph_counts( MATCH (n) WHERE n.graph_scope = $scope UNWIND labels(n) AS label - RETURN label, count(*) AS count + RETURN + label, + count(*) AS count ORDER BY label """, scope=GRAPH_SCOPE, @@ -295,16 +358,25 @@ def graph_counts( return nodes, relationships -def work_quality_counts( +def titled_entity_counts( session, + label: str, ) -> tuple[int, int]: + if label not in { + "Thesis", + "Work", + }: + raise ValueError( + f"Unsupported label: {label}" + ) + record = session.run( - """ - MATCH (w:Work) - WHERE w.graph_scope = $scope + f""" + MATCH (n:{label}) + WHERE n.graph_scope = $scope RETURN - count(w) AS total, - count(w.title) AS with_title + count(n) AS total, + count(n.title) AS with_title """, scope=GRAPH_SCOPE, ).single() @@ -330,10 +402,12 @@ def print_payload_summary( for key in ( "people", - "authors", + "document_authors", "documents", "categories", "topics", + "keywords", + "theses", "works", ): print( @@ -341,6 +415,25 @@ def print_payload_summary( f"{len(payload[key]):>6}" ) + print() + print("Relationships") + + for key in ( + "person_documents", + "document_author_documents", + "document_categories", + "document_keywords", + "document_describes_topics", + "person_theses", + "person_works", + "thesis_documents", + "work_documents", + ): + print( + f"{key:<30}" + f"{len(payload[key]):>6}" + ) + print("=" * 60) @@ -374,16 +467,31 @@ def print_database_summary( f"{count:>6}" ) - total, with_title = ( - work_quality_counts( - session + thesis_total, thesis_with_title = ( + titled_entity_counts( + session, + "Thesis", + ) + ) + + work_total, work_with_title = ( + titled_entity_counts( + session, + "Work", ) ) print() + print( + "Theses with extracted title: " + f"{thesis_with_title}/" + f"{thesis_total}" + ) + print( "Works with extracted title: " - f"{with_title}/{total}" + f"{work_with_title}/" + f"{work_total}" ) print("=" * 60) @@ -427,6 +535,10 @@ def main() -> None: with driver.session( database=settings.database ) as session: + drop_legacy_constraints( + session + ) + create_constraints( session ) diff --git a/tools/graphrag/graph_search.py b/tools/graphrag/graph_search.py index 72da71e..320611e 100644 --- a/tools/graphrag/graph_search.py +++ b/tools/graphrag/graph_search.py @@ -12,8 +12,8 @@ DEFAULT_LIMIT = 40 DEFAULT_MAX_ENTITIES = 8 RELATIONSHIP_TYPES = ( - "HAS_DOCUMENT|AUTHORED|IN_CATEGORY|HAS_TOPIC|" - "DESCRIBES|HAS_WORK|EVIDENCED_BY" + "HAS_DOCUMENT|AUTHORED|IN_CATEGORY|HAS_KEYWORD|" + "DESCRIBES|HAS_WORK|EVIDENCED_BY|AUTHORED_DOCUMENT" ) CATALOG_QUERY = """ @@ -21,19 +21,23 @@ MATCH (n) WHERE n.graph_scope = $scope AND ( n:Person + OR n:Thesis + OR n:Work + OR n:Keyword OR n:Topic OR n:Category - OR n:Author - OR n:Work + OR n:DocumentAuthor ) RETURN elementId(n) AS element_id, CASE WHEN n:Person THEN 'Person' + WHEN n:Thesis THEN 'Thesis' + WHEN n:Work THEN 'Work' + WHEN n:Keyword THEN 'Keyword' WHEN n:Topic THEN 'Topic' WHEN n:Category THEN 'Category' - WHEN n:Author THEN 'Author' - WHEN n:Work THEN 'Work' + WHEN n:DocumentAuthor THEN 'DocumentAuthor' ELSE 'Other' END AS kind, n.id AS id, @@ -42,7 +46,10 @@ RETURN n.heading AS heading, n.start_year AS start_year, n.year AS year, + n.thesis_type AS thesis_type, + n.thesis_type_label AS thesis_type_label, n.work_type AS work_type, + n.work_type_label AS work_type_label, n.person_kind AS person_kind, n.source_document AS source_document ORDER BY kind, name, title, heading @@ -58,12 +65,17 @@ WHERE all(node IN nodes(p) WHERE node.graph_scope = $scope) 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) + ('Person' IN $target_kinds AND node:Person) + OR ('Thesis' IN $target_kinds AND node:Thesis) OR ('Work' IN $target_kinds AND node:Work) + OR ('Keyword' IN $target_kinds AND node:Keyword) + OR ('Topic' IN $target_kinds AND node:Topic) + OR ('Category' IN $target_kinds AND node:Category) + OR ( + 'DocumentAuthor' IN $target_kinds + AND node:DocumentAuthor + ) OR ('Document' IN $target_kinds AND node:Document) - OR ('Person' IN $target_kinds AND node:Person) ) ) RETURN @@ -71,11 +83,13 @@ RETURN element_id: elementId(node), kind: CASE WHEN node:Person THEN 'Person' - WHEN node:Document THEN 'Document' + WHEN node:Thesis THEN 'Thesis' WHEN node:Work THEN 'Work' + WHEN node:Document THEN 'Document' + WHEN node:Keyword THEN 'Keyword' WHEN node:Topic THEN 'Topic' WHEN node:Category THEN 'Category' - WHEN node:Author THEN 'Author' + WHEN node:DocumentAuthor THEN 'DocumentAuthor' ELSE 'Other' END, id: node.id, @@ -85,11 +99,17 @@ RETURN path: node.path, start_year: node.start_year, year: node.year, + thesis_type: node.thesis_type, + thesis_type_label: node.thesis_type_label, work_type: node.work_type, + work_type_label: node.work_type_label, + person_kind: node.person_kind, source_document: node.source_document, source_url: node.source_url }}] AS nodes, - [relationship IN relationships(p) | type(relationship)] AS relationships, + [relationship IN relationships(p) | + type(relationship) + ] AS relationships, length(p) AS hops ORDER BY hops LIMIT $limit @@ -104,19 +124,18 @@ 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:Thesis THEN 'Thesis' WHEN node:Work THEN 'Work' + WHEN node:Document THEN 'Document' + WHEN node:Keyword THEN 'Keyword' WHEN node:Topic THEN 'Topic' WHEN node:Category THEN 'Category' - WHEN node:Author THEN 'Author' + WHEN node:DocumentAuthor THEN 'DocumentAuthor' ELSE 'Other' END, id: node.id, @@ -126,18 +145,132 @@ RETURN path: node.path, start_year: node.start_year, year: node.year, + thesis_type: node.thesis_type, + thesis_type_label: node.thesis_type_label, work_type: node.work_type, + work_type_label: node.work_type_label, + person_kind: node.person_kind, source_document: node.source_document, source_url: node.source_url }}] AS nodes, - [relationship IN relationships(p) | type(relationship)] AS relationships, + [relationship IN relationships(p) | + type(relationship) + ] AS relationships, length(p) AS hops ORDER BY hops LIMIT $limit """ +THESIS_FACTS_QUERY = """ +MATCH + (person:Person) + -[:AUTHORED]-> + (thesis:Thesis) + -[:EVIDENCED_BY]-> + (document:Document) +WHERE person.graph_scope = $scope + AND thesis.graph_scope = $scope + AND document.graph_scope = $scope + AND ( + size($person_ids) = 0 + OR elementId(person) IN $person_ids + ) + AND ( + size($thesis_ids) = 0 + OR elementId(thesis) IN $thesis_ids + ) + AND ( + size($years) = 0 + OR thesis.year IN $years + ) + AND ( + size($thesis_types) = 0 + OR thesis.thesis_type IN $thesis_types + ) +RETURN + elementId(person) AS person_element_id, + person.name AS person, + person.start_year AS start_year, + elementId(thesis) AS thesis_element_id, + thesis.id AS thesis_id, + thesis.title AS title, + thesis.heading AS heading, + thesis.year AS year, + thesis.thesis_type AS thesis_type, + thesis.thesis_type_label AS thesis_type_label, + document.path AS document_path, + document.source_url AS source_url, + thesis.source_chunk_ids AS source_chunk_ids +ORDER BY person, year, thesis_type, title +LIMIT $limit +""" + +THESIS_PATHS_QUERY = """ +MATCH p= + (person:Person) + -[:AUTHORED]-> + (thesis:Thesis) + -[:EVIDENCED_BY]-> + (document:Document) +WHERE person.graph_scope = $scope + AND thesis.graph_scope = $scope + AND document.graph_scope = $scope + AND ( + size($person_ids) = 0 + OR elementId(person) IN $person_ids + ) + AND ( + size($thesis_ids) = 0 + OR elementId(thesis) IN $thesis_ids + ) + AND ( + size($years) = 0 + OR thesis.year IN $years + ) + AND ( + size($thesis_types) = 0 + OR thesis.thesis_type IN $thesis_types + ) +RETURN + [node IN nodes(p) | { + element_id: elementId(node), + kind: CASE + WHEN node:Person THEN 'Person' + WHEN node:Thesis THEN 'Thesis' + WHEN node:Document THEN 'Document' + 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, + thesis_type: node.thesis_type, + thesis_type_label: node.thesis_type_label, + 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, + thesis.year, + thesis.thesis_type, + thesis.title +LIMIT $limit +""" + WORK_FACTS_QUERY = """ -MATCH (person:Person)-[:HAS_WORK]->(work:Work)-[:EVIDENCED_BY]->(document:Document) +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 @@ -175,30 +308,13 @@ 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) +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 @@ -223,8 +339,8 @@ RETURN element_id: elementId(node), kind: CASE WHEN node:Person THEN 'Person' - WHEN node:Document THEN 'Document' WHEN node:Work THEN 'Work' + WHEN node:Document THEN 'Document' ELSE 'Other' END, id: node.id, @@ -235,26 +351,180 @@ RETURN start_year: node.start_year, year: node.year, work_type: node.work_type, + work_type_label: node.work_type_label, source_document: node.source_document, source_url: node.source_url }] AS nodes, - [relationship IN relationships(p) | type(relationship)] AS relationships, + [relationship IN relationships(p) | + type(relationship) + ] AS relationships, length(p) AS hops -ORDER BY person.name, work.year, work.work_type, work.title +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) +START_YEAR_FACTS_QUERY = """ +MATCH (person:Person) +WHERE person.graph_scope = $scope + AND ( + size($person_ids) = 0 + OR elementId(person) IN $person_ids + ) + AND ( + size($years) = 0 + OR person.start_year IN $years + ) +OPTIONAL MATCH (document:Document) +WHERE document.graph_scope = $scope + AND document.path = person.source_document +RETURN + elementId(person) AS person_element_id, + person.name AS person, + person.start_year AS start_year, + person.source_document AS document_path, + document.source_url AS source_url +ORDER BY person +LIMIT $limit +""" + +PERSON_KEYWORD_FACTS_QUERY = """ +MATCH + (person:Person) + -[:HAS_DOCUMENT]-> + (document:Document) + -[:HAS_KEYWORD]-> + (keyword:Keyword) +WHERE person.graph_scope = $scope + AND document.graph_scope = $scope + AND keyword.graph_scope = $scope + AND ( + size($person_ids) = 0 + OR elementId(person) IN $person_ids + ) + AND ( + size($keyword_ids) = 0 + OR elementId(keyword) IN $keyword_ids + ) +WITH + person, + keyword, + collect(DISTINCT document) AS documents +WITH + person, + keyword, + head( + [ + document IN documents + WHERE document.path = person.source_document + ] + + documents + ) AS document +RETURN + elementId(person) AS person_element_id, + person.name AS person, + elementId(keyword) AS keyword_element_id, + keyword.id AS keyword_id, + keyword.name AS keyword, + document.path AS document_path, + document.source_url AS source_url +ORDER BY person, keyword +LIMIT $limit +""" + +PERSON_CATEGORY_FACTS_QUERY = """ +MATCH + (person:Person) + -[:HAS_DOCUMENT]-> + (document:Document) + -[:IN_CATEGORY]-> + (category:Category) +WHERE person.graph_scope = $scope + AND document.graph_scope = $scope + AND category.graph_scope = $scope + AND ( + size($person_ids) = 0 + OR elementId(person) IN $person_ids + ) + AND ( + size($category_ids) = 0 + OR elementId(category) IN $category_ids + ) +WITH + person, + category, + collect(DISTINCT document) AS documents +WITH + person, + category, + head( + [ + document IN documents + WHERE document.path = person.source_document + ] + + documents + ) AS document +RETURN + elementId(person) AS person_element_id, + person.name AS person, + elementId(category) AS category_element_id, + category.name AS category, + category.category_type AS category_type, + category.year AS category_year, + document.path AS document_path, + document.source_url AS source_url +ORDER BY person, category +LIMIT $limit +""" + +SHARED_KEYWORD_FACTS_QUERY = """ +MATCH + (left:Person) + -[:HAS_DOCUMENT]-> + (left_document:Document) + -[:HAS_KEYWORD]-> + (keyword:Keyword) +MATCH + (right:Person) + -[:HAS_DOCUMENT]-> + (right_document:Document) + -[:HAS_KEYWORD]-> + (keyword) 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 keyword.graph_scope = $scope AND elementId(left) IN $person_ids AND elementId(right) IN $person_ids AND elementId(left) < elementId(right) +WITH + left, + right, + keyword, + collect(DISTINCT left_document) AS left_documents, + collect(DISTINCT right_document) AS right_documents +WITH + left, + right, + keyword, + head( + [ + document IN left_documents + WHERE document.path = left.source_document + ] + + left_documents + ) AS left_document, + head( + [ + document IN right_documents + WHERE document.path = right.source_document + ] + + right_documents + ) AS right_document RETURN elementId(left) AS left_person_element_id, left.name AS left_person, @@ -264,15 +534,140 @@ RETURN 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 + elementId(keyword) AS keyword_element_id, + keyword.id AS keyword_id, + keyword.name AS keyword +ORDER BY left_person, right_person, keyword +LIMIT $limit +""" + +SHARED_CATEGORY_FACTS_QUERY = """ +MATCH + (left:Person) + -[:HAS_DOCUMENT]-> + (left_document:Document) + -[:IN_CATEGORY]-> + (category:Category) +MATCH + (right:Person) + -[:HAS_DOCUMENT]-> + (right_document:Document) + -[:IN_CATEGORY]-> + (category) +WHERE left.graph_scope = $scope + AND right.graph_scope = $scope + AND left_document.graph_scope = $scope + AND right_document.graph_scope = $scope + AND category.graph_scope = $scope + AND elementId(left) IN $person_ids + AND elementId(right) IN $person_ids + AND elementId(left) < elementId(right) +WITH + left, + right, + category, + collect(DISTINCT left_document) AS left_documents, + collect(DISTINCT right_document) AS right_documents +WITH + left, + right, + category, + head( + [ + document IN left_documents + WHERE document.path = left.source_document + ] + + left_documents + ) AS left_document, + head( + [ + document IN right_documents + WHERE document.path = right.source_document + ] + + right_documents + ) AS right_document +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(category) AS category_element_id, + category.name AS category, + category.category_type AS category_type, + category.year AS category_year +ORDER BY left_person, right_person, category +LIMIT $limit +""" + +SHARED_DOCUMENT_AUTHOR_FACTS_QUERY = """ +MATCH + (author:DocumentAuthor) + -[:AUTHORED_DOCUMENT]-> + (left_document:Document) + <-[:HAS_DOCUMENT]- + (left:Person) +MATCH + (author) + -[:AUTHORED_DOCUMENT]-> + (right_document:Document) + <-[:HAS_DOCUMENT]- + (right:Person) +WHERE author.graph_scope = $scope + AND left.graph_scope = $scope + AND right.graph_scope = $scope + AND left_document.graph_scope = $scope + AND right_document.graph_scope = $scope + AND elementId(left) IN $person_ids + AND elementId(right) IN $person_ids + AND elementId(left) < elementId(right) +WITH + left, + right, + author, + collect(DISTINCT left_document) AS left_documents, + collect(DISTINCT right_document) AS right_documents +WITH + left, + right, + author, + head( + [ + document IN left_documents + WHERE document.path = left.source_document + ] + + left_documents + ) AS left_document, + head( + [ + document IN right_documents + WHERE document.path = right.source_document + ] + + right_documents + ) AS right_document +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(author) AS document_author_element_id, + author.id AS document_author_id, + author.name AS document_author +ORDER BY left_person, right_person, document_author LIMIT $limit """ -def normalize_for_match(value: str) -> str: +def normalize_for_match( + value: str, +) -> str: decomposed = unicodedata.normalize( "NFKD", str(value), @@ -281,7 +676,9 @@ def normalize_for_match(value: str) -> str: without_marks = "".join( character for character in decomposed - if not unicodedata.combining(character) + if not unicodedata.combining( + character + ) ) normalized = re.sub( @@ -297,8 +694,12 @@ def normalize_for_match(value: str) -> str: ).strip() -def normalized_tokens(value: str) -> list[str]: - normalized = normalize_for_match(value) +def normalized_tokens( + value: str, +) -> list[str]: + normalized = normalize_for_match( + value + ) if not normalized: return [] @@ -311,7 +712,9 @@ def has_stem( stem: str, ) -> bool: return any( - token.startswith(stem) + token.startswith( + stem + ) for token in tokens ) @@ -319,9 +722,14 @@ def has_stem( def extract_query_constraints( query: str, ) -> dict[str, Any]: - normalized = normalize_for_match(query) + normalized = normalize_for_match( + query + ) + tokens = normalized.split() - token_set = set(tokens) + token_set = set( + tokens + ) years = sorted( { @@ -333,8 +741,6 @@ def extract_query_constraints( } ) - work_types: list[str] = [] - has_work_word = has_stem( tokens, "prac", @@ -345,30 +751,56 @@ def extract_query_constraints( "projekt", ) + thesis_types: list[str] = [] + work_types: list[str] = [] + if ( ( - has_stem(tokens, "bakalar") + has_stem( + tokens, + "bakalar", + ) and has_work_word ) or "bp" in token_set ): - work_types.append( + thesis_types.append( "bachelor_thesis" ) if ( ( - has_stem(tokens, "diplom") + has_stem( + tokens, + "diplom", + ) and has_work_word ) or "dp" in token_set ): - work_types.append( + thesis_types.append( "diploma_thesis" ) if ( - has_stem(tokens, "diplom") + ( + has_stem( + tokens, + "dizert", + ) + and has_work_word + ) + or "phd" in token_set + ): + thesis_types.append( + "doctoral_thesis" + ) + + if ( + has_stem( + tokens, + "diplom", + ) and has_project_word ): work_types.append( @@ -376,7 +808,10 @@ def extract_query_constraints( ) if ( - has_stem(tokens, "timov") + has_stem( + tokens, + "timov", + ) and has_project_word ): work_types.append( @@ -384,24 +819,16 @@ def extract_query_constraints( ) if ( - has_stem(tokens, "vyskum") + 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" @@ -412,77 +839,239 @@ def extract_query_constraints( "research_project" ) + thesis_types = list( + dict.fromkeys( + thesis_types + ) + ) + work_types = list( - dict.fromkeys(work_types) + dict.fromkeys( + work_types + ) ) mentions_start_year = ( - has_stem(tokens, "zaciat") - and has_stem(tokens, "stud") + has_stem( + tokens, + "zaciat", + ) + and has_stem( + tokens, + "stud", + ) ) - mentions_topic = ( - has_stem(tokens, "tem") + mentions_keyword = ( + has_stem( + tokens, + "kluc", + ) or "tag" in token_set or "tagy" in token_set ) - mentions_category = has_stem( - tokens, - "kategori", + mentions_topic = ( + has_stem( + tokens, + "tem", + ) ) - mentions_author = ( - has_stem(tokens, "autor") - or has_stem(tokens, "veduc") + mentions_category = ( + has_stem( + tokens, + "kategori", + ) ) - mentions_title = has_stem( - tokens, - "nazov", + mentions_document_author = ( + has_stem( + tokens, + "autor", + ) + and ( + has_stem( + tokens, + "dokument", + ) + or has_stem( + tokens, + "strank", + ) + or "wiki" in token_set + ) + ) + + mentions_title = ( + has_stem( + tokens, + "nazov", + ) ) mentions_shared = ( - has_stem(tokens, "spoloc") - or has_stem(tokens, "rovnak") + has_stem( + tokens, + "spoloc", + ) + or has_stem( + tokens, + "rovnak", + ) + ) + + mentions_person_answer = ( + "kto" in token_set + or ( + has_stem( + tokens, + "ktor", + ) + and ( + has_stem( + tokens, + "student", + ) + or has_stem( + tokens, + "osob", + ) + or has_stem( + tokens, + "clov", + ) + ) + ) + ) + + mentions_thesis = ( + bool( + thesis_types + ) + or ( + has_stem( + tokens, + "zaverec", + ) + and has_work_word + ) + ) + + mentions_work = bool( + work_types ) target_kinds: list[str] = [] + if mentions_keyword: + target_kinds.append( + "Keyword" + ) + if mentions_topic: - target_kinds.append("Topic") + target_kinds.extend( + [ + "Keyword", + "Topic", + ] + ) if mentions_category: - target_kinds.append("Category") + target_kinds.append( + "Category" + ) - if mentions_author: - target_kinds.append("Author") + if mentions_document_author: + target_kinds.append( + "DocumentAuthor" + ) + + if mentions_thesis: + target_kinds.append( + "Thesis" + ) + + if mentions_work: + target_kinds.append( + "Work" + ) if ( mentions_title - or work_types - or ( - years - and not mentions_start_year - ) + and not mentions_thesis + and not mentions_work ): - target_kinds.append("Work") + target_kinds.extend( + [ + "Thesis", + "Work", + ] + ) + + if ( + years + and not mentions_start_year + and not mentions_thesis + and not mentions_work + ): + target_kinds.extend( + [ + "Thesis", + "Work", + ] + ) + + if ( + mentions_person_answer + and not mentions_shared + ): + target_kinds = [ + "Person" + ] return { "years": years, - "work_types": work_types, + "thesis_types": ( + thesis_types + ), + "work_types": ( + work_types + ), "mentions_start_year": ( mentions_start_year ), - "mentions_topic": mentions_topic, + "mentions_keyword": ( + mentions_keyword + ), + "mentions_topic": ( + mentions_topic + ), "mentions_category": ( mentions_category ), - "mentions_author": mentions_author, - "mentions_title": mentions_title, - "mentions_shared": mentions_shared, + "mentions_document_author": ( + mentions_document_author + ), + "mentions_title": ( + mentions_title + ), + "mentions_shared": ( + mentions_shared + ), + "mentions_person_answer": ( + mentions_person_answer + ), + "mentions_thesis": ( + mentions_thesis + ), + "mentions_work": ( + mentions_work + ), "target_kinds": list( - dict.fromkeys(target_kinds) + dict.fromkeys( + target_kinds + ) ), } @@ -491,7 +1080,9 @@ def entity_aliases( entity: dict[str, Any], ) -> list[str]: kind = str( - entity.get("kind") + entity.get( + "kind" + ) or "" ) @@ -499,28 +1090,44 @@ def entity_aliases( if kind in { "Person", - "Author", + "DocumentAuthor", }: values.append( - entity.get("name") + entity.get( + "name" + ) ) - elif kind == "Topic": + elif kind in { + "Keyword", + "Topic", + }: values.extend( [ - entity.get("name"), - entity.get("id"), + entity.get( + "name" + ), + entity.get( + "id" + ), ] ) elif kind == "Category": values.append( - entity.get("name") + entity.get( + "name" + ) ) - elif kind == "Work": + elif kind in { + "Thesis", + "Work", + }: values.append( - entity.get("title") + entity.get( + "title" + ) ) aliases: list[str] = [] @@ -534,10 +1141,14 @@ def entity_aliases( ) if normalized: - aliases.append(normalized) + aliases.append( + normalized + ) return list( - dict.fromkeys(aliases) + dict.fromkeys( + aliases + ) ) @@ -575,40 +1186,57 @@ def alias_match_score( query_tokens: set[str], alias: str, ) -> float: - alias_tokens = alias.split() + alias_tokens = ( + alias.split() + ) if not alias_tokens: return 0.0 - if len(alias_tokens) == 1: + if len( + alias_tokens + ) == 1: if alias not in query_tokens: return 0.0 - return 50.0 + min( - len(alias), - 20, - ) / 100.0 + return ( + 50.0 + + min( + len(alias), + 20, + ) + / 100.0 + ) padded_query = ( - " " - + normalized_query - + " " + f" {normalized_query} " ) padded_alias = ( - " " - + alias - + " " + f" {alias} " ) - if padded_alias in padded_query: + if ( + padded_alias + in padded_query + ): score = ( 100.0 - + len(alias_tokens) * 5.0 - + min(len(alias), 100) / 100.0 + + len( + alias_tokens + ) + * 5.0 + + min( + len(alias), + 100, + ) + / 100.0 ) - if normalized_query == alias: + if ( + normalized_query + == alias + ): score += 50.0 return score @@ -619,34 +1247,155 @@ def alias_match_score( alias_token, query_token, ) - for query_token in query_tokens + for query_token + in query_tokens ) - for alias_token in alias_tokens + for alias_token + in alias_tokens ): return ( 95.0 - + len(alias_tokens) * 5.0 - + min(len(alias), 100) / 100.0 + + len( + alias_tokens + ) + * 5.0 + + min( + len(alias), + 100, + ) + / 100.0 ) return 0.0 +def suppress_nested_entity_matches( + matches: list[ + dict[str, Any] + ], +) -> list[dict[str, Any]]: + result: list[ + dict[str, Any] + ] = [] + + for item in matches: + kind = str( + item.get( + "kind" + ) + or "" + ) + + if kind not in { + "Thesis", + "Work", + }: + result.append( + item + ) + continue + + alias = str( + item.get( + "matched_alias" + ) + or "" + ) + + score = float( + item.get( + "match_score" + ) + or 0.0 + ) + + dominated = False + + for other in matches: + if other is item: + continue + + if ( + other.get( + "kind" + ) + != kind + ): + continue + + other_alias = str( + other.get( + "matched_alias" + ) + or "" + ) + + other_score = float( + other.get( + "match_score" + ) + or 0.0 + ) + + if ( + not alias + or not other_alias + or alias == other_alias + ): + continue + + if ( + alias in other_alias + and len( + other_alias + ) + > len( + alias + ) + and other_score + > score + ): + dominated = True + break + + if not dominated: + result.append( + item + ) + + return result + + def resolve_entities( query: str, - catalog: list[dict[str, Any]], + catalog: list[ + dict[str, Any] + ], *, - max_entities: int = DEFAULT_MAX_ENTITIES, + max_entities: int = ( + DEFAULT_MAX_ENTITIES + ), + preferred_kinds: Iterable[ + str + ] = (), ) -> list[dict[str, Any]]: - normalized_query = normalize_for_match( - query + normalized_query = ( + normalize_for_match( + query + ) ) query_tokens = set( normalized_query.split() ) - matches: list[dict[str, Any]] = [] + preferred = set( + preferred_kinds + ) + + matches: list[ + dict[str, Any] + ] = [] for entity in catalog: best_score = 0.0 @@ -655,12 +1404,23 @@ def resolve_entities( for alias in entity_aliases( entity ): - score = alias_match_score( - normalized_query, - query_tokens, - alias, + score = ( + alias_match_score( + normalized_query, + query_tokens, + alias, + ) ) + if ( + score > 0 + and entity.get( + "kind" + ) + in preferred + ): + score += 2.0 + if score > best_score: best_score = score best_alias = alias @@ -668,53 +1428,104 @@ def resolve_entities( if best_score <= 0.0: continue - item = dict(entity) - item["matched_alias"] = best_alias - item["match_score"] = round( + item = dict( + entity + ) + + item[ + "matched_alias" + ] = best_alias + + item[ + "match_score" + ] = round( best_score, 3, ) - matches.append(item) + matches.append( + item + ) person_aliases = { - str(item.get("matched_alias") or "") + str( + item.get( + "matched_alias" + ) + or "" + ) for item in matches - if item.get("kind") == "Person" + if item.get( + "kind" + ) + == "Person" } matches = [ item for item in matches if not ( - item.get("kind") == "Author" + item.get( + "kind" + ) + == "DocumentAuthor" and str( - item.get("matched_alias") + item.get( + "matched_alias" + ) or "" ) in person_aliases ) ] + matches = ( + suppress_nested_entity_matches( + matches + ) + ) + matches.sort( key=lambda item: ( - -float(item["match_score"]), - str(item.get("kind") or ""), + -float( + item[ + "match_score" + ] + ), str( - item.get("name") - or item.get("title") - or item.get("id") + 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() + selected: list[ + dict[str, Any] + ] = [] + + seen_ids: set[ + str + ] = set() for item in matches: element_id = str( - item.get("element_id") + item.get( + "element_id" + ) or "" ) @@ -724,10 +1535,18 @@ def resolve_entities( if element_id in seen_ids: continue - seen_ids.add(element_id) - selected.append(item) + seen_ids.add( + element_id + ) - if len(selected) >= max_entities: + selected.append( + item + ) + + if ( + len(selected) + >= max_entities + ): break return selected @@ -736,12 +1555,19 @@ def resolve_entities( def rows_from_result( result: Any, ) -> list[dict[str, Any]]: - rows: list[dict[str, Any]] = [] + rows: list[ + dict[str, Any] + ] = [] for record in result: - if isinstance(record, dict): + if isinstance( + record, + dict, + ): rows.append( - dict(record) + dict( + record + ) ) continue @@ -751,14 +1577,20 @@ def rows_from_result( None, ) - if callable(data): + if callable( + data + ): rows.append( - dict(data()) + dict( + data() + ) ) continue rows.append( - dict(record) + dict( + record + ) ) return rows @@ -778,10 +1610,17 @@ def run_rows( def deduplicate_paths( - paths: list[dict[str, Any]], + paths: list[ + dict[str, Any] + ], ) -> list[dict[str, Any]]: - result: list[dict[str, Any]] = [] - seen: set[str] = set() + result: list[ + dict[str, Any] + ] = [] + + seen: set[ + str + ] = set() for path in paths: key = json.dumps( @@ -790,9 +1629,11 @@ def deduplicate_paths( "nodes", [], ), - "relationships": path.get( - "relationships", - [], + "relationships": ( + path.get( + "relationships", + [], + ) ), }, ensure_ascii=False, @@ -803,16 +1644,70 @@ def deduplicate_paths( if key in seen: continue - seen.add(key) - result.append(path) + seen.add( + key + ) + + result.append( + path + ) + + return result + + +def deduplicate_rows( + rows: list[ + dict[str, Any] + ], + key_fields: tuple[ + str, + ..., + ], +) -> list[dict[str, Any]]: + result: list[ + dict[str, Any] + ] = [] + + seen: set[ + tuple[str, ...] + ] = set() + + for row in rows: + key = tuple( + str( + row.get( + field + ) + or "" + ) + for field + in key_fields + ) + + if key in seen: + continue + + seen.add( + key + ) + + result.append( + row + ) return result def collect_documents( - resolved_entities: list[dict[str, Any]], - paths: list[dict[str, Any]], - facts: list[dict[str, Any]], + resolved_entities: list[ + dict[str, Any] + ], + paths: list[ + dict[str, Any] + ], + facts: list[ + dict[str, Any] + ], ) -> list[dict[str, Any]]: documents: dict[ str, @@ -831,12 +1726,16 @@ def collect_documents( if not value: return - current = documents.setdefault( - value, - { - "document_path": value, - "source_url": None, - }, + current = ( + documents.setdefault( + value, + { + "document_path": ( + value + ), + "source_url": None, + }, + ) ) if ( @@ -845,11 +1744,15 @@ def collect_documents( "source_url" ) ): - current["source_url"] = str( + current[ + "source_url" + ] = str( source_url ) - for entity in resolved_entities: + for entity in ( + resolved_entities + ): add_document( entity.get( "source_document" @@ -867,10 +1770,19 @@ def collect_documents( ): continue - if node.get("kind") == "Document": + if ( + node.get( + "kind" + ) + == "Document" + ): add_document( - node.get("path"), - node.get("source_url"), + node.get( + "path" + ), + node.get( + "source_url" + ), ) add_document( @@ -884,7 +1796,9 @@ def collect_documents( fact.get( "document_path" ), - fact.get("source_url"), + fact.get( + "source_url" + ), ) add_document( @@ -907,56 +1821,235 @@ def collect_documents( return sorted( documents.values(), - key=lambda item: item[ - "document_path" - ], + 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]], + thesis_rows: list[ + dict[str, Any] + ], + work_rows: list[ + dict[str, Any] + ], + start_year_rows: list[ + dict[str, Any] + ], + person_keyword_rows: list[ + dict[str, Any] + ], + person_category_rows: list[ + dict[str, Any] + ], + shared_keyword_rows: list[ + dict[str, Any] + ], + shared_category_rows: list[ + dict[str, Any] + ], + shared_document_author_rows: list[ + dict[str, Any] + ], *, mentions_shared: bool, ) -> list[dict[str, Any]]: - facts: list[dict[str, Any]] = [] + facts: list[ + dict[str, Any] + ] = [] + + for row in thesis_rows: + fact = dict( + row + ) + + fact[ + "fact_type" + ] = "person_thesis" + + facts.append( + fact + ) for row in work_rows: - fact = dict(row) - fact["fact_type"] = ( - "person_work" + fact = dict( + row ) - facts.append(fact) - for row in start_year_rows: - fact = dict(row) - fact["fact_type"] = ( - "person_start_year" - ) - facts.append(fact) + fact[ + "fact_type" + ] = "person_work" - for row in shared_topic_rows: - fact = dict(row) - fact["fact_type"] = ( - "shared_topic" + facts.append( + fact + ) + + unique_start_rows = ( + deduplicate_rows( + start_year_rows, + ( + "person_element_id", + ), + ) + ) + + for row in unique_start_rows: + fact = dict( + row + ) + + fact[ + "fact_type" + ] = "person_start_year" + + facts.append( + fact + ) + + unique_person_keywords = ( + deduplicate_rows( + person_keyword_rows, + ( + "person_element_id", + "keyword_element_id", + ), + ) + ) + + for row in unique_person_keywords: + fact = dict( + row + ) + + fact[ + "fact_type" + ] = "person_keyword" + + facts.append( + fact + ) + + unique_person_categories = ( + deduplicate_rows( + person_category_rows, + ( + "person_element_id", + "category_element_id", + ), + ) + ) + + for row in unique_person_categories: + fact = dict( + row + ) + + fact[ + "fact_type" + ] = "person_category" + + facts.append( + fact + ) + + unique_shared_keywords = ( + deduplicate_rows( + shared_keyword_rows, + ( + "left_person_element_id", + "right_person_element_id", + "keyword_element_id", + ), + ) + ) + + for row in unique_shared_keywords: + fact = dict( + row + ) + + fact[ + "fact_type" + ] = "shared_keyword" + + facts.append( + fact + ) + + unique_shared_categories = ( + deduplicate_rows( + shared_category_rows, + ( + "left_person_element_id", + "right_person_element_id", + "category_element_id", + ), + ) + ) + + for row in unique_shared_categories: + fact = dict( + row + ) + + fact[ + "fact_type" + ] = "shared_category" + + facts.append( + fact + ) + + unique_shared_authors = ( + deduplicate_rows( + shared_document_author_rows, + ( + "left_person_element_id", + "right_person_element_id", + "document_author_element_id", + ), + ) + ) + + for row in unique_shared_authors: + fact = dict( + row + ) + + fact[ + "fact_type" + ] = ( + "shared_document_author" + ) + + facts.append( + fact ) - facts.append(fact) if ( mentions_shared - and len(start_year_rows) >= 2 + and len( + unique_start_rows + ) + >= 2 ): ordered = sorted( - start_year_rows, + unique_start_rows, key=lambda row: str( - row.get("person") + row.get( + "person" + ) or "" ), ) years = [ - row.get("start_year") + row.get( + "start_year" + ) for row in ordered ] @@ -967,9 +2060,17 @@ def build_facts( ] same = ( - len(known_years) - == len(years) - and len(set(known_years)) + len( + known_years + ) + == len( + years + ) + and len( + set( + known_years + ) + ) == 1 ) @@ -980,21 +2081,27 @@ def build_facts( ), "people": [ str( - row.get("person") + row.get( + "person" + ) or "" ) for row in ordered ], "start_years": { str( - row.get("person") + row.get( + "person" + ) or "" ): row.get( "start_year" ) for row in ordered }, - "same_start_year": same, + "same_start_year": ( + same + ), "shared_start_year": ( known_years[0] if same @@ -1010,8 +2117,12 @@ def search_graph_with_session( session: Any, query: str, *, - limit: int = DEFAULT_LIMIT, - max_entities: int = DEFAULT_MAX_ENTITIES, + limit: int = ( + DEFAULT_LIMIT + ), + max_entities: int = ( + DEFAULT_MAX_ENTITIES + ), ) -> dict[str, Any]: clean_query = str( query @@ -1033,8 +2144,10 @@ def search_graph_with_session( "max_entities musí byť aspoň 1" ) - constraints = extract_query_constraints( - clean_query + constraints = ( + extract_query_constraints( + clean_query + ) ) catalog = run_rows( @@ -1043,65 +2156,300 @@ def search_graph_with_session( scope=GRAPH_SCOPE, ) - resolved_entities = resolve_entities( - clean_query, - catalog, - max_entities=max_entities, + resolved_entities = ( + resolve_entities( + clean_query, + catalog, + max_entities=( + max_entities + ), + preferred_kinds=( + constraints[ + "target_kinds" + ] + ), + ) ) start_ids = [ - str(entity["element_id"]) - for entity in resolved_entities - if entity.get("element_id") + 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 + str( + entity[ + "element_id" + ] + ) + for entity + in resolved_entities if ( - entity.get("element_id") - and entity.get("kind") + entity.get( + "element_id" + ) + and entity.get( + "kind" + ) == "Person" ) ] - work_ids = [ - str(entity["element_id"]) - for entity in resolved_entities + thesis_ids = [ + str( + entity[ + "element_id" + ] + ) + for entity + in resolved_entities if ( - entity.get("element_id") - and entity.get("kind") + entity.get( + "element_id" + ) + and entity.get( + "kind" + ) + == "Thesis" + ) + ] + + work_ids = [ + str( + entity[ + "element_id" + ] + ) + for entity + in resolved_entities + if ( + entity.get( + "element_id" + ) + and entity.get( + "kind" + ) == "Work" ) ] + keyword_ids = [ + str( + entity[ + "element_id" + ] + ) + for entity + in resolved_entities + if ( + entity.get( + "element_id" + ) + and entity.get( + "kind" + ) + == "Keyword" + ) + ] + + category_ids = [ + str( + entity[ + "element_id" + ] + ) + for entity + in resolved_entities + if ( + entity.get( + "element_id" + ) + and entity.get( + "kind" + ) + == "Category" + ) + ] + + year_specific = bool( + constraints[ + "years" + ] + and not constraints[ + "mentions_start_year" + ] + ) + + title_specific = ( + constraints[ + "mentions_title" + ] + ) + + thesis_specific = bool( + constraints[ + "mentions_thesis" + ] + or thesis_ids + or ( + title_specific + and not constraints[ + "mentions_work" + ] + ) + or ( + year_specific + and not constraints[ + "mentions_work" + ] + ) + ) + work_specific = bool( - constraints["work_types"] - or constraints["mentions_title"] + constraints[ + "mentions_work" + ] or work_ids or ( - constraints["years"] + title_specific and not constraints[ - "mentions_start_year" + "mentions_thesis" + ] + ) + or ( + year_specific + and not constraints[ + "mentions_thesis" ] ) ) - shared_relation_query = bool( - constraints["mentions_shared"] - and len(person_ids) >= 2 + shared_keyword_query = bool( + constraints[ + "mentions_shared" + ] + and len( + person_ids + ) + >= 2 and ( - constraints["mentions_topic"] - or constraints[ - "mentions_category" + constraints[ + "mentions_keyword" ] or constraints[ - "mentions_author" + "mentions_topic" ] ) ) - paths: list[dict[str, Any]] = [] + shared_category_query = bool( + constraints[ + "mentions_shared" + ] + and len( + person_ids + ) + >= 2 + and constraints[ + "mentions_category" + ] + ) + + shared_document_author_query = ( + bool( + constraints[ + "mentions_shared" + ] + and len( + person_ids + ) + >= 2 + and constraints[ + "mentions_document_author" + ] + ) + ) + + keyword_specific = bool( + not constraints[ + "mentions_shared" + ] + and ( + keyword_ids + or ( + person_ids + and constraints[ + "mentions_keyword" + ] + ) + ) + ) + + category_specific = bool( + not constraints[ + "mentions_shared" + ] + and ( + category_ids + or ( + person_ids + and constraints[ + "mentions_category" + ] + ) + ) + ) + + structured_query = ( + thesis_specific + or work_specific + or constraints[ + "mentions_start_year" + ] + or keyword_specific + or category_specific + or shared_keyword_query + or shared_category_query + or shared_document_author_query + ) + + paths: list[ + dict[str, Any] + ] = [] + + if thesis_specific: + paths.extend( + run_rows( + session, + THESIS_PATHS_QUERY, + scope=GRAPH_SCOPE, + person_ids=( + person_ids + ), + thesis_ids=( + thesis_ids + ), + years=( + constraints[ + "years" + ] + ), + thesis_types=( + constraints[ + "thesis_types" + ] + ), + limit=limit, + ) + ) if work_specific: paths.extend( @@ -1109,9 +2457,17 @@ def search_graph_with_session( session, WORK_PATHS_QUERY, scope=GRAPH_SCOPE, - person_ids=person_ids, - work_ids=work_ids, - years=constraints["years"], + person_ids=( + person_ids + ), + work_ids=( + work_ids + ), + years=( + constraints[ + "years" + ] + ), work_types=( constraints[ "work_types" @@ -1121,60 +2477,97 @@ def search_graph_with_session( ) ) - 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, + if ( + not structured_query + and start_ids + ): + if ( + constraints[ + "mentions_shared" + ] + and len( + start_ids ) + >= 2 + ): + paths.extend( + run_rows( + session, + PAIR_PATH_QUERY, + scope=GRAPH_SCOPE, + start_ids=( + start_ids + ), + limit=limit, + ) + ) + + else: + paths.extend( + run_rows( + session, + NEIGHBORHOOD_QUERY, + scope=GRAPH_SCOPE, + start_ids=( + start_ids + ), + target_kinds=( + constraints[ + "target_kinds" + ] + ), + limit=limit, + ) + ) + + thesis_rows: list[ + dict[str, Any] + ] = [] + + if thesis_specific: + thesis_rows = run_rows( + session, + THESIS_FACTS_QUERY, + scope=GRAPH_SCOPE, + person_ids=( + person_ids + ), + thesis_ids=( + thesis_ids + ), + years=( + constraints[ + "years" + ] + ), + thesis_types=( + constraints[ + "thesis_types" + ] + ), + 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]] = [] + 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"], + person_ids=( + person_ids + ), + work_ids=( + work_ids + ), + years=( + constraints[ + "years" + ] + ), work_types=( constraints[ "work_types" @@ -1190,30 +2583,112 @@ def search_graph_with_session( 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, + 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[ + person_keyword_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, + if keyword_specific: + person_keyword_rows = ( + run_rows( + session, + PERSON_KEYWORD_FACTS_QUERY, + scope=GRAPH_SCOPE, + person_ids=( + person_ids + ), + keyword_ids=( + keyword_ids + ), + limit=limit, + ) + ) + + person_category_rows: list[ + dict[str, Any] + ] = [] + + if category_specific: + person_category_rows = ( + run_rows( + session, + PERSON_CATEGORY_FACTS_QUERY, + scope=GRAPH_SCOPE, + person_ids=( + person_ids + ), + category_ids=( + category_ids + ), + limit=limit, + ) + ) + + shared_keyword_rows: list[ + dict[str, Any] + ] = [] + + if shared_keyword_query: + shared_keyword_rows = ( + run_rows( + session, + SHARED_KEYWORD_FACTS_QUERY, + scope=GRAPH_SCOPE, + person_ids=( + person_ids + ), + limit=limit, + ) + ) + + shared_category_rows: list[ + dict[str, Any] + ] = [] + + if shared_category_query: + shared_category_rows = ( + run_rows( + session, + SHARED_CATEGORY_FACTS_QUERY, + scope=GRAPH_SCOPE, + person_ids=( + person_ids + ), + limit=limit, + ) + ) + + shared_document_author_rows: list[ + dict[str, Any] + ] = [] + + if shared_document_author_query: + shared_document_author_rows = ( + run_rows( + session, + SHARED_DOCUMENT_AUTHOR_FACTS_QUERY, + scope=GRAPH_SCOPE, + person_ids=( + person_ids + ), + limit=limit, + ) ) paths = deduplicate_paths( @@ -1221,9 +2696,14 @@ def search_graph_with_session( ) facts = build_facts( + thesis_rows, work_rows, start_year_rows, - shared_topic_rows, + person_keyword_rows, + person_category_rows, + shared_keyword_rows, + shared_category_rows, + shared_document_author_rows, mentions_shared=( constraints[ "mentions_shared" @@ -1237,7 +2717,9 @@ def search_graph_with_session( facts, ) - strategies: list[str] = [] + strategies: list[ + str + ] = [] if resolved_entities: strategies.append( @@ -1249,6 +2731,11 @@ def search_graph_with_session( "graph_traversal" ) + if thesis_rows: + strategies.append( + "thesis_constraints" + ) + if work_rows: strategies.append( "work_constraints" @@ -1259,9 +2746,29 @@ def search_graph_with_session( "start_year_constraints" ) - if shared_topic_rows: + if person_keyword_rows: strategies.append( - "shared_topic_constraints" + "person_keyword_constraints" + ) + + if person_category_rows: + strategies.append( + "person_category_constraints" + ) + + if shared_keyword_rows: + strategies.append( + "shared_keyword_constraints" + ) + + if shared_category_rows: + strategies.append( + "shared_category_constraints" + ) + + if shared_document_author_rows: + strategies.append( + "shared_document_author_constraints" ) return { @@ -1280,15 +2787,21 @@ def search_graph_with_session( def search_graph( query: str, *, - limit: int = DEFAULT_LIMIT, - max_entities: int = DEFAULT_MAX_ENTITIES, + limit: int = ( + DEFAULT_LIMIT + ), + max_entities: int = ( + DEFAULT_MAX_ENTITIES + ), ) -> dict[str, Any]: from .graph_db import ( create_neo4j_driver, get_neo4j_settings, ) - settings = get_neo4j_settings() + settings = ( + get_neo4j_settings() + ) with create_neo4j_driver( settings @@ -1296,21 +2809,29 @@ def search_graph( driver.verify_connectivity() with driver.session( - database=settings.database + database=( + settings.database + ) ) as session: - return search_graph_with_session( - session, - query, - limit=limit, - max_entities=max_entities, + return ( + search_graph_with_session( + session, + query, + limit=limit, + max_entities=( + max_entities + ), + ) ) def main() -> None: - parser = argparse.ArgumentParser() + parser = ( + argparse.ArgumentParser() + ) parser.add_argument( - "query", + "query" ) parser.add_argument( @@ -1322,7 +2843,9 @@ def main() -> None: parser.add_argument( "--max-entities", type=int, - default=DEFAULT_MAX_ENTITIES, + default=( + DEFAULT_MAX_ENTITIES + ), ) args = parser.parse_args() @@ -1330,7 +2853,9 @@ def main() -> None: result = search_graph( args.query, limit=args.limit, - max_entities=args.max_entities, + max_entities=( + args.max_entities + ), ) print( diff --git a/tools/graphrag/graphrag.py b/tools/graphrag/graphrag.py index d32d3df..3821829 100644 --- a/tools/graphrag/graphrag.py +++ b/tools/graphrag/graphrag.py @@ -9,6 +9,12 @@ from typing import Any GRAPH_SCOPE = "zpwiki" +THESIS_TYPES = { + "bachelor_thesis", + "diploma_thesis", + "doctoral_thesis", +} + WORK_TYPE_LABELS = { "bachelor_thesis": "Bakalárska práca", "diploma_thesis": "Diplomová práca", @@ -32,42 +38,92 @@ TITLE_LABELS = { "nazov prace", "nazov bakalarskej prace", "nazov diplomovej prace", + "predbezny nazov", "tema", + "navrh na temu", "navrh na nazov", "navrh na nazov bp", "navrh na nazov dp", "navrh na nazov bakalarskej prace", + "navrh na nazov bakalarkej prace", "navrh na nazov diplomovej prace", } +WEAK_TITLE_LABELS = { + "navrh na zadanie", + "navrh na zadanie bp", + "navrh na zadanie dp", + "navrh na zadanie prace", + "navrh na zadanie bakalarskej prace", + "navrh na zadanie diplomovej prace", +} + FALLBACK_TITLE_LABELS = { "bachelor_thesis": { "navrh na nazov bp", "navrh na nazov bakalarskej prace", + "navrh na nazov bakalarkej prace", + "navrh na temu", }, "diploma_thesis": { "navrh na nazov dp", "navrh na nazov diplomovej prace", + "navrh na temu", }, } -SKIP_LINE_PREFIXES = ( - "dokument:", - "sekcia:", - "rok zaciatku studia", - "meno veduceho", +TITLE_BLOCK_BOUNDARY_PREFIXES = ( "zadanie", "navrh na zadanie", "navrh na zadamie", - "ciel:", - "ciele:", - "ulohy:", - "zasobnik uloh", + "ciel", + "ciele", + "ulohy", + "uloha", + "zasobnik", "stretnutie", - "stav:", + "sttetnutie", + "virtualne stretnutie", + "stav", + "myslienky", + "napady", + "namety", + "stare zadanie", + "analyza ulohy", +) + +NON_TITLE_PREFIXES = ( + "dokument", + "sekcia", + "rok zaciatku studia", + "meno veduceho", + "veduci", + "konzultant", + "spolupraca", "repozitar", "git repozitar", "crzp", + "literatura", + "suvisiace materialy", + "praca suvisi", + "do buducnosti", + "do dalsieho", + "do buduceho", + "poznamky", + "kriteria", + "projektove stranky", + "pisomna praca", + "vystupy", + "hotove", + "ukoncil predcasne", + "praca uspesne obhajena", + "projekt zakonceny", + "dp sa preklada", + "na maise", + "je potrebne", + "todo", + "to-do", + "v roku", ) @@ -141,6 +197,69 @@ def document_kind(path: str) -> str: return "other" +def person_profile_path( + path: str, +) -> str | None: + normalized = path.replace("\\", "/") + parts = normalized.split("/") + + if ( + len(parts) >= 4 + and parts[0] == "pages" + and parts[1] == "students" + and re.fullmatch( + r"20\d{2}", + parts[2], + ) + and parts[3] + ): + return ( + "pages/students/" + + parts[2] + + "/" + + parts[3] + + "/README.md" + ) + + if ( + len(parts) >= 3 + and parts[0] == "pages" + and parts[1] == "interns" + and parts[2] + ): + return ( + "pages/interns/" + + parts[2] + + "/README.md" + ) + + return None + + +def person_kind_from_profile_path( + path: str, +) -> str | None: + normalized = path.replace("\\", "/") + profile_path = person_profile_path( + normalized + ) + + if profile_path != normalized: + return None + + if normalized.startswith( + "pages/students/" + ): + return "student" + + if normalized.startswith( + "pages/interns/" + ): + return "intern" + + return None + + def extract_full_years(value: str) -> list[int]: years = [ int(match) @@ -296,55 +415,139 @@ def clean_markdown_line( return value.strip() -def next_title_line( +def strip_list_prefix( + line: str, +) -> str: + value = clean_markdown_line( + line + ) + + value = re.sub( + r"^[-+•]\s+(?:\[[ xX-]\]\s*)?", + "", + value, + ) + + return value.strip() + + +def is_numbered_or_bulleted_line( + line: str, +) -> bool: + value = clean_markdown_line( + line + ) + + if re.match( + r"^\d+[\.\)]\s*", + value, + ): + return True + + return bool( + re.match( + r"^[-+•]\s+", + value, + ) + ) + + +def is_title_candidate( + line: str, +) -> bool: + value = strip_list_prefix( + line + ) + + if not value: + return False + + normalized = normalize_text( + value + ) + + if normalized.startswith( + NON_TITLE_PREFIXES + ): + return False + + if value.startswith( + ( + "http://", + "https://", + "[", + "./", + "../", + ) + ): + return False + + if ( + "http://" in value + or "https://" in value + ): + return False + + if re.match( + r"^\d+[\.\)]\s*", + value, + ): + return False + + if value.endswith( + ":" + ): + return False + + if len(value) < 6: + return False + + if len(value) > 180: + return False + + return True + + +def immediate_title_after_label( lines: list[str], start_index: int, ) -> str | None: for raw_line in lines[start_index:]: - line = clean_markdown_line( + if not raw_line.strip(): + continue + + if is_numbered_or_bulleted_line( + raw_line + ): + return None + + candidate = strip_list_prefix( raw_line ) - if not line: - continue - - normalized = normalize_text(line) - - if normalized.startswith( - SKIP_LINE_PREFIXES + if not is_title_candidate( + candidate ): - continue + return None - if line.startswith( - ("http://", "https://", "[") - ): - continue - - if re.match( - r"^[-•]\s+", - line, - ): - continue - - if re.match( - r"^\d+[\.\)]\s*", - line, - ): - continue - - return line + return candidate.rstrip( + " ," + ) return None -def extract_work_title( +def extract_labeled_title( texts: list[str], + labels: set[str], ) -> str | None: for text in texts: lines = text.splitlines() - for index, raw_line in enumerate(lines): - line = clean_markdown_line( + for index, raw_line in enumerate( + lines + ): + line = strip_list_prefix( raw_line ) @@ -352,76 +555,157 @@ def extract_work_title( continue before, separator, after = ( - line.partition(":") + line.partition( + ":" + ) ) - normalized_before = normalize_text( - before + if not separator: + continue + + normalized_before = ( + normalize_text( + before + ) ) - if ( - separator - and normalized_before - in TITLE_LABELS - ): - title = after.strip() + if normalized_before not in labels: + continue - if title: - return title + title = after.strip() - candidate = next_title_line( + if title: + if is_title_candidate( + title + ): + return title.rstrip( + " ," + ) + + continue + + candidate = ( + immediate_title_after_label( lines, index + 1, ) + ) - if candidate: - return candidate + if candidate: + return candidate + return None + + +def extract_leading_title( + texts: list[str], +) -> str | None: for text in texts: - lines = text.splitlines() + candidates: list[str] = [] + saw_content = False + + for raw_line in text.splitlines(): + if not raw_line.strip(): + continue - for raw_line in lines: line = clean_markdown_line( raw_line ) - - if not line: - continue - + stripped = strip_list_prefix( + raw_line + ) normalized = normalize_text( - line + stripped ) if normalized.startswith( - SKIP_LINE_PREFIXES + ( + "dokument:", + "sekcia:", + "rok zaciatku studia", + ) ): continue - if line.startswith( - ("http://", "https://", "[") + if ( + stripped.startswith( + "[" + ) + and "](" in stripped ): continue - if re.match( - r"^[-•]\s+", - line, + if normalized.startswith( + ( + "spolupraca", + "suvisiaca praca", + "suvisiaca dizertacna praca", + ) ): continue - if re.match( - r"^\d+[\.\)]\s*", - line, + saw_content = True + + if is_numbered_or_bulleted_line( + raw_line + ): + if candidates: + return candidates[-1] + + return None + + if normalized.startswith( + TITLE_BLOCK_BOUNDARY_PREFIXES + ): + if candidates: + return candidates[-1] + + return None + + if not is_title_candidate( + line ): continue - if len(line) < 8: - continue + candidates.append( + stripped.rstrip( + " ," + ) + ) - return line + if candidates: + return candidates[-1] + + if saw_content: + return None return None +def extract_work_title( + texts: list[str], +) -> str | None: + title = extract_labeled_title( + texts, + TITLE_LABELS, + ) + + if title: + return title + + title = extract_labeled_title( + texts, + WEAK_TITLE_LABELS, + ) + + if title: + return title + + return extract_leading_title( + texts + ) + + def extract_explicit_proposal_title( work_type: str, chunks: list[dict[str, Any]], @@ -459,8 +743,10 @@ def extract_explicit_proposal_title( lines = text.splitlines() - for index, raw_line in enumerate(lines): - line = clean_markdown_line( + for index, raw_line in enumerate( + lines + ): + line = strip_list_prefix( raw_line ) @@ -468,14 +754,18 @@ def extract_explicit_proposal_title( continue before, separator, after = ( - line.partition(":") + line.partition( + ":" + ) ) if not separator: continue - normalized_before = normalize_text( - before + normalized_before = ( + normalize_text( + before + ) ) if normalized_before not in labels: @@ -483,22 +773,32 @@ def extract_explicit_proposal_title( title = after.strip() - if not title: - title = next_title_line( - lines, - index + 1, + if title: + if not is_title_candidate( + title + ): + continue + + title = title.rstrip( + " ," + ) + else: + title = ( + immediate_title_after_label( + lines, + index + 1, + ) ) if not title: continue - normalized_title = normalize_text( - title + normalized_title = ( + normalize_text( + title + ) ) - if not normalized_title: - continue - candidate = candidates.setdefault( normalized_title, { @@ -597,7 +897,20 @@ def topic_key_from_document_path( ) -def topic_display_name( +def is_topic_root_document( + path: str, +) -> bool: + normalized = path.replace("\\", "/") + + return bool( + re.fullmatch( + r"pages/topics/[^/]+/README\.md", + normalized, + ) + ) + + +def keyword_display_name( value: str, ) -> str: aliases = { @@ -607,10 +920,12 @@ def topic_display_name( "qa": "QA", "ie": "IE", "lm": "LM", + "llm": "LLM", "nmt": "NMT", "pos": "POS", "gpu": "GPU", "spacy": "Spacy", + "db": "DB", } key = normalize_key(value) @@ -651,23 +966,35 @@ def extract_works( [], ) - candidates = [] + chunk_headings: dict[ + tuple[str, str], + WorkHeading, + ] = {} for path in heading_paths: - if isinstance(path, list): - candidates.extend( - str(value) - for value in path + values = ( + path + if isinstance( + path, + list, ) - elif path: - candidates.append( - str(path) + else [path] + ) + + heading = None + + for raw_heading in reversed( + values + ): + if not raw_heading: + continue + + heading = find_work_heading( + str(raw_heading) ) - for raw_heading in candidates: - heading = find_work_heading( - raw_heading - ) + if heading is not None: + break if heading is None: continue @@ -679,6 +1006,11 @@ def extract_works( ), ) + chunk_headings[key] = heading + + for key, heading in ( + chunk_headings.items() + ): entry = grouped.setdefault( key, { @@ -806,12 +1138,17 @@ def build_graph_payload( path ].append(chunk) + documents_by_path = { + str(document["path"]): document + for document in documents + } + people: dict[ str, dict[str, Any], ] = {} - authors: dict[ + document_authors: dict[ str, dict[str, Any], ] = {} @@ -831,17 +1168,29 @@ def build_graph_payload( dict[str, Any], ] = {} + keywords: dict[ + str, + dict[str, Any], + ] = {} + + theses: dict[ + str, + dict[str, Any], + ] = {} + works: dict[ str, dict[str, Any], ] = {} person_documents = set() - author_documents = set() + document_author_documents = set() document_categories = set() - document_topics = set() + document_keywords = set() document_describes_topics = set() + person_theses = set() person_works = set() + thesis_documents = set() work_documents = set() for document in documents: @@ -919,15 +1268,18 @@ def build_graph_payload( author_name ) - authors[author_id] = { - "id": author_id, - "name": author_name, - "graph_scope": ( - GRAPH_SCOPE - ), - } + if author_id not in document_authors: + document_authors[ + author_id + ] = { + "id": author_id, + "name": author_name, + "graph_scope": ( + GRAPH_SCOPE + ), + } - author_documents.add( + document_author_documents.add( ( author_id, path, @@ -960,19 +1312,18 @@ def build_graph_payload( ) for tag in doc_tags: - topic_id = normalize_key( + keyword_id = normalize_key( tag ) - current = topics.get( - topic_id - ) + if not keyword_id: + continue - if current is None: - topics[topic_id] = { - "id": topic_id, + if keyword_id not in keywords: + keywords[keyword_id] = { + "id": keyword_id, "name": ( - topic_display_name( + keyword_display_name( tag ) ), @@ -981,10 +1332,10 @@ def build_graph_payload( ), } - document_topics.add( + document_keywords.add( ( path, - topic_id, + keyword_id, ) ) @@ -995,15 +1346,21 @@ def build_graph_payload( ) if described_topic: - topics[ - described_topic - ] = { - "id": described_topic, - "name": title, - "graph_scope": ( - GRAPH_SCOPE - ), - } + if ( + described_topic not in topics + or is_topic_root_document( + path + ) + ): + topics[ + described_topic + ] = { + "id": described_topic, + "name": title, + "graph_scope": ( + GRAPH_SCOPE + ), + } document_describes_topics.add( ( @@ -1012,40 +1369,68 @@ def build_graph_payload( ) ) - if kind not in { - "student", - "intern", - }: + profile_path = person_profile_path( + path + ) + + if profile_path is None: + continue + + profile_document = documents_by_path.get( + profile_path + ) + + if profile_document is None: + continue + + person_kind = ( + person_kind_from_profile_path( + profile_path + ) + ) + + if person_kind is None: continue person_id = ( - f"person:{path}" + f"person:{profile_path}" ) - document_chunks = ( + profile_title = str( + profile_document.get( + "title", + profile_path, + ) + ) + + profile_chunks = ( chunks_by_document.get( - path, + profile_path, [], ) ) - people[person_id] = { - "id": person_id, - "name": title, - "person_kind": kind, - "start_year": ( - extract_start_year( - path, - document_chunks, - ) - if kind == "student" - else None - ), - "source_document": path, - "graph_scope": ( - GRAPH_SCOPE - ), - } + if person_id not in people: + people[person_id] = { + "id": person_id, + "name": profile_title, + "person_kind": person_kind, + "start_year": ( + extract_start_year( + profile_path, + profile_chunks, + ) + if person_kind + == "student" + else None + ), + "source_document": ( + profile_path + ), + "graph_scope": ( + GRAPH_SCOPE + ), + } person_documents.add( ( @@ -1054,42 +1439,34 @@ def build_graph_payload( ) ) - if kind != "student": + if ( + person_kind != "student" + or path != profile_path + ): continue + document_chunks = ( + chunks_by_document.get( + profile_path, + [], + ) + ) + for work in extract_works( document, document_chunks, ): - work_id = ( - "work:" - + path - + ":" - + normalize_key( - work["heading"] - ) - ) - - display_name = ( - work["title"] - or work["heading"] - ) - - works[work_id] = { - "id": work_id, - "name": display_name, + common_data = { + "name": ( + work["title"] + or work["heading"] + ), "title": work[ "title" ], "heading": work[ "heading" ], - "work_type": work[ - "work_type" - ], - "work_type_label": work[ - "work_type_label" - ], "year": work[ "year" ], @@ -1102,6 +1479,65 @@ def build_graph_payload( ), } + if work["work_type"] in THESIS_TYPES: + thesis_id = ( + "thesis:" + + path + + ":" + + normalize_key( + work["heading"] + ) + ) + + theses[thesis_id] = { + "id": thesis_id, + "thesis_type": work[ + "work_type" + ], + "thesis_type_label": ( + work[ + "work_type_label" + ] + ), + **common_data, + } + + person_theses.add( + ( + person_id, + thesis_id, + ) + ) + + thesis_documents.add( + ( + thesis_id, + path, + ) + ) + + continue + + work_id = ( + "work:" + + path + + ":" + + normalize_key( + work["heading"] + ) + ) + + works[work_id] = { + "id": work_id, + "work_type": work[ + "work_type" + ], + "work_type_label": work[ + "work_type_label" + ], + **common_data, + } + person_works.add( ( person_id, @@ -1136,8 +1572,8 @@ def build_graph_payload( people.values(), key=lambda row: row["id"], ), - "authors": sorted( - authors.values(), + "document_authors": sorted( + document_authors.values(), key=lambda row: row["id"], ), "documents": sorted( @@ -1156,6 +1592,14 @@ def build_graph_payload( topics.values(), key=lambda row: row["id"], ), + "keywords": sorted( + keywords.values(), + key=lambda row: row["id"], + ), + "theses": sorted( + theses.values(), + key=lambda row: row["id"], + ), "works": sorted( works.values(), key=lambda row: row["id"], @@ -1165,9 +1609,9 @@ def build_graph_payload( "person_id", "document_path", ), - "author_documents": relation_rows( - author_documents, - "author_id", + "document_author_documents": relation_rows( + document_author_documents, + "document_author_id", "document_path", ), "document_categories": relation_rows( @@ -1175,10 +1619,10 @@ def build_graph_payload( "document_path", "category_name", ), - "document_topics": relation_rows( - document_topics, + "document_keywords": relation_rows( + document_keywords, "document_path", - "topic_id", + "keyword_id", ), "document_describes_topics": ( relation_rows( @@ -1187,11 +1631,21 @@ def build_graph_payload( "topic_id", ) ), + "person_theses": relation_rows( + person_theses, + "person_id", + "thesis_id", + ), "person_works": relation_rows( person_works, "person_id", "work_id", ), + "thesis_documents": relation_rows( + thesis_documents, + "thesis_id", + "document_path", + ), "work_documents": relation_rows( work_documents, "work_id", diff --git a/tools/graphrag/tests/test_graph_search.py b/tools/graphrag/tests/test_graph_search.py index 44c7bc7..b385901 100644 --- a/tools/graphrag/tests/test_graph_search.py +++ b/tools/graphrag/tests/test_graph_search.py @@ -49,6 +49,7 @@ def sample_catalog() -> list[dict[str, Any]]: "heading": None, "start_year": 2016, "year": None, + "thesis_type": None, "work_type": None, "person_kind": "student", "source_document": ( @@ -65,6 +66,7 @@ def sample_catalog() -> list[dict[str, Any]]: "heading": None, "start_year": 2016, "year": None, + "thesis_type": None, "work_type": None, "person_kind": "student", "source_document": ( @@ -73,35 +75,68 @@ def sample_catalog() -> list[dict[str, Any]]: ), }, { - "element_id": "t1", - "kind": "Topic", + "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": "t2", - "kind": "Topic", + "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": "w1", - "kind": "Work", - "id": "work:jakub:dp2021", + "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 " @@ -112,10 +147,65 @@ def sample_catalog() -> list[dict[str, Any]]: "pomenovaných entít v " "slovenskom jazyku." ), - "heading": "Diplomová práca 2021", + "heading": ( + "Diplomová práca 2021" + ), "start_year": None, "year": 2021, - "work_type": "diploma_thesis", + "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/" @@ -142,47 +232,174 @@ def test_extract_query_constraints_for_diploma_thesis_year() -> None: ) ) - assert result["years"] == [ - 2021 - ] assert result[ - "work_types" + "years" + ] == [2021] + + assert result[ + "thesis_types" ] == [ "diploma_thesis" ] + + assert result[ + "work_types" + ] == [] + assert result[ "mentions_title" ] is True - assert "Work" in result[ + + 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?" + "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(), +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 = { @@ -197,170 +414,179 @@ def test_resolve_entities_finds_two_people_and_topic() -> None: "Person", "Maroš Harahus", ) in resolved + assert ( "Person", "Jakub Maruniak", ) in resolved + assert ( - "Topic", + "Keyword", "Spacy", ) in resolved + assert ( + "Topic", + "Spracovanie prirodzeného jazyka", + ) 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, - } - ] +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, + result = ( + graph_search.resolve_entities( + ( + "Čomu sa venuje " + "Maroš Harahus?" + ), + catalog, + ) ) matching = [ item for item in result - if item.get("matched_alias") + if item.get( + "matched_alias" + ) == "maros harahus" ] - assert len(matching) == 1 - assert matching[0]["kind"] == "Person" + assert len( + matching + ) == 1 -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 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" + 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(), +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"] == "Work" - and item["element_id"] == "w1" + item["kind"] + == "Thesis" + and item[ + "element_id" + ] + == "th1" 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" +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?" ), - "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, + sample_catalog(), + ) ) - work_results = [ - item + thesis_ids = [ + item[ + "element_id" + ] for item in result - if item["kind"] == "Work" + if item[ + "kind" + ] + == "Thesis" ] - assert work_results[0]["element_id"] == "w1" + assert thesis_ids == [ + "th1" + ] -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, + +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( @@ -368,9 +594,9 @@ def test_search_graph_with_session_uses_topic_path_for_two_people() -> None: graph_search.CATALOG_QUERY: ( sample_catalog() ), - graph_search.NEIGHBORHOOD_QUERY: [], - graph_search.PAIR_PATH_QUERY: [ - topic_path + graph_search.THESIS_PATHS_QUERY: [], + graph_search.THESIS_FACTS_QUERY: [ + thesis_fact ], } ) @@ -378,45 +604,116 @@ def test_search_graph_with_session_uses_topic_path_for_two_people() -> None: result = ( graph_search.search_graph_with_session( session, - "Akú spoločnú tému majú " - "Maroš Harahus a Jakub Maruniak?", + ( + "Aký bol názov diplomovej " + "práce Jakuba Maruniaka " + "v roku 2021?" + ), ) ) - assert result[ - "constraints" - ]["mentions_topic"] is True - assert len( - result["paths"] - ) == 1 - assert { - item["document_path"] - for item in result[ - "documents" + thesis_facts = [ + fact + for fact in result[ + "facts" ] - } == { - ( - "pages/students/2016/" - "maros_harahus/README.md" - ), - ( - "pages/students/2016/" - "jakub_maruniak/README.md" - ), - } + if fact[ + "fact_type" + ] + == "person_thesis" + ] - pair_calls = [ + 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.PAIR_PATH_QUERY + == graph_search.THESIS_FACTS_QUERY ] - assert len(pair_calls) == 1 - assert pair_calls[0][ - "require_topic" - ] is True + 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: @@ -424,29 +721,27 @@ def test_search_graph_with_session_uses_work_constraints() -> None: "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." + "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" ), - "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" - ) - ], + "source_chunk_ids": [], } session = FakeSession( @@ -454,7 +749,7 @@ def test_search_graph_with_session_uses_work_constraints() -> None: graph_search.CATALOG_QUERY: ( sample_catalog() ), - graph_search.NEIGHBORHOOD_QUERY: [], + graph_search.WORK_PATHS_QUERY: [], graph_search.WORK_FACTS_QUERY: [ work_fact ], @@ -464,68 +759,418 @@ def test_search_graph_with_session_uses_work_constraints() -> None: result = ( graph_search.search_graph_with_session( session, - "Aký bol názov diplomovej práce " - "Jakuba Maruniaka v roku 2021?", + ( + "Aký diplomový projekt mal " + "Jakub Maruniak " + "v roku 2020?" + ), ) ) assert any( - fact["fact_type"] + fact[ + "fact_type" + ] == "person_work" - and fact["year"] == 2021 - and fact["work_type"] - == "diploma_thesis" - for fact in result["facts"] + and fact[ + "year" + ] + == 2020 + and fact[ + "work_type" + ] + == "diploma_project" + for fact in result[ + "facts" + ] ) - work_calls = [ - parameters - for query, parameters - in session.calls - if query - == graph_search.WORK_FACTS_QUERY + +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, + }, ] - 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.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", + "person": ( + "Maroš Harahus" + ), "start_year": 2016, "document_path": ( - "pages/students/2016/" - "maros_harahus/README.md" + "maros.md" ), "source_url": None, }, { "person_element_id": "p2", - "person": "Jakub Maruniak", + "person": ( + "Jakub Maruniak" + ), "start_year": 2016, "document_path": ( - "pages/students/2016/" - "jakub_maruniak/README.md" + "jakub.md" + ), + "source_url": None, + }, + { + "person_element_id": "p2", + "person": ( + "Jakub Maruniak" + ), + "start_year": 2016, + "document_path": ( + "duplicate.md" ), "source_url": None, }, @@ -536,29 +1181,142 @@ def test_search_graph_with_session_uses_start_year_facts() -> None: result = ( graph_search.search_graph_with_session( session, - "Mali Maroš Harahus a Jakub " - "Maruniak rovnaký rok začiatku " - "štúdia?", + ( + "Mali Maroš Harahus a " + "Jakub Maruniak rovnaký " + "rok začiatku štúdia?" + ), ) ) - values = { - fact["person"]: fact[ - "start_year" + person_facts = [ + fact + for fact in result[ + "facts" + ] + if fact[ + "fact_type" ] - 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 values == { - "Maroš Harahus": 2016, - "Jakub Maruniak": 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({}) + session = FakeSession( + {} + ) with pytest.raises( ValueError, @@ -568,217 +1326,3 @@ def test_search_graph_rejects_empty_query() -> None: 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"] diff --git a/tools/graphrag/tests/test_graphrag.py b/tools/graphrag/tests/test_graphrag.py index 0d1c3d3..c2e7ec7 100644 --- a/tools/graphrag/tests/test_graphrag.py +++ b/tools/graphrag/tests/test_graphrag.py @@ -4,6 +4,8 @@ from tools.graphrag.graphrag import ( build_graph_payload, extract_explicit_proposal_title, extract_works, + person_kind_from_profile_path, + person_profile_path, ) @@ -12,12 +14,19 @@ JAKUB_PATH = ( "jakub_maruniak/README.md" ) +JAKUB_NESTED_PATH = ( + "pages/students/2016/" + "jakub_maruniak/dp2021/README.md" +) + def make_document( *, path: str = JAKUB_PATH, title: str = "Jakub Maruniak", categories: list[str] | None = None, + tags: list[str] | None = None, + author: str = "Daniel Hladek", ) -> dict: return { "path": path, @@ -25,13 +34,20 @@ def make_document( "categories": ( categories if categories is not None - else ["dp2021", "bp2019"] + else [ + "dp2021", + "bp2019", + ] ), - "tags": [ - "spacy", - "ner", - ], - "author": "Daniel Hladek", + "tags": ( + tags + if tags is not None + else [ + "spacy", + "ner", + ] + ), + "author": author, "published": True, } @@ -42,13 +58,14 @@ def make_chunk( text: str, *, path: str = JAKUB_PATH, + title: str = "Jakub Maruniak", ) -> dict: return { "chunk_id": ( f"{path}::chunk-{index}" ), "document_path": path, - "title": "Jakub Maruniak", + "title": title, "categories": [ "dp2021", "bp2019", @@ -70,6 +87,70 @@ def make_chunk( } +def test_person_profile_path_for_student_profile() -> None: + assert ( + person_profile_path( + JAKUB_PATH + ) + == JAKUB_PATH + ) + + +def test_person_profile_path_for_nested_student_document() -> None: + assert ( + person_profile_path( + JAKUB_NESTED_PATH + ) + == JAKUB_PATH + ) + + +def test_person_profile_path_for_non_person_document() -> None: + assert ( + person_profile_path( + "pages/topics/nlp/README.md" + ) + is None + ) + + +def test_person_kind_accepts_only_profile_root() -> None: + assert ( + person_kind_from_profile_path( + JAKUB_PATH + ) + == "student" + ) + + assert ( + person_kind_from_profile_path( + JAKUB_NESTED_PATH + ) + is None + ) + + +def test_person_kind_for_intern_profile() -> None: + path = ( + "pages/interns/" + "john_doe/README.md" + ) + + assert ( + person_profile_path( + path + ) + == path + ) + + assert ( + person_kind_from_profile_path( + path + ) + == "intern" + ) + + def test_explicit_dp_proposal_is_found() -> None: chunks = [ make_chunk( @@ -152,13 +233,16 @@ def test_extract_works_uses_dp_proposal_fallback() -> None: work = diploma_works[0] assert work["year"] == 2021 + assert work["title"] == ( "Anotácia a rozpoznávanie " "pomenovaných entít v " "slovenskom jazyku." ) - assert work["source_chunk_ids"] == [ + assert work[ + "source_chunk_ids" + ] == [ f"{JAKUB_PATH}::chunk-0", f"{JAKUB_PATH}::chunk-1", ] @@ -207,7 +291,9 @@ def test_existing_work_title_has_priority() -> None: "Finálny názov práce" ) - assert work["source_chunk_ids"] == [ + assert work[ + "source_chunk_ids" + ] == [ f"{JAKUB_PATH}::chunk-1" ] @@ -359,14 +445,17 @@ def test_proposal_is_not_shared_between_multiple_dp_works() -> None: == "diploma_thesis" ] - assert len(diploma_works) == 2 + assert len( + diploma_works + ) == 2 + assert all( work["title"] is None for work in diploma_works ) -def test_graph_payload_contains_fallback_title() -> None: +def test_graph_payload_contains_thesis_with_fallback_title() -> None: document = make_document() chunks = [ @@ -401,22 +490,556 @@ def test_graph_payload_contains_fallback_title() -> None: chunks, ) - diploma_work = next( - work - for work in payload["works"] - if work["work_type"] + diploma_thesis = next( + thesis + for thesis + in payload["theses"] + if thesis["thesis_type"] == "diploma_thesis" ) - assert diploma_work["title"] == ( + assert diploma_thesis[ + "title" + ] == ( "Anotácia a rozpoznávanie " "pomenovaných entít v " "slovenskom jazyku." ) - assert diploma_work[ + assert diploma_thesis[ "source_chunk_ids" ] == [ f"{JAKUB_PATH}::chunk-0", f"{JAKUB_PATH}::chunk-1", ] + + assert payload[ + "works" + ] == [] + + +def test_tags_become_keywords_not_topics() -> None: + document = make_document( + tags=[ + "spacy", + "ner", + "rag", + ] + ) + + payload = build_graph_payload( + [document], + [], + ) + + keyword_ids = { + keyword["id"] + for keyword + in payload["keywords"] + } + + assert keyword_ids == { + "spacy", + "ner", + "rag", + } + + assert payload[ + "topics" + ] == [] + + assert len( + payload[ + "document_keywords" + ] + ) == 3 + + +def test_topic_is_created_from_topic_document() -> None: + path = ( + "pages/topics/nlp/README.md" + ) + + document = make_document( + path=path, + title=( + "Nástroje a zdroje pre " + "spracovanie prirodzeného jazyka" + ), + categories=[ + "info", + ], + tags=[], + ) + + payload = build_graph_payload( + [document], + [], + ) + + assert payload["topics"] == [ + { + "id": "nlp", + "name": ( + "Nástroje a zdroje pre " + "spracovanie prirodzeného jazyka" + ), + "graph_scope": "zpwiki", + } + ] + + assert payload[ + "document_describes_topics" + ] == [ + { + "document_path": path, + "topic_id": "nlp", + } + ] + + +def test_metadata_author_becomes_document_author() -> None: + document = make_document( + author="Daniel Hladek" + ) + + payload = build_graph_payload( + [document], + [], + ) + + assert payload[ + "document_authors" + ] == [ + { + "id": "daniel-hladek", + "name": "Daniel Hladek", + "graph_scope": "zpwiki", + } + ] + + assert payload[ + "document_author_documents" + ] == [ + { + "document_author_id": ( + "daniel-hladek" + ), + "document_path": ( + JAKUB_PATH + ), + } + ] + + +def test_nested_document_does_not_create_second_person() -> None: + profile = make_document() + + nested = make_document( + path=JAKUB_NESTED_PATH, + title="Diplomová práca 2021", + categories=[ + "dp2021", + ], + tags=[ + "spacy", + ], + ) + + payload = build_graph_payload( + [ + profile, + nested, + ], + [], + ) + + assert len( + payload["people"] + ) == 1 + + person = payload[ + "people" + ][0] + + assert person["name"] == ( + "Jakub Maruniak" + ) + + assert person[ + "source_document" + ] == JAKUB_PATH + + assert payload[ + "person_documents" + ] == [ + { + "person_id": ( + f"person:{JAKUB_PATH}" + ), + "document_path": ( + JAKUB_PATH + ), + }, + { + "person_id": ( + f"person:{JAKUB_PATH}" + ), + "document_path": ( + JAKUB_NESTED_PATH + ), + }, + ] + + +def test_thesis_and_other_work_are_separated() -> None: + document = make_document( + categories=[ + "dp2021", + "tp2020", + ] + ) + + chunks = [ + make_chunk( + 0, + [ + "Jakub Maruniak", + "Diplomová práca 2021", + ], + ( + "Dokument: Jakub Maruniak\n" + "Sekcia: Diplomová práca 2021\n\n" + "Názov: Diplomová práca" + ), + ), + make_chunk( + 1, + [ + "Jakub Maruniak", + "Tímový projekt 2020", + ], + ( + "Dokument: Jakub Maruniak\n" + "Sekcia: Tímový projekt 2020\n\n" + "Názov: Tímový projekt" + ), + ), + ] + + payload = build_graph_payload( + [document], + chunks, + ) + + assert len( + payload["theses"] + ) == 1 + + assert len( + payload["works"] + ) == 1 + + thesis = payload[ + "theses" + ][0] + + work = payload[ + "works" + ][0] + + assert thesis[ + "thesis_type" + ] == "diploma_thesis" + + assert work[ + "work_type" + ] == "team_project" + + assert payload[ + "person_theses" + ] == [ + { + "person_id": ( + f"person:{JAKUB_PATH}" + ), + "thesis_id": thesis["id"], + } + ] + + assert payload[ + "person_works" + ] == [ + { + "person_id": ( + f"person:{JAKUB_PATH}" + ), + "work_id": work["id"], + } + ] + + +def test_bulleted_title_label_is_used() -> None: + document = make_document( + categories=["bp2020"] + ) + + chunks = [ + make_chunk( + 0, + ["Bakalárska práca 2020"], + ( + "Dokument: Jakub Maruniak\n" + "Sekcia: Bakalárska práca 2020\n\n" + "- Názov: Jednoduchý chatbot v jazyku Python,\n\n" + "Zadanie:\n\n" + "1. Prvá úloha." + ), + ) + ] + + work = extract_works( + document, + chunks, + )[0] + + assert work["title"] == ( + "Jednoduchý chatbot v jazyku Python" + ) + + +def test_preliminary_title_is_used() -> None: + document = make_document( + categories=["bp2025"] + ) + + chunks = [ + make_chunk( + 0, + ["Bakalárska práca 2025"], + ( + "Dokument: Jakub Maruniak\n" + "Sekcia: Bakalárska práca 2025\n\n" + "Vedúci: doc.\n\n" + "Predbežný názov:\n\n" + "Testovanie hybridného klaudu" + ), + ) + ] + + work = extract_works( + document, + chunks, + )[0] + + assert work["title"] == ( + "Testovanie hybridného klaudu" + ) + + +def test_proposed_topic_is_used() -> None: + document = make_document( + categories=["dp2027"] + ) + + chunks = [ + make_chunk( + 0, + ["Diplomová práca 2027"], + ( + "Dokument: Jakub Maruniak\n" + "Sekcia: Diplomová práca 2027\n\n" + "Návrh na tému:\n\n" + "Prepis reči pre zdravotný záznam\n\n" + "Ciele:\n\n" + "- Vytvoriť systém." + ), + ) + ] + + work = extract_works( + document, + chunks, + )[0] + + assert work["title"] == ( + "Prepis reči pre zdravotný záznam" + ) + + +def test_leading_title_skips_status_lines() -> None: + document = make_document( + categories=["bp2020"] + ) + + chunks = [ + make_chunk( + 0, + ["Bakalárska práca 2020"], + ( + "Dokument: Jakub Maruniak\n" + "Sekcia: Bakalárska práca 2020\n\n" + "Práca úspešne obhájená.\n\n" + "Odovzdaný klaster mal zlé káble.\n\n" + "Pokusný Kubernetes klaster\n\n" + "1. Prvá úloha." + ), + ) + ] + + work = extract_works( + document, + chunks, + )[0] + + assert work["title"] == ( + "Pokusný Kubernetes klaster" + ) + + +def test_goal_is_not_promoted_to_title() -> None: + document = make_document( + categories=["dp2025"] + ) + + chunks = [ + make_chunk( + 0, + ["Diplomová práca 2025"], + ( + "Dokument: Jakub Maruniak\n" + "Sekcia: Diplomová práca 2025\n\n" + "Cieľ:\n\n" + "Overiť možnosti generovania testov\n\n" + "Zadanie:\n\n" + "1. Prvá úloha." + ), + ) + ] + + work = extract_works( + document, + chunks, + )[0] + + assert work["title"] is None + + +def test_deepest_work_heading_owns_nested_chunk() -> None: + document = make_document( + categories=[ + "dp2027", + "bp2025", + ] + ) + + chunks = [ + make_chunk( + 0, + ["Diplomová práca 2027"], + ( + "Dokument: Jakub Maruniak\n" + "Sekcia: Diplomová práca 2027\n\n" + "Téma:\n\n" + "Diplomová téma" + ), + ), + make_chunk( + 1, + [ + "Diplomová práca 2027", + "Bakalárska práca 2025", + ], + ( + "Dokument: Jakub Maruniak\n" + "Sekcia: Diplomová práca 2027 > " + "Bakalárska práca 2025\n\n" + "Téma:\n\n" + "Bakalárska téma" + ), + ), + ] + + works = extract_works( + document, + chunks, + ) + + diploma = next( + work + for work in works + if work["work_type"] + == "diploma_thesis" + ) + + bachelor = next( + work + for work in works + if work["work_type"] + == "bachelor_thesis" + ) + + assert diploma["title"] == ( + "Diplomová téma" + ) + + assert bachelor["title"] == ( + "Bakalárska téma" + ) + + +def test_nested_document_does_not_create_academic_work() -> None: + profile = make_document( + categories=[] + ) + + nested = make_document( + path=JAKUB_NESTED_PATH, + title="Diplomový projekt", + categories=[], + tags=[], + ) + + nested_chunk = make_chunk( + 0, + ["Diplomový projekt 2020"], + ( + "Dokument: Diplomový projekt\n" + "Sekcia: Diplomový projekt 2020\n\n" + "Názov: Falošná samostatná práca" + ), + path=JAKUB_NESTED_PATH, + title="Diplomový projekt", + ) + + payload = build_graph_payload( + [ + profile, + nested, + ], + [nested_chunk], + ) + + assert payload["works"] == [] + assert payload["theses"] == [] + + assert payload[ + "person_documents" + ] == [ + { + "person_id": ( + f"person:{JAKUB_PATH}" + ), + "document_path": ( + JAKUB_PATH + ), + }, + { + "person_id": ( + f"person:{JAKUB_PATH}" + ), + "document_path": ( + JAKUB_NESTED_PATH + ), + }, + ]