argo-search 1.0.1

package/scripts/health_check.py

@@ -0,0 +1,437 @@
+#!/usr/bin/env python3
+"""
+health_check.py — 引擎健康检查
+
+职责：
+  - 遍历三层引擎注册表所有引擎
+  - 标记可用/降级/不可用
+  - 按层级筛选可用引擎
+  - 输出健康报告供路由决策使用
+
+调用方式：
+  - 内部：from health_check import health_check_all, get_available_engines_by_tier
+  - CLI：python3 health_check.py [--tier T1] [--json]
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+# 添加 scripts 目录到 import 路径
+SCRIPT_DIR = Path(__file__).resolve().parent.parent / "scripts"
+sys.path.insert(0, str(SCRIPT_DIR))
+
+try:
+    import yaml
+except ImportError:
+    yaml = None  # type: ignore
+
+
+REGISTRY_PATH = Path(__file__).resolve().parent.parent / "backends" / "engine_registry.yaml"
+
+# 本地搜索模块路径（优雅降级，缺失时不阻断）
+# 默认路径为项目内 local-search/ 子目录，可通过环境变量 LOCAL_SEARCH_PATH 覆盖
+_LOCAL_SEARCH_V3_PATH = Path(os.environ.get(
+    "LOCAL_SEARCH_PATH",
+    str(Path(__file__).resolve().parent.parent / "local-search" / "search_v3.py")
+))
+_SEARXNG_BRIDGE_PATH = Path(os.environ.get(
+    "SEARXNG_BRIDGE_PATH",
+    str(Path(__file__).resolve().parent.parent / "local-search" / "searxng_bridge.py")
+))
+
+
+def _load_registry() -> list[dict[str, Any]]:
+    """加载引擎注册表 YAML。"""
+    if not REGISTRY_PATH.exists():
+        return []
+
+    if yaml is None:
+        # 无 PyYAML 时尝试简易解析（仅处理基本结构）
+        return _parse_yaml_light()
+
+    try:
+        with open(REGISTRY_PATH, "r", encoding="utf-8") as f:
+            data = yaml.safe_load(f)
+    except Exception:
+        return []
+
+    if not isinstance(data, dict):
+        return []
+
+    return data.get("engines", [])
+
+
+def _parse_yaml_light() -> list[dict[str, Any]]:
+    """简易 YAML 解析器回退 — 仅提取 engines 列表。"""
+    try:
+        text = REGISTRY_PATH.read_text(encoding="utf-8")
+        # 查找 engines: 段落后面的条目
+        # 这是一个非常基础的解析，PyYAML 不可用时才使用
+        return []
+    except Exception:
+        return []
+
+
+def _check_t1_engine(name: str, timeout: float = 5) -> dict[str, Any]:
+    """检查 T1 引擎健康状态。
+
+    通过 engines.py 的 search() 函数快速探测。
+    """
+    result = {
+        "engine": name,
+        "tier": "T1",
+        "available": False,
+        "latency_ms": 0,
+        "error": None,
+        "status": "unavailable",
+    }
+
+    try:
+        from engines import search, get_registry
+        registry = get_registry()
+        if name not in registry:
+            result["error"] = f"引擎 {name} 未在 engines.py 注册表中"
+            result["status"] = "unavailable"
+            return result
+
+        fn = registry[name]
+        t0 = time.time()
+        # 快速探测：只请求 1 条结果
+        try:
+            results = fn("test", n=1, _timeout=timeout)
+        except TypeError:
+            results = fn("test", n=1)
+
+        elapsed = round((time.time() - t0) * 1000)
+        result["latency_ms"] = elapsed
+
+        if results and isinstance(results, list) and len(results) > 0:
+            result["available"] = True
+            result["status"] = "ok"
+            result["sample"] = results[0]
+        else:
+            result["error"] = "返回空结果"
+            result["status"] = "degraded"
+
+    except ImportError:
+        result["error"] = "engines.py 不可导入"
+        result["status"] = "unavailable"
+    except Exception as e:
+        result["error"] = f"{type(e).__name__}: {e}"
+        result["status"] = "unavailable"
+
+    return result
+
+
+def _check_t2_engine(name: str, timeout: float = 5) -> dict[str, Any]:
+    """检查 T2 引擎健康状态。
+
+    通过 search_v3.check_engine_health() 探测。
+    引擎名去掉 local/ 前缀后传给 check_engine_health。
+    """
+    local_name = name.replace("local/", "") if name.startswith("local/") else name
+
+    result = {
+        "engine": name,
+        "tier": "T2",
+        "available": False,
+        "latency_ms": 0,
+        "error": None,
+        "status": "unavailable",
+    }
+
+    try:
+        sys.path.insert(0, str(_LOCAL_SEARCH_V3_PATH.parent))
+        from search_v3 import check_engine_health
+        health = check_engine_health(local_name, timeout=timeout)
+        result["available"] = health.get("available", False)
+        result["latency_ms"] = health.get("latency_ms", 0)
+        result["error"] = health.get("error")
+        result["status"] = health.get("status", "unavailable")
+        if health.get("sample"):
+            result["sample"] = health["sample"]
+    except ImportError:
+        result["error"] = "search_v3.py 不可导入"
+        result["status"] = "unavailable"
+    except Exception as e:
+        result["error"] = f"{type(e).__name__}: {e}"
+        result["status"] = "unavailable"
+
+    return result
+
+
+def _check_t3_engine(name: str, timeout: float = 5) -> dict[str, Any]:
+    """检查 T3 引擎健康状态。
+
+    通过 searxng_bridge 检查 SearXNG 是否可达，然后检查引擎是否在列表中。
+    """
+    result = {
+        "engine": name,
+        "tier": "T3",
+        "available": False,
+        "latency_ms": 0,
+        "error": None,
+        "status": "unavailable",
+    }
+
+    try:
+        sys.path.insert(0, str(_SEARXNG_BRIDGE_PATH.parent))
+        from searxng_bridge import fetch_searxng_engine_list
+
+        t0 = time.time()
+        engine_data = fetch_searxng_engine_list(timeout=timeout)
+        elapsed = round((time.time() - t0) * 1000)
+        result["latency_ms"] = elapsed
+
+        if engine_data.get("error"):
+            result["error"] = engine_data["error"]
+            result["status"] = "unavailable"
+            return result
+
+        # 检查指定引擎是否在垂直可用列表中
+        searxng_engine_name = name.replace("searxng/", "")
+        available_verticals = engine_data.get("available_vertical", [])
+        available_names = {e.get("name", "").lower() for e in available_verticals}
+
+        if searxng_engine_name.lower() in available_names or \
+           searxng_engine_name.lower().replace("_", " ") in available_names:
+            result["available"] = True
+            result["status"] = "ok"
+        else:
+            # 也检查是否在 category_map 中
+            cmap = engine_data.get("category_map", {})
+            found = False
+            for cat_engines in cmap.values():
+                if searxng_engine_name.lower() in [e.lower() for e in cat_engines]:
+                    found = True
+                    break
+            if found:
+                result["available"] = True
+                result["status"] = "ok"
+            else:
+                result["error"] = f"引擎 {searxng_engine_name} 不在 SearXNG 可用垂直列表中"
+                result["status"] = "degraded"
+
+    except ImportError:
+        result["error"] = "searxng_bridge.py 不可导入"
+        result["status"] = "unavailable"
+    except Exception as e:
+        result["error"] = f"{type(e).__name__}: {e}"
+        result["status"] = "unavailable"
+
+    return result
+
+
+# 健康检查映射
+_CHECKERS = {
+    "T1": _check_t1_engine,
+    "T2": _check_t2_engine,
+    "T3": _check_t3_engine,
+}
+
+
+def health_check_all(
+    timeout: float = 5,
+    tiers: list[str] | None = None,
+    fast: bool = False,
+) -> dict[str, Any]:
+    """遍历注册表所有引擎，标记可用/降级/不可用。
+
+    Args:
+        timeout: 单引擎探测超时（秒）
+        tiers: 要检查的层级列表（None = 全部）
+        fast: 快速模式 — 仅检查推荐引擎
+
+    Returns:
+        {
+            "ok": [...],
+            "degraded": [...],
+            "unavailable": [...],
+            "by_tier": {"T1": [...], "T2": [...], "T3": [...]},
+            "summary": {"total": N, "ok_count": N, "degraded_count": N, ...},
+            "elapsed_ms": N,
+        }
+    """
+    t0_total = time.time()
+
+    registry = _load_registry()
+    if not registry:
+        return {
+            "ok": [], "degraded": [], "unavailable": [],
+            "by_tier": {}, "summary": {"total": 0, "ok_count": 0,
+                                       "degraded_count": 0, "unavailable_count": 0},
+            "elapsed_ms": 0, "error": "引擎注册表为空或不可读",
+        }
+
+    target_tiers = set(tiers) if tiers else {"T1", "T2", "T3"}
+
+    ok_list: list[dict] = []
+    degraded_list: list[dict] = []
+    unavailable_list: list[dict] = []
+    by_tier: dict[str, list[dict]] = {"T1": [], "T2": [], "T3": []}
+
+    for engine in registry:
+        name = engine.get("name", "")
+        tier = engine.get("tier", "")
+        status = engine.get("status", "unknown")
+
+        if tier not in target_tiers:
+            continue
+
+        # 快速模式：仅检查推荐引擎 + 状态为 ok 的引擎
+        if fast:
+            is_recommended = engine.get("recommended", False)
+            is_ok = status == "ok"
+            if not (is_recommended or is_ok):
+                engine["_checked"] = False
+                engine["_skip_reason"] = "非推荐引擎（快速模式）"
+                continue
+
+        # 如果注册表已标记为 disabled，直接跳过
+        if status == "disabled":
+            engine["_checked"] = False
+            engine["_skip_reason"] = "已禁用"
+            continue
+
+        # 对于 degraded 状态的引擎（如 Baidu），不做探测直接使用注册表状态
+        if status == "degraded":
+            check_result = {
+                "engine": name,
+                "tier": tier,
+                "available": False,
+                "latency_ms": engine.get("latency_ms", 0),
+                "error": engine.get("note", "已降级"),
+                "status": "degraded",
+                "_checked": True,
+                "_from_registry": True,
+            }
+            degraded_list.append(check_result)
+            by_tier.setdefault(tier, []).append(check_result)
+            continue
+
+        # 执行健康检查
+        checker = _CHECKERS.get(tier)
+        if checker:
+            check_result = checker(name, timeout=timeout)
+            check_result["_checked"] = True
+        else:
+            check_result = {
+                "engine": name,
+                "tier": tier,
+                "available": False,
+                "latency_ms": 0,
+                "error": f"未知层级: {tier}",
+                "status": "unavailable",
+                "_checked": True,
+            }
+
+        # 分类
+        st = check_result.get("status", "unavailable")
+        if st == "ok":
+            ok_list.append(check_result)
+        elif st in ("degraded", "slow"):
+            degraded_list.append(check_result)
+        else:
+            unavailable_list.append(check_result)
+
+        by_tier.setdefault(tier, []).append(check_result)
+
+    elapsed = round((time.time() - t0_total) * 1000)
+
+    return {
+        "ok": ok_list,
+        "degraded": degraded_list,
+        "unavailable": unavailable_list,
+        "by_tier": by_tier,
+        "summary": {
+            "total": len(ok_list) + len(degraded_list) + len(unavailable_list),
+            "ok_count": len(ok_list),
+            "degraded_count": len(degraded_list),
+            "unavailable_count": len(unavailable_list),
+        },
+        "elapsed_ms": elapsed,
+    }
+
+
+def get_available_engines_by_tier(
+    tier: str,
+    include_degraded: bool = False,
+) -> list[str]:
+    """按层级筛选可用引擎名称列表。
+
+    Args:
+        tier: 层级名（"T1" / "T2" / "T3"）
+        include_degraded: 是否包含降级但仍可用的引擎
+
+    Returns:
+        引擎名称列表（仅为 ok 状态，若 include_degraded 则包含 degraded）
+    """
+    registry = _load_registry()
+    names: list[str] = []
+
+    for engine in registry:
+        if engine.get("tier") != tier:
+            continue
+        status = engine.get("status", "unknown")
+        if status == "ok":
+            names.append(engine.get("name", ""))
+        elif include_degraded and status in ("degraded", "slow"):
+            names.append(engine.get("name", ""))
+
+    return names
+
+
+# ── CLI 调试用 ─────────────────────────────────────────────────────────────────
+def _cli():
+    import argparse
+
+    parser = argparse.ArgumentParser(description="引擎健康检查")
+    parser.add_argument("--tier", "-t", choices=["T1", "T2", "T3"],
+                        help="仅检查指定层级")
+    parser.add_argument("--json", action="store_true", help="JSON 格式输出")
+    parser.add_argument("--fast", action="store_true", help="快速模式（仅检查推荐引擎）")
+    parser.add_argument("--list", action="store_true", help="列出注册表所有引擎（不探测）")
+    args = parser.parse_args()
+
+    if args.list:
+        registry = _load_registry()
+        if args.tier:
+            registry = [e for e in registry if e.get("tier") == args.tier]
+        for e in registry:
+            print(f"[{e.get('tier', '?')}] {e.get('name', '?')} "
+                  f"({e.get('status', '?')}) — {e.get('desc', '')}")
+        return
+
+    tiers = [args.tier] if args.tier else None
+    result = health_check_all(timeout=5, tiers=tiers, fast=args.fast)
+
+    if args.json:
+        print(json.dumps(result, ensure_ascii=False, indent=2))
+    else:
+        print(f"健康检查完成 — 耗时 {result['elapsed_ms']}ms")
+        summary = result["summary"]
+        print(f"  可用: {summary['ok_count']}  降级: {summary['degraded_count']}  "
+              f"不可用: {summary['unavailable_count']}  总计: {summary['total']}")
+        print()
+        if result["ok"]:
+            print("✅ 可用:")
+            for e in result["ok"]:
+                print(f"  [{e['tier']}] {e['engine']} ({e['latency_ms']}ms)")
+        if result["degraded"]:
+            print("⚠️  降级:")
+            for e in result["degraded"]:
+                print(f"  [{e['tier']}] {e['engine']} — {e.get('error', '')}")
+        if result["unavailable"]:
+            print("❌ 不可用:")
+            for e in result["unavailable"]:
+                print(f"  [{e['tier']}] {e['engine']} — {e.get('error', '')}")
+
+
+if __name__ == "__main__":
+    _cli()

package/scripts/health_probe.py

@@ -0,0 +1,218 @@
+#!/usr/bin/env python3
+"""
+health_probe.py — 引擎健康探针（v2.5 新增）
+
+对 HTTP 引擎做可达性检查，结果写入 SQLite 供 route.py 查询。
+- 启动时全量探测
+- 每 5 分钟后台自动刷新
+- 连续 2 次失败标记 unavailable，1 次成功恢复
+
+用法：
+  python3 health_probe.py          # 单次探测
+  python3 health_probe.py --watch  # 后台持续监控（每 5 分钟）
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sqlite3
+import sys
+import threading
+import time
+import urllib.request
+from pathlib import Path
+from typing import Any
+
+SCRIPT_DIR = Path(__file__).parent
+sys.path.insert(0, str(SCRIPT_DIR))
+
+try:
+    from config import load_config, get_engines
+except ImportError:
+    from config import load_config, get_engines
+
+
+# ── 路径 ──────────────────────────────────────────────────────────────────────
+
+DB_PATH = Path.home() / ".cache" / "unified-search" / "health.db"
+PROBE_INTERVAL = 300  # 5 分钟
+PROBE_TIMEOUT = 1.5  # 单次探测超时（快，不能拖慢启动）
+MAX_FAILURES = 2  # 连续失败次数阈值
+
+
+# ── SQLite 存储 ───────────────────────────────────────────────────────────────
+
+def _connect():
+    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(str(DB_PATH), timeout=5)
+    conn.execute("PRAGMA journal_mode=WAL")
+    return conn
+
+
+def _init_db():
+    with _connect() as conn:
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS engine_health (
+                engine TEXT PRIMARY KEY,
+                status TEXT NOT NULL DEFAULT 'unknown',
+                consecutive_failures INTEGER DEFAULT 0,
+                last_probe_at REAL DEFAULT 0,
+                last_latency_ms REAL DEFAULT 0,
+                last_error TEXT DEFAULT ''
+            )
+        """)
+
+
+def _probe_http(url: str, timeout: float = PROBE_TIMEOUT) -> tuple[bool, float, str]:
+    """探测 HTTP 引擎：发起 HEAD 请求检查可达性。"""
+    try:
+        req = urllib.request.Request(url, method="HEAD")
+        t0 = time.time()
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            elapsed = (time.time() - t0) * 1000
+            return True, elapsed, ""
+    except Exception as e:
+        return False, 0, str(e)
+
+
+def _probe_cli(cmd: list[str], timeout: float = 2) -> tuple[bool, float, str]:
+    """探测 CLI 引擎：检查命令是否存在。"""
+    import subprocess
+    if not cmd:
+        return False, 0, "empty command"
+    exe = cmd[0]
+    if isinstance(exe, str) and exe.startswith("~"):
+        exe = str(Path(exe).expanduser())
+    try:
+        result = subprocess.run(["which", exe], capture_output=True, text=True, timeout=timeout)
+        if result.returncode == 0:
+            return True, 0, ""
+        return False, 0, f"command not found: {exe}"
+    except Exception as e:
+        return False, 0, str(e)
+
+
+def probe_all_engines():
+    """探测所有已启用的引擎。"""
+    _init_db()
+    cfg = load_config()
+    engines = get_engines(cfg)
+    results = {}
+
+    for name, spec in engines.items():
+        engine_type = spec.get("type", "")
+        url = spec.get("url", "")
+        cmd = spec.get("cmd", [])
+
+        if engine_type == "http" and url:
+            ok, latency, err = _probe_http(url)
+        elif engine_type == "cli" and cmd:
+            ok, latency, err = _probe_cli(cmd)
+        else:
+            continue
+
+        # 更新数据库
+        with _connect() as conn:
+            row = conn.execute(
+                "SELECT consecutive_failures FROM engine_health WHERE engine = ?",
+                (name,),
+            ).fetchone()
+
+            if ok:
+                new_status = "available"
+                new_failures = 0
+            else:
+                prev_failures = row[0] if row else 0
+                new_failures = prev_failures + 1
+                new_status = "unavailable" if new_failures >= MAX_FAILURES else "degraded"
+
+            conn.execute(
+                """INSERT OR REPLACE INTO engine_health
+                   (engine, status, consecutive_failures, last_probe_at, last_latency_ms, last_error)
+                   VALUES (?, ?, ?, ?, ?, ?)""",
+                (name, new_status, new_failures, time.time(), latency, err),
+            )
+            conn.commit()
+
+        results[name] = {
+            "status": "available" if ok else "unavailable",
+            "latency_ms": round(latency, 1),
+            "error": err,
+        }
+
+    return results
+
+
+def get_engine_status(engine: str) -> dict[str, Any]:
+    """查询单个引擎的健康状态。"""
+    _init_db()
+    with _connect() as conn:
+        row = conn.execute(
+            "SELECT status, consecutive_failures, last_probe_at, last_latency_ms FROM engine_health WHERE engine = ?",
+            (engine,),
+        ).fetchone()
+    if not row:
+        return {"status": "unknown", "consecutive_failures": 0, "available": True}
+    return {
+        "status": row[0],
+        "consecutive_failures": row[1],
+        "last_probe_at": row[2],
+        "last_latency_ms": row[3],
+        "available": row[0] != "unavailable",
+    }
+
+
+def get_all_status() -> dict[str, dict[str, Any]]:
+    """获取所有引擎的健康状态。"""
+    _init_db()
+    with _connect() as conn:
+        rows = conn.execute(
+            "SELECT engine, status, consecutive_failures, last_probe_at, last_latency_ms FROM engine_health"
+        ).fetchall()
+    return {
+        row[0]: {
+            "status": row[1],
+            "consecutive_failures": row[2],
+            "last_probe_at": row[3],
+            "last_latency_ms": row[4],
+            "available": row[1] != "unavailable",
+        }
+        for row in rows
+    }
+
+
+def watch_background():
+    """后台持续监控模式。在后台线程中运行，不阻塞 MCP 启动。"""
+    def _loop():
+        # 首次探测延迟 10 秒，让 MCP 先完成 initialize 握手
+        time.sleep(10)
+        try:
+            probe_all_engines()
+        except Exception:
+            pass
+        while True:
+            time.sleep(PROBE_INTERVAL)
+            try:
+                probe_all_engines()
+            except Exception:
+                pass
+
+    t = threading.Thread(target=_loop, daemon=True)
+    t.start()
+
+
+# ── CLI ───────────────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    if "--watch" in sys.argv:
+        print("启动后台健康监控（每 5 分钟）...", file=sys.stderr)
+        watch_background()
+        try:
+            while True:
+                time.sleep(60)
+        except KeyboardInterrupt:
+            pass
+    else:
+        results = probe_all_engines()
+        print(json.dumps(results, ensure_ascii=False, indent=2))