726 lines
15 KiB
Python
726 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from tools.rag.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 COALESCE(published, 1) = 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 COALESCE(published, 1) = 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
|