241 lines
5.0 KiB
Python
241 lines
5.0 KiB
Python
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
# Tieto importy sú zámerne verejné.
|
|
# Zachovávajú spätnú kompatibilitu pre testy
|
|
# a evaluačné skripty, ktoré ich importujú
|
|
# zo scripts.search_utils.
|
|
from scripts.search_core import (
|
|
ANY_TERM_RRF_WEIGHT,
|
|
BM25_SQL,
|
|
BM25_WEIGHTS,
|
|
DEFAULT_CANDIDATE_MULTIPLIER,
|
|
FTS_RRF_WEIGHT,
|
|
MIN_CANDIDATES,
|
|
MIN_STEM_PREFIX_LENGTH,
|
|
RRF_K,
|
|
STRATEGY_PRIORITY,
|
|
VECTOR_RRF_WEIGHT,
|
|
WORD_RE,
|
|
add_fts_metadata,
|
|
add_vector_metadata,
|
|
build_match_queries,
|
|
database_bool,
|
|
diversify_results,
|
|
embedding_index_info,
|
|
exact_match_bonus,
|
|
fuse_hybrid_results,
|
|
load_labels,
|
|
load_result_labels,
|
|
make_source_url,
|
|
normalize_for_compare,
|
|
parse_heading_paths,
|
|
query_contains_title,
|
|
query_tokens,
|
|
quote_fts_token,
|
|
run_fts_query,
|
|
run_vector_query,
|
|
tokens_match_for_title,
|
|
verify_search_schema,
|
|
)
|
|
|
|
|
|
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(
|
|
"Databáza neexistuje: "
|
|
f"{db_file}"
|
|
)
|
|
|
|
clean_query = query.strip()
|
|
|
|
if not clean_query:
|
|
return {
|
|
"engine": (
|
|
"hybrid_fts5_embeddings"
|
|
),
|
|
"strategies": [],
|
|
"results": [],
|
|
}
|
|
|
|
match_queries = (
|
|
build_match_queries(
|
|
clean_query
|
|
)
|
|
)
|
|
|
|
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
|
|
)
|
|
|
|
# -------------------------
|
|
# FTS5 kandidáti
|
|
# -------------------------
|
|
|
|
fts_candidates: list[
|
|
dict[str, Any]
|
|
] = []
|
|
|
|
used_strategies: list[
|
|
str
|
|
] = []
|
|
|
|
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
|
|
|
|
fts_candidates = rows
|
|
|
|
used_strategies = [
|
|
strategy
|
|
]
|
|
|
|
# Používa sa prvá stratégia,
|
|
# ktorá vráti výsledky.
|
|
break
|
|
|
|
fts_results = (
|
|
add_fts_metadata(
|
|
conn,
|
|
clean_query,
|
|
fts_candidates,
|
|
)
|
|
)
|
|
|
|
# -------------------------
|
|
# Embedding kandidáti
|
|
# -------------------------
|
|
|
|
vector_candidates = (
|
|
run_vector_query(
|
|
conn,
|
|
clean_query,
|
|
candidate_limit,
|
|
published_only,
|
|
)
|
|
)
|
|
|
|
vector_results = (
|
|
add_vector_metadata(
|
|
conn,
|
|
vector_candidates,
|
|
)
|
|
)
|
|
|
|
# -------------------------
|
|
# Presné FTS výsledky
|
|
# -------------------------
|
|
#
|
|
# Pri all_terms alebo prefix_terms
|
|
# embeddingy iba preradia chunky,
|
|
# ktoré už našiel FTS.
|
|
#
|
|
# Tým sa zabráni pridávaniu
|
|
# sémanticky podobných, ale
|
|
# lexikálne nesúvisiacich výsledkov.
|
|
|
|
if (
|
|
used_strategies
|
|
and used_strategies[0]
|
|
in {
|
|
"all_terms",
|
|
"prefix_terms",
|
|
}
|
|
):
|
|
fts_chunk_ids = {
|
|
item[
|
|
"chunk_id"
|
|
]
|
|
for item
|
|
in fts_results
|
|
}
|
|
|
|
vector_results = [
|
|
item
|
|
for item
|
|
in vector_results
|
|
if item[
|
|
"chunk_id"
|
|
]
|
|
in fts_chunk_ids
|
|
]
|
|
|
|
# -------------------------
|
|
# Hybrid fusion
|
|
# -------------------------
|
|
|
|
hybrid_results = (
|
|
fuse_hybrid_results(
|
|
fts_results,
|
|
vector_results,
|
|
)
|
|
)
|
|
|
|
# -------------------------
|
|
# Document diversification
|
|
# -------------------------
|
|
|
|
final_results = (
|
|
diversify_results(
|
|
hybrid_results,
|
|
limit,
|
|
max_per_document,
|
|
)
|
|
)
|
|
|
|
return {
|
|
"engine": (
|
|
"hybrid_fts5_embeddings"
|
|
),
|
|
"strategies": (
|
|
used_strategies
|
|
),
|
|
"results": (
|
|
final_results
|
|
),
|
|
}
|