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()
|