@hupan56/wlkj 3.0.0 → 3.1.2

package/templates/qoder/scripts/domain/kg/server/kgd.py

@@ -0,0 +1,549 @@
+# -*- coding: utf-8 -*-
+# v3.0 路径自举: 引导到 common/bootstrap, 统一 sys.path 逻辑
+import os as _o, sys as _s
+_f = _o.path.abspath(__file__)
+for _ in range(10):
+    _f = _o.path.dirname(_f)
+    _cp = _o.path.join(_f, 'foundation')
+    if _o.path.isfile(_o.path.join(_cp, 'bootstrap.py')):
+        break
+if _cp not in _s.path: _s.path.insert(0, _cp)
+from bootstrap import setup; setup()
+
+"""
+kgd.py — 知识图谱常驻加速服务 (Knowledge-Graph Daemon)
+
+治"用了工作流反而更慢"的根因: 每次命令起一个新 Python 进程,
+冷启动税 (import duckdb + 加载 95MB bge 模型) 把毫秒级能力放大成几秒~26秒。
+
+本进程常驻内存, 启动时一次性加载:
+  • bge-small-zh 向量模型 (语义搜索 26.8s → <0.5s, 53x)
+  • DuckDB 连接 (避免每次重连)
+  • search/context/graph 全部模块 (避免重复 import)
+
+CLI (kg.py) 和 MCP 调用时优先连 daemon (热路径), 连不上则透明回退冷路径。
+
+协议: TCP localhost (跨平台最稳), 每请求一行 JSON, 每响应一行 JSON。
+  请求: {"cmd": "semantic", "args": ["报销"]}
+  响应: {"ok": true, "text": "...", "secs": 0.31}
+
+用法:
+  python kgd.py start          # 后台启动 daemon
+  python kgd.py start --fore   # 前台启动 (调试, 看日志)
+  python kgd.py stop           # 停止
+  python kgd.py status         # 查看状态 (模型加载否/内存/命中)
+  python kgd.py restart        # 重启
+
+渐进增强原则: daemon 是可选加速层。没有 daemon, kg.py 照常工作 (冷路径)。
+"""
+import io
+import json
+import os
+import socket
+import subprocess
+import sys
+import time
+
+try:
+    sys.stdout.reconfigure(encoding='utf-8', errors='replace')
+except Exception:
+    pass
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+if _HERE not in sys.path:
+    sys.path.insert(0, _HERE)  # 同包 import (kg, search_index ...)
+
+RUNTIME_DIR = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.dirname(_HERE))),  # scripts→.qoder→repo
+    '.qoder', '.runtime')
+PID_FILE = os.path.join(RUNTIME_DIR, 'kgd.pid')
+PORT_FILE = os.path.join(RUNTIME_DIR, 'kgd.port')
+LOG_FILE = os.path.join(RUNTIME_DIR, 'kgd.log')
+
+# 端口范围: 避开常用端口, 在动态区间内选
+PORT_MIN = 49152
+PORT_MAX = 49200
+CONNECT_TIMEOUT = 0.3   # client 探活超时 (秒) — 超过即回退冷路径
+READ_TIMEOUT = 120      # 单请求最长等待 (语义搜索首次需加载)
+
+# ── 这些 handler 在 daemon 里会被直接调用 (复用 kg.py 的实现) ──
+# 哪些子命令值得走 daemon? 必须是"进程内直接执行"的重量级命令
+# (import 慢 / 模型重 / 计算重), daemon 常驻后能省掉冷启动。
+#
+# ⚠️ 排除"subprocess 再起进程"类命令 (fill-prototype / design-system):
+# 它们内部用 subprocess.run 另起 Python 跑脚本, daemon 常驻的优势用不上,
+# 反而多一层 IPC 转发开销 (实测 daemon 比冷路径慢 ~1s)。
+# 这类命令应走冷路径, daemon 不要接管。
+DAEMON_COMMANDS = {
+    'search', 'api', 'prd', 'impact', 'context360', 'context',
+    'prefetch', 'coverage', 'feature', 'workflow', 'hop',
+    'wiki', 'semantic',
+    # 排除: 'fill-prototype', 'design-system' (subprocess 再起, daemon 帮倒忙)
+}
+# 其中 semantic 收益最大 (26.8s→0.5s), context/impact/prefetch 次之。
+
+
+# ============================================================
+# Daemon 服务端
+# ============================================================
+class KgDaemon:
+    """常驻服务进程。启动时预加载, 之后监听 TCP 转发请求。"""
+
+    def __init__(self):
+        self.handlers = {}        # cmd_name -> handler_fn
+        self.model_loaded = False
+        self.started_at = time.time()
+        self.requests = 0
+        self.cache_hits = 0
+        self._db_mtime = 0
+        self._db_path = ''
+
+    # ── 预加载 (启动时跑一次, 后续全热) ──────────────────────────
+    def preload(self):
+        """一次性加载所有重资源。"""
+        log('preload: 导入 kg.py 模块 (cwd=%s, _HERE=%s) ...' % (os.getcwd(), _HERE))
+        # 注意: kg 既是包名(有 __init__.py) 又含 kg.py CLI 模块。
+        # `import kg` 会命中包而非 kg.py → 必须按文件路径加载 kg.py 本身。
+        import importlib.util
+        kg_file = os.path.join(_HERE, 'kg.py')
+        _spec = importlib.util.spec_from_file_location('_kgd_kg', kg_file)
+        _kg = importlib.util.module_from_spec(_spec)
+        _spec.loader.exec_module(_kg)
+        cmds = getattr(_kg, 'COMMANDS', [])
+        log('preload: kg.py = %s, COMMANDS = %d' % (kg_file, len(cmds)))
+        for name, fn, _desc in cmds:
+            if name in DAEMON_COMMANDS:
+                self.handlers[name] = fn
+        log('preload: %d handlers 注册' % len(self.handlers))
+
+        # 预加载 DuckDB 连接 (各 handler 内部会自建连接, 这里只是 warm-up import)
+        try:
+            import duckdb
+            self._db_path = os.path.join(
+                os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(_HERE)))),
+                'data', 'index', 'kg.duckdb')
+            if os.path.isfile(self._db_path):
+                self._db_mtime = os.path.getmtime(self._db_path)
+                con = duckdb.connect(self._db_path, read_only=True)
+                con.execute('SELECT COUNT(*) FROM entities').fetchone()
+                con.close()
+                log('preload: DuckDB warm-up ok (entities 表可读)')
+        except Exception as e:
+            log('preload: DuckDB warm-up 跳过 (%s)' % str(e)[:60])
+
+        # 预加载向量模型 (头号收益: semantic 26.8s→0.5s)
+        # 放最后, 因为它最重 (95MB, ~25s)。即使失败也不影响其他 handler。
+        # ⚠️ 关键: 强制离线加载。sentence_transformers 默认会联网校验模型元数据,
+        #    在 DETACHED daemon 进程里网络不通会崩溃。设 HF_HUB_OFFLINE 后纯本地,
+        #    既根治联网不稳定, 又避免每次 ~25s 的网络等待 (真正的冷加载只需几秒)。
+        log('preload: 预加载 bge 向量模型 (强制离线, 本地缓存) ...')
+        os.environ['HF_HUB_OFFLINE'] = '1'
+        os.environ['TRANSFORMERS_OFFLINE'] = '1'
+        try:
+            from domain.kg.graph import kg_semantic
+            m = kg_semantic._load_model()
+            self.model_loaded = m is not None
+            if self.model_loaded:
+                log('preload: ✓ bge 模型已常驻, semantic 将秒回')
+            else:
+                log('preload: bge 模型未安装 (sentence-transformers 缺失), semantic 仍走冷路径')
+        except Exception as e:
+            log('preload: 模型预加载失败 (%s)' % str(e)[:80])
+        log('preload 完成, 用时 %.1fs' % (time.time() - self.started_at))
+
+    def _db_changed(self):
+        """检测 kg.duckdb 是否被重建 (mtime 变化)。变了需 handler 内部自重连。"""
+        try:
+            if self._db_path and os.path.isfile(self._db_path):
+                return os.path.getmtime(self._db_path) != self._db_mtime
+        except Exception:
+            pass
+        return False
+
+    # ── 请求处理 ──────────────────────────────────────────────────
+    # ── 内部命令 (不走 handler 表) ────────────────────────────────
+    def _internal(self, cmd):
+        """处理 __ping / __status 等内部探测命令。返回 dict 或 None。"""
+        if cmd == '__ping':
+            return {'ok': True, 'text': 'pong', 'secs': 0}
+        if cmd == '__status':
+            return {'ok': True, 'text': json.dumps({
+                'handlers': len(self.handlers),
+                'model': self.model_loaded,
+                'requests': self.requests,
+                'uptime_min': round((time.time() - self.started_at) / 60, 1),
+                'db_changed': self._db_changed(),
+            }, ensure_ascii=False), 'secs': 0}
+        return None  # 不是内部命令
+
+    def handle(self, cmd, args):
+        """执行一个子命令, 返回 dict 响应。"""
+        # 内部探测命令优先
+        internal = self._internal(cmd)
+        if internal is not None:
+            return internal
+
+        self.requests += 1
+        if cmd not in self.handlers:
+            return {'ok': False, 'error': 'daemon 不支持: %s' % cmd,
+                    'fallback': True}  # 让 client 回退冷路径
+
+        # 捕获 handler 的输出 + 返回值。
+        # ⚠️ sys.stdout 是进程全局, 多线程并发会串台 → 必须用线程本地隔离。
+        # 同时 handler 有两种风格: cmd_semantic 直接 return 字符串,
+        # cmd_search 通过 _capture 捕获 stdout → 两者都要收集。
+        import threading
+        if not hasattr(self, '_tls'):
+            self._tls = threading.local()
+        buf = io.StringIO()
+        real_stdout = sys.stdout
+        self._tls.buf = buf  # 线程私有
+        # 把当前线程的 sys.stdout 指向 buf (只影响本线程调用的代码)
+        # 注意: sys.stdout 是全局的, 这里靠"单请求处理很快+GIL"近似隔离;
+        # 更严格的隔离需改 handler, 但当前 kg.py 的 handler 都是瞬时完成的, 风险低。
+        sys.stdout = buf
+        t0 = time.time()
+        ret = None
+        try:
+            ret = self.handlers[cmd](args or [])
+            text = buf.getvalue().strip()
+            # handler 优先用返回值, 没有返回值才用 stdout 捕获
+            if ret:
+                text = (text + '\n' + str(ret)).strip() if text else str(ret).strip()
+            return {'ok': True, 'text': text, 'secs': round(time.time() - t0, 3)}
+        except Exception as e:
+            return {'ok': False, 'error': str(e)[:200],
+                    'fallback': True}  # 出错让 client 回退
+        finally:
+            sys.stdout = real_stdout
+
+    def serve(self, port):
+        """TCP 监听主循环。"""
+        srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        srv.bind(('127.0.0.1', port))
+        srv.listen(8)
+        srv.settimeout(1.0)  # 1s 唤醒一次, 便于优雅退出
+        log('kgd 监听 127.0.0.1:%d (PID %d)' % (port, os.getpid()))
+
+        running = [True]
+        try:
+            while running[0]:
+                try:
+                    conn, _ = srv.accept()
+                except socket.timeout:
+                    if _should_exit():
+                        break
+                    continue
+                except OSError:
+                    break
+                # 每连接一个线程 (请求都很短, 用 daemon 线程即可)
+                import threading
+                def serve_conn(c):
+                    try:
+                        c.settimeout(READ_TIMEOUT)
+                        data = b''
+                        while b'\n' not in data:
+                            chunk = c.recv(8192)
+                            if not chunk:
+                                break
+                            data += chunk
+                        line = data.decode('utf-8', 'replace').strip()
+                        if not line:
+                            return
+                        req = json.loads(line)
+                        resp = self.handle(req.get('cmd', ''), req.get('args', []))
+                        c.sendall((json.dumps(resp, ensure_ascii=False) + '\n').encode('utf-8'))
+                    except Exception as e:
+                        try:
+                            err = json.dumps({'ok': False, 'error': str(e)[:200],
+                                              'fallback': True})
+                            c.sendall((err + '\n').encode('utf-8'))
+                        except Exception:
+                            pass
+                    finally:
+                        c.close()
+                threading.Thread(target=serve_conn, args=(conn,), daemon=True).start()
+        finally:
+            srv.close()
+            log('kgd 已停止')
+
+
+def _should_exit():
+    """检测停止信号 (删除 pidfile 即视为要求退出)。"""
+    return not os.path.isfile(PID_FILE)
+
+
+def log(msg):
+    """写日志 (追加, 带时间戳)。"""
+    try:
+        os.makedirs(RUNTIME_DIR, exist_ok=True)
+        with open(LOG_FILE, 'a', encoding='utf-8') as f:
+            f.write('[%s] %s\n' % (time.strftime('%H:%M:%S'), msg))
+    except Exception:
+        pass
+
+
+def _find_free_port():
+    """在 PORT_MIN..PORT_MAX 找一个可用端口。"""
+    for p in range(PORT_MIN, PORT_MAX + 1):
+        try:
+            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            s.bind(('127.0.0.1', p))
+            s.close()
+            return p
+        except OSError:
+            continue
+    return 0
+
+
+def _read_pid_port():
+    """读 pidfile+portfile, 返回 (pid, port) 或 (None, None)。"""
+    try:
+        if not os.path.isfile(PID_FILE):
+            return None, None
+        with open(PID_FILE) as f:
+            pid = int(f.read().strip())
+        port = None
+        if os.path.isfile(PORT_FILE):
+            with open(PORT_FILE) as f:
+                port = int(f.read().strip())
+        return pid, port
+    except Exception:
+        return None, None
+
+
+def _is_alive(pid):
+    """进程是否存活 (跨平台)。"""
+    if not pid:
+        return False
+    try:
+        if os.name == 'nt':
+            # Windows: kernel32 OpenProcess
+            import ctypes
+            k32 = ctypes.windll.kernel32
+            PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
+            h = k32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid)
+            if h:
+                k32.CloseHandle(h)
+                return True
+            return False
+        else:
+            os.kill(pid, 0)  # signal 0 = 探活
+            return True
+    except Exception:
+        return False
+
+
+# ============================================================
+# Client (供 kg.py / MCP / perf_bench 调用)
+# ============================================================
+class _KgdClient:
+    """daemon 客户端。所有方法在 daemon 不可用时返回 None (触发调用方回退)。"""
+
+    def ping(self, timeout=CONNECT_TIMEOUT):
+        """探活。daemon 在线返回 True。"""
+        pid, port = _read_pid_port()
+        if not pid or not port or not _is_alive(pid):
+            return False
+        try:
+            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            s.settimeout(timeout)
+            s.connect(('127.0.0.1', port))
+            # 发一个 no-op ping (status 查询)
+            s.sendall((json.dumps({'cmd': '__ping', 'args': []}) + '\n').encode('utf-8'))
+            data = b''
+            while b'\n' not in data:
+                chunk = s.recv(1024)
+                if not chunk:
+                    break
+                data += chunk
+            s.close()
+            return b'\n' in data
+        except Exception:
+            return False
+
+    def call(self, cmd, args, timeout=READ_TIMEOUT):
+        """转发请求到 daemon。返回文本结果, 或 None (连不上/出错 → 调用方回退)。
+
+        Args:
+            cmd: 子命令名 (如 'semantic')
+            args: 参数列表 (如 ['报销'])
+        Returns:
+            (text, secs) 或 None
+        """
+        pid, port = _read_pid_port()
+        if not pid or not port or not _is_alive(pid):
+            return None
+        try:
+            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            s.settimeout(CONNECT_TIMEOUT)
+            s.connect(('127.0.0.1', port))
+            s.settimeout(timeout)
+            req = json.dumps({'cmd': cmd, 'args': args or []}, ensure_ascii=False)
+            s.sendall((req + '\n').encode('utf-8'))
+            data = b''
+            while b'\n' not in data:
+                chunk = s.recv(65536)
+                if not chunk:
+                    break
+                data += chunk
+            s.close()
+            if b'\n' not in data:
+                return None
+            resp = json.loads(data.decode('utf-8', 'replace').strip())
+            if resp.get('ok'):
+                return resp.get('text', ''), resp.get('secs', 0)
+            # daemon 要求回退 (不支持的命令 / 出错)
+            return None
+        except Exception:
+            return None
+
+
+kgd_client = _KgdClient()
+
+
+# ============================================================
+# 命令行入口: start / stop / status / restart
+# ============================================================
+def cmd_start(foreground=False):
+    """启动 daemon。"""
+    pid, port = _read_pid_port()
+    if pid and _is_alive(pid):
+        print('kgd 已在运行 (PID %d, port %s)' % (pid, port or '?'))
+        return 0
+
+    port = _find_free_port()
+    if not port:
+        print('✗ 无可用端口')
+        return 1
+
+    if foreground:
+        # 前台模式: 直接跑 (调试用)
+        os.makedirs(RUNTIME_DIR, exist_ok=True)
+        with open(PID_FILE, 'w') as f:
+            f.write(str(os.getpid()))
+        with open(PORT_FILE, 'w') as f:
+            f.write(str(port))
+        try:
+            d = KgDaemon()
+            d.preload()
+            d.serve(port)
+        finally:
+            for p in (PID_FILE, PORT_FILE):
+                try:
+                    os.remove(p)
+                except OSError:
+                    pass
+        return 0
+
+    # 后台模式: spawn 自己 --fore (detach)
+    env = dict(os.environ)
+    env['PYTHONIOENCODING'] = 'utf-8'
+    creationflags = 0
+    if os.name == 'nt':
+        creationflags = subprocess.CREATE_NEW_PROCESS_GROUP | 0x00000008  # DETACHED
+    try:
+        subprocess.Popen(
+            [sys.executable, os.path.abspath(__file__), 'start', '--fore'],
+            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+            stdin=subprocess.DEVNULL, cwd=os.getcwd(), env=env,
+            creationflags=creationflags)
+    except Exception as e:
+        print('✗ 启动失败: %s' % e)
+        return 1
+
+    # 等待 daemon 就绪 (轮询 ping, 最多等 preload 完)
+    print('kgd 启动中 (预加载模型 ~25s) ...', end='', flush=True)
+    for _ in range(60):  # 最多等 60s
+        time.sleep(1)
+        print('.', end='', flush=True)
+        if kgd_client.ping(timeout=0.5):
+            _, p = _read_pid_port()
+            print(' 就绪 (port %s)' % p)
+            return 0
+    print(' 超时, 查日志: %s' % LOG_FILE)
+    return 1
+
+
+def cmd_stop():
+    """停止 daemon。"""
+    pid, port = _read_pid_port()
+    if not pid:
+        print('kgd 未运行')
+        return 0
+    if not _is_alive(pid):
+        for p in (PID_FILE, PORT_FILE):
+            try:
+                os.remove(p)
+            except OSError:
+                pass
+        print('kgd 未运行 (清理了残留 pidfile)')
+        return 0
+    try:
+        if os.name == 'nt':
+            subprocess.run(['taskkill', '/PID', str(pid), '/F'],
+                           capture_output=True)
+        else:
+            os.kill(pid, 15)
+        time.sleep(0.5)
+    except Exception:
+        pass
+    for p in (PID_FILE, PORT_FILE):
+        try:
+            os.remove(p)
+        except OSError:
+            pass
+    print('kgd 已停止 (PID %d)' % pid)
+    return 0
+
+
+def cmd_status():
+    """查状态。"""
+    pid, port = _read_pid_port()
+    alive = _is_alive(pid) if pid else False
+    if not alive:
+        print('kgd: 未运行')
+        print('  启动: python .qoder/scripts/domain/kg/server/kgd.py start')
+        return 0
+    # 在线, 查详情
+    uptime = '?'
+    res = kgd_client.call('__status', [], timeout=2)
+    print('kgd: 运行中 (PID %d, port %s)' % (pid, port))
+    if res:
+        text, secs = res
+        if text:
+            try:
+                info = json.loads(text)
+                print('  handlers: %d' % info.get('handlers', '?'))
+                print('  模型加载: %s' % ('✓' if info.get('model') else '✗'))
+                print('  请求数: %d' % info.get('requests', '?'))
+                print('  运行时长: %.0f 分钟' % info.get('uptime_min', 0))
+            except Exception:
+                print('  %s' % text[:200])
+    print('  日志: %s' % LOG_FILE)
+    return 0
+
+
+def main():
+    args = sys.argv[1:]
+    if not args or args[0] in ('-h', '--help', 'help'):
+        print(__doc__)
+        return 0
+    cmd = args[0]
+    if cmd == 'start':
+        return cmd_start(foreground='--fore' in args)
+    if cmd == 'stop':
+        return cmd_stop()
+    if cmd == 'status':
+        return cmd_status()
+    if cmd == 'restart':
+        cmd_stop()
+        time.sleep(1)
+        return cmd_start(foreground='--fore' in args)
+    print('未知命令: %s (start/stop/status/restart)' % cmd)
+    return 1
+
+
+if __name__ == '__main__':
+    sys.exit(main())