PlantAI-rag-new/sdk/app/upstream.py

"""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]]]