547 lines
25 KiB
Python
547 lines
25 KiB
Python
"""Local Chinese lexical recall, reciprocal-rank fusion and scoped query caches.
|
||
|
||
The keyword channel has its own persisted FTS5 index, so it can find documents
|
||
outside vector top-k. Ranking is lightweight lexical ranking, not a paid/model
|
||
reranker. Derived indexes contain only the files in the requested scope.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from collections import OrderedDict
|
||
from contextlib import contextmanager
|
||
from copy import deepcopy
|
||
try:
|
||
import fcntl
|
||
except ImportError: # Windows
|
||
fcntl = None
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
import math
|
||
from pathlib import Path
|
||
import re
|
||
import sqlite3
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
|
||
from ..config import DATA_DIR
|
||
from . import store
|
||
|
||
logger = logging.getLogger("yawen.rag.hybrid")
|
||
INDEX_DIR = DATA_DIR / "rag_hybrid"
|
||
VERSION = "zh-bigram-fts5-rrf-v2"
|
||
CACHE_TTL = 300
|
||
CACHE_MAX = 192
|
||
CACHE_MAX_CHARS = 8_000_000
|
||
_cache: OrderedDict = OrderedDict()
|
||
_cache_lock = threading.RLock()
|
||
_build_locks: dict[str, threading.Lock] = {}
|
||
# Index construction is deliberately kept out of request threads. A single
|
||
# process can have several libraries become ready at once (for example after a
|
||
# bulk import), so use a small bounded pool and coalesce duplicate requests by
|
||
# revision. The persisted SQLite index remains protected by _build_lock.
|
||
_prewarm_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="rag-prewarm")
|
||
_prewarm_state_lock = threading.RLock()
|
||
_prewarm_pending: set[tuple[int, str]] = set()
|
||
_lexical_metric_local = threading.local()
|
||
_fusion_metric_local = threading.local()
|
||
_fts_pool_lock = threading.RLock()
|
||
_fts_connection_pool: OrderedDict[str, tuple[tuple, sqlite3.Connection, threading.RLock]] = OrderedDict()
|
||
_FTS_CONNECTION_MAX = 32
|
||
_STOP = {"如何", "怎么", "什么", "为什么", "为何", "可以", "需要", "我们", "你们", "他们", "这个", "那个", "哪些", "以及", "一个", "进行", "方法"}
|
||
_RUNS = re.compile(r"[\u3400-\u9fff]+|[a-z0-9]+(?:[._-][a-z0-9]+)*", re.I)
|
||
|
||
|
||
def tokens(value: str) -> list[str]:
|
||
"""Overlapping Chinese 2/3-grams retain names/numbers without a dictionary."""
|
||
result = []
|
||
for run in _RUNS.findall(str(value or "").lower()):
|
||
if '\u3400' <= run[0] <= '\u9fff':
|
||
if len(run) == 1:
|
||
result.append(run)
|
||
for size in (2, 3):
|
||
result.extend(run[i:i + size] for i in range(len(run) - size + 1)
|
||
if run[i:i + size] not in _STOP)
|
||
else:
|
||
result.append(run)
|
||
return result
|
||
|
||
|
||
def _digest(value) -> str:
|
||
return hashlib.sha256(json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode()).hexdigest()
|
||
|
||
|
||
def revision(kb_id: int) -> str:
|
||
"""Store revisions are persisted so invalidation also reaches query workers."""
|
||
return store.kb_revision(int(kb_id))
|
||
|
||
|
||
def warm_state(org: int, kb_id: int, rev: str | None = None,
|
||
files: set[str] | None = None) -> dict:
|
||
"""Return a lightweight, read-only FTS warm state for diagnostics."""
|
||
try:
|
||
library_id, current = int(kb_id), str(rev or revision(kb_id))
|
||
ready = index_available(int(org or 0), library_id, current, files)
|
||
with _prewarm_state_lock:
|
||
pending = (library_id, current) in _prewarm_pending
|
||
return {"state": "ready" if ready else ("warming" if pending else "cold"),
|
||
"ready": bool(ready), "pending": bool(pending),
|
||
"revision": current}
|
||
except Exception as exc:
|
||
return {"state": "error", "ready": False, "pending": False,
|
||
"error": type(exc).__name__}
|
||
|
||
|
||
def last_lexical_metric() -> dict:
|
||
return dict(getattr(_lexical_metric_local, "metric", {}) or {})
|
||
|
||
|
||
def last_fusion_metric() -> dict:
|
||
return dict(getattr(_fusion_metric_local, "metric", {}) or {})
|
||
|
||
|
||
def _scope_key(org: int, kb_id: int, rev: str, files: set[str] | None) -> str:
|
||
return _digest([VERSION, int(org), int(kb_id), rev, None if files is None else sorted(files)])
|
||
|
||
|
||
@contextmanager
|
||
def _build_lock(path: Path):
|
||
with _cache_lock:
|
||
lock = _build_locks.setdefault(str(path), threading.Lock())
|
||
with lock:
|
||
if fcntl is None:
|
||
yield
|
||
return
|
||
with path.with_suffix('.lock').open('a') as handle:
|
||
fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
|
||
try:
|
||
yield
|
||
finally:
|
||
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
||
|
||
|
||
def _valid_index(path: Path, signature: str) -> bool:
|
||
if not path.exists():
|
||
return False
|
||
try:
|
||
with sqlite3.connect(f"file:{path}?mode=ro", uri=True) as db:
|
||
row = db.execute("SELECT value FROM settings WHERE key='signature'").fetchone()
|
||
return bool(row and row[0] == signature)
|
||
except sqlite3.Error:
|
||
return False
|
||
|
||
|
||
def _fts_connection(path: Path) -> tuple[sqlite3.Connection, threading.RLock, bool]:
|
||
"""Return a shared read-only FTS connection for ``path``.
|
||
|
||
FTS indexes are immutable between atomic replacements. Reusing a
|
||
connection avoids repeatedly opening the index on the external volume;
|
||
the inode/size/mtime signature forces a reconnect whenever a new revision
|
||
replaces the file, so a worker never serves an old index indefinitely.
|
||
"""
|
||
stat = path.stat()
|
||
signature = (str(path), int(getattr(stat, "st_ino", 0)), int(stat.st_size),
|
||
int(getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1_000_000_000))))
|
||
key = str(path)
|
||
with _fts_pool_lock:
|
||
holder = _fts_connection_pool.get(key)
|
||
if holder and holder[0] == signature:
|
||
_fts_connection_pool.move_to_end(key)
|
||
return holder[1], holder[2], True
|
||
# A revision is published with an atomic rename. Serialize rotation
|
||
# with any query using the previous handle before closing it.
|
||
gate = holder[2] if holder else threading.RLock()
|
||
with gate:
|
||
uri = f"file:{path}?mode=ro"
|
||
# The bounded recall pool may hand the same KB to different
|
||
# workers across requests. A per-handle gate serializes access;
|
||
# disable sqlite's thread identity check so the handle can be
|
||
# safely reused by those workers.
|
||
connection = sqlite3.connect(uri, uri=True, check_same_thread=False)
|
||
# These indexes are read-only and query-heavy. A larger per-
|
||
# connection page cache and mmap keep repeated Chinese MATCH
|
||
# queries off the external volume.
|
||
for pragma in (
|
||
"PRAGMA query_only=ON",
|
||
"PRAGMA temp_store=MEMORY",
|
||
"PRAGMA cache_size=-32768",
|
||
"PRAGMA mmap_size=268435456",
|
||
):
|
||
try:
|
||
connection.execute(pragma)
|
||
except sqlite3.Error:
|
||
# Older SQLite builds may reject mmap_size; the query
|
||
# remains correct and still benefits from other settings.
|
||
if pragma != "PRAGMA mmap_size=268435456":
|
||
connection.close()
|
||
raise
|
||
if holder:
|
||
try:
|
||
holder[1].close()
|
||
except sqlite3.Error:
|
||
pass
|
||
_fts_connection_pool[key] = (signature, connection, gate)
|
||
_fts_connection_pool.move_to_end(key)
|
||
while len(_fts_connection_pool) > _FTS_CONNECTION_MAX:
|
||
old_key, (_, old_connection, old_gate) = _fts_connection_pool.popitem(last=False)
|
||
if old_key == key:
|
||
# Defensive guard; the just-created handle should never
|
||
# be evicted, but restore it if the limit is misconfigured.
|
||
_fts_connection_pool[key] = (signature, connection, gate)
|
||
break
|
||
with old_gate:
|
||
try:
|
||
old_connection.close()
|
||
except sqlite3.Error:
|
||
pass
|
||
return connection, gate, False
|
||
|
||
|
||
def _drop_fts_connection() -> None:
|
||
# The failed query may have come from any recall worker. Drop all handles
|
||
# so the next request reopens the affected index with a clean connection.
|
||
with _fts_pool_lock:
|
||
for _, connection, connection_gate in list(_fts_connection_pool.values()):
|
||
with connection_gate:
|
||
try:
|
||
connection.close()
|
||
except sqlite3.Error:
|
||
pass
|
||
_fts_connection_pool.clear()
|
||
|
||
|
||
def _index(org: int, kb_id: int, rev: str, files: set[str] | None) -> Path:
|
||
signature = _scope_key(org, kb_id, rev, files)
|
||
INDEX_DIR.mkdir(parents=True, exist_ok=True)
|
||
# One derived file per scope; revisions atomically replace it rather than
|
||
# retaining a full text copy of every historical library revision.
|
||
scope_name = _scope_key(org, kb_id, "", files)
|
||
path = INDEX_DIR / f"kb_{kb_id}_{scope_name}.sqlite3"
|
||
if _valid_index(path, signature):
|
||
return path
|
||
with _build_lock(path):
|
||
if _valid_index(path, signature):
|
||
return path
|
||
stage = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
|
||
started = time.perf_counter()
|
||
count = 0
|
||
try:
|
||
with sqlite3.connect(str(stage)) as db:
|
||
db.execute('PRAGMA journal_mode=OFF')
|
||
db.execute('PRAGMA synchronous=OFF')
|
||
db.execute('CREATE TABLE settings(key TEXT PRIMARY KEY, value TEXT)')
|
||
db.execute('CREATE TABLE sources(rowid INTEGER PRIMARY KEY, chunk_id TEXT NOT NULL, source_text TEXT NOT NULL, metadata TEXT NOT NULL)')
|
||
# Contentless FTS stores postings, not another expanded n-gram copy
|
||
# of every source document. The original text lives once in sources.
|
||
db.execute('CREATE VIRTUAL TABLE docs USING fts5(body, heading, content="", tokenize="unicode61")')
|
||
batch = []
|
||
for item in store.iter_chunks(kb_id, file_names=files):
|
||
meta = item.get('meta') or {}
|
||
# Defense in depth, including future store implementations.
|
||
if files is not None and meta.get('file') not in files:
|
||
raise RuntimeError('关键词索引读取到了所选分类以外的文件')
|
||
source = str(item.get('text') or '')
|
||
if not source.strip() or not item.get('id'):
|
||
continue
|
||
count += 1
|
||
batch.append((count, ' '.join(tokens(source)), ' '.join(tokens(meta.get('heading', ''))),
|
||
str(item['id']), source, json.dumps(meta, ensure_ascii=False)))
|
||
if len(batch) >= 128:
|
||
db.executemany('INSERT INTO docs(rowid,body,heading) VALUES(?,?,?)', (row[:3] for row in batch))
|
||
db.executemany('INSERT INTO sources VALUES(?,?,?,?)', ((row[0], *row[3:]) for row in batch))
|
||
batch.clear()
|
||
if batch:
|
||
db.executemany('INSERT INTO docs(rowid,body,heading) VALUES(?,?,?)', (row[:3] for row in batch))
|
||
db.executemany('INSERT INTO sources VALUES(?,?,?,?)', ((row[0], *row[3:]) for row in batch))
|
||
if revision(kb_id) != rev:
|
||
raise RuntimeError('知识库索引正在更新，请稍后重新检索')
|
||
db.execute('INSERT INTO settings VALUES(?,?)', ('signature', signature))
|
||
db.execute("INSERT INTO docs(docs) VALUES('optimize')")
|
||
stage.replace(path)
|
||
logger.info('中文关键词索引就绪：知识库 %s，%s 段，%.3f 秒', kb_id, count, time.perf_counter() - started)
|
||
finally:
|
||
stage.unlink(missing_ok=True)
|
||
return path
|
||
|
||
|
||
def _index_path(org: int, kb_id: int, rev: str, files: set[str] | None) -> Path:
|
||
"""Return the persisted index path without opening or building it."""
|
||
scope_name = _scope_key(org, kb_id, "", files)
|
||
INDEX_DIR.mkdir(parents=True, exist_ok=True)
|
||
return INDEX_DIR / f"kb_{int(kb_id)}_{scope_name}.sqlite3"
|
||
|
||
|
||
def index_available(org: int, kb_id: int, rev: str, files: set[str] | None = None) -> bool:
|
||
"""Return whether a derived FTS index already exists for this scope.
|
||
|
||
Building a first index can require scanning tens of thousands of chunks. It
|
||
must never happen synchronously on a content generation request: callers can
|
||
use the vector channel immediately and let ``schedule_prewarm`` finish the
|
||
lexical index in the background.
|
||
"""
|
||
try:
|
||
signature = _scope_key(int(org or 0), int(kb_id), str(rev), files)
|
||
path = _index_path(int(org or 0), int(kb_id), str(rev), files)
|
||
return _valid_index(path, signature)
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def schedule_prewarm(org: int, kb_id: int, rev: str | None = None) -> bool:
|
||
"""Build the full-library lexical index in the background.
|
||
|
||
Upload/rebuild callers invoke this only after their transaction has
|
||
committed. It is intentionally best-effort: a failed prewarm must never
|
||
make an otherwise successful upload fail; the normal first query can still
|
||
build the index under the same cross-process lock. Requests for the same
|
||
library revision are coalesced so repeated progress callbacks do not spawn
|
||
unbounded workers.
|
||
"""
|
||
try:
|
||
library_id = int(kb_id)
|
||
current_revision = str(rev or revision(library_id))
|
||
key = (library_id, current_revision)
|
||
except Exception:
|
||
return False
|
||
with _prewarm_state_lock:
|
||
if key in _prewarm_pending:
|
||
return False
|
||
_prewarm_pending.add(key)
|
||
|
||
def _worker():
|
||
started = time.perf_counter()
|
||
try:
|
||
# A newer upload may have committed while this task was queued;
|
||
# skip stale work and let the newer revision be scheduled instead.
|
||
if revision(library_id) != current_revision:
|
||
return
|
||
_index(int(org or 0), library_id, current_revision, None)
|
||
logger.info("知识库全文索引已后台预热：知识库 %s，%.3f 秒",
|
||
library_id, time.perf_counter() - started)
|
||
except Exception:
|
||
# Best effort only. Keep the exception in server logs for
|
||
# diagnosis, while preserving the successful import/rebuild API.
|
||
logger.exception("知识库全文索引后台预热失败：知识库 %s", library_id)
|
||
finally:
|
||
with _prewarm_state_lock:
|
||
_prewarm_pending.discard(key)
|
||
|
||
try:
|
||
_prewarm_executor.submit(_worker)
|
||
return True
|
||
except RuntimeError:
|
||
# Interpreter shutdown or executor teardown; requests can fall back to
|
||
# lazy construction on the next search.
|
||
with _prewarm_state_lock:
|
||
_prewarm_pending.discard(key)
|
||
return False
|
||
|
||
|
||
def lexical_search(org: int, kb_id: int, rev: str, query: str, limit: int,
|
||
files: set[str] | None = None) -> list[dict]:
|
||
started = time.perf_counter()
|
||
metric = {"kb_id": int(kb_id), "fts_ms": 0.0, "metadata_filter_ms": 0.0,
|
||
"warm_state": warm_state(org, kb_id, rev, files)}
|
||
_lexical_metric_local.metric = metric
|
||
if files is not None and not files:
|
||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||
return []
|
||
terms = list(dict.fromkeys(tokens(query)))[:192]
|
||
if not terms:
|
||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||
return []
|
||
# A category/file scope can be very large (tens of thousands of chunks).
|
||
# Building a dedicated FTS database synchronously in the generation
|
||
# request used to block the RAG worker until its 45s deadline. Full
|
||
# library indexes are prewarmed after import, so use one as a safe fallback
|
||
# and filter its metadata after the match. The vector channel still runs
|
||
# concurrently when a full index is being prepared.
|
||
post_filter_files = None
|
||
scope_path = _index_path(org, kb_id, rev, files)
|
||
scope_signature = _scope_key(org, kb_id, rev, files)
|
||
if files is not None and not _valid_index(scope_path, scope_signature):
|
||
full_path = _index_path(org, kb_id, rev, None)
|
||
full_signature = _scope_key(org, kb_id, rev, None)
|
||
if not _valid_index(full_path, full_signature):
|
||
# Do not construct a potentially multi-hundred-MB index on the
|
||
# request thread. Import/rebuild schedules this work in the
|
||
# background; returning no lexical hits is preferable to timing
|
||
# out the whole generation request.
|
||
schedule_prewarm(org, kb_id, rev)
|
||
logger.info("知识库 %s 的全文关键词索引尚未就绪，已后台预热；本次仅使用向量召回", kb_id)
|
||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||
return []
|
||
path = full_path
|
||
post_filter_files = files
|
||
else:
|
||
# The full-library index is normally ready from the post-import
|
||
# prewarm. If a revision changed and the replacement is still being
|
||
# built, keep this request bounded as well instead of synchronously
|
||
# scanning a large collection in the RAG worker.
|
||
if files is None:
|
||
full_path = _index_path(org, kb_id, rev, None)
|
||
full_signature = _scope_key(org, kb_id, rev, None)
|
||
if not _valid_index(full_path, full_signature):
|
||
schedule_prewarm(org, kb_id, rev)
|
||
logger.info("知识库 %s 的全文关键词索引尚未就绪，已后台预热；本次仅使用向量召回", kb_id)
|
||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||
return []
|
||
path = full_path
|
||
else:
|
||
path = scope_path
|
||
# Every token is quoted; user input cannot introduce FTS operators.
|
||
match = ' OR '.join('"' + term.replace('"', '""') + '"' for term in terms)
|
||
try:
|
||
fts_started = time.perf_counter()
|
||
db, db_gate, reused = _fts_connection(path)
|
||
metric["connection_reused"] = bool(reused)
|
||
# Pull a bounded oversample when filtering a full-library index so
|
||
# that relevant files are not lost behind unrelated matches.
|
||
fetch_limit = min(2000, max(limit, limit * 8)) if post_filter_files is not None else limit
|
||
with db_gate:
|
||
rows = db.execute('SELECT sources.chunk_id,sources.source_text,sources.metadata,bm25(docs,1.0,2.2) AS rank FROM docs JOIN sources ON sources.rowid=docs.rowid WHERE docs MATCH ? ORDER BY rank,sources.chunk_id LIMIT ?',
|
||
(match, fetch_limit)).fetchall()
|
||
metric["fts_ms"] = round((time.perf_counter() - fts_started) * 1000, 2)
|
||
except sqlite3.Error as exc:
|
||
_drop_fts_connection()
|
||
raise RuntimeError(f'知识库 {kb_id} 关键词检索失败：{exc}') from exc
|
||
result = []
|
||
for row in rows:
|
||
filter_started = time.perf_counter()
|
||
metadata = json.loads(row[2])
|
||
if post_filter_files is not None and metadata.get('file') not in post_filter_files:
|
||
metric["metadata_filter_ms"] += (time.perf_counter() - filter_started) * 1000
|
||
continue
|
||
result.append({'id': row[0], 'text': row[1], 'meta': metadata,
|
||
'lexical_score': max(0.0, -float(row[3]))})
|
||
if len(result) >= limit:
|
||
metric["metadata_filter_ms"] += (time.perf_counter() - filter_started) * 1000
|
||
break
|
||
metric["metadata_filter_ms"] += (time.perf_counter() - filter_started) * 1000
|
||
metric["metadata_filter_ms"] = round(metric["metadata_filter_ms"], 2)
|
||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||
return result
|
||
|
||
|
||
def cache_get(key):
|
||
with _cache_lock:
|
||
cached = _cache.get(key)
|
||
if cached is None:
|
||
return None
|
||
if time.monotonic() - cached[0] > CACHE_TTL:
|
||
_cache.pop(key, None)
|
||
return None
|
||
_cache.move_to_end(key)
|
||
return deepcopy(cached[1])
|
||
|
||
|
||
def cache_put(key, value):
|
||
with _cache_lock:
|
||
size = len(json.dumps(value, ensure_ascii=False))
|
||
if size > CACHE_MAX_CHARS:
|
||
return
|
||
_cache[key] = (time.monotonic(), deepcopy(value), size)
|
||
_cache.move_to_end(key)
|
||
while len(_cache) > CACHE_MAX or sum(row[2] for row in _cache.values()) > CACHE_MAX_CHARS:
|
||
_cache.popitem(last=False)
|
||
|
||
|
||
def clear_cache():
|
||
with _cache_lock:
|
||
_cache.clear()
|
||
|
||
|
||
def prewarm_index(org: int, kb_id: int, files: set[str] | None = None) -> None:
|
||
"""Build the FTS scope outside the user's request path.
|
||
|
||
Import/rebuild jobs call this after the vector index is committed. The
|
||
existing file lock makes concurrent warmups idempotent and safe across
|
||
API and worker processes.
|
||
"""
|
||
_index(int(org or 1), int(kb_id), revision(int(kb_id)), files)
|
||
|
||
|
||
def candidate_key(org: int, user: int, scopes: list[tuple], query: str, limit: int,
|
||
model: str, *, mode: str = "hybrid", top_k: int | None = None,
|
||
permissions=None, index_version: str | None = None) -> str:
|
||
"""Build an isolated retrieval candidate-pool cache key.
|
||
|
||
Query vectors are cached in :mod:`embed` and deliberately do not use this
|
||
key. Candidate pools depend on every scope and authorization dimension;
|
||
callers should pass a stable permission signature when one is available.
|
||
Defaults preserve compatibility with older six-argument callers.
|
||
"""
|
||
# ``scopes`` already carries KB id, index revision and category/file scope.
|
||
# Sorting permission values makes equivalent grant sets map to one key while
|
||
# still preventing cross-user or cross-permission pool reuse.
|
||
permission_scope = permissions
|
||
if isinstance(permission_scope, (set, frozenset, list, tuple)):
|
||
permission_scope = sorted(str(value) for value in permission_scope)
|
||
return _digest([VERSION, int(org or 0), int(user or 0), scopes,
|
||
str(query).strip(), int(limit), int(top_k or limit),
|
||
str(model), str(mode or "hybrid"), permission_scope,
|
||
str(index_version or "")])
|
||
|
||
|
||
def fuse(vector: list[dict], lexical: list[dict], query: str) -> list[dict]:
|
||
"""RRF balances independent recalls; token coverage refines bounded candidates."""
|
||
started = time.perf_counter()
|
||
metric = {"merge_ms": 0.0, "dedup_ms": 0.0, "rrf_ms": 0.0}
|
||
_fusion_metric_local.metric = metric
|
||
merge_started = time.perf_counter()
|
||
merged = {}
|
||
for channel, rows in (('vector', vector), ('lexical', lexical)):
|
||
for rank, row in enumerate(rows, 1):
|
||
key = (int(row['kb_id']), str(row['id']))
|
||
current = merged.setdefault(key, {**deepcopy(row), 'retrieval_channels': [], '_rrf': 0.0})
|
||
current['retrieval_channels'].append(channel)
|
||
current['_rrf'] += 1.0 / (60 + rank)
|
||
if channel == 'lexical':
|
||
current['lexical_score'] = row.get('lexical_score', 0)
|
||
metric["merge_ms"] = round((time.perf_counter() - merge_started) * 1000, 2)
|
||
metric["dedup_ms"] = metric["merge_ms"]
|
||
rrf_started = time.perf_counter()
|
||
qterms = set(tokens(query))
|
||
# Candidate-local inverse document frequency makes rare, exact names useful.
|
||
df = {}
|
||
term_sets = {}
|
||
for key, row in merged.items():
|
||
terms = set(tokens(row.get('text', '') + ' ' + str((row.get('meta') or {}).get('heading', ''))))
|
||
term_sets[key] = terms
|
||
for term in qterms & terms:
|
||
df[term] = df.get(term, 0) + 1
|
||
weights = {term: math.log(1 + (len(merged) + 1) / (df.get(term, 0) + 1)) for term in qterms}
|
||
total_weight = sum(weights.values()) or 1
|
||
for key, row in merged.items():
|
||
coverage = sum(weights[term] for term in qterms & term_sets[key]) / total_weight
|
||
row['score'] = row.pop('_rrf') + 0.045 * coverage
|
||
row['keyword_coverage'] = round(coverage, 4)
|
||
result = sorted(merged.values(), key=lambda row: (-row['score'], int(row['kb_id']), str(row['id'])))
|
||
metric["rrf_ms"] = round((time.perf_counter() - rrf_started) * 1000, 2)
|
||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||
return result
|
||
|
||
|
||
def finalize(candidates: list[dict], k: int, excluded: dict[int, set[str]],
|
||
allowed: dict[int, set[str] | None]) -> list[dict]:
|
||
"""Always run after cache lookup: consumption and access scopes are live."""
|
||
output, seen = [], set()
|
||
for row in candidates:
|
||
kid, meta = int(row['kb_id']), row.get('meta') or {}
|
||
if kid not in allowed or row['id'] in excluded.get(kid, set()):
|
||
continue
|
||
files = allowed[kid]
|
||
if files is not None and meta.get('file') not in files:
|
||
continue
|
||
# Same heading/body in one KB is a repeat; numbers and distinct headings survive.
|
||
fingerprint = (kid, re.sub(r'\s+', ' ', str(meta.get('heading') or '')).strip(),
|
||
re.sub(r'\s+', ' ', row.get('text', '')).strip())
|
||
if fingerprint in seen:
|
||
continue
|
||
seen.add(fingerprint)
|
||
output.append(row)
|
||
if len(output) >= k:
|
||
break
|
||
return output
|