dp-zp-agent/scripts/graphrag.py
2026-09-29 18:42:17 +02:00

1049 lines
22 KiB
Python

from __future__ import annotations
import re
import unicodedata
from collections import defaultdict
from dataclasses import dataclass
from typing import Any
GRAPH_SCOPE = "zpwiki"
WORK_TYPE_LABELS = {
"bachelor_thesis": "Bakalárska práca",
"diploma_thesis": "Diplomová práca",
"diploma_project": "Diplomový projekt",
"team_project": "Tímový projekt",
"research_project": "Výskumný projekt",
"doctoral_thesis": "Dizertačná práca",
}
WORK_HEADING_PATTERNS = (
("bachelor_thesis", "bakalarska praca", "bp"),
("diploma_thesis", "diplomova praca", "dp"),
("diploma_project", "diplomovy projekt", None),
("team_project", "timovy projekt", "tp"),
("research_project", "vyskumny projekt", "vp"),
("doctoral_thesis", "dizertacna praca", "phd"),
)
TITLE_LABELS = {
"nazov",
"nazov prace",
"nazov bakalarskej prace",
"nazov diplomovej prace",
"tema",
"navrh na nazov",
"navrh na nazov bp",
"navrh na nazov dp",
"navrh na nazov bakalarskej prace",
"navrh na nazov diplomovej prace",
}
SKIP_LINE_PREFIXES = (
"dokument:",
"sekcia:",
"rok zaciatku studia",
"meno veduceho",
"zadanie",
"navrh na zadanie",
"navrh na zadamie",
"ciel:",
"ciele:",
"ulohy:",
"zasobnik uloh",
"stretnutie",
"stav:",
"repozitar",
"git repozitar",
"crzp",
)
@dataclass(frozen=True)
class WorkHeading:
work_type: str
heading: str
category_prefix: str | None
def normalize_text(value: str) -> str:
decomposed = unicodedata.normalize(
"NFKD",
str(value),
)
without_marks = "".join(
char
for char in decomposed
if not unicodedata.combining(char)
)
return re.sub(
r"\s+",
" ",
without_marks.lower(),
).strip()
def normalize_key(value: str) -> str:
normalized = normalize_text(value)
normalized = re.sub(
r"[^a-z0-9]+",
"-",
normalized,
)
return normalized.strip("-")
def source_url_from_path(path: str) -> str:
normalized = path.replace("\\", "/")
if normalized.startswith("pages/"):
normalized = normalized[len("pages/"):]
if normalized.endswith("/README.md"):
normalized = normalized[:-len("/README.md")]
elif normalized.endswith(".md"):
normalized = normalized[:-3]
return (
"https://zp.kemt.fei.tuke.sk/"
+ normalized.strip("/")
)
def document_kind(path: str) -> str:
normalized = path.replace("\\", "/")
if normalized.startswith("pages/students/"):
return "student"
if normalized.startswith("pages/interns/"):
return "intern"
if normalized.startswith("pages/topics/"):
return "topic"
return "other"
def extract_full_years(value: str) -> list[int]:
years = [
int(match)
for match in re.findall(
r"\b(20\d{2})\b",
value,
)
]
for match in re.finditer(
r"\b(20\d{2})\s*/\s*(\d{2})\b",
value,
):
first = int(match.group(1))
second_short = int(match.group(2))
second = (
(first // 100) * 100
+ second_short
)
if second not in years:
years.append(second)
return sorted(set(years))
def parse_category(
category: str,
) -> tuple[str, int | None]:
match = re.fullmatch(
r"(bp|dp|tp|vp|phd)(20\d{2})",
category.lower(),
)
if not match:
return "other", None
prefix = match.group(1)
year = int(match.group(2))
category_types = {
"bp": "bachelor_thesis",
"dp": "diploma_thesis",
"tp": "team_project",
"vp": "research_project",
"phd": "doctoral_thesis",
}
return category_types[prefix], year
def find_work_heading(
heading: str,
) -> WorkHeading | None:
normalized = normalize_text(heading)
for (
work_type,
phrase,
category_prefix,
) in WORK_HEADING_PATTERNS:
if phrase in normalized:
return WorkHeading(
work_type=work_type,
heading=heading,
category_prefix=category_prefix,
)
return None
def category_years(
categories: list[str],
prefix: str | None,
) -> list[int]:
if prefix is None:
return []
pattern = re.compile(
rf"^{re.escape(prefix)}(20\d{{2}})$",
flags=re.IGNORECASE,
)
result = []
for category in categories:
match = pattern.fullmatch(category)
if match:
result.append(
int(match.group(1))
)
return sorted(set(result))
def resolve_work_year(
heading: WorkHeading,
categories: list[str],
) -> int | None:
heading_years = extract_full_years(
heading.heading
)
category_candidates = category_years(
categories,
heading.category_prefix,
)
if heading.work_type == "diploma_project":
if heading_years:
return heading_years[-1]
return None
common = sorted(
set(heading_years)
& set(category_candidates)
)
if common:
return common[-1]
if len(heading_years) > 1:
if len(category_candidates) == 1:
return category_candidates[0]
return heading_years[-1]
if heading_years:
return heading_years[-1]
if category_candidates:
return category_candidates[-1]
return None
def clean_markdown_line(
line: str,
) -> str:
value = line.strip()
value = value.replace(
"*",
"",
).replace(
"`",
"",
)
return value.strip()
def next_title_line(
lines: list[str],
start_index: int,
) -> str | None:
for raw_line in lines[start_index:]:
line = clean_markdown_line(
raw_line
)
if not line:
continue
normalized = normalize_text(line)
if normalized.startswith(
SKIP_LINE_PREFIXES
):
continue
if line.startswith(
("http://", "https://", "[")
):
continue
if re.match(
r"^[-•]\s+",
line,
):
continue
if re.match(
r"^\d+[\.\)]\s*",
line,
):
continue
return line
return None
def extract_work_title(
texts: list[str],
) -> str | None:
for text in texts:
lines = text.splitlines()
for index, raw_line in enumerate(lines):
line = clean_markdown_line(
raw_line
)
if not line:
continue
before, separator, after = (
line.partition(":")
)
normalized_before = normalize_text(
before
)
if (
separator
and normalized_before
in TITLE_LABELS
):
title = after.strip()
if title:
return title
candidate = next_title_line(
lines,
index + 1,
)
if candidate:
return candidate
for text in texts:
lines = text.splitlines()
for raw_line in lines:
line = clean_markdown_line(
raw_line
)
if not line:
continue
normalized = normalize_text(
line
)
if normalized.startswith(
SKIP_LINE_PREFIXES
):
continue
if line.startswith(
("http://", "https://", "[")
):
continue
if re.match(
r"^[-•]\s+",
line,
):
continue
if re.match(
r"^\d+[\.\)]\s*",
line,
):
continue
if len(line) < 8:
continue
return line
return None
def extract_start_year(
path: str,
chunks: list[dict[str, Any]],
) -> int | None:
pattern = re.compile(
r"rok\s+zaciatku\s+studia"
r"\s*\*?\s*:\s*\*?\s*"
r"(20\d{2})"
)
for chunk in chunks:
normalized = normalize_text(
str(chunk.get("text", ""))
)
match = pattern.search(
normalized
)
if match:
return int(
match.group(1)
)
match = re.search(
r"pages/students/(20\d{2})/",
path.replace("\\", "/"),
)
if match:
return int(
match.group(1)
)
return None
def topic_key_from_document_path(
path: str,
) -> str | None:
normalized = path.replace("\\", "/")
if not normalized.startswith(
"pages/topics/"
):
return None
value = normalized[
len("pages/topics/"):
]
if value.endswith("/README.md"):
value = value[:-len("/README.md")]
elif value.endswith(".md"):
value = value[:-3]
if not value:
return None
return normalize_key(
value.split("/")[0]
)
def topic_display_name(
value: str,
) -> str:
aliases = {
"nlp": "NLP",
"rag": "RAG",
"ner": "NER",
"qa": "QA",
"ie": "IE",
"lm": "LM",
"nmt": "NMT",
"pos": "POS",
"gpu": "GPU",
"spacy": "Spacy",
}
key = normalize_key(value)
if key in aliases:
return aliases[key]
return str(value).strip()
def extract_works(
document: dict[str, Any],
chunks: list[dict[str, Any]],
) -> list[dict[str, Any]]:
grouped: dict[
tuple[str, str],
dict[str, Any],
] = {}
categories = list(
document.get(
"categories",
[],
)
)
for chunk in sorted(
chunks,
key=lambda item: int(
item.get(
"chunk_index",
0,
)
),
):
heading_paths = chunk.get(
"heading_paths",
[],
)
candidates = []
for path in heading_paths:
if isinstance(path, list):
candidates.extend(
str(value)
for value in path
)
elif path:
candidates.append(
str(path)
)
for raw_heading in candidates:
heading = find_work_heading(
raw_heading
)
if heading is None:
continue
key = (
heading.work_type,
normalize_text(
heading.heading
),
)
entry = grouped.setdefault(
key,
{
"heading": heading,
"texts": [],
"chunk_ids": [],
},
)
text = str(
chunk.get(
"text",
"",
)
)
if text:
entry["texts"].append(
text
)
chunk_id = chunk.get(
"chunk_id"
)
if chunk_id:
entry["chunk_ids"].append(
str(chunk_id)
)
result = []
for entry in grouped.values():
heading = entry["heading"]
year = resolve_work_year(
heading,
categories,
)
title = extract_work_title(
entry["texts"]
)
result.append(
{
"work_type": heading.work_type,
"work_type_label": (
WORK_TYPE_LABELS[
heading.work_type
]
),
"heading": heading.heading,
"year": year,
"title": title,
"source_chunk_ids": sorted(
set(
entry[
"chunk_ids"
]
)
),
}
)
return sorted(
result,
key=lambda item: (
item["year"] or 0,
item["work_type"],
item["heading"],
),
)
def build_graph_payload(
documents: list[dict[str, Any]],
chunks: list[dict[str, Any]],
) -> dict[str, list[dict[str, Any]]]:
chunks_by_document: dict[
str,
list[dict[str, Any]],
] = defaultdict(list)
for chunk in chunks:
path = str(
chunk.get(
"document_path",
"",
)
)
if path:
chunks_by_document[
path
].append(chunk)
people: dict[
str,
dict[str, Any],
] = {}
authors: dict[
str,
dict[str, Any],
] = {}
document_nodes: dict[
str,
dict[str, Any],
] = {}
categories: dict[
str,
dict[str, Any],
] = {}
topics: dict[
str,
dict[str, Any],
] = {}
works: dict[
str,
dict[str, Any],
] = {}
person_documents = set()
author_documents = set()
document_categories = set()
document_topics = set()
document_describes_topics = set()
person_works = set()
work_documents = set()
for document in documents:
path = str(
document["path"]
)
title = str(
document.get(
"title",
path,
)
)
kind = document_kind(
path
)
doc_categories = [
str(value)
for value in document.get(
"categories",
[],
)
]
doc_tags = [
str(value)
for value in document.get(
"tags",
[],
)
]
author = document.get(
"author"
)
document_nodes[path] = {
"path": path,
"name": title,
"title": title,
"source_url": (
source_url_from_path(
path
)
),
"document_kind": kind,
"published": bool(
document.get(
"published",
False,
)
),
"categories": (
doc_categories
),
"tags": doc_tags,
"author": (
str(author)
if author
else None
),
"graph_scope": (
GRAPH_SCOPE
),
}
if author:
author_name = str(
author
)
author_id = normalize_key(
author_name
)
authors[author_id] = {
"id": author_id,
"name": author_name,
"graph_scope": (
GRAPH_SCOPE
),
}
author_documents.add(
(
author_id,
path,
)
)
for category in doc_categories:
category_type, year = (
parse_category(
category
)
)
categories[category] = {
"name": category,
"category_type": (
category_type
),
"year": year,
"graph_scope": (
GRAPH_SCOPE
),
}
document_categories.add(
(
path,
category,
)
)
for tag in doc_tags:
topic_id = normalize_key(
tag
)
current = topics.get(
topic_id
)
if current is None:
topics[topic_id] = {
"id": topic_id,
"name": (
topic_display_name(
tag
)
),
"graph_scope": (
GRAPH_SCOPE
),
}
document_topics.add(
(
path,
topic_id,
)
)
described_topic = (
topic_key_from_document_path(
path
)
)
if described_topic:
topics[
described_topic
] = {
"id": described_topic,
"name": title,
"graph_scope": (
GRAPH_SCOPE
),
}
document_describes_topics.add(
(
path,
described_topic,
)
)
if kind not in {
"student",
"intern",
}:
continue
person_id = (
f"person:{path}"
)
document_chunks = (
chunks_by_document.get(
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
),
}
person_documents.add(
(
person_id,
path,
)
)
if kind != "student":
continue
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,
"title": work[
"title"
],
"heading": work[
"heading"
],
"work_type": work[
"work_type"
],
"work_type_label": work[
"work_type_label"
],
"year": work[
"year"
],
"source_document": path,
"source_chunk_ids": work[
"source_chunk_ids"
],
"graph_scope": (
GRAPH_SCOPE
),
}
person_works.add(
(
person_id,
work_id,
)
)
work_documents.add(
(
work_id,
path,
)
)
def relation_rows(
pairs: set[tuple[str, str]],
source_name: str,
target_name: str,
) -> list[dict[str, str]]:
return [
{
source_name: source,
target_name: target,
}
for source, target in sorted(
pairs
)
]
return {
"people": sorted(
people.values(),
key=lambda row: row["id"],
),
"authors": sorted(
authors.values(),
key=lambda row: row["id"],
),
"documents": sorted(
document_nodes.values(),
key=lambda row: row[
"path"
],
),
"categories": sorted(
categories.values(),
key=lambda row: row[
"name"
],
),
"topics": sorted(
topics.values(),
key=lambda row: row["id"],
),
"works": sorted(
works.values(),
key=lambda row: row["id"],
),
"person_documents": relation_rows(
person_documents,
"person_id",
"document_path",
),
"author_documents": relation_rows(
author_documents,
"author_id",
"document_path",
),
"document_categories": relation_rows(
document_categories,
"document_path",
"category_name",
),
"document_topics": relation_rows(
document_topics,
"document_path",
"topic_id",
),
"document_describes_topics": (
relation_rows(
document_describes_topics,
"document_path",
"topic_id",
)
),
"person_works": relation_rows(
person_works,
"person_id",
"work_id",
),
"work_documents": relation_rows(
work_documents,
"work_id",
"document_path",
),
}