Mutual_Fund/service/advisor/agent_client.py

"""投顾Agent HTTP 客户端 + 错误码映射层（工作台 → Agent 的唯一出口）。

职责：
1. 通过 httpx 调用投顾Agent 独立服务（统一前缀 /api/advisor-agent，Agent 文档 §5）；
2. 透传上游投顾 JWT 与 X-Trace-Id；
3. 快接口超时/重试（仅幂等 GET 重试，写操作不盲目重试）；
4. 将 Agent 自有错误码（0/40001/40020/40030/40401/50001/50002）映射为工作台异常，
   绝不把 Agent 码透传给上层调用方（common_const §6 错误码域边界）。
"""
from __future__ import annotations

import httpx

from common_const import (
    AGENT_ERR_MESSAGE,
    ERR_CODE_DRAFT_NOT_FOUND,
    ERR_CODE_FORBIDDEN_CUSTOMER,
    ERR_CODE_GRAPH_ERROR,
    ERR_CODE_LLM_ERROR,
    ERR_CODE_NOT_SIGNED_REBALANCE,
    ERR_CODE_OK,
    ERR_CODE_SUITABILITY_INVALID,
)
from config.settings import settings
from utils.exceptions import (
    ForbiddenError,
    LLMFailError,
    NotFoundError,
    NotSuitableError,
    ParamError,
)

# Agent 统一前缀（Agent 文档 §5）
_AGENT_PREFIX = "/api/advisor-agent"


def translate_agent_error(code: int, message: str | None = None) -> str | None:
    """把 Agent 业务码映射为工作台结果。

    - code == 0：正常，返回 None；
    - code == 50002：降级但成功，返回告警文案（不抛异常）；
    - 其余：抛出映射后的 ApiError（工作台码域，不透传 Agent 码）。
    """
    if code == ERR_CODE_OK:
        return None
    if code == ERR_CODE_GRAPH_ERROR:
        return message or AGENT_ERR_MESSAGE.get(code, "图谱查询异常，已降级返回部分结果")
    default = message or AGENT_ERR_MESSAGE.get(code, "AI 服务调用异常，请稍后重试")
    if code == ERR_CODE_FORBIDDEN_CUSTOMER:
        raise ForbiddenError(default)
    if code == ERR_CODE_SUITABILITY_INVALID:
        raise NotSuitableError(default)
    if code == ERR_CODE_NOT_SIGNED_REBALANCE:
        raise ParamError(default)
    if code == ERR_CODE_DRAFT_NOT_FOUND:
        raise NotFoundError(default)
    if code == ERR_CODE_LLM_ERROR:
        raise LLMFailError(default)
    # 未知 Agent 业务码：统一按 AI 服务异常兜底，不透传
    raise LLMFailError(default)


class AdvisorAgentClient:
    """投顾Agent 客户端（模块级单例 get_agent_client() 获取）。"""

    def __init__(self, base_url: str, timeout: float, retry: int):
        self.base_url = (base_url or "").rstrip("/")
        self.timeout = timeout
        self.retry = max(0, retry)

    @property
    def configured(self) -> bool:
        return bool(self.base_url)

    def _url(self, path: str) -> str:
        return f"{self.base_url}{_AGENT_PREFIX}{path}"

    def _headers(self, auth_header: str, trace_id: str) -> dict:
        return {
            "Authorization": auth_header,
            "X-Trace-Id": trace_id,
            "Content-Type": "application/json",
        }

    async def _request(
        self,
        method: str,
        path: str,
        *,
        auth_header: str,
        trace_id: str,
        params: dict | None = None,
        json: dict | None = None,
        allow_retry: bool = False,
    ) -> dict:
        """统一请求：校验配置 → 超时重试 → 解析返回体 → 映射业务码。

        返回 {"data": ..., "warning": ...}；业务失败抛出映射后的 ApiError。
        """
        if not self.configured:
            raise LLMFailError("投顾Agent 服务未配置，请联系管理员")
        url = self._url(path)
        headers = self._headers(auth_header, trace_id)
        attempts = self.retry + 1 if allow_retry else 1
        for attempt in range(attempts):
            try:
                async with httpx.AsyncClient(timeout=self.timeout) as client:
                    resp = await client.request(
                        method, url, headers=headers, params=params, json=json
                    )
            except (httpx.TimeoutException, httpx.TransportError) as exc:
                if attempt < attempts - 1:
                    continue
                raise LLMFailError("AI 服务调用异常，请稍后重试") from exc

            if resp.status_code != 200:
                if attempt < attempts - 1:
                    continue
                raise LLMFailError("AI 服务调用异常，请稍后重试")

            try:
                body = resp.json()
            except ValueError as exc:
                raise LLMFailError("AI 服务返回格式异常，请稍后重试") from exc

            code = body.get("code", ERR_CODE_LLM_ERROR)
            warning = translate_agent_error(code, body.get("message"))
            return {"data": body.get("data"), "warning": warning}
        raise LLMFailError("AI 服务调用异常，请稍后重试")  # 理论不可达，防御

    # ---- 各 Agent 接口（字段与 Agent 文档 §5 对齐） ----
    async def draft_list(
        self,
        *,
        auth_header: str,
        trace_id: str,
        advisor_id: int,
        customer_id: int | None = None,
        status: str | None = None,
        page: int = 1,
        page_size: int = 10,
    ) -> dict:
        params = {"advisor_id": advisor_id, "page": page, "page_size": page_size}
        if customer_id is not None:
            params["customer_id"] = customer_id
        if status:
            params["status"] = status
        return await self._request(
            "GET", "/draft/list", auth_header=auth_header, trace_id=trace_id, params=params,
            allow_retry=True,
        )

    async def draft_detail(self, draft_id: str, *, auth_header: str, trace_id: str) -> dict:
        return await self._request(
            "GET", f"/draft/{draft_id}", auth_header=auth_header, trace_id=trace_id,
            allow_retry=True,
        )

    async def draft_save(
        self, draft_id: str, payload: dict, *, auth_header: str, trace_id: str
    ) -> dict:
        return await self._request(
            "PUT", f"/draft/{draft_id}/save", auth_header=auth_header,
            trace_id=trace_id, json=payload,
        )

    async def draft_operate(
        self, draft_id: str, action: str, *, auth_header: str, trace_id: str
    ) -> dict:
        return await self._request(
            "POST", f"/draft/{draft_id}/operate", auth_header=auth_header,
            trace_id=trace_id, json={"operation": action},
        )

    async def rebalance_run(
        self, customer_id: int, *, auth_header: str, trace_id: str
    ) -> dict:
        return await self._request(
            "POST", "/rebalance/run", auth_header=auth_header,
            trace_id=trace_id, json={"query": f"请为客户{customer_id}执行调仓再平衡"},
        )

    async def generate_talk_script(
        self, customer_id: int, scene_type: str, *, auth_header: str, trace_id: str
    ) -> dict:
        return await self._request(
            "POST", "/generate-talk-script", auth_header=auth_header,
            trace_id=trace_id,
            json={"query": f"请为客户{customer_id}生成{scene_type}沟通话术"},
        )

    async def data_query(
        self,
        payload: dict,
        *,
        auth_header: str,
        trace_id: str,
    ) -> dict:
        """代理当前投顾选中客户的数据查询，不向工作台暴露 SQL。"""
        return await self._request(
            "POST", "/data-query", auth_header=auth_header,
            trace_id=trace_id, json=payload,
        )


_client: AdvisorAgentClient | None = None


def get_agent_client() -> AdvisorAgentClient:
    """取全局客户端单例（按 config.settings 组装）。"""
    global _client
    if _client is None:
        cfg = settings.advisor_agent
        _client = AdvisorAgentClient(cfg.base_url, cfg.timeout, cfg.retry)
    return _client