644 lines
12 KiB
Python
644 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
import sqlite3
|
|
import unicodedata
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
WORD_RE = re.compile(
|
|
r"[^\W_]+",
|
|
re.UNICODE,
|
|
)
|
|
|
|
|
|
# Poradie zodpovedá stĺpcom v chunks_fts:
|
|
# chunk_id, title, author, document_path,
|
|
# tags, categories, text
|
|
BM25_WEIGHTS = (
|
|
0.0,
|
|
10.0,
|
|
7.0,
|
|
5.0,
|
|
8.0,
|
|
4.0,
|
|
1.0,
|
|
)
|
|
|
|
BM25_SQL = ", ".join(
|
|
str(value)
|
|
for value in BM25_WEIGHTS
|
|
)
|
|
|
|
|
|
DEFAULT_CANDIDATE_MULTIPLIER = 8
|
|
MIN_CANDIDATES = 50
|
|
MIN_STEM_PREFIX_LENGTH = 5
|
|
|
|
|
|
STRATEGY_PRIORITY = {
|
|
"all_terms": 3,
|
|
"prefix_terms": 2,
|
|
"any_term": 1,
|
|
}
|
|
|
|
|
|
def normalize_for_compare(
|
|
text: str,
|
|
) -> str:
|
|
"""Normalizácia pre pomocné bonusové skóre."""
|
|
text = unicodedata.normalize(
|
|
"NFKD",
|
|
text.casefold(),
|
|
)
|
|
|
|
text = "".join(
|
|
character
|
|
for character in text
|
|
if not unicodedata.combining(character)
|
|
)
|
|
|
|
return " ".join(
|
|
WORD_RE.findall(text)
|
|
)
|
|
|
|
|
|
def query_tokens(
|
|
query: str,
|
|
) -> list[str]:
|
|
"""Vytvorí bezpečné tokeny pre FTS5."""
|
|
tokens: list[str] = []
|
|
seen: set[str] = set()
|
|
|
|
for token in WORD_RE.findall(query):
|
|
normalized = normalize_for_compare(
|
|
token
|
|
)
|
|
|
|
if not normalized:
|
|
continue
|
|
|
|
if normalized in seen:
|
|
continue
|
|
|
|
tokens.append(token)
|
|
seen.add(normalized)
|
|
|
|
return tokens
|
|
|
|
|
|
def quote_fts_token(
|
|
token: str,
|
|
*,
|
|
use_prefix: bool = True,
|
|
shorten: bool = False,
|
|
) -> str:
|
|
value = token
|
|
|
|
if (
|
|
shorten
|
|
and len(value) > MIN_STEM_PREFIX_LENGTH
|
|
):
|
|
value = value[:MIN_STEM_PREFIX_LENGTH]
|
|
|
|
escaped = value.replace(
|
|
'"',
|
|
'""',
|
|
)
|
|
|
|
suffix = (
|
|
"*"
|
|
if use_prefix and len(value) >= 4
|
|
else ""
|
|
)
|
|
|
|
return f'"{escaped}"{suffix}'
|
|
|
|
|
|
def build_match_queries(
|
|
query: str,
|
|
) -> list[tuple[str, str]]:
|
|
"""
|
|
Vráti stratégie od najpresnejšej:
|
|
|
|
all_terms -> prefix_terms -> any_term
|
|
"""
|
|
tokens = query_tokens(query)
|
|
|
|
if not tokens:
|
|
return []
|
|
|
|
full_terms = [
|
|
quote_fts_token(token)
|
|
for token in tokens
|
|
]
|
|
|
|
all_terms_query = " AND ".join(
|
|
full_terms
|
|
)
|
|
|
|
queries = [
|
|
(
|
|
"all_terms",
|
|
all_terms_query,
|
|
)
|
|
]
|
|
|
|
shortened_terms = [
|
|
quote_fts_token(
|
|
token,
|
|
shorten=True,
|
|
)
|
|
for token in tokens
|
|
]
|
|
|
|
shortened_query = " AND ".join(
|
|
shortened_terms
|
|
)
|
|
|
|
if shortened_query != all_terms_query:
|
|
queries.append(
|
|
(
|
|
"prefix_terms",
|
|
shortened_query,
|
|
)
|
|
)
|
|
|
|
if len(full_terms) > 1:
|
|
queries.append(
|
|
(
|
|
"any_term",
|
|
" OR ".join(full_terms),
|
|
)
|
|
)
|
|
|
|
return queries
|
|
|
|
|
|
def verify_search_schema(
|
|
conn: sqlite3.Connection,
|
|
) -> None:
|
|
row = conn.execute(
|
|
"""
|
|
SELECT 1
|
|
FROM sqlite_master
|
|
WHERE type = 'table'
|
|
AND name = 'chunks_fts'
|
|
"""
|
|
).fetchone()
|
|
|
|
if row is None:
|
|
raise RuntimeError(
|
|
"FTS5 index v databáze chýba. "
|
|
"Spusti python scripts/rebuild_index.py."
|
|
)
|
|
|
|
|
|
def make_source_url(
|
|
document_path: str,
|
|
) -> str:
|
|
clean_path = document_path
|
|
|
|
if clean_path.startswith("pages/"):
|
|
clean_path = clean_path[
|
|
len("pages/"):
|
|
]
|
|
|
|
if clean_path.endswith("/README.md"):
|
|
clean_path = clean_path[
|
|
:-len("/README.md")
|
|
]
|
|
|
|
return (
|
|
"https://zp.kemt.fei.tuke.sk/"
|
|
f"{clean_path}"
|
|
)
|
|
|
|
|
|
def load_labels(
|
|
conn: sqlite3.Connection,
|
|
table: str,
|
|
column: str,
|
|
chunk_ids: list[str],
|
|
) -> dict[str, list[str]]:
|
|
if not chunk_ids:
|
|
return {}
|
|
|
|
placeholders = ",".join(
|
|
"?"
|
|
for _ in chunk_ids
|
|
)
|
|
|
|
rows = conn.execute(
|
|
f"""
|
|
SELECT chunk_id, {column}
|
|
FROM {table}
|
|
WHERE chunk_id IN ({placeholders})
|
|
ORDER BY chunk_id, {column}
|
|
""",
|
|
chunk_ids,
|
|
).fetchall()
|
|
|
|
values: dict[str, list[str]] = defaultdict(
|
|
list
|
|
)
|
|
|
|
for chunk_id, value in rows:
|
|
values[chunk_id].append(value)
|
|
|
|
return dict(values)
|
|
|
|
|
|
def run_fts_query(
|
|
conn: sqlite3.Connection,
|
|
match_query: str,
|
|
candidate_limit: int,
|
|
published_only: bool,
|
|
) -> list[dict[str, Any]]:
|
|
rows = conn.execute(
|
|
f"""
|
|
SELECT
|
|
chunks.chunk_id,
|
|
chunks.document_path,
|
|
chunks.title,
|
|
chunks.author,
|
|
chunks.published,
|
|
chunks.chunk_index,
|
|
chunks.heading_paths_json,
|
|
chunks.text,
|
|
chunks.text_length,
|
|
chunks.token_count,
|
|
chunks.content_hash,
|
|
chunks_fts.rank AS bm25_score,
|
|
snippet(
|
|
chunks_fts,
|
|
6,
|
|
'',
|
|
'',
|
|
' … ',
|
|
36
|
|
) AS snippet
|
|
FROM chunks_fts
|
|
JOIN chunks
|
|
ON chunks.id = chunks_fts.rowid
|
|
WHERE chunks_fts MATCH ?
|
|
AND chunks_fts.rank MATCH
|
|
'bm25({BM25_SQL})'
|
|
AND (
|
|
? = 0
|
|
OR chunks.published = 1
|
|
)
|
|
ORDER BY
|
|
chunks_fts.rank ASC,
|
|
chunks.id ASC
|
|
LIMIT ?
|
|
""",
|
|
(
|
|
match_query,
|
|
1 if published_only else 0,
|
|
candidate_limit,
|
|
),
|
|
).fetchall()
|
|
|
|
return [
|
|
dict(row)
|
|
for row in rows
|
|
]
|
|
|
|
|
|
def exact_match_bonus(
|
|
query: str,
|
|
item: dict[str, Any],
|
|
tags: list[str],
|
|
categories: list[str],
|
|
) -> float:
|
|
normalized_query = normalize_for_compare(
|
|
query
|
|
)
|
|
|
|
if not normalized_query:
|
|
return 0.0
|
|
|
|
title = normalize_for_compare(
|
|
item.get("title") or ""
|
|
)
|
|
|
|
author = normalize_for_compare(
|
|
item.get("author") or ""
|
|
)
|
|
|
|
path = normalize_for_compare(
|
|
item.get("document_path") or ""
|
|
)
|
|
|
|
text = normalize_for_compare(
|
|
item.get("text") or ""
|
|
)
|
|
|
|
normalized_tags = [
|
|
normalize_for_compare(value)
|
|
for value in tags
|
|
]
|
|
|
|
normalized_categories = [
|
|
normalize_for_compare(value)
|
|
for value in categories
|
|
]
|
|
|
|
bonus = 0.0
|
|
|
|
if title == normalized_query:
|
|
bonus += 6.0
|
|
|
|
elif normalized_query in title:
|
|
bonus += 3.0
|
|
|
|
if author == normalized_query:
|
|
bonus += 5.0
|
|
|
|
elif normalized_query in author:
|
|
bonus += 2.0
|
|
|
|
if normalized_query in path:
|
|
bonus += 2.0
|
|
|
|
if normalized_query in normalized_tags:
|
|
bonus += 4.0
|
|
|
|
if normalized_query in normalized_categories:
|
|
bonus += 3.0
|
|
|
|
if normalized_query in text:
|
|
bonus += 1.5
|
|
|
|
return bonus
|
|
|
|
|
|
def database_bool(
|
|
value: Any,
|
|
) -> bool | None:
|
|
"""Prevedie SQLite 0/1 na API boolean."""
|
|
if value is None:
|
|
return None
|
|
|
|
return bool(value)
|
|
|
|
|
|
def add_labels_and_scores(
|
|
conn: sqlite3.Connection,
|
|
query: str,
|
|
candidates: list[dict[str, Any]],
|
|
) -> list[dict[str, Any]]:
|
|
chunk_ids = [
|
|
item["chunk_id"]
|
|
for item in candidates
|
|
]
|
|
|
|
tags_by_chunk = load_labels(
|
|
conn,
|
|
"chunk_tags",
|
|
"tag",
|
|
chunk_ids,
|
|
)
|
|
|
|
categories_by_chunk = load_labels(
|
|
conn,
|
|
"chunk_categories",
|
|
"category",
|
|
chunk_ids,
|
|
)
|
|
|
|
results: list[dict[str, Any]] = []
|
|
|
|
for item in candidates:
|
|
chunk_id = item["chunk_id"]
|
|
|
|
tags = tags_by_chunk.get(
|
|
chunk_id,
|
|
[],
|
|
)
|
|
|
|
categories = categories_by_chunk.get(
|
|
chunk_id,
|
|
[],
|
|
)
|
|
|
|
bm25_score = float(
|
|
item.pop("bm25_score")
|
|
)
|
|
|
|
strategy = item.pop("strategy")
|
|
|
|
base_score = max(
|
|
0.0,
|
|
-bm25_score,
|
|
)
|
|
|
|
score = (
|
|
base_score
|
|
+ exact_match_bonus(
|
|
query,
|
|
item,
|
|
tags,
|
|
categories,
|
|
)
|
|
)
|
|
|
|
try:
|
|
heading_paths = json.loads(
|
|
item.pop(
|
|
"heading_paths_json"
|
|
)
|
|
or "[]"
|
|
)
|
|
|
|
except json.JSONDecodeError:
|
|
heading_paths = []
|
|
|
|
item["published"] = database_bool(
|
|
item.get("published")
|
|
)
|
|
|
|
item["_strategy_priority"] = (
|
|
STRATEGY_PRIORITY[strategy]
|
|
)
|
|
|
|
item.update(
|
|
{
|
|
"heading_paths": heading_paths,
|
|
"tags": tags,
|
|
"categories": categories,
|
|
"score": round(
|
|
score,
|
|
6,
|
|
),
|
|
"bm25_score": round(
|
|
bm25_score,
|
|
6,
|
|
),
|
|
"match_strategy": strategy,
|
|
"source_url": make_source_url(
|
|
item["document_path"]
|
|
),
|
|
}
|
|
)
|
|
|
|
results.append(item)
|
|
|
|
results.sort(
|
|
key=lambda item: (
|
|
-item["_strategy_priority"],
|
|
-item["score"],
|
|
item["bm25_score"],
|
|
item["document_path"],
|
|
item["chunk_index"],
|
|
)
|
|
)
|
|
|
|
for item in results:
|
|
item.pop(
|
|
"_strategy_priority",
|
|
None,
|
|
)
|
|
|
|
return results
|
|
|
|
|
|
def diversify_results(
|
|
results: list[dict[str, Any]],
|
|
limit: int,
|
|
max_per_document: int,
|
|
) -> list[dict[str, Any]]:
|
|
if max_per_document <= 0:
|
|
return results[:limit]
|
|
|
|
selected: list[dict[str, Any]] = []
|
|
|
|
document_counts: dict[str, int] = (
|
|
defaultdict(int)
|
|
)
|
|
|
|
for item in results:
|
|
document_path = item[
|
|
"document_path"
|
|
]
|
|
|
|
if (
|
|
document_counts[document_path]
|
|
>= max_per_document
|
|
):
|
|
continue
|
|
|
|
selected.append(item)
|
|
|
|
document_counts[
|
|
document_path
|
|
] += 1
|
|
|
|
if len(selected) >= limit:
|
|
break
|
|
|
|
return selected
|
|
|
|
|
|
def search_database(
|
|
db_file: Path,
|
|
query: str,
|
|
limit: int = 10,
|
|
published_only: bool = False,
|
|
max_per_document: int = 3,
|
|
) -> dict[str, Any]:
|
|
if not db_file.exists():
|
|
raise FileNotFoundError(
|
|
f"Databáza neexistuje: {db_file}"
|
|
)
|
|
|
|
clean_query = query.strip()
|
|
|
|
if not clean_query:
|
|
return {
|
|
"engine": "sqlite_fts5",
|
|
"strategies": [],
|
|
"results": [],
|
|
}
|
|
|
|
match_queries = build_match_queries(
|
|
clean_query
|
|
)
|
|
|
|
if not match_queries:
|
|
return {
|
|
"engine": "sqlite_fts5",
|
|
"strategies": [],
|
|
"results": [],
|
|
}
|
|
|
|
candidate_limit = max(
|
|
MIN_CANDIDATES,
|
|
limit * DEFAULT_CANDIDATE_MULTIPLIER,
|
|
)
|
|
|
|
with sqlite3.connect(
|
|
db_file,
|
|
timeout=5.0,
|
|
) as conn:
|
|
conn.row_factory = sqlite3.Row
|
|
conn.execute(
|
|
"PRAGMA query_only = ON"
|
|
)
|
|
|
|
verify_search_schema(conn)
|
|
|
|
candidates: list[
|
|
dict[str, Any]
|
|
] = []
|
|
|
|
used_strategies: list[str] = []
|
|
|
|
# Použije sa iba prvá stratégia,
|
|
# ktorá nájde aspoň jeden výsledok:
|
|
#
|
|
# all_terms -> prefix_terms -> any_term
|
|
#
|
|
# any_term teda nedopĺňa presné
|
|
# výsledky nerelevantným obsahom.
|
|
for strategy, match_query in match_queries:
|
|
rows = run_fts_query(
|
|
conn,
|
|
match_query,
|
|
candidate_limit,
|
|
published_only,
|
|
)
|
|
|
|
if not rows:
|
|
continue
|
|
|
|
for row in rows:
|
|
row["strategy"] = strategy
|
|
|
|
candidates = rows
|
|
used_strategies = [
|
|
strategy
|
|
]
|
|
|
|
break
|
|
|
|
results = add_labels_and_scores(
|
|
conn,
|
|
clean_query,
|
|
candidates,
|
|
)
|
|
|
|
return {
|
|
"engine": "sqlite_fts5",
|
|
"strategies": used_strategies,
|
|
"results": diversify_results(
|
|
results,
|
|
limit,
|
|
max_per_document,
|
|
),
|
|
}
|