116 lines
4.5 KiB
Python
116 lines
4.5 KiB
Python
|
"""常驻检索进程。标准输出只用于 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()
|