Mutual_Fund/tool/llm.py

"""统一 LLM 客户端：本地 Ollama / OpenAI 兼容 API（api_key）双后端，按优先级降级。

设计：两个后端的底层都是 OpenAI 兼容 /v1 端点，业务代码一律调用 `chat()` / `embed()`，
不感知底层。后端链由 .env 决定（LLM_MODE）：
  - auto（默认）：本地 Ollama 参数齐全则优先本地，LLM_API_KEY 齐全则作为其后的降级后端；
  - ollama / api：强制只用该后端。
chat 降级链：后端内 主模型 → fallback_chat_model（每模型指数退避重试）→ 下一后端 → 全部失败抛
LLMFailError（§1001），调用方兜底话术见 FALLBACK_REPLY。
embed 只走首选后端、不跨后端降级：不同模型的向量空间不兼容，混入 Milvus 会污染检索。
"""
from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass

import httpx

from config.settings import LLMCfg, settings
from utils.exceptions import LLMFailError

logger = logging.getLogger("tool.llm")

# 兜底话术：所有后端 + 备用模型均失败时返回
FALLBACK_REPLY = "抱歉，服务暂时不可用，请稍后再试，或拨打客服热线 400-XXX-XXXX。"


@dataclass(frozen=True)
class Backend:
    """一个 OpenAI 兼容端点及其模型；name 为 ollama / api。"""

    name: str
    base_url: str  # 以 /v1 结尾
    chat_model: str
    embed_model: str
    api_key: str = ""

    @property
    def is_ollama(self) -> bool:
        return self.name == "ollama"

    @property
    def headers(self) -> dict:
        h = {"Content-Type": "application/json"}
        if self.api_key:
            h["Authorization"] = f"Bearer {self.api_key}"
        return h

    def client(self, timeout: float) -> httpx.AsyncClient:
        """ollama 走本机端点，必须绕过系统代理：httpx 会读取 Windows 系统代理，
        但不识别其 ProxyOverride 白名单，localhost 请求会被转发到代理并返回 502。
        API 后端保留 trust_env，远端接口可能依赖代理。"""
        return httpx.AsyncClient(timeout=timeout, trust_env=not self.is_ollama)


def resolve_backends(cfg: LLMCfg) -> list[Backend]:
    """按 LLM_MODE 与 .env 完整度解析后端链，首个为首选后端。"""
    ollama = (
        Backend(
            name="ollama",
            base_url=cfg.ollama_base.rstrip("/") + "/v1",
            chat_model=cfg.ollama_chat_model,
            embed_model=cfg.ollama_embed_model,
        )
        if cfg.ollama_base and cfg.ollama_chat_model
        else None
    )
    api = (
        Backend(
            name="api",
            base_url=cfg.api_base.rstrip("/"),
            chat_model=cfg.api_chat_model,
            embed_model=cfg.api_embed_model,
            api_key=cfg.api_key,
        )
        if cfg.api_key and cfg.api_base and cfg.api_chat_model
        else None
    )
    mode = cfg.mode.lower()
    if mode == "ollama":
        chain = [ollama]
    elif mode == "api":
        chain = [api]
    elif mode == "auto":
        chain = [ollama, api]  # 本地优先，API 兜底
    else:
        raise ValueError(f"LLM_MODE 无效: {cfg.mode!r}，应为 auto / ollama / api")
    chain = [b for b in chain if b is not None]
    if not chain:
        raise ValueError(
            f"LLM_MODE={cfg.mode} 下没有可用后端：本地需 LLM_OLLAMA_BASE + LLM_OLLAMA_CHAT_MODEL，"
            "API 需 LLM_API_KEY + LLM_API_BASE + LLM_API_CHAT_MODEL"
        )
    return chain


class LLMClient:
    def __init__(self, cfg: LLMCfg | None = None):
        self.cfg = cfg or settings.llm
        self.backends = resolve_backends(self.cfg)

    # ---- 首选后端（健康检查/日志/embed 使用） -----------------------------
    @property
    def primary(self) -> Backend:
        return self.backends[0]

    @property
    def is_ollama(self) -> bool:
        return self.primary.is_ollama

    @property
    def base_url(self) -> str:
        return self.primary.base_url

    @property
    def chat_model(self) -> str:
        return self.primary.chat_model

    @property
    def embed_model(self) -> str:
        return self.primary.embed_model

    def describe(self) -> str:
        """健康检查/日志用：后端链（首选在前）+ 各自模型。"""
        chain = " -> ".join(
            f"{b.name}[{b.base_url}] chat={b.chat_model} embed={b.embed_model or '-'}"
            for b in self.backends
        )
        return f"mode={self.cfg.mode} {chain}"

    # ---- 对话 -----------------------------------------------------------
    async def chat(
        self,
        messages: list[dict],
        *,
        temperature: float | None = None,
        max_tokens: int | None = None,
        model: str | None = None,
    ) -> str:
        """chat/completions，返回 assistant 文本。

        降级链：首选后端 [model > chat_model > fallback_chat_model] → 后续后端 [chat_model >
        fallback_chat_model]。显式 model 只作用于首选后端——模型名与后端绑定，跨后端无意义。
        """
        last_err: Exception | None = None
        for index, backend in enumerate(self.backends):
            models = [(model or backend.chat_model) if index == 0 else backend.chat_model]
            if self.cfg.fallback_chat_model and self.cfg.fallback_chat_model not in models:
                models.append(self.cfg.fallback_chat_model)
            for m in models:
                try:
                    return await self._chat_once(backend, m, messages, temperature, max_tokens)
                except Exception as e:  # noqa: BLE001 重试/切换模型或后端，异常向上收敛
                    logger.warning(
                        "chat backend=%s model=%r failed: %s: %s",
                        backend.name, m, type(e).__name__, e,
                    )
                    last_err = e
            if index < len(self.backends) - 1:
                logger.warning("backend=%s exhausted, falling back to %s",
                               backend.name, self.backends[index + 1].name)
        raise LLMFailError() from last_err

    async def _chat_once(
        self,
        backend: Backend,
        model: str,
        messages: list[dict],
        temperature: float | None,
        max_tokens: int | None,
    ) -> str:
        payload = {
            "model": model,
            "messages": messages,
            "temperature": self.cfg.temperature if temperature is None else temperature,
            "max_tokens": self.cfg.max_tokens if max_tokens is None else max_tokens,
            "stream": False,
        }
        url = f"{backend.base_url}/chat/completions"
        # 指数退避：retry_backoff_sec → 翻倍 → …，最多 max_retries 次
        for attempt in range(self.cfg.max_retries):
            try:
                async with backend.client(self.cfg.timeout) as client:
                    r = await client.post(url, headers=backend.headers, json=payload)
                    r.raise_for_status()
                return r.json()["choices"][0]["message"]["content"]
            except Exception:
                if attempt == self.cfg.max_retries - 1:
                    raise
                await asyncio.sleep(self.cfg.retry_backoff_sec * (2**attempt))
        raise LLMFailError()  # 理论不可达，防御

    # ---- Embedding（RAG 入口的统一向量化） -------------------------------
    async def embed(self, texts: list[str]) -> list[list[float]]:
        """批量向量化，返回与 input 一一对应的向量列表。

        只走首选后端（见模块说明）。按 OpenAI 兼容协议传 `dimensions`（来自
        LLM_EMBED_DIMENSIONS），MRL 模型（qwen3-embedding 等）会截断并重新归一化到目标维度。
        """
        backend = self.primary
        if not backend.embed_model:
            raise LLMFailError(f"backend={backend.name} 未配置 embed 模型")
        url = f"{backend.base_url}/embeddings"
        payload = {
            "model": backend.embed_model,
            "input": texts,
            # "dimensions": self.cfg.embed_dimensions,
        }
        async with backend.client(self.cfg.timeout) as client:
            r = await client.post(url, headers=backend.headers, json=payload)
            r.raise_for_status()
        data = r.json()
        return [item["embedding"] for item in data["data"]]

    async def embed_one(self, text: str) -> list[float]:
        return (await self.embed([text]))[0]

    # ---- 就绪探测（启动/健康检查可选接入） --------------------------------
    async def check_health(self) -> None:
        """验证首选后端端点可达。ollama 走 /api/tags，API 走 /models（需鉴权，401 也算可达）。"""
        backend = self.primary
        if backend.is_ollama:
            url = backend.base_url.removesuffix("/v1") + "/api/tags"
        else:
            url = f"{backend.base_url}/models"
        async with backend.client(min(self.cfg.timeout, 10)) as client:
            r = await client.get(url, headers=backend.headers)
            if r.status_code >= 500:
                r.raise_for_status()


# 全局单例：Agent 统一引入
llm = LLMClient()