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

server/app/rag/rag_worker.py

@@ -0,0 +1,115 @@
+"""常驻检索进程。标准输出只用于 JSON 协议，初始化和检索诊断写标准错误。"""
+import sys
+import json
+import time
+from contextlib import redirect_stdout
+from pathlib import Path
+
+# 从任意 cwd 启动都能找到 app 包（uvicorn 工作目录不保证是 server/）
+_SERVER_ROOT = Path(__file__).resolve().parents[2]
+if str(_SERVER_ROOT) not in sys.path:
+    sys.path.insert(0, str(_SERVER_ROOT))
+
+
+def initialize():
+    """每个 worker 仅初始化一次，避免每次查询重复执行迁移和历史表扫描。"""
+    try:
+        import fcntl
+    except ImportError:
+        fcntl = None
+    from app import db as db_mod
+    from app.config import DATA_DIR
+    # 多个首次查询同时启动进程时，不并发执行相同的数据库迁移。
+    with (DATA_DIR / 'rag-worker-init.lock').open('a') as lock:
+        if fcntl:
+            fcntl.flock(lock, fcntl.LOCK_EX)
+        db_mod.init_db()
+
+
+def handle(payload: dict) -> list:
+    from sqlmodel import Session
+    from app import db as db_mod
+    from app.rag import pipeline as rag_pipeline
+
+    with Session(db_mod.engine) as session:
+        results = rag_pipeline.search_kbs(
+            session,
+            payload["kb_ids"],
+            payload["query"],
+            payload.get("k", 6),
+            excluded_chunk_ids={int(k): set(v) for k, v in (payload.get("excluded") or {}).items()},
+            campaign_id=payload.get("campaign_id", 0),
+            trace_id=payload.get("trace_id", ""),
+            user_id=payload.get("user_id", 0),
+            organization_id=payload.get("organization_id", 1),
+            file_whitelists={int(k): set(v) for k, v in (payload.get("file_whitelists") or {}).items()},
+        )
+    # 透传库名、召回来源、重排分数和各阶段耗时，不裁掉检索诊断信息。
+    return results
+
+
+def handle_warmup() -> dict:
+    """Warm the resident worker without issuing a billable retrieval."""
+    from app.rag import embed, store, rerank
+    started = time.monotonic()
+    embed.prewarm()
+    # Optional local Cross-Encoder is loaded in the resident worker so the
+    # first real retrieval does not pay model initialization latency.
+    rerank.prewarm()
+    # Construct the persistent vector client once. Collections remain opened
+    # lazily per KB, while the client and embedding runtime stay resident.
+    store.get_client()
+    return {"warm_state": "warm", "warmup_ms": round((time.monotonic() - started) * 1000, 2)}
+
+
+def main():
+    protocol_out = sys.stdout
+    try:
+        with redirect_stdout(sys.stderr):
+            initialize()
+    except Exception as exc:
+        from app.system_log import redact
+        protocol_out.write(json.dumps({'ok': False, 'error': '检索进程初始化失败：' + redact(exc, 400)}, ensure_ascii=False) + '\n')
+        protocol_out.flush()
+        return
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            started = time.monotonic()
+            with redirect_stdout(sys.stderr):
+                payload = json.loads(line)
+                if payload.get("op") == "warmup":
+                    warm = handle_warmup()
+                    protocol_out.write(json.dumps({"ok": True, **warm}, ensure_ascii=False) + "\n")
+                    protocol_out.flush()
+                    continue
+                out = handle(payload)
+            from app.rag import pipeline as _pipeline
+            metric = _pipeline.last_retrieval_metric()
+            envelope = {"ok": True, "results": out,
+                        "worker_timings": {"search_ms": round((time.monotonic() - started) * 1000, 2)},
+                        "retrieval_metric": metric}
+            serialization_started = time.perf_counter()
+            json.dumps(envelope, ensure_ascii=False)
+            metric["serialization_ms"] = round((time.perf_counter() - serialization_started) * 1000, 2)
+            envelope["retrieval_metric"] = metric
+            protocol_out.write(json.dumps(envelope, ensure_ascii=False) + "\n")
+        except Exception as exc:
+            from app.system_log import redact
+            metric = {}
+            try:
+                from app.rag import pipeline as _pipeline
+                metric = _pipeline.last_retrieval_metric()
+            except Exception:
+                pass
+            if not metric:
+                metric = {"error": "retrieval_failed"}
+            protocol_out.write(json.dumps({"ok": False, "error": redact(exc, 400),
+                "retrieval_metric": metric}, ensure_ascii=False) + "\n")
+        protocol_out.flush()
+
+
+if __name__ == "__main__":
+    main()