941 lines
20 KiB
Python
941 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
import sqlite3
|
|
import unicodedata
|
|
from collections import defaultdict
|
|
from typing import Any
|
|
|
|
from tools.rag.embedding_utils import (
|
|
blob_to_vector,
|
|
cosine_similarity,
|
|
embed_query,
|
|
model_profile,
|
|
)
|
|
|
|
WORD_RE = re.compile(
|
|
r"[^\W_]+",
|
|
re.UNICODE,
|
|
)
|
|
|
|
QUOTED_PHRASE_RE = re.compile(
|
|
r'[„“”"]([^„“”"]+)[„“”"]',
|
|
re.UNICODE,
|
|
)
|
|
|
|
|
|
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
|
|
|
|
RRF_K = 60
|
|
FTS_RRF_WEIGHT = 1.0
|
|
VECTOR_RRF_WEIGHT = 1.5
|
|
ANY_TERM_RRF_WEIGHT = 0.25
|
|
|
|
# Ak FTS kandidát obsahuje veľmi silný
|
|
# lexikálny signál, napríklad meno osoby
|
|
# alebo presný názov práce, nesmie byť
|
|
# potlačený iba embeddingovým výsledkom.
|
|
LEXICAL_ANCHOR_MIN_BONUS = 4.5
|
|
LEXICAL_ANCHOR_RRF_WEIGHT = 2.0
|
|
|
|
|
|
STRATEGY_PRIORITY = {
|
|
"all_terms": 3,
|
|
"prefix_terms": 2,
|
|
"any_term": 1,
|
|
}
|
|
|
|
|
|
def normalize_for_compare(
|
|
text: str,
|
|
) -> str:
|
|
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]:
|
|
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 quoted_query_phrases(
|
|
query: str,
|
|
) -> list[str]:
|
|
phrases: list[str] = []
|
|
seen: set[str] = set()
|
|
|
|
for match in QUOTED_PHRASE_RE.finditer(query):
|
|
phrase = normalize_for_compare(match.group(1))
|
|
|
|
if not phrase:
|
|
continue
|
|
|
|
# Jednoslovné výrazy ako "RAG"
|
|
# alebo "chatbot" sú príliš široké
|
|
# na silný lexikálny anchor.
|
|
if len(phrase.split()) < 2:
|
|
continue
|
|
|
|
if phrase in seen:
|
|
continue
|
|
|
|
phrases.append(phrase)
|
|
|
|
seen.add(phrase)
|
|
|
|
return phrases
|
|
|
|
|
|
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]]:
|
|
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:
|
|
fts_row = conn.execute("""
|
|
SELECT 1
|
|
FROM sqlite_master
|
|
WHERE type = 'table'
|
|
AND name = 'chunks_fts'
|
|
""").fetchone()
|
|
|
|
if fts_row is None:
|
|
raise RuntimeError(
|
|
"FTS5 index v databáze chýba. " "Spusti python -m tools.rag.rebuild_index."
|
|
)
|
|
|
|
embedding_row = conn.execute("""
|
|
SELECT 1
|
|
FROM sqlite_master
|
|
WHERE type = 'table'
|
|
AND name = 'chunk_embeddings'
|
|
""").fetchone()
|
|
|
|
if embedding_row is None:
|
|
raise RuntimeError(
|
|
"Embedding index v databáze chýba. " "Spusti python -m tools.rag.rebuild_index."
|
|
)
|
|
|
|
|
|
def embedding_index_info(
|
|
conn: sqlite3.Connection,
|
|
) -> tuple[str, int]:
|
|
row = conn.execute("""
|
|
SELECT
|
|
model,
|
|
dimensions
|
|
FROM chunk_embeddings
|
|
LIMIT 1
|
|
""").fetchone()
|
|
|
|
if row is None:
|
|
raise RuntimeError("Embedding index je prázdny")
|
|
if conn.execute("SELECT 1 FROM sqlite_master WHERE name='embedding_model_info'").fetchone():
|
|
metadata = conn.execute("SELECT metadata_json FROM embedding_model_info").fetchone()
|
|
if (
|
|
metadata
|
|
and json.loads(metadata[0])["fingerprint"]
|
|
!= model_profile(str(row["model"]))["fingerprint"]
|
|
):
|
|
raise RuntimeError("Verzia embeddingového modelu nezodpovedá indexu; obnovte index.")
|
|
|
|
return (
|
|
str(row["model"]),
|
|
int(row["dimensions"]),
|
|
)
|
|
|
|
|
|
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 COALESCE(chunks.published, 1) = 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 run_vector_query(
|
|
conn: sqlite3.Connection,
|
|
query: str,
|
|
candidate_limit: int,
|
|
published_only: bool,
|
|
) -> list[dict[str, Any]]:
|
|
model_name, dimensions = embedding_index_info(conn)
|
|
|
|
query_vector = embed_query(
|
|
query,
|
|
model_name=model_name,
|
|
)
|
|
|
|
if int(query_vector.shape[0]) != dimensions:
|
|
raise RuntimeError("Rozmer query embeddingu " "sa nezhoduje s indexom")
|
|
|
|
rows = conn.execute(
|
|
"""
|
|
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,
|
|
chunk_embeddings.embedding,
|
|
chunk_embeddings.dimensions
|
|
FROM chunk_embeddings
|
|
JOIN chunks
|
|
ON chunks.chunk_id
|
|
= chunk_embeddings.chunk_id
|
|
WHERE (
|
|
? = 0
|
|
OR COALESCE(chunks.published, 1) = 1
|
|
)
|
|
""",
|
|
(1 if published_only else 0,),
|
|
).fetchall()
|
|
|
|
candidates: list[dict[str, Any]] = []
|
|
|
|
for row in rows:
|
|
item = dict(row)
|
|
|
|
blob = item.pop("embedding")
|
|
|
|
stored_dimensions = int(item.pop("dimensions"))
|
|
|
|
vector = blob_to_vector(
|
|
blob,
|
|
stored_dimensions,
|
|
)
|
|
|
|
similarity = cosine_similarity(
|
|
query_vector,
|
|
vector,
|
|
)
|
|
|
|
item["vector_score"] = round(
|
|
similarity,
|
|
6,
|
|
)
|
|
|
|
candidates.append(item)
|
|
|
|
candidates.sort(
|
|
key=lambda item: (
|
|
-item["vector_score"],
|
|
item["document_path"],
|
|
item["chunk_index"],
|
|
)
|
|
)
|
|
|
|
# A chunk can have several lossless embedding segments. Rank by its best
|
|
# matching segment, and never spend several result slots on the same chunk.
|
|
unique = {}
|
|
for item in candidates:
|
|
unique.setdefault(item["chunk_id"], item)
|
|
return list(unique.values())[:candidate_limit]
|
|
|
|
|
|
def tokens_match_for_title(
|
|
title_token: str,
|
|
query_token: str,
|
|
) -> bool:
|
|
if not title_token or not query_token:
|
|
return False
|
|
|
|
if title_token == query_token:
|
|
return True
|
|
|
|
shorter_length = min(
|
|
len(title_token),
|
|
len(query_token),
|
|
)
|
|
|
|
# Pri veľmi krátkych tokenoch nechceme
|
|
# agresívne prefixové zhody.
|
|
if shorter_length < 3:
|
|
return False
|
|
|
|
return title_token.startswith(query_token) or query_token.startswith(title_token)
|
|
|
|
|
|
def query_contains_title(
|
|
query: str,
|
|
title: str,
|
|
) -> bool:
|
|
normalized_query = normalize_for_compare(query)
|
|
|
|
normalized_title = normalize_for_compare(title)
|
|
|
|
if not normalized_query or not normalized_title:
|
|
return False
|
|
|
|
title_tokens = normalized_title.split()
|
|
|
|
query_token_values = normalized_query.split()
|
|
|
|
# Jednoslovné názvy by mohli vytvárať
|
|
# príliš veľa falošných boostov.
|
|
if len(title_tokens) < 2:
|
|
return False
|
|
|
|
return all(
|
|
any(
|
|
tokens_match_for_title(
|
|
title_token,
|
|
query_token,
|
|
)
|
|
for query_token in query_token_values
|
|
)
|
|
for title_token in title_tokens
|
|
)
|
|
|
|
|
|
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
|
|
|
|
raw_title = item.get("title") or ""
|
|
|
|
title = normalize_for_compare(raw_title)
|
|
|
|
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
|
|
|
|
# Pri otázkach typu:
|
|
#
|
|
# Kto sa venuje práci s názvom
|
|
# „Systém získavania informácií
|
|
# v slovenskom jazyku“?
|
|
#
|
|
# je celý názov práce veľmi silný
|
|
# lexikálny signál, aj keď title
|
|
# dokumentu je meno študenta.
|
|
for phrase in quoted_query_phrases(query):
|
|
if phrase in title or phrase in text:
|
|
bonus += 6.0
|
|
break
|
|
|
|
if title == normalized_query:
|
|
bonus += 6.0
|
|
|
|
elif query_contains_title(
|
|
query,
|
|
raw_title,
|
|
):
|
|
bonus += 5.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:
|
|
if value is None:
|
|
return None
|
|
|
|
return bool(value)
|
|
|
|
|
|
def parse_heading_paths(
|
|
item: dict[str, Any],
|
|
) -> list:
|
|
try:
|
|
return json.loads(item.pop("heading_paths_json") or "[]")
|
|
|
|
except json.JSONDecodeError:
|
|
return []
|
|
|
|
|
|
def load_result_labels(
|
|
conn: sqlite3.Connection,
|
|
candidates: list[dict[str, Any]],
|
|
) -> tuple[
|
|
dict[str, list[str]],
|
|
dict[str, list[str]],
|
|
]:
|
|
chunk_ids = [item["chunk_id"] for item in candidates]
|
|
|
|
tags = load_labels(
|
|
conn,
|
|
"chunk_tags",
|
|
"tag",
|
|
chunk_ids,
|
|
)
|
|
|
|
categories = load_labels(
|
|
conn,
|
|
"chunk_categories",
|
|
"category",
|
|
chunk_ids,
|
|
)
|
|
|
|
return (
|
|
tags,
|
|
categories,
|
|
)
|
|
|
|
|
|
def add_fts_metadata(
|
|
conn: sqlite3.Connection,
|
|
query: str,
|
|
candidates: list[dict[str, Any]],
|
|
) -> list[dict[str, Any]]:
|
|
(
|
|
tags_by_chunk,
|
|
categories_by_chunk,
|
|
) = load_result_labels(
|
|
conn,
|
|
candidates,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
match_bonus = exact_match_bonus(
|
|
query,
|
|
item,
|
|
tags,
|
|
categories,
|
|
)
|
|
|
|
score = base_score + match_bonus
|
|
|
|
heading_paths = parse_heading_paths(item)
|
|
|
|
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 add_vector_metadata(
|
|
conn: sqlite3.Connection,
|
|
candidates: list[dict[str, Any]],
|
|
) -> list[dict[str, Any]]:
|
|
(
|
|
tags_by_chunk,
|
|
categories_by_chunk,
|
|
) = load_result_labels(
|
|
conn,
|
|
candidates,
|
|
)
|
|
|
|
results: list[dict[str, Any]] = []
|
|
|
|
for item in candidates:
|
|
chunk_id = item["chunk_id"]
|
|
|
|
item["published"] = database_bool(item.get("published"))
|
|
|
|
item["heading_paths"] = parse_heading_paths(item)
|
|
|
|
item["tags"] = tags_by_chunk.get(
|
|
chunk_id,
|
|
[],
|
|
)
|
|
|
|
item["categories"] = categories_by_chunk.get(
|
|
chunk_id,
|
|
[],
|
|
)
|
|
|
|
item["source_url"] = make_source_url(item["document_path"])
|
|
|
|
text = item.get("text") or ""
|
|
|
|
item["snippet"] = text[:320].strip()
|
|
|
|
results.append(item)
|
|
|
|
return results
|
|
|
|
|
|
def lexical_bonus_from_fts_result(
|
|
item: dict[str, Any],
|
|
) -> float:
|
|
fts_score = item.get("score")
|
|
|
|
bm25_score = item.get("bm25_score")
|
|
|
|
if fts_score is None or bm25_score is None:
|
|
return 0.0
|
|
|
|
try:
|
|
score_value = float(fts_score)
|
|
|
|
bm25_value = float(bm25_score)
|
|
|
|
except (
|
|
TypeError,
|
|
ValueError,
|
|
):
|
|
return 0.0
|
|
|
|
base_score = max(
|
|
0.0,
|
|
-bm25_value,
|
|
)
|
|
|
|
return max(
|
|
0.0,
|
|
score_value - base_score,
|
|
)
|
|
|
|
|
|
def fuse_hybrid_results(
|
|
fts_results: list[dict[str, Any]],
|
|
vector_results: list[dict[str, Any]],
|
|
) -> list[dict[str, Any]]:
|
|
merged: dict[
|
|
str,
|
|
dict[str, Any],
|
|
] = {}
|
|
|
|
scores: dict[
|
|
str,
|
|
float,
|
|
] = defaultdict(float)
|
|
|
|
fts_ranks: dict[
|
|
str,
|
|
int,
|
|
] = {}
|
|
|
|
vector_ranks: dict[
|
|
str,
|
|
int,
|
|
] = {}
|
|
|
|
for rank, item in enumerate(
|
|
fts_results,
|
|
start=1,
|
|
):
|
|
chunk_id = item["chunk_id"]
|
|
|
|
merged[chunk_id] = dict(item)
|
|
|
|
fts_ranks[chunk_id] = rank
|
|
|
|
strategy = item.get("match_strategy")
|
|
|
|
fts_weight = ANY_TERM_RRF_WEIGHT if strategy == "any_term" else FTS_RRF_WEIGHT
|
|
|
|
lexical_bonus = lexical_bonus_from_fts_result(item)
|
|
|
|
# Silný lexikálny anchor má prednosť
|
|
# pred čisto sémantickou podobnosťou.
|
|
#
|
|
# Typické prípady:
|
|
# - meno osoby je priamo title dokumentu,
|
|
# - názov práce je uvedený v úvodzovkách,
|
|
# - query obsahuje celý názov dokumentu.
|
|
#
|
|
# Toto rieši najmä situáciu, keď
|
|
# any_term FTS nájde správneho Bogdana
|
|
# na 1. mieste, ale pôvodná váha 0.25
|
|
# ho po RRF vytlačila mimo Top 5.
|
|
if lexical_bonus >= LEXICAL_ANCHOR_MIN_BONUS:
|
|
fts_weight = max(
|
|
fts_weight,
|
|
LEXICAL_ANCHOR_RRF_WEIGHT,
|
|
)
|
|
|
|
scores[chunk_id] += fts_weight / (RRF_K + rank)
|
|
|
|
for rank, item in enumerate(
|
|
vector_results,
|
|
start=1,
|
|
):
|
|
chunk_id = item["chunk_id"]
|
|
|
|
vector_ranks[chunk_id] = rank
|
|
|
|
scores[chunk_id] += VECTOR_RRF_WEIGHT / (RRF_K + rank)
|
|
|
|
if chunk_id not in merged:
|
|
merged[chunk_id] = dict(item)
|
|
|
|
else:
|
|
merged[chunk_id]["vector_score"] = item["vector_score"]
|
|
|
|
results: list[dict[str, Any]] = []
|
|
|
|
for chunk_id, item in merged.items():
|
|
item["fts_score"] = item.get("score")
|
|
|
|
item["fts_rank"] = fts_ranks.get(chunk_id)
|
|
|
|
item["vector_rank"] = vector_ranks.get(chunk_id)
|
|
|
|
item.setdefault(
|
|
"vector_score",
|
|
None,
|
|
)
|
|
|
|
item.setdefault(
|
|
"bm25_score",
|
|
None,
|
|
)
|
|
|
|
item.setdefault(
|
|
"match_strategy",
|
|
None,
|
|
)
|
|
|
|
hybrid_score = scores[chunk_id]
|
|
|
|
item["hybrid_score"] = round(
|
|
hybrid_score,
|
|
8,
|
|
)
|
|
|
|
item["score"] = round(
|
|
hybrid_score,
|
|
8,
|
|
)
|
|
|
|
results.append(item)
|
|
|
|
results.sort(
|
|
key=lambda item: (
|
|
-item["hybrid_score"],
|
|
item["document_path"],
|
|
item["chunk_index"],
|
|
)
|
|
)
|
|
|
|
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
|