PlantAI-rag-new/sdk/app/routers/qa.py

"""routers/qa.py —— /api/qa/* ：7920 代理 + 归一化 + SSE 中继（A2 归属，API.md §3/§4）。

端点：
- GET  /api/qa/info         上游 /info 精简（name/version/pipeline/config）
- POST /api/qa/chat         上游 /chat，300s 读超时，invalid_llm_citation 自动重试一次
- POST /api/qa/chat/stream  SSE 中继（app.services.sse_relay.relay_stream）
- POST /api/qa/search       上游 /search，强制 use_llm:false，180s 读超时，不重试，
                            响应 mode:"search"、answer:""

请求体一律透传（前端字段集见 API.md §4，BFF 不增删字段）。
"""
from __future__ import annotations

import json
import logging
from typing import Any, Optional

import httpx
from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse, StreamingResponse

from ..services import sse_relay

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/qa", tags=["qa"])

# 超时阶梯（TODO(A2→A1 对齐点)：A1 的 app.config.settings 就绪后改为配置读取）
CHAT_READ_TIMEOUT = 300.0
SEARCH_READ_TIMEOUT = 180.0
INFO_READ_TIMEOUT = 5.0
CONNECT_TIMEOUT = 10.0

# ---------------------------------------------------------------------------
# 错误信封（TODO(A2→A1 对齐点)：HTTP JSON 信封换用 app.upstream.envelope()）
# ---------------------------------------------------------------------------
MESSAGES = {
    "UPSTREAM_DOWN": "知识问答服务当前不可用，请稍后重试",
    "UPSTREAM_ERROR": "知识问答服务返回错误",
    "TIMEOUT": "检索超时",
    "INVALID_CITATION": "引用校验失败，请重试",
    "VALIDATION": "请求参数不合法",
    "INTERNAL": "内部错误",
}


def _envelope(code: str, status: int, message: Optional[str] = None, degraded: Optional[bool] = None) -> JSONResponse:
    if degraded is None:
        degraded = code == "UPSTREAM_DOWN"
    return JSONResponse(
        status_code=status,
        content={"ok": False, "error": {"code": code, "message": message or MESSAGES[code], "degraded": degraded}},
    )


# ---------------------------------------------------------------------------
# 上游调用
# ---------------------------------------------------------------------------


def _upstream_path(endpoint: str) -> str:
    return f"{sse_relay.QA_UPSTREAM_PREFIX}/{endpoint}"


def _read_timeout(seconds: float) -> httpx.Timeout:
    return httpx.Timeout(connect=CONNECT_TIMEOUT, read=seconds, write=30.0, pool=CONNECT_TIMEOUT)


async def _parse_json(text: Optional[bytes]) -> Any:
    if not text:
        return None
    try:
        return json.loads(text.decode("utf-8", errors="replace"))
    except Exception:
        return None


def _error_from_status(status: int, payload: Any) -> tuple[str, int, Optional[str]]:
    """上游 HTTP≥400 → (code, bff_http, message)。HTTP 按 API.md §1 固定映射，不透传上游状态码。"""
    if sse_relay.is_invalid_citation(payload):
        return "INVALID_CITATION", 502, None
    if status in (408, 504):
        return "TIMEOUT", 504, None
    detail = payload.get("detail") if isinstance(payload, dict) else None
    if isinstance(detail, dict):
        message = detail.get("message") or detail.get("error")
        return "UPSTREAM_ERROR", 502, message
    if isinstance(payload, dict):
        message = payload.get("message") or payload.get("error")
        return "UPSTREAM_ERROR", 502, message
    return "UPSTREAM_ERROR", 502, None


async def _proxy(
    endpoint: str,
    body: Optional[dict] = None,
    *,
    read_timeout: float,
    retry_invalid_citation: bool,
    mode: str,
    normalizer=None,
) -> Any:
    """非流式上游代理：重试、错误信封、归一化。返回 (JSONResponse)。"""
    client = sse_relay.get_qa_client()
    normalizer = normalizer or sse_relay.get_qa_normalizer()
    url = f"{sse_relay.qa_base_url()}{_upstream_path(endpoint)}"
    timeout = _read_timeout(read_timeout)
    json_body = body if body is not None else {}

    for attempt in (1, 2):
        try:
            resp = await client.request("POST", url, json=json_body, timeout=timeout)
        except httpx.ReadTimeout:
            return _envelope("TIMEOUT", 504)
        except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as exc:
            logger.warning("上游 %s 连接失败: %s", endpoint, exc)
            return _envelope("UPSTREAM_DOWN", 503)
        except httpx.HTTPError as exc:
            logger.warning("上游 %s 请求失败: %s", endpoint, exc)
            return _envelope("UPSTREAM_ERROR", 502)
        payload = await _parse_json(resp.content)
        if resp.status_code >= 400:
            if (
                retry_invalid_citation
                and attempt == 1
                and sse_relay.is_invalid_citation(payload)
            ):
                logger.warning(
                    "上游 %s 返回 %d invalid_llm_citation，自动重试（第 1 次）", endpoint, resp.status_code
                )
                continue
            code, bff_status, message = _error_from_status(resp.status_code, payload)
            return _envelope(code, bff_status, message)
        try:
            normalized = normalizer(payload if isinstance(payload, dict) else {})
        except Exception:
            logger.exception("归一化上游 %s 响应失败", endpoint)
            return _envelope("INTERNAL", 500)
        normalized["mode"] = mode
        if mode == "search":
            normalized["answer"] = ""
        # 图文并茂（2026-08-28）：同步 /chat 同样注入图片（SSE 路径在 sse_relay）
        if mode == "chat":
            try:
                from ..services.sse_relay import _inject_figures_async

                normalized = await _inject_figures_async(normalized)
            except Exception as exc:
                logger.warning("图片注入失败（降级为纯文本答案）: %s", exc)
        return JSONResponse(content=normalized)
    # 不可达（重试循环内必然 return）；防御性兜底
    return _envelope("UPSTREAM_ERROR", 502)


# ---------------------------------------------------------------------------
# 端点
# ---------------------------------------------------------------------------


@router.get("/info")
async def qa_info():
    """上游 /info 精简（API.md §3）。"""
    client = sse_relay.get_qa_client()
    url = f"{sse_relay.qa_base_url()}{_upstream_path('info')}"
    try:
        resp = await client.get(url, timeout=_read_timeout(INFO_READ_TIMEOUT))
    except httpx.ReadTimeout:
        return _envelope("TIMEOUT", 504)
    except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as exc:
        logger.warning("上游 /info 连接失败: %s", exc)
        return _envelope("UPSTREAM_DOWN", 503)
    except httpx.HTTPError as exc:
        logger.warning("上游 /info 请求失败: %s", exc)
        return _envelope("UPSTREAM_ERROR", 502)
    payload = await _parse_json(resp.content)
    if resp.status_code >= 400 or not isinstance(payload, dict):
        return _envelope("UPSTREAM_ERROR", 502)
    info = payload.get("info") if isinstance(payload.get("info"), dict) else payload
    return JSONResponse(
        content={
            "name": info.get("name") or payload.get("name"),
            "version": info.get("version") or payload.get("version"),
            "pipeline": info.get("pipeline") or payload.get("pipeline") or [],
            "config": info.get("config") or payload.get("config") or {},
        }
    )


@router.post("/chat")
async def qa_chat(request: Request):
    """非流式问答：上游 /chat + 归一化 + invalid_llm_citation 重试一次。"""
    body = await request.json()
    return await _proxy(
        "chat", body, read_timeout=CHAT_READ_TIMEOUT, retry_invalid_citation=True, mode="chat"
    )


@router.post("/chat/stream")
async def qa_chat_stream(request: Request):
    """SSE 中继（API.md §4）。

    evidence_first（2026-08-26 方案 3）：两级体验开关——BFF 并发打上游
    /search 先出证据帧；该字段是 BFF 契约字段，上游 QARequest 无此键，
    必须剥离后再转发，否则上游 422。
    """
    try:
        body = await request.json()
    except Exception:
        return _envelope("VALIDATION", 422)
    evidence_first = bool(body.pop("evidence_first", False))
    return StreamingResponse(
        sse_relay.relay_stream(body, evidence_first=evidence_first),
        media_type="text/event-stream",
        headers=sse_relay.sse_headers,
    )


@router.post("/search")
async def qa_search(request: Request):
    """检索模式：强制 use_llm:false，180s 读超时，不重试。"""
    body = await request.json()
    body = {**body, "use_llm": False}
    return await _proxy(
        "search", body, read_timeout=SEARCH_READ_TIMEOUT, retry_invalid_citation=False, mode="search"
    )