zhenxun_bot/zhenxun/services/ai/llm/system/network.py

"""
LLM 核心基础设施模块

包含执行 LLM 请求所需的底层组件，如 HTTP 客户端、API Key 存储和智能重试逻辑。
"""

import asyncio
import json
import os
from pathlib import Path
import time
from typing import Any

import aiofiles
import httpx
import nonebot

from zhenxun.configs.config import BotConfig
from zhenxun.configs.path_config import DATA_PATH
from zhenxun.services.ai.config import ProviderConfig
from zhenxun.services.ai.core.exceptions import (
    AuthenticationException,
    ConfigurationException,
    LocationNotSupportedException,
    QuotaExceededException,
    RateLimitException,
)
from zhenxun.services.ai.utils.logger import log_llm as logger
from zhenxun.utils.pydantic_compat import model_dump, parse_as
from zhenxun.utils.user_agent import get_user_agent

from .models import (
    CircuitBreakerPolicy,
    GlobalHealthState,
    HttpClientConfig,
    KeyHealthStatus,
    ProviderHealthStatus,
    RouteHealthState,
    RouteHealthStatus,
)

driver = nonebot.get_driver()


class LLMHttpClient:
    """[内部 API] LLM 服务专用异步 HTTP 客户端封装。"""

    def __init__(self, config: HttpClientConfig | None = None):
        """初始化 LLM 服务专用 HTTP 客户端"""
        self.config = config or HttpClientConfig()
        self._client: httpx.AsyncClient | None = None
        self._active_requests = 0
        self._lock = asyncio.Lock()

    async def _ensure_client_initialized(self) -> httpx.AsyncClient:
        """确保底层的 AsyncClient 已完成初始化"""
        if self._client is None or self._client.is_closed:
            async with self._lock:
                if self._client is None or self._client.is_closed:
                    logger.debug(
                        f"LLMHttpClient: 正在初始化新的 httpx.AsyncClient "
                        f"配置: {self.config}"
                    )
                    headers = get_user_agent()
                    limits = httpx.Limits(
                        max_connections=self.config.max_connections,
                        max_keepalive_connections=self.config.max_keepalive_connections,
                    )
                    timeout = httpx.Timeout(self.config.timeout)

                    client_kwargs = {}
                    if BotConfig.system_proxy:
                        try:
                            version_parts = httpx.__version__.split(".")
                            major = int(
                                "".join(c for c in version_parts[0] if c.isdigit())
                            )
                            minor = (
                                int("".join(c for c in version_parts[1] if c.isdigit()))
                                if len(version_parts) > 1
                                else 0
                            )
                            if (major, minor) >= (0, 28):
                                client_kwargs["proxy"] = BotConfig.system_proxy
                            else:
                                client_kwargs["proxies"] = BotConfig.system_proxy
                        except (ValueError, IndexError):
                            client_kwargs["proxies"] = BotConfig.system_proxy
                            logger.warning(
                                f"无法解析 httpx version '{httpx.__version__}'，"
                                "LLM模块将默认使用旧版 'proxies' 参数语法。"
                            )

                    self._client = httpx.AsyncClient(
                        headers=headers,
                        limits=limits,
                        timeout=timeout,
                        follow_redirects=True,
                        **client_kwargs,
                    )
        if self._client is None:
            raise ConfigurationException(
                "HTTP 客户端初始化失败。",
            )
        return self._client

    async def request(self, method: str, url: str, **kwargs: Any) -> httpx.Response:
        """发送异步 HTTP 请求"""
        client = await self._ensure_client_initialized()
        async with self._lock:
            self._active_requests += 1
        try:
            return await client.request(method, url, **kwargs)
        finally:
            async with self._lock:
                self._active_requests -= 1

    async def post(self, url: str, **kwargs: Any) -> httpx.Response:
        """发送异步 POST 请求"""
        return await self.request("POST", url, **kwargs)

    async def close(self):
        """安全关闭 HTTP 客户端并释放连接池"""
        async with self._lock:
            if self._client and not self._client.is_closed:
                logger.debug(
                    f"LLMHttpClient: 正在关闭，配置: {self.config}. "
                    f"活跃请求数: {self._active_requests}"
                )
                if self._active_requests > 0:
                    logger.warning(
                        f"LLMHttpClient: 关闭时仍有 {self._active_requests} "
                        f"个请求处于活跃状态。"
                    )
                await self._client.aclose()
            self._client = None
        logger.debug(f"配置为 {self.config} 的 LLMHttpClient 已完全关闭。")

    @property
    def is_closed(self) -> bool:
        """检查底层客户端是否已关闭"""
        return self._client is None or self._client.is_closed


class LLMHttpClientManager:
    """[内部 API] 负责管理与复用 LLMHttpClient 连接池。"""

    def __init__(self):
        """初始化客户端管理器"""
        self._clients: dict[tuple[str, str, int], LLMHttpClient] = {}
        self._lock = asyncio.Lock()

    def _get_client_key(self, provider_config: ProviderConfig) -> tuple[str, str, int]:
        """获取客户端唯一缓存标识"""
        api_base = provider_config.api_base or ""
        return (provider_config.api_type, api_base, provider_config.timeout)

    async def get_client(self, provider_config: ProviderConfig) -> LLMHttpClient:
        """获取或创建指定配置的 HTTP 客户端实例"""
        key = self._get_client_key(provider_config)
        async with self._lock:
            client = self._clients.get(key)
            if client and not client.is_closed:
                logger.debug(
                    f"LLMHttpClientManager: 复用现有的 LLMHttpClient 密钥: {key}"
                )
                return client

            if client and client.is_closed:
                logger.debug(
                    f"LLMHttpClientManager: 发现密钥 {key} 对应的客户端已关闭。"
                    f"正在创建新的客户端。"
                )

            logger.debug(f"LLMHttpClientManager: 为密钥 {key} 创建新的 LLMHttpClient")
            http_client_config = HttpClientConfig(timeout=provider_config.timeout)
            new_client = LLMHttpClient(config=http_client_config)
            self._clients[key] = new_client
            return new_client

    async def shutdown(self):
        """关闭所有托管的客户端连接池"""
        async with self._lock:
            logger.info(
                f"LLMHttpClientManager: 正在关闭。关闭 {len(self._clients)} 个客户端。"
            )
            close_tasks = [
                client.close()
                for client in self._clients.values()
                if client and not client.is_closed
            ]
            if close_tasks:
                await asyncio.gather(*close_tasks, return_exceptions=True)
            self._clients.clear()
        logger.info("LLMHttpClientManager: 关闭完成。")


http_client_manager = LLMHttpClientManager()


async def create_llm_http_client(
    timeout: int = 180,
) -> LLMHttpClient:
    """创建并返回一个新的 HTTP 客户端"""
    config = HttpClientConfig(timeout=timeout)
    return LLMHttpClient(config)


class HealthStatePersister:
    """后台异步持久化管理器"""

    def __init__(self, state: GlobalHealthState, file_path: Path):
        """初始化持久化管理器"""
        self.state = state
        self.file_path = file_path
        self._is_dirty = False
        self._lock = asyncio.Lock()
        self._watchdog_task: asyncio.Task | None = None
        self._stop_event = asyncio.Event()

    def start(self):
        """启动后台异步保存定时任务"""
        if self._watchdog_task is None or self._watchdog_task.done():
            self._stop_event.clear()
            self._watchdog_task = asyncio.create_task(self._watchdog_loop())

    def mark_dirty(self):
        """标记内存状态已脏，需要存盘"""
        self._is_dirty = True

    async def _watchdog_loop(self):
        """后台循环检测并持久化脏数据"""
        while not self._stop_event.is_set():
            await asyncio.sleep(5)
            if self._is_dirty:
                await self.force_save()

    async def force_save(self):
        """强制将内存中的状态同步写入磁盘文件"""
        if not self._is_dirty:
            return
        async with self._lock:
            self._is_dirty = False

            data_to_save = model_dump(self.state)
            try:
                self.file_path.parent.mkdir(parents=True, exist_ok=True)
                temp_path = self.file_path.with_suffix(".json.tmp")
                async with aiofiles.open(temp_path, "w", encoding="utf-8") as f:
                    await f.write(
                        json.dumps(data_to_save, ensure_ascii=False, indent=2)
                    )
                if self.file_path.exists():
                    self.file_path.unlink()
                os.rename(temp_path, self.file_path)
            except Exception as e:
                logger.error(f"保存密钥状态到文件失败: {e}", e=e)

    async def stop(self):
        """停止后台任务并保存所有脏状态"""
        self._stop_event.set()
        if self._watchdog_task:
            self._watchdog_task.cancel()
            try:
                await self._watchdog_task
            except asyncio.CancelledError:
                pass
        await self.force_save()


class KeyRotationManager:
    """专门负责 API Key 的负载均衡与冷却状态维护"""

    def __init__(self, state: GlobalHealthState):
        """初始化 API Key 轮询管理器"""
        self.state = state
        self._provider_key_index: dict[str, int] = {}

    def get_next_available_key(
        self,
        provider_name: str,
        api_keys: list[str],
        exclude_keys: set[str] | None = None,
        strict_mode: bool = False,
    ) -> str | None:
        """轮询策略获取下一个健康可用的 API Key"""
        if not api_keys:
            return None

        exclude_keys = exclude_keys or set()

        provider_state = self.state.providers.setdefault(
            provider_name, ProviderHealthStatus()
        )
        for key in api_keys:
            if key not in provider_state.api_keys:
                provider_state.api_keys[key] = KeyHealthStatus()

        now = time.time()
        available_keys = [
            key
            for key in api_keys
            if key not in exclude_keys
            and provider_state.api_keys[key].cooldown_until <= now
        ]

        if not available_keys:
            if strict_mode:
                return None
            return api_keys[0]

        current_index = self._provider_key_index.get(provider_name, 0)
        selected_key = available_keys[current_index % len(available_keys)]
        self._provider_key_index[provider_name] = current_index + 1

        stats = provider_state.api_keys[selected_key]
        total_usage = stats.successes + stats.failures
        logger.debug(f"轮询选择API密钥 (使用次数: {total_usage})")
        return selected_key

    def record_key_success(self, provider_name: str, api_key: str):
        """记录指定 API Key 调用成功状态"""
        provider_state = self.state.providers.setdefault(
            provider_name, ProviderHealthStatus()
        )
        stats = provider_state.api_keys.setdefault(api_key, KeyHealthStatus())
        stats.cooldown_until = 0.0
        stats.successes += 1
        stats.status = "HEALTHY"
        stats.last_error = None

    def record_key_failure(
        self,
        provider_name: str,
        api_key: str,
        exception: Exception,
        policy: CircuitBreakerPolicy,
    ):
        """记录并处理指定 API Key 的调用失败冷却"""
        now = time.time()
        cooldown_duration = 0

        error_message = str(exception)

        if isinstance(exception, LocationNotSupportedException):
            provider_state = self.state.providers.setdefault(
                provider_name, ProviderHealthStatus()
            )
            stats = provider_state.api_keys.setdefault(api_key, KeyHealthStatus())
            stats.failures += 1
            stats.last_error = error_message[:256]
            return

        if isinstance(exception, QuotaExceededException):
            cooldown_duration = policy.quota_error_cooldown
        elif isinstance(exception, AuthenticationException):
            cooldown_duration = policy.auth_error_cooldown
        elif isinstance(exception, RateLimitException):
            cooldown_duration = policy.rate_limit_cooldown

        provider_state = self.state.providers.setdefault(
            provider_name, ProviderHealthStatus()
        )
        stats = provider_state.api_keys.setdefault(api_key, KeyHealthStatus())
        if cooldown_duration > 0:
            stats.cooldown_until = now + cooldown_duration
            stats.status = (
                "COOLDOWN"
                if cooldown_duration < policy.auth_error_cooldown
                else "DISABLED"
            )
        stats.failures += 1
        stats.last_error = error_message[:256]

    def reset_key_status(self, provider_name: str, api_key: str):
        """重置 API Key 状态为健康"""
        provider_state = self.state.providers.setdefault(
            provider_name, ProviderHealthStatus()
        )
        stats = provider_state.api_keys.setdefault(api_key, KeyHealthStatus())
        stats.cooldown_until = 0.0
        stats.last_error = None
        stats.status = "HEALTHY"


class CircuitBreakerManager:
    """专门负责模型路由级别的熔断与探活 (无锁化设计)"""

    def __init__(self, state: GlobalHealthState):
        """初始化路由级熔断管理器"""
        self.state = state

    def is_route_healthy(self, route_name: str, strict_mode: bool = True) -> bool:
        """检查路由节点的健康与熔断状态"""
        stats = self.state.routes.get(route_name)
        if not stats:
            return True

        if stats.state == RouteHealthState.CLOSED:
            return True

        if stats.state == RouteHealthState.OPEN:
            if time.time() >= stats.cooldown_until:
                stats.state = RouteHealthState.HALF_OPEN
                logger.info(
                    f"🔄 [Route-Level] 节点 '{route_name}' "
                    "冷却期结束，进入 HALF_OPEN 半开试探状态。"
                )
                return True
            if not strict_mode:
                logger.debug(
                    f"👉 [Route-Level] 节点 '{route_name}' 处于熔断状态(OPEN)，"
                    "但因非严格模式(单模型直调)，强制放行执行探活。"
                )
                return True
            return False

        if stats.state == RouteHealthState.HALF_OPEN:
            if not strict_mode:
                return True
            return False

        return True

    def record_route_success(self, route_name: str, latency: float):
        """记录路由请求成功并尝试闭合熔断器"""
        stats = self.state.routes.setdefault(route_name, RouteHealthStatus())
        stats.successes += 1
        total = stats.successes + stats.failures
        stats.success_rate = (stats.successes / total) * 100

        if stats.latency_ema == 0.0:
            stats.latency_ema = latency
        else:
            stats.latency_ema = 0.2 * latency + 0.8 * stats.latency_ema

        if stats.state != RouteHealthState.CLOSED:
            logger.info(
                f"✅ [Route-Level] 节点 '{route_name}' "
                "试探成功！已完全恢复健康状态 (CLOSED)。"
            )
            stats.state = RouteHealthState.CLOSED
            stats.cooldown_until = 0.0

        stats.last_error = None

    def record_route_failure(
        self, route_name: str, exception: Exception, policy: CircuitBreakerPolicy
    ):
        """记录路由失败并开启熔断状态"""
        now = time.time()
        cooldown_duration = policy.server_error_cooldown

        stats = self.state.routes.setdefault(route_name, RouteHealthStatus())
        stats.failures += 1
        total = stats.successes + stats.failures
        stats.success_rate = (stats.successes / total) * 100

        stats.state = RouteHealthState.OPEN
        stats.cooldown_until = now + cooldown_duration
        stats.last_error = str(exception)[:256]

    def get_best_fallback_route(self, route_names: list[str]) -> str:
        """选择处于熔断冷却最少或最健康的备选路由"""

        def get_cooldown(name: str) -> float:
            stats = self.state.routes.get(name)
            return stats.cooldown_until if stats else 0.0

        return sorted(route_names, key=get_cooldown)[0]


class HealthManager:
    """全局 AI 健康与遥测门面"""

    def __init__(self):
        """初始化健康监测与遥测门面"""
        self.state = GlobalHealthState()
        self._file_path = DATA_PATH / "ai" / "api_key.json"
        self._persister: HealthStatePersister | None = None
        self._key_manager = KeyRotationManager(self.state)
        self._circuit_manager = CircuitBreakerManager(self.state)
        self.policy = CircuitBreakerPolicy()

    async def initialize(self):
        """从本地文件异步加载遥测状态"""
        if not self._file_path.exists():
            logger.debug("未找到遥测状态文件，将使用内存状态启动。")
        else:
            try:
                import aiofiles

                logger.debug(f"正在从 {self._file_path} 加载密钥状态...")
                async with aiofiles.open(self._file_path, encoding="utf-8") as f:
                    content = await f.read()
                    if content:
                        self.state = parse_as(GlobalHealthState, json.loads(content))
                        self._key_manager.state = self.state
                        self._circuit_manager.state = self.state
                total_keys = sum(
                    len(provider.api_keys) for provider in self.state.providers.values()
                )
                logger.debug(f"成功加载 {total_keys} 个密钥的状态。")
            except json.JSONDecodeError:
                logger.error(f"遥测状态文件 {self._file_path} 格式错误，无法解析。")
            except Exception as e:
                logger.error(f"加载遥测状态文件时发生错误: {e}", e=e)

        self._persister = HealthStatePersister(self.state, self._file_path)
        self._persister.start()

    async def shutdown(self):
        """在应用关闭时安全地持久化健康状态"""
        if self._persister:
            await self._persister.stop()
        logger.debug("HealthManager 已在关闭前保存遥测状态。")

    async def get_next_available_key(
        self,
        provider_name: str,
        api_keys: list[str],
        exclude_keys: set[str] | None = None,
        strict_mode: bool = False,
    ) -> str | None:
        """路由获取下一个可用 API Key"""
        return self._key_manager.get_next_available_key(
            provider_name, api_keys, exclude_keys, strict_mode
        )

    def is_route_healthy(self, route_name: str, strict_mode: bool = True) -> bool:
        """路由判断指定模型路由是否健康"""
        return self._circuit_manager.is_route_healthy(route_name, strict_mode)

    async def record_route_success(self, route_name: str, latency: float):
        """路由记录模型请求成功"""
        self._circuit_manager.record_route_success(route_name, latency)
        if self._persister:
            self._persister.mark_dirty()

    async def record_route_failure(self, route_name: str, exception: Exception):
        """路由记录模型请求失败并记录熔断"""
        self._circuit_manager.record_route_failure(route_name, exception, self.policy)
        if self._persister:
            self._persister.mark_dirty()
        logger.warning(
            f"🚨 [Route-Level] 节点 '{route_name}' 发生服务端故障，"
            f"已触发熔断 (OPEN)。错误: {exception}"
        )

    def get_best_fallback_route(self, route_names: list[str]) -> str:
        """路由获取最佳备选健康节点"""
        return self._circuit_manager.get_best_fallback_route(route_names)

    async def record_key_success(self, provider_name: str, api_key: str):
        """路由记录 API Key 成功状态"""
        self._key_manager.record_key_success(provider_name, api_key)
        if self._persister:
            self._persister.mark_dirty()

    async def record_key_failure(
        self,
        provider_name: str,
        api_key: str,
        exception: Exception,
    ):
        """路由记录 API Key 失败冷却状态"""
        self._key_manager.record_key_failure(
            provider_name, api_key, exception, self.policy
        )
        if self._persister:
            self._persister.mark_dirty()
        key_id = self._get_key_id(api_key)
        logger.debug(f"API Key {key_id} 发生失败: {exception}")

    async def reset_key_status(self, provider_name: str, api_key: str):
        """路由重置 API Key 的健康状态"""
        self._key_manager.reset_key_status(provider_name, api_key)
        if self._persister:
            self._persister.mark_dirty()
        logger.info(f"重置API密钥状态: {self._get_key_id(api_key)}")

    def _get_key_id(self, api_key: str) -> str:
        """脱敏获取用于日志的 API Key 摘要 ID"""
        if len(api_key) <= 8:
            return api_key
        return f"{api_key[:4]}...{api_key[-4:]}"


health_manager = HealthManager()


@driver.on_shutdown
async def _shutdown_health_manager():
    await health_manager.shutdown()