"""upstream.py —— 上游健康探测(httpx 单例 + 5s TTL 缓存)与统一错误信封。 - `probe_all()`:并发探测 qa(/health 或 /info)、rag(/health)、ragAlt(/health)、 Neo4j(延迟导入 `app.services.neo4j.probe_neo4j()`,A3 提供;若模块缺失则置 `{ok:false}`)。探测超时 1s(settings.timeout_health),结果缓存 5s。 - `envelope(code, message, degraded, status)`:API.md §1 错误信封构造器,供所有 router 用。 """ from __future__ import annotations import asyncio import logging import time from typing import Any, Awaitable, Callable import httpx from .config import get_settings logger = logging.getLogger(__name__) _client: httpx.AsyncClient | None = None _httpx_lock = asyncio.Lock() async def get_client() -> httpx.AsyncClient: """httpx.AsyncClient 单例(供 qa/kg/extraction 各代理层共用)。""" global _client if _client is None or _client.is_closed: async with _httpx_lock: if _client is None or _client.is_closed: settings = get_settings() _client = httpx.AsyncClient( timeout=httpx.Timeout( connect=settings.timeout_health, read=settings.timeout_qa_chat, write=10.0, pool=5.0, ), follow_redirects=True, ) return _client async def close_client() -> None: """app shutdown 钩子。""" global _client if _client is not None and not _client.is_closed: await _client.aclose() _client = None # --------------------------------------------------------------------------- # 错误信封(API.md §1)——供所有 router 用 # --------------------------------------------------------------------------- _ERROR_MESSAGES: dict[str, str] = { "UPSTREAM_DOWN": "知识问答服务当前不可用,请稍后重试", "UPSTREAM_ERROR": "上游服务返回错误", "TIMEOUT": "检索超时", "INVALID_CITATION": "引用校验失败", "NOT_FOUND": "资源不存在", "VALIDATION": "请求参数不合法", "INTERNAL": "服务内部错误", } def envelope( code: str, message: str | None = None, degraded: bool = False, status: int | None = None, ) -> tuple[dict[str, Any], int]: """构造 `{"ok": false, "error": {code, message, degraded}}` 信封 + HTTP 状态。 status 缺省时按 code 推导(API.md §1 表): UPSTREAM_DOWN=503, UPSTREAM_ERROR=502, TIMEOUT=504, INVALID_CITATION=502, NOT_FOUND=404, VALIDATION=422, INTERNAL=500。 """ if status is None: status = { "UPSTREAM_DOWN": 503, "UPSTREAM_ERROR": 502, "TIMEOUT": 504, "INVALID_CITATION": 502, "NOT_FOUND": 404, "VALIDATION": 422, "INTERNAL": 500, }.get(code, 500) if message is None: message = _ERROR_MESSAGES.get(code, "未知错误") return { "ok": False, "error": {"code": code, "message": message, "degraded": degraded}, }, status # --------------------------------------------------------------------------- # 健康探测 # --------------------------------------------------------------------------- _probe_lock = asyncio.Lock() _probe_at: float = 0.0 _probe_result: dict[str, Any] | None = None async def _get_json(url: str) -> tuple[bool, Any]: """1s 超时 GET + JSON 解析;任何失败返回 (False, None),不抛错。""" try: client = await get_client() resp = await client.get(url, timeout=httpx.Timeout(get_settings().timeout_health)) if resp.status_code >= 400: return False, None return True, resp.json() except Exception: return False, None async def _probe_qa() -> dict[str, Any]: settings = get_settings() base = settings.qa_base_url.rstrip("/") # 按 DISCOVERY.md:health 为可靠探测端点;其不可用时回落 info ok, data = await _get_json(f"{base}{settings.qa_api_path}/health") if not ok: ok, data = await _get_json(f"{base}{settings.qa_api_path}/info") if not ok or not isinstance(data, dict): return {"ok": False} version = data.get("version") or data.get("app_version") return {"ok": True, "version": version} async def _probe_rag(base_url: str) -> dict[str, Any]: ok, data = await _get_json(f"{base_url.rstrip('/')}/health") if not ok: return {"ok": False} return {"ok": True} async def _probe_neo4j() -> dict[str, Any]: """Neo4j 状态由 A3 提供(app.services.neo4j.probe_neo4j)。 延迟导入:A3 未实现时 catch ImportError,置 {ok:false}。 """ try: from .services.neo4j import probe_neo4j except ImportError: return {"ok": False} try: result = await probe_neo4j() except Exception: return {"ok": False} if not isinstance(result, dict): return {"ok": False} out: dict[str, Any] = {"ok": bool(result.get("ok"))} if result.get("active"): out["active"] = result["active"] if result.get("detail"): out["detail"] = result["detail"] return out async def _probe_all_uncached() -> dict[str, Any]: settings = get_settings() qa_t, rag_t, alt_t, neo_t = await asyncio.gather( _probe_qa(), _probe_rag(settings.rag_base_url), _probe_rag(settings.rag_alt_base_url), _probe_neo4j(), ) return {"qa": qa_t, "rag": rag_t, "ragAlt": alt_t, "neo4j": neo_t} async def probe_all(force: bool = False) -> dict[str, Any]: """并发探测全部上游;结果缓存 settings.health_cache_ttl 秒(默认 5)。""" global _probe_at, _probe_result now = time.monotonic() if not force and _probe_result is not None: if now - _probe_at < get_settings().health_cache_ttl: return _probe_result async with _probe_lock: now = time.monotonic() if not force and _probe_result is not None: if now - _probe_at < get_settings().health_cache_ttl: return _probe_result result = await _probe_all_uncached() _probe_at = time.monotonic() _probe_result = result return result async def probe_all_refresh() -> dict[str, Any]: """强制刷新(无缓存)。""" return await probe_all(force=True) # 供 A2/A3 复用的类型别名 UpstreamProbe = Callable[..., Awaitable[dict[str, Any]]]