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

2026-09-16 19:06:20 +08:00
parent 72fc5b47c7
commit d4b5b3f43a
122 changed files with 31906 additions and 0 deletions
@@ -0,0 +1,237 @@
+"""Small isolated RAG process pool with bounded admission and wall-clock deadlines."""
+from __future__ import annotations
+
+import atexit
+import json
+import math
+import os
+from pathlib import Path
+import queue
+import selectors
+import subprocess
+import sys
+import threading
+import time
+
+
+class RagWorkerError(RuntimeError):
+    pass
+
+
+class RagWorkerBusy(RagWorkerError):
+    pass
+
+
+class RagWorkerTimeout(RagWorkerError):
+    pass
+
+
+def _number(name: str, default: float, minimum: float, maximum: float) -> float:
+    try:
+        value = float(os.environ.get(name, default))
+        return min(maximum, max(minimum, value)) if math.isfinite(value) else default
+    except (TypeError, ValueError):
+        return default
+
+
+class _Worker:
+    def __init__(self, command: list[str], stderr_path: Path | None):
+        self.command = command
+        self.stderr_path = stderr_path
+        self.process: subprocess.Popen | None = None
+        self.last_start_ms = 0.0
+        self.last_reused = False
+
+    def start(self):
+        if self.process is not None and self.process.poll() is None:
+            self.last_start_ms = 0.0
+            self.last_reused = True
+            return self.process
+        self.close()
+        started = time.perf_counter()
+        stderr = subprocess.DEVNULL
+        if self.stderr_path:
+            self.stderr_path.parent.mkdir(parents=True, exist_ok=True)
+            stderr = self.stderr_path.open('ab')
+        try:
+            self.process = subprocess.Popen(
+                self.command, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+                stderr=stderr, bufsize=0)
+        finally:
+            if hasattr(stderr, 'close'):
+                stderr.close()
+        os.set_blocking(self.process.stdin.fileno(), False)
+        os.set_blocking(self.process.stdout.fileno(), False)
+        self.last_start_ms = round((time.perf_counter() - started) * 1000, 2)
+        self.last_reused = False
+        return self.process
+
+    def close(self):
+        process, self.process = self.process, None
+        if process is None:
+            return
+        try:
+            if process.poll() is None:
+                process.kill()
+            process.wait(timeout=0.2)
+        except (OSError, subprocess.TimeoutExpired):
+            pass
+        finally:
+            for pipe in (process.stdin, process.stdout):
+                if pipe:
+                    pipe.close()
+
+    def request(self, payload: dict, deadline: float) -> dict:
+        if time.monotonic() >= deadline:
+            raise RagWorkerTimeout('知识库检索超时，请缩小资料范围后重试。')
+        process = self.start()
+        encoded = memoryview((json.dumps(payload, ensure_ascii=False) + '\n').encode('utf-8'))
+        buffer = bytearray()
+        with selectors.DefaultSelector() as selector:
+            selector.register(process.stdin, selectors.EVENT_WRITE)
+            selector.register(process.stdout, selectors.EVENT_READ)
+            while True:
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise RagWorkerTimeout('知识库检索超时，请缩小资料范围后重试。')
+                events = selector.select(remaining)
+                if not events:
+                    raise RagWorkerTimeout('知识库检索超时，请缩小资料范围后重试。')
+                for key, _ in events:
+                    if key.fileobj is process.stdin:
+                        try:
+                            sent = os.write(key.fd, encoded[:65536])
+                        except BlockingIOError:
+                            continue
+                        encoded = encoded[sent:]
+                        if not encoded:
+                            selector.unregister(process.stdin)
+                    else:
+                        try:
+                            part = os.read(key.fd, 65536)
+                        except BlockingIOError:
+                            continue
+                        if not part:
+                            raise RagWorkerError('知识库检索进程意外退出，请重试。')
+                        buffer.extend(part)
+                        if len(buffer) > 32 * 1024 * 1024:
+                            raise RagWorkerError('知识库检索返回内容过大，请减少检索片段数。')
+                        if b'\n' in buffer:
+                            line, trailing = buffer.split(b'\n', 1)
+                            if trailing.strip():
+                                raise RagWorkerError('知识库检索进程返回格式异常，请重试。')
+                            try:
+                                response = json.loads(line)
+                            except (ValueError, UnicodeDecodeError) as exc:
+                                raise RagWorkerError('知识库检索进程返回格式异常，请重试。') from exc
+                            if not isinstance(response, dict):
+                                raise RagWorkerError('知识库检索进程返回格式异常，请重试。')
+                            return response
+
+
+class RagWorkerPool:
+    """At most three active workers, a small waiting room, no in-process fallback."""
+    def __init__(self, *, size: int = 2, total_timeout: float = 45,
+                 queue_timeout: float = 3, max_pending: int | None = None,
+                 command: list[str] | None = None, stderr_path: Path | None = None):
+        self.size = min(3, max(1, int(size)))
+        self.total_timeout = max(0.01, float(total_timeout))
+        self.queue_timeout = max(0.0, min(float(queue_timeout), self.total_timeout))
+        self._closed = False
+        self._workers = [_Worker(command or [sys.executable, str(Path(__file__).with_name('rag_worker.py'))],
+                                 stderr_path) for _ in range(self.size)]
+        self._available = queue.Queue(maxsize=self.size)
+        for worker in self._workers:
+            self._available.put_nowait(worker)
+        self._admission = threading.BoundedSemaphore(self.size + (max_pending if max_pending is not None else self.size * 2))
+
+    def request(self, payload: dict) -> dict:
+        started = time.monotonic()
+        deadline = started + self.total_timeout
+        if self._closed:
+            raise RagWorkerError('知识库检索服务正在关闭，请稍后重试。')
+        if not self._admission.acquire(blocking=False):
+            raise RagWorkerBusy('知识库检索任务较多，等待队列已满，请稍后重试。')
+        worker = None
+        try:
+            try:
+                worker = self._available.get(timeout=min(self.queue_timeout, max(0, deadline - time.monotonic())))
+            except queue.Empty as exc:
+                raise RagWorkerBusy('知识库检索排队超时，请稍后重试。') from exc
+            queue_ms = round((time.monotonic() - started) * 1000, 2)
+            try:
+                response = worker.request(payload, deadline)
+                if response.get('ok') is not True:
+                    from ..system_log import redact
+                    detail = redact(response.get('error') or '未知错误', 400)
+                    error = RagWorkerError(f'知识库检索失败：{detail}')
+                    setattr(error, 'retrieval_metric', response.get('retrieval_metric') or {"error": detail})
+                    raise error
+                if payload.get('op') == 'warmup':
+                    return response
+                if not isinstance(response.get('results'), list):
+                    raise RagWorkerError('知识库检索返回格式异常，请重试。')
+            except Exception as exc:
+                # A timed-out response may never be consumed by a subsequent request.
+                worker.close()
+                if isinstance(exc, RagWorkerError):
+                    raise
+                raise RagWorkerError('知识库检索进程通信失败，请重试。') from exc
+            response['transport_timings'] = {
+                'queue_ms': queue_ms, 'total_ms': round((time.monotonic() - started) * 1000, 2),
+                'worker_ms': round((time.monotonic() - started) * 1000 - queue_ms, 2),
+                'process_start_ms': worker.last_start_ms,
+                'process_reused': worker.last_reused,
+                'worker_pid': int(worker.process.pid) if worker.process else 0,
+            }
+            # Child search timing is reported in the response envelope.  The
+            # residual after spawn and search is the parent/child pipe and JSON
+            # transport overhead, useful for spotting IPC regressions.
+            child_ms = float((response.get('worker_timings') or {}).get('search_ms') or 0.0)
+            response['transport_timings']['ipc_ms'] = round(max(
+                0.0,
+                response['transport_timings']['worker_ms']
+                - child_ms
+                - float(response['transport_timings']['process_start_ms'] or 0.0)), 2)
+            return response
+        finally:
+            if worker is not None:
+                self._available.put_nowait(worker)
+            self._admission.release()
+
+    def prewarm(self) -> dict:
+        """Start one resident worker and warm its model/client in the background."""
+        return self.request({'op': 'warmup'})
+
+    def close(self):
+        self._closed = True
+        for worker in self._workers:
+            worker.close()
+
+
+_pool: RagWorkerPool | None = None
+_pool_lock = threading.Lock()
+
+
+def get_pool() -> RagWorkerPool:
+    global _pool
+    with _pool_lock:
+        if _pool is None:
+            from ..config import DATA_DIR
+            _pool = RagWorkerPool(
+                # The local embedding model is loaded once per worker. Keep one
+                # resident worker by default on the single-host deployment to
+                # avoid duplicating several hundred MB and triggering macOS
+                # memory pressure; operators can raise RAG_WORKER_POOL_SIZE
+                # explicitly when concurrent retrieval is required.
+                size=int(_number('RAG_WORKER_POOL_SIZE', 1, 1, 3)),
+                total_timeout=_number('RAG_SEARCH_TIMEOUT_SECONDS', 45, 1, 180),
+                # Content tasks commonly fan out several segment requests at
+                # once. Give queued requests enough time to use the single
+                # memory-efficient worker instead of surfacing a spurious
+                # timeout while an earlier segment is still being searched.
+                queue_timeout=_number('RAG_QUEUE_TIMEOUT_SECONDS', 10, 0, 30),
+                max_pending=int(_number('RAG_MAX_PENDING', 8, 0, 32)),
+                stderr_path=DATA_DIR / 'logs' / 'rag-worker.log')
+            atexit.register(_pool.close)
+        return _pool