yawan-workbench/server/app/rag/rerank.py

"""Optional local Cross-Encoder candidate reranking.

The first-stage vector/FTS recall remains the default.  A Cross-Encoder is
only enabled when ``RAG_RERANKER_MODEL`` points at a local model directory (or
``RAG_RERANKER_ALLOW_DOWNLOAD=1`` is explicitly set).  This keeps production
search deterministic and offline when the optional model is absent, while
allowing operators to add a Chinese model such as ``BAAI/bge-reranker-base``
without changing the retrieval contract.
"""
from __future__ import annotations

import logging
import os
from pathlib import Path
import threading
import time

logger = logging.getLogger("yawen.rag.rerank")

_lock = threading.Lock()
_metric_lock = threading.Lock()
_model = None
_model_name = ""
_model_error = ""
_last_metric: dict = {
    "enabled": False,
    "model": "",
    "model_loaded": False,
    "model_load_ms": 0.0,
    "inference_ms": 0.0,
    "candidate_count": 0,
    "reranked_count": 0,
    "fallback": True,
}


def _configured_model() -> str:
    return str(os.environ.get("RAG_RERANKER_MODEL", "")).strip()


def _allow_download() -> bool:
    return str(os.environ.get("RAG_RERANKER_ALLOW_DOWNLOAD", "0")).lower() in {
        "1", "true", "yes", "on"
    }


def _set_metric(**values) -> None:
    global _last_metric
    with _metric_lock:
        row = dict(_last_metric)
        row.update(values)
        for key in ("model_load_ms", "inference_ms", "total_ms"):
            if key in row:
                try:
                    row[key] = round(float(row[key]), 2)
                except (TypeError, ValueError):
                    row[key] = 0.0
        _last_metric = row


def last_metric() -> dict:
    with _metric_lock:
        return dict(_last_metric)


def _local_model_path(name: str) -> str | None:
    """Resolve a configured model without triggering a network download."""
    candidate = Path(name).expanduser()
    if candidate.is_dir() and (candidate / "config.json").exists():
        return str(candidate)
    try:
        from ..config import MODELS_DIR
        cached = MODELS_DIR / name.replace("/", "--")
        if cached.is_dir() and (cached / "config.json").exists():
            return str(cached)
    except Exception:
        pass
    return None


def _get_model():
    """Lazily load a local sentence-transformers CrossEncoder."""
    global _model, _model_name, _model_error
    name = _configured_model()
    if not name:
        return None
    with _lock:
        if _model is not None and _model_name == name:
            return _model
        if _model_error and _model_name == name:
            return None
        started = time.perf_counter()
        local = _local_model_path(name)
        if local is None and not _allow_download():
            _model_name = name
            _model_error = "模型目录不存在；为避免请求线程下载，已回退基础排序"
            _set_metric(enabled=False, model=name, model_loaded=False,
                        fallback=True, error=_model_error)
            logger.warning("RAG 重排模型不可用：%s（%s）", name, _model_error)
            return None
        try:
            from sentence_transformers import CrossEncoder
            _model = CrossEncoder(local or name)
            _model_name = name
            _model_error = ""
            _set_metric(enabled=True, model=name, model_loaded=True,
                        model_load_ms=(time.perf_counter() - started) * 1000,
                        fallback=False)
            return _model
        except Exception as exc:  # optional feature must never break RAG
            _model = None
            _model_name = name
            _model_error = f"{type(exc).__name__}: {exc}"
            _set_metric(enabled=False, model=name, model_loaded=False,
                        model_load_ms=(time.perf_counter() - started) * 1000,
                        fallback=True, error=_model_error)
            logger.exception("RAG Cross-Encoder 加载失败，回退基础排序：%s", name)
            return None


def rerank(candidates: list[dict], query: str, *, limit: int | None = None) -> list[dict]:
    """Rerank candidates in place-compatible copies and return stable results.

    Only the first ``limit`` candidates are scored; the tail keeps its fused
    order.  If the model is not configured/available, candidates are returned
    unchanged and no exception escapes the retrieval path.
    """
    configured = _configured_model()
    if not candidates or not configured:
        _set_metric(enabled=False, model=configured, candidate_count=len(candidates),
                    reranked_count=0, fallback=True, inference_ms=0.0)
        return candidates

    count = min(len(candidates), max(1, min(int(limit or os.environ.get("RAG_RERANKER_TOP_K", 32)), 64)))
    model = _get_model()
    if model is None:
        _set_metric(candidate_count=len(candidates), reranked_count=0, fallback=True)
        return candidates
    pairs = [(str(query), str(row.get("text") or "")) for row in candidates[:count]]
    started = time.perf_counter()
    try:
        scores = model.predict(pairs, batch_size=min(32, count), show_progress_bar=False)
        scored = []
        for index, (row, score) in enumerate(zip(candidates[:count], scores)):
            item = dict(row)
            item["rerank_score"] = float(score)
            scored.append((item, float(score), index))
        scored.sort(key=lambda value: (-value[1], value[2]))
        result = [item for item, _, _ in scored] + [dict(row) for row in candidates[count:]]
        _set_metric(candidate_count=len(candidates), reranked_count=count,
                    inference_ms=(time.perf_counter() - started) * 1000,
                    fallback=False)
        return result
    except Exception as exc:
        logger.exception("RAG Cross-Encoder 推理失败，回退基础排序")
        _set_metric(candidate_count=len(candidates), reranked_count=0,
                    inference_ms=(time.perf_counter() - started) * 1000,
                    fallback=True, error=f"{type(exc).__name__}: {exc}")
        return candidates


def prewarm() -> bool:
    """Load the optional reranker in the resident RAG worker at startup."""
    if not _configured_model():
        return False
    return _get_model() is not None