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