1180 lines
20 KiB
Python
1180 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 scripts.embedding_utils import (
|
|
blob_to_vector,
|
|
cosine_similarity,
|
|
embed_query,
|
|
)
|
|
|
|
|
|
WORD_RE = re.compile(
|
|
r"[^\W_]+",
|
|
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
|
|
|
|
|
|
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 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 scripts/rebuild_index.py."
|
|
)
|
|
|
|
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 scripts/rebuild_index.py."
|
|
)
|
|
|
|
|
|
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"
|
|
)
|
|
|
|
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 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 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 chunks.published = 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"],
|
|
)
|
|
)
|
|
|
|
return candidates[
|
|
: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
|
|
|
|
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,
|
|
)
|
|
|
|
score = (
|
|
base_score
|
|
+ exact_match_bonus(
|
|
query,
|
|
item,
|
|
tags,
|
|
categories,
|
|
)
|
|
)
|
|
|
|
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 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
|
|
)
|
|
|
|
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
|