PlantAI-rag-new/sdk/app/services/docgen/client.py

"""services/docgen/client.py —— vstore HTTP 客户端（BFF → vstore 全经此处）。

- httpx.AsyncClient（连接池复用）+ 超时/重试（指数退避 0.5/1.0/2.0s）
- 方法：ingest / search_documents / get_document / get_style / delete_document /
  update_metadata / rebuild_index / health
- 所有失败路径不抛未捕获异常：网络/超时抛 DocgenUnavailable，业务错误
  返回 {"ok": False, "error": {...}} 结构供上层兜底
"""
from __future__ import annotations

import logging
from typing import Any

import httpx

from .settings import get_paths

logger = logging.getLogger(__name__)

DEFAULT_VSTORE_BASE_URL = "http://127.0.0.1:8094"
RETRY_BACKOFF = (0.5, 1.0, 2.0)  # 重试退避（秒）
MAX_RETRIES = 3
DEFAULT_TIMEOUT = 600.0  # 2026-09-04 上传无限制：大文件解析+嵌入可超 60s，与前端 600s 兜底对齐


class DocgenUnavailable(Exception):
    """vstore 不可达（dev2 路由层兜底捕获：UPSTREAM_DOWN）。"""

    def __init__(self, message: str, *, cause: BaseException | None = None) -> None:
        super().__init__(message)
        self.cause = cause


def _vstore_base_url() -> str:
    """vstore 基地址：VSTORE_BASE_URL env（dev2 会在 config.py 提供），
    兜底 http://127.0.0.1:8094。"""
    import os

    return (os.environ.get("VSTORE_BASE_URL") or DEFAULT_VSTORE_BASE_URL).rstrip("/")


_client: httpx.AsyncClient | None = None


def get_client() -> httpx.AsyncClient:
    """全局 AsyncClient（app 关闭时 close_client）。"""
    global _client
    if _client is None:
        _client = httpx.AsyncClient(
            base_url=_vstore_base_url(),
            timeout=httpx.Timeout(DEFAULT_TIMEOUT, connect=5.0),
        )
    return _client


def set_client(client: httpx.AsyncClient) -> None:
    """注入 client（测试用 MockTransport）。"""
    global _client
    _client = client


async def close_client() -> None:
    global _client
    if _client is not None:
        await _client.aclose()
        _client = None


async def _request(
    method: str,
    path: str,
    *,
    retries: int = MAX_RETRIES,
    **kwargs: Any,
) -> httpx.Response:
    """带重试的请求。网络/超时/5xx 重试；重试尽 → DocgenUnavailable。"""
    client = get_client()
    last_exc: BaseException | None = None
    for attempt in range(retries + 1):
        try:
            resp = await client.request(method, path, **kwargs)
            if resp.status_code < 500:
                return resp
            last_exc = RuntimeError(f"vstore 返回 {resp.status_code}")
        except (httpx.HTTPError, httpx.TimeoutException) as exc:
            last_exc = exc
        if attempt < retries:
            import asyncio

            await asyncio.sleep(RETRY_BACKOFF[min(attempt, len(RETRY_BACKOFF) - 1)])
    raise DocgenUnavailable(f"vstore 不可达（{path}）：{last_exc}", cause=last_exc)


def _check_ok(resp: httpx.Response) -> dict[str, Any]:
    """响应 → dict；ok=False 返回错误结构（不抛），HTTP 错误抛 DocgenUnavailable。"""
    if resp.status_code >= 400:
        raise DocgenUnavailable(f"vstore 返回错误状态 {resp.status_code}")
    try:
        payload = resp.json()
    except ValueError as exc:
        raise DocgenUnavailable(f"vstore 响应非 JSON：{exc}") from exc
    return payload if isinstance(payload, dict) else {"ok": False, "error": {"code": "BAD_RESPONSE", "message": "响应形态异常"}}


# ---------------------------------------------------------------------------
# dev2 契约导出（签名与 sdk/app/routers/docgen.py 桩段对齐）
# ---------------------------------------------------------------------------

async def ingest_bytes(
    filename: str,
    content: bytes,
    *,
    name: str,
    category: str = "",
    tags: str = "",
    description: str = "",
    as_template: bool = True,
) -> dict[str, Any]:
    """POST /api/vstore/ingest（multipart）。返回 {document_id, name, enc_path, ...}。"""
    resp = await _request(
        "POST", "/api/vstore/ingest",
        params={"as_template": bool(as_template)},
        data={"name": name, "category": category, "tags": tags, "description": description},
        files={"file": (filename, content)},
        retries=MAX_RETRIES,
    )
    payload = _check_ok(resp)
    if not payload.get("ok"):
        return payload  # 错误结构供上层兜底
    item = payload.get("item") or {}
    if not item.get("document_id"):
        return {"ok": False, "error": {"code": "BAD_RESPONSE", "message": "vstore 未返回 document_id"}}
    return {**item, "ok": True}


async def search_documents(
    query: str,
    *,
    as_template: bool | None = None,
    top_k: int = 20,
    category: str = "",
) -> list[dict[str, Any]]:
    """GET /api/vstore/search。检索失败抛 DocgenUnavailable（router 兜底 → 空模版集）。"""
    params: dict[str, Any] = {"q": query, "top_k": top_k}
    if as_template is not None:
        params["as_template"] = as_template
    if category:
        params["category"] = category
    resp = await _request("GET", "/api/vstore/search", params=params)
    payload = _check_ok(resp)
    if not payload.get("ok"):
        return []  # 检索失败按空结果（不阻断生成）
    return [item for item in payload.get("items", []) if isinstance(item, dict)]


async def search_chunks(
    query: str,
    *,
    top_k: int = 8,
    min_score: float = 0.30,
    timeout: float = 12.0,
) -> dict[str, Any]:
    """GET /api/vstore/search-chunks（FR-PRIV-004）。失败抛 DocgenUnavailable。"""
    resp = await _request(
        "GET",
        "/api/vstore/search-chunks",
        params={"q": query, "top_k": top_k, "min_score": min_score},
        timeout=timeout,
    )
    payload = _check_ok(resp)
    if not isinstance(payload, dict) or not payload.get("ok"):
        raise DocgenUnavailable("vstore search-chunks 返回异常")
    return payload


async def get_document(document_id: str) -> dict[str, Any] | None:
    """GET /api/vstore/documents/{id}（明文全文 + outline + meta）。

    404 → None（资料缺失占位提示，不阻断）；不可达 → DocgenUnavailable。
    """
    client = get_client()
    try:
        resp = await client.get(f"/api/vstore/documents/{document_id}", timeout=DEFAULT_TIMEOUT)
    except (httpx.HTTPError, httpx.TimeoutException) as exc:
        raise DocgenUnavailable(f"vstore 不可达（documents/{document_id}）：{exc}", cause=exc) from exc
    if resp.status_code == 404:
        return None
    if resp.status_code >= 400:
        raise DocgenUnavailable(f"vstore 返回错误状态 {resp.status_code}")
    payload = _check_ok(resp)
    if not payload.get("ok"):
        return None
    doc = payload.get("document")
    return doc if isinstance(doc, dict) else None


async def get_style(document_id: str) -> dict[str, Any] | None:
    """GET /api/vstore/documents/{id}/style（docx 样式提取）。404/失败 → None。"""
    client = get_client()
    try:
        resp = await client.get(f"/api/vstore/documents/{document_id}/style", timeout=DEFAULT_TIMEOUT)
    except (httpx.HTTPError, httpx.TimeoutException):
        return None
    if resp.status_code >= 400:
        return None
    try:
        payload = resp.json()
    except ValueError:
        return None
    hints = payload.get("style_hints") if isinstance(payload, dict) else None
    return hints if isinstance(hints, dict) else None


async def delete_document(document_id: str) -> dict[str, Any]:
    """DELETE /api/vstore/documents/{id}。"""
    resp = await _request("DELETE", f"/api/vstore/documents/{document_id}")
    return _check_ok(resp)


async def update_metadata(document_id: str, metadata: dict[str, Any]) -> dict[str, Any]:
    """PATCH 元数据。vstore 未提供 PATCH 端点 → 转发后返回错误结构（容错）。"""
    # vstore V1 未提供 PATCH；meta 以 BFF doc_library 落库为准，
    # 透传失败返回错误结构，由 router 容错忽略
    return {"ok": False, "error": {"code": "NOT_IMPLEMENTED", "message": "vstore 无 meta 更新端点（以 BFF 落库为准）"}}


async def rebuild_index() -> dict[str, Any]:
    """POST /api/vstore/rebuild-index。"""
    resp = await _request("POST", "/api/vstore/rebuild-index")
    return _check_ok(resp)


async def health() -> dict[str, Any]:
    """GET /api/vstore/health。失败返回 {ok: False} 结构（不抛）。"""
    try:
        resp = await _request("GET", "/api/vstore/health", retries=0)
        return _check_ok(resp)
    except DocgenUnavailable as exc:
        return {"ok": False, "error": {"code": "UPSTREAM_DOWN", "message": str(exc)}}