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
|