feat(server): add workbench API, DSH planner worker and RAG pipeline

server/app/rag/rerank.py

@@ -0,0 +1,166 @@
+"""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