1079 lines
19 KiB
Python
1079 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from scripts.search_utils import (
|
|
make_source_url,
|
|
normalize_for_compare,
|
|
)
|
|
|
|
|
|
YEAR_RE = re.compile(
|
|
r"\b(?:19|20)\d{2}\b"
|
|
)
|
|
|
|
|
|
SECTION_INTENTS: dict[
|
|
str,
|
|
tuple[str, ...],
|
|
] = {
|
|
"diploma_thesis": (
|
|
"diplomov",
|
|
"prac",
|
|
),
|
|
"bachelor_thesis": (
|
|
"bakalarsk",
|
|
"prac",
|
|
),
|
|
"doctoral_thesis": (
|
|
"dizertacn",
|
|
"prac",
|
|
),
|
|
"diploma_project": (
|
|
"diplomov",
|
|
"projekt",
|
|
),
|
|
"team_project": (
|
|
"timov",
|
|
"projekt",
|
|
),
|
|
}
|
|
|
|
|
|
WORK_INTENTS = {
|
|
"diploma_thesis",
|
|
"bachelor_thesis",
|
|
"doctoral_thesis",
|
|
}
|
|
|
|
|
|
QUERY_STOPWORDS = {
|
|
"a",
|
|
"aj",
|
|
"aka",
|
|
"ake",
|
|
"aky",
|
|
"je",
|
|
"k",
|
|
"kde",
|
|
"ktora",
|
|
"ktore",
|
|
"ktory",
|
|
"ma",
|
|
"mal",
|
|
"mala",
|
|
"malo",
|
|
"na",
|
|
"o",
|
|
"osoba",
|
|
"osobe",
|
|
"osoby",
|
|
"pri",
|
|
"pre",
|
|
"sa",
|
|
"s",
|
|
"so",
|
|
"su",
|
|
"ta",
|
|
"ten",
|
|
"to",
|
|
"u",
|
|
"v",
|
|
"vo",
|
|
"z",
|
|
"za",
|
|
}
|
|
|
|
|
|
def normalized_tokens(
|
|
value: str,
|
|
) -> list[str]:
|
|
normalized = (
|
|
normalize_for_compare(
|
|
value
|
|
)
|
|
)
|
|
|
|
if not normalized:
|
|
return []
|
|
|
|
return normalized.split()
|
|
|
|
|
|
def has_prefix(
|
|
tokens: list[str],
|
|
prefix: str,
|
|
) -> bool:
|
|
return any(
|
|
token.startswith(
|
|
prefix
|
|
)
|
|
for token in tokens
|
|
)
|
|
|
|
|
|
def detect_section_intent(
|
|
query: str,
|
|
) -> str | None:
|
|
tokens = normalized_tokens(
|
|
query
|
|
)
|
|
|
|
if not tokens:
|
|
return None
|
|
|
|
# Diplomový projekt musí byť pred
|
|
# diplomovou prácou, pretože oba
|
|
# začínajú tokenom "diplomov".
|
|
ordered_intents = (
|
|
"diploma_project",
|
|
"team_project",
|
|
"diploma_thesis",
|
|
"bachelor_thesis",
|
|
"doctoral_thesis",
|
|
)
|
|
|
|
for intent in ordered_intents:
|
|
prefixes = (
|
|
SECTION_INTENTS[
|
|
intent
|
|
]
|
|
)
|
|
|
|
if all(
|
|
has_prefix(
|
|
tokens,
|
|
prefix,
|
|
)
|
|
for prefix in prefixes
|
|
):
|
|
return intent
|
|
|
|
return None
|
|
|
|
|
|
def extract_years(
|
|
value: str,
|
|
) -> list[int]:
|
|
return [
|
|
int(
|
|
match
|
|
)
|
|
for match in YEAR_RE.findall(
|
|
value
|
|
)
|
|
]
|
|
|
|
|
|
def parse_heading_paths(
|
|
value: Any,
|
|
) -> list[Any]:
|
|
if isinstance(
|
|
value,
|
|
list,
|
|
):
|
|
return value
|
|
|
|
if not value:
|
|
return []
|
|
|
|
try:
|
|
parsed = json.loads(
|
|
str(
|
|
value
|
|
)
|
|
)
|
|
|
|
except json.JSONDecodeError:
|
|
return []
|
|
|
|
if not isinstance(
|
|
parsed,
|
|
list,
|
|
):
|
|
return []
|
|
|
|
return parsed
|
|
|
|
|
|
def flatten_heading_paths(
|
|
heading_paths: list[Any],
|
|
) -> str:
|
|
parts: list[str] = []
|
|
|
|
for item in heading_paths:
|
|
if isinstance(
|
|
item,
|
|
str,
|
|
):
|
|
value = (
|
|
item.strip()
|
|
)
|
|
|
|
if value:
|
|
parts.append(
|
|
value
|
|
)
|
|
|
|
continue
|
|
|
|
if isinstance(
|
|
item,
|
|
(
|
|
list,
|
|
tuple,
|
|
),
|
|
):
|
|
for part in item:
|
|
value = str(
|
|
part
|
|
).strip()
|
|
|
|
if value:
|
|
parts.append(
|
|
value
|
|
)
|
|
|
|
return " ".join(
|
|
parts
|
|
)
|
|
|
|
|
|
def section_matches_intent(
|
|
value: str,
|
|
intent: str,
|
|
) -> bool:
|
|
prefixes = (
|
|
SECTION_INTENTS.get(
|
|
intent
|
|
)
|
|
)
|
|
|
|
if not prefixes:
|
|
return False
|
|
|
|
tokens = normalized_tokens(
|
|
value
|
|
)
|
|
|
|
return all(
|
|
has_prefix(
|
|
tokens,
|
|
prefix,
|
|
)
|
|
for prefix in prefixes
|
|
)
|
|
|
|
|
|
def is_primary_student_document(
|
|
document_path: str,
|
|
) -> bool:
|
|
parts = Path(
|
|
document_path
|
|
).parts
|
|
|
|
return (
|
|
len(
|
|
parts
|
|
)
|
|
== 5
|
|
and parts[0]
|
|
== "pages"
|
|
and parts[1]
|
|
== "students"
|
|
and parts[-1]
|
|
== "README.md"
|
|
)
|
|
|
|
|
|
def title_occurs_in_query(
|
|
title: str,
|
|
query: str,
|
|
) -> bool:
|
|
normalized_title = (
|
|
normalize_for_compare(
|
|
title
|
|
)
|
|
)
|
|
|
|
normalized_query = (
|
|
normalize_for_compare(
|
|
query
|
|
)
|
|
)
|
|
|
|
if (
|
|
not normalized_title
|
|
or not normalized_query
|
|
):
|
|
return False
|
|
|
|
# Jednoslovné názvy sú príliš široké
|
|
# na bezpečné exact-document rozšírenie.
|
|
if len(
|
|
normalized_title.split()
|
|
) < 2:
|
|
return False
|
|
|
|
return (
|
|
f" {normalized_title} "
|
|
in f" {normalized_query} "
|
|
)
|
|
|
|
|
|
def sqlite_table_exists(
|
|
conn: sqlite3.Connection,
|
|
table_name: str,
|
|
) -> bool:
|
|
row = conn.execute(
|
|
"""
|
|
SELECT 1
|
|
FROM sqlite_master
|
|
WHERE type = 'table'
|
|
AND name = ?
|
|
LIMIT 1
|
|
""",
|
|
(
|
|
table_name,
|
|
),
|
|
).fetchone()
|
|
|
|
return row is not None
|
|
|
|
|
|
def find_exact_student_document(
|
|
conn: sqlite3.Connection,
|
|
query: str,
|
|
*,
|
|
published_only: bool,
|
|
) -> dict[str, Any] | None:
|
|
rows = conn.execute(
|
|
"""
|
|
SELECT
|
|
path,
|
|
title,
|
|
author,
|
|
published
|
|
FROM documents
|
|
WHERE (
|
|
? = 0
|
|
OR published = 1
|
|
)
|
|
ORDER BY
|
|
LENGTH(title) DESC,
|
|
path ASC
|
|
""",
|
|
(
|
|
1
|
|
if published_only
|
|
else 0,
|
|
),
|
|
).fetchall()
|
|
|
|
candidates: list[
|
|
dict[str, Any]
|
|
] = []
|
|
|
|
for row in rows:
|
|
document_path = str(
|
|
row[
|
|
"path"
|
|
]
|
|
or ""
|
|
).strip()
|
|
|
|
title = str(
|
|
row[
|
|
"title"
|
|
]
|
|
or ""
|
|
).strip()
|
|
|
|
if (
|
|
not document_path
|
|
or not title
|
|
):
|
|
continue
|
|
|
|
if not is_primary_student_document(
|
|
document_path
|
|
):
|
|
continue
|
|
|
|
if not title_occurs_in_query(
|
|
title,
|
|
query,
|
|
):
|
|
continue
|
|
|
|
candidates.append(
|
|
{
|
|
"document_path": (
|
|
document_path
|
|
),
|
|
"title": title,
|
|
"author": row[
|
|
"author"
|
|
],
|
|
"published": (
|
|
row[
|
|
"published"
|
|
]
|
|
),
|
|
}
|
|
)
|
|
|
|
if not candidates:
|
|
return None
|
|
|
|
candidates.sort(
|
|
key=lambda item: (
|
|
-len(
|
|
normalized_tokens(
|
|
str(
|
|
item[
|
|
"title"
|
|
]
|
|
)
|
|
)
|
|
),
|
|
-len(
|
|
str(
|
|
item[
|
|
"title"
|
|
]
|
|
)
|
|
),
|
|
str(
|
|
item[
|
|
"document_path"
|
|
]
|
|
),
|
|
)
|
|
)
|
|
|
|
return candidates[
|
|
0
|
|
]
|
|
|
|
|
|
def query_content_tokens(
|
|
query: str,
|
|
title: str,
|
|
) -> list[str]:
|
|
query_tokens = (
|
|
normalized_tokens(
|
|
query
|
|
)
|
|
)
|
|
|
|
title_tokens = set(
|
|
normalized_tokens(
|
|
title
|
|
)
|
|
)
|
|
|
|
result: list[str] = []
|
|
|
|
for token in query_tokens:
|
|
if token in title_tokens:
|
|
continue
|
|
|
|
if token in QUERY_STOPWORDS:
|
|
continue
|
|
|
|
if len(token) < 3:
|
|
continue
|
|
|
|
result.append(
|
|
token
|
|
)
|
|
|
|
return result
|
|
|
|
|
|
def local_chunk_score(
|
|
*,
|
|
query: str,
|
|
title: str,
|
|
intent: str,
|
|
heading_paths: list[Any],
|
|
text: str,
|
|
chunk_index: int,
|
|
) -> tuple[
|
|
int,
|
|
int,
|
|
int,
|
|
]:
|
|
heading_text = (
|
|
flatten_heading_paths(
|
|
heading_paths
|
|
)
|
|
)
|
|
|
|
heading_normalized = (
|
|
normalize_for_compare(
|
|
heading_text
|
|
)
|
|
)
|
|
|
|
text_normalized = (
|
|
normalize_for_compare(
|
|
text
|
|
)
|
|
)
|
|
|
|
score = 0
|
|
|
|
heading_matches = (
|
|
section_matches_intent(
|
|
heading_text,
|
|
intent,
|
|
)
|
|
)
|
|
|
|
text_matches = (
|
|
section_matches_intent(
|
|
text,
|
|
intent,
|
|
)
|
|
)
|
|
|
|
if heading_matches:
|
|
score += 1000
|
|
|
|
if text_matches:
|
|
score += 500
|
|
|
|
# Ak chunk vôbec nepatrí k požadovanému
|
|
# typu sekcie, nemá byť lokálnym víťazom.
|
|
if (
|
|
not heading_matches
|
|
and not text_matches
|
|
):
|
|
return (
|
|
-1,
|
|
-1,
|
|
-chunk_index,
|
|
)
|
|
|
|
query_years = extract_years(
|
|
query
|
|
)
|
|
|
|
heading_years = extract_years(
|
|
heading_text
|
|
)
|
|
|
|
text_years = extract_years(
|
|
text
|
|
)
|
|
|
|
if query_years:
|
|
if any(
|
|
year in heading_years
|
|
for year in query_years
|
|
):
|
|
score += 1500
|
|
|
|
elif any(
|
|
year in text_years
|
|
for year in query_years
|
|
):
|
|
score += 750
|
|
|
|
else:
|
|
score -= 500
|
|
|
|
content_tokens = (
|
|
query_content_tokens(
|
|
query,
|
|
title,
|
|
)
|
|
)
|
|
|
|
for token in content_tokens:
|
|
if token in heading_normalized:
|
|
score += 30
|
|
|
|
if token in text_normalized:
|
|
score += 8
|
|
|
|
if intent in WORK_INTENTS:
|
|
if (
|
|
"nazov diplomovej prace"
|
|
in text_normalized
|
|
or "nazov bakalarskej prace"
|
|
in text_normalized
|
|
or "nazov dizertacnej prace"
|
|
in text_normalized
|
|
):
|
|
score += 120
|
|
|
|
elif "nazov" in text_normalized:
|
|
score += 80
|
|
|
|
if "tema" in text_normalized:
|
|
score += 50
|
|
|
|
if has_prefix(
|
|
normalized_tokens(
|
|
query
|
|
),
|
|
"rok",
|
|
):
|
|
if heading_years:
|
|
score += 100
|
|
|
|
# Ak rok nie je v dopyte a dokument má
|
|
# viac rovnakých typov prác, preferujeme
|
|
# najnovšiu sekciu.
|
|
latest_year = (
|
|
max(
|
|
heading_years
|
|
)
|
|
if heading_years
|
|
else -1
|
|
)
|
|
|
|
return (
|
|
score,
|
|
latest_year,
|
|
-chunk_index,
|
|
)
|
|
|
|
|
|
def load_best_document_section_chunk(
|
|
conn: sqlite3.Connection,
|
|
*,
|
|
query: str,
|
|
document: dict[str, Any],
|
|
intent: str,
|
|
published_only: bool,
|
|
) -> dict[str, Any] | None:
|
|
document_path = str(
|
|
document[
|
|
"document_path"
|
|
]
|
|
)
|
|
|
|
rows = conn.execute(
|
|
"""
|
|
SELECT
|
|
chunk_id,
|
|
document_path,
|
|
title,
|
|
author,
|
|
published,
|
|
chunk_index,
|
|
heading_paths_json,
|
|
text
|
|
FROM chunks
|
|
WHERE document_path = ?
|
|
AND (
|
|
? = 0
|
|
OR published = 1
|
|
)
|
|
ORDER BY
|
|
chunk_index ASC,
|
|
id ASC
|
|
""",
|
|
(
|
|
document_path,
|
|
(
|
|
1
|
|
if published_only
|
|
else 0
|
|
),
|
|
),
|
|
).fetchall()
|
|
|
|
best_item: dict[
|
|
str,
|
|
Any,
|
|
] | None = None
|
|
|
|
best_score: tuple[
|
|
int,
|
|
int,
|
|
int,
|
|
] | None = None
|
|
|
|
for row in rows:
|
|
heading_paths = (
|
|
parse_heading_paths(
|
|
row[
|
|
"heading_paths_json"
|
|
]
|
|
)
|
|
)
|
|
|
|
text = str(
|
|
row[
|
|
"text"
|
|
]
|
|
or ""
|
|
).strip()
|
|
|
|
if not text:
|
|
continue
|
|
|
|
chunk_index = int(
|
|
row[
|
|
"chunk_index"
|
|
]
|
|
)
|
|
|
|
score = local_chunk_score(
|
|
query=query,
|
|
title=str(
|
|
document[
|
|
"title"
|
|
]
|
|
),
|
|
intent=intent,
|
|
heading_paths=(
|
|
heading_paths
|
|
),
|
|
text=text,
|
|
chunk_index=(
|
|
chunk_index
|
|
),
|
|
)
|
|
|
|
if score[
|
|
0
|
|
] < 0:
|
|
continue
|
|
|
|
if (
|
|
best_score is not None
|
|
and score <= best_score
|
|
):
|
|
continue
|
|
|
|
best_score = score
|
|
|
|
best_item = {
|
|
"chunk_id": str(
|
|
row[
|
|
"chunk_id"
|
|
]
|
|
),
|
|
"document_path": (
|
|
document_path
|
|
),
|
|
"title": (
|
|
row[
|
|
"title"
|
|
]
|
|
or document[
|
|
"title"
|
|
]
|
|
),
|
|
"author": (
|
|
row[
|
|
"author"
|
|
]
|
|
or document.get(
|
|
"author"
|
|
)
|
|
),
|
|
"published": (
|
|
bool(
|
|
row[
|
|
"published"
|
|
]
|
|
)
|
|
if row[
|
|
"published"
|
|
]
|
|
is not None
|
|
else None
|
|
),
|
|
"chunk_index": (
|
|
chunk_index
|
|
),
|
|
"heading_paths": (
|
|
heading_paths
|
|
),
|
|
"text": text,
|
|
"source_url": (
|
|
make_source_url(
|
|
document_path
|
|
)
|
|
),
|
|
}
|
|
|
|
return best_item
|
|
|
|
|
|
def exact_document_expansion_metadata(
|
|
*,
|
|
applied: bool,
|
|
document_path: str | None,
|
|
chunk_id: str | None,
|
|
chunk_index: int | None,
|
|
added_source: bool,
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"strategy": (
|
|
"exact_document_section"
|
|
),
|
|
"applied": applied,
|
|
"document_path": (
|
|
document_path
|
|
),
|
|
"chunk_id": (
|
|
chunk_id
|
|
),
|
|
"chunk_index": (
|
|
chunk_index
|
|
),
|
|
"added_source": (
|
|
added_source
|
|
),
|
|
}
|
|
|
|
|
|
def expand_results_with_exact_document_section(
|
|
db_path: Path,
|
|
query: str,
|
|
results: list[
|
|
dict[str, Any]
|
|
],
|
|
*,
|
|
published_only: bool,
|
|
limit: int,
|
|
) -> list[dict[str, Any]]:
|
|
if not results:
|
|
base_results: list[
|
|
dict[str, Any]
|
|
] = []
|
|
|
|
else:
|
|
base_results = [
|
|
dict(
|
|
result
|
|
)
|
|
for result in results
|
|
]
|
|
|
|
intent = detect_section_intent(
|
|
query
|
|
)
|
|
|
|
if intent is None:
|
|
return base_results
|
|
|
|
if not db_path.exists():
|
|
return base_results
|
|
|
|
with sqlite3.connect(
|
|
db_path,
|
|
timeout=5.0,
|
|
) as conn:
|
|
conn.row_factory = (
|
|
sqlite3.Row
|
|
)
|
|
|
|
conn.execute(
|
|
"PRAGMA query_only = ON"
|
|
)
|
|
|
|
# Niektoré unit testy používajú zámerne
|
|
# minimalistickú SQLite databázu iba
|
|
# s tabuľkou chunks.
|
|
#
|
|
# Exact-document expansion je voliteľná
|
|
# RAG-only nadstavba. Ak potrebná schéma
|
|
# nie je dostupná, musí zostať bezpečný
|
|
# pôvodný výsledok retrievalu.
|
|
if (
|
|
not sqlite_table_exists(
|
|
conn,
|
|
"documents",
|
|
)
|
|
or not sqlite_table_exists(
|
|
conn,
|
|
"chunks",
|
|
)
|
|
):
|
|
return base_results
|
|
|
|
document = (
|
|
find_exact_student_document(
|
|
conn,
|
|
query,
|
|
published_only=(
|
|
published_only
|
|
),
|
|
)
|
|
)
|
|
|
|
if document is None:
|
|
return base_results
|
|
|
|
local_chunk = (
|
|
load_best_document_section_chunk(
|
|
conn,
|
|
query=query,
|
|
document=document,
|
|
intent=intent,
|
|
published_only=(
|
|
published_only
|
|
),
|
|
)
|
|
)
|
|
|
|
if local_chunk is None:
|
|
return base_results
|
|
|
|
document_path = str(
|
|
document[
|
|
"document_path"
|
|
]
|
|
)
|
|
|
|
existing_index: int | None = None
|
|
|
|
for index, result in enumerate(
|
|
base_results
|
|
):
|
|
if str(
|
|
result.get(
|
|
"document_path"
|
|
)
|
|
or ""
|
|
) == document_path:
|
|
existing_index = (
|
|
index
|
|
)
|
|
|
|
break
|
|
|
|
if existing_index is not None:
|
|
item = dict(
|
|
base_results[
|
|
existing_index
|
|
]
|
|
)
|
|
|
|
local_chunk_id = str(
|
|
local_chunk[
|
|
"chunk_id"
|
|
]
|
|
)
|
|
|
|
primary_chunk_id = str(
|
|
item.get(
|
|
"chunk_id"
|
|
)
|
|
or ""
|
|
)
|
|
|
|
if (
|
|
local_chunk_id
|
|
!= primary_chunk_id
|
|
):
|
|
item[
|
|
"exact_document_text"
|
|
] = local_chunk[
|
|
"text"
|
|
]
|
|
|
|
item[
|
|
"exact_document_heading_paths"
|
|
] = local_chunk[
|
|
"heading_paths"
|
|
]
|
|
|
|
item[
|
|
"document_expansion"
|
|
] = (
|
|
exact_document_expansion_metadata(
|
|
applied=(
|
|
local_chunk_id
|
|
!= primary_chunk_id
|
|
),
|
|
document_path=(
|
|
document_path
|
|
),
|
|
chunk_id=(
|
|
local_chunk_id
|
|
),
|
|
chunk_index=int(
|
|
local_chunk[
|
|
"chunk_index"
|
|
]
|
|
),
|
|
added_source=False,
|
|
)
|
|
)
|
|
|
|
del base_results[
|
|
existing_index
|
|
]
|
|
|
|
# Presne zhodný dokument má byť
|
|
# v RAG kontexte pred dokumentmi,
|
|
# kde je meno iba v author poli.
|
|
base_results.insert(
|
|
0,
|
|
item,
|
|
)
|
|
|
|
return (
|
|
base_results[
|
|
:limit
|
|
]
|
|
if limit > 0
|
|
else base_results
|
|
)
|
|
|
|
synthetic = {
|
|
**local_chunk,
|
|
"match_strategy": (
|
|
"exact_document_section"
|
|
),
|
|
"fts_rank": None,
|
|
"vector_rank": None,
|
|
"vector_score": None,
|
|
"hybrid_score": None,
|
|
"document_expansion": (
|
|
exact_document_expansion_metadata(
|
|
applied=True,
|
|
document_path=(
|
|
document_path
|
|
),
|
|
chunk_id=str(
|
|
local_chunk[
|
|
"chunk_id"
|
|
]
|
|
),
|
|
chunk_index=int(
|
|
local_chunk[
|
|
"chunk_index"
|
|
]
|
|
),
|
|
added_source=True,
|
|
)
|
|
),
|
|
}
|
|
|
|
expanded = [
|
|
synthetic,
|
|
*base_results,
|
|
]
|
|
|
|
if limit > 0:
|
|
return expanded[
|
|
:limit
|
|
]
|
|
|
|
return expanded
|