✨ feat(llm): 新增LLM模型管理插件并增强API密钥管理 (#1972)

🔧 新增功能: - LLM模型管理插件 (builtin_plugins/llm_manager/) • llm list - 查看可用模型列表 (图片格式) • llm info - 查看模型详细信息 (Markdown图片) • llm default - 管理全局默认模型 • llm test - 测试模型连通性 • llm keys - 查看API Key状态 (表格图片，含健康度/成功率/延迟) • llm reset-key - 重置API Key失败状态 🏗️ 架构重构: - 会话管理: AI/AIConfig 类迁移至独立的 session.py - 类型定义: TaskType 枚举移至 types/enums.py - API增强: • chat() 函数返回完整 LLMResponse，支持工具调用 • 新增 generate() 函数用于一次性响应生成 • 统一API调用核心方法 _perform_api_call，返回使用的API密钥 🚀 密钥管理增强: - 详细状态跟踪: 健康度、成功率、平均延迟、错误信息、建议操作 - 状态持久化: 启动时加载，关闭时自动保存密钥状态 - 智能冷却策略: 根据错误类型设置不同冷却时间 - 延迟监控: with_smart_retry 记录API调用延迟并更新统计 Co-authored-by: webjoin111 <455457521@qq.com> Co-authored-by: HibiKier <45528451+HibiKier@users.noreply.github.com>
2026-10-19 19:20:03 +08:00 · 2025-07-14 22:39:17 +08:00
parent 8649aaaa54
commit 46a0768a45
14 changed files with 1423 additions and 682 deletions
@@ -5,17 +5,27 @@ LLM 核心基础设施模块
 """

 import asyncio
+from dataclasses import asdict, dataclass
+from enum import IntEnum
+import json
+import os
+import time
 from typing import Any

+import aiofiles
 import httpx
+import nonebot
 from pydantic import BaseModel

+from zhenxun.configs.path_config import DATA_PATH
 from zhenxun.services.log import logger
 from zhenxun.utils.user_agent import get_user_agent

 from .types import ProviderConfig
 from .types.exceptions import LLMErrorCode, LLMException

+driver = nonebot.get_driver()
+

 class HttpClientConfig(BaseModel):
    """HTTP客户端配置"""
@@ -194,6 +204,82 @@ async def create_llm_http_client(
    return LLMHttpClient(config)


+class KeyStatus(IntEnum):
+    """用于排序和展示的密钥状态枚举"""
+
+    DISABLED = 0
+    ERROR = 1
+    COOLDOWN = 2
+    WARNING = 3
+    HEALTHY = 4
+    UNUSED = 5
+
+
+@dataclass
+class KeyStats:
+    """单个API Key的详细状态和统计信息"""
+
+    cooldown_until: float = 0.0
+    success_count: int = 0
+    failure_count: int = 0
+    total_latency: float = 0.0
+    last_error_info: str | None = None
+
+    @property
+    def is_available(self) -> bool:
+        """检查Key当前是否可用"""
+        return time.time() >= self.cooldown_until
+
+    @property
+    def avg_latency(self) -> float:
+        """计算平均延迟"""
+        return (
+            self.total_latency / self.success_count if self.success_count > 0 else 0.0
+        )
+
+    @property
+    def success_rate(self) -> float:
+        """计算成功率"""
+        total = self.success_count + self.failure_count
+        return self.success_count / total * 100 if total > 0 else 100.0
+
+    @property
+    def status(self) -> KeyStatus:
+        """根据当前统计数据动态计算状态"""
+        now = time.time()
+        cooldown_left = max(0, self.cooldown_until - now)
+
+        if cooldown_left > 31536000 - 60:
+            return KeyStatus.DISABLED
+        if cooldown_left > 0:
+            return KeyStatus.COOLDOWN
+
+        total_calls = self.success_count + self.failure_count
+        if total_calls == 0:
+            return KeyStatus.UNUSED
+
+        if self.success_rate < 80:
+            return KeyStatus.ERROR
+
+        if total_calls >= 5 and self.avg_latency > 15000:
+            return KeyStatus.WARNING
+
+        return KeyStatus.HEALTHY
+
+    @property
+    def suggested_action(self) -> str:
+        """根据状态给出建议操作"""
+        status_actions = {
+            KeyStatus.DISABLED: "更换Key",
+            KeyStatus.ERROR: "检查网络/重置",
+            KeyStatus.COOLDOWN: "等待/重置",
+            KeyStatus.WARNING: "观察",
+            KeyStatus.HEALTHY: "-",
+            KeyStatus.UNUSED: "-",
+        }
+        return status_actions.get(self.status, "未知")
+
+
 class RetryConfig:
    """重试配置"""

@@ -236,26 +322,38 @@ async def with_smart_retry(
    last_exception: Exception | None = None
    failed_keys: set[str] = set()

+    model_instance = next((arg for arg in args if hasattr(arg, "api_keys")), None)
+    all_provider_keys = model_instance.api_keys if model_instance else []
+
    for attempt in range(config.max_retries + 1):
        try:
            if config.key_rotation and "failed_keys" in func.__code__.co_varnames:
                kwargs["failed_keys"] = failed_keys

-            return await func(*args, **kwargs)
+            start_time = time.monotonic()
+            result = await func(*args, **kwargs)
+            latency = (time.monotonic() - start_time) * 1000
+
+            if key_store and isinstance(result, tuple) and len(result) == 2:
+                final_result, api_key_used = result
+                if api_key_used:
+                    await key_store.record_success(api_key_used, latency)
+                return final_result
+            else:
+                return result

        except LLMException as e:
            last_exception = e
+            api_key_in_use = e.details.get("api_key")

-            if e.code in [
-                LLMErrorCode.API_KEY_INVALID,
-                LLMErrorCode.API_QUOTA_EXCEEDED,
-            ]:
-                if hasattr(e, "details") and e.details and "api_key" in e.details:
-                    failed_keys.add(e.details["api_key"])
-                    if key_store and provider_name:
-                        await key_store.record_failure(
-                            e.details["api_key"], e.details.get("status_code")
-                        )
+            if api_key_in_use:
+                failed_keys.add(api_key_in_use)
+                if key_store and provider_name and len(all_provider_keys) > 1:
+                    status_code = e.details.get("status_code")
+                    error_message = f"({e.code.name}) {e.message}"
+                    await key_store.record_failure(
+                        api_key_in_use, status_code, error_message
+                    )

            should_retry = _should_retry_llm_error(e, attempt, config.max_retries)
            if not should_retry:
@@ -267,7 +365,7 @@ async def with_smart_retry(
                if config.exponential_backoff:
                    wait_time *= 2**attempt
                logger.warning(
-                    f"请求失败，{wait_time}秒后重试 (第{attempt + 1}次): {e}"
+                    f"请求失败，{wait_time:.2f}秒后重试 (第{attempt + 1}次): {e}"
                )
                await asyncio.sleep(wait_time)
            else:
@@ -325,14 +423,66 @@ def _should_retry_llm_error(


 class KeyStatusStore:
-    """API Key 状态管理存储 - 优化版本，支持轮询和负载均衡"""
+    """API Key 状态管理存储 - 支持持久化"""

    def __init__(self):
-        self._key_status: dict[str, bool] = {}
-        self._key_usage_count: dict[str, int] = {}
-        self._key_last_used: dict[str, float] = {}
+        self._key_stats: dict[str, KeyStats] = {}
        self._provider_key_index: dict[str, int] = {}
        self._lock = asyncio.Lock()
+        self._file_path = DATA_PATH / "llm" / "key_status.json"
+
+    async def initialize(self):
+        """从文件异步加载密钥状态，在应用启动时调用"""
+        async with self._lock:
+            if not self._file_path.exists():
+                logger.info("未找到密钥状态文件，将使用内存状态启动。")
+                return
+
+            try:
+                logger.info(f"正在从 {self._file_path} 加载密钥状态...")
+                async with aiofiles.open(self._file_path, encoding="utf-8") as f:
+                    content = await f.read()
+                    if not content:
+                        logger.warning("密钥状态文件为空。")
+                        return
+                    data = json.loads(content)
+
+                for key, stats_dict in data.items():
+                    self._key_stats[key] = KeyStats(**stats_dict)
+
+                logger.info(f"成功加载 {len(self._key_stats)} 个密钥的状态。")
+
+            except json.JSONDecodeError:
+                logger.error(f"密钥状态文件 {self._file_path} 格式错误，无法解析。")
+            except Exception as e:
+                logger.error(f"加载密钥状态文件时发生错误: {e}", e=e)
+
+    async def _save_to_file_internal(self):
+        """
+        [内部方法] 将当前密钥状态安全地写入JSON文件。
+        假定调用方已持有锁。
+        """
+        data_to_save = {key: asdict(stats) for key, stats in self._key_stats.items()}
+
+        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)
+            logger.debug("密钥状态已成功持久化到文件。")
+        except Exception as e:
+            logger.error(f"保存密钥状态到文件失败: {e}", e=e)
+
+    async def shutdown(self):
+        """在应用关闭时安全地保存状态"""
+        async with self._lock:
+            await self._save_to_file_internal()
+        logger.info("KeyStatusStore 已在关闭前保存状态。")

    async def get_next_available_key(
        self,
@@ -355,88 +505,122 @@ class KeyStatusStore:
            return None

        exclude_keys = exclude_keys or set()
-        available_keys = [
-            key
-            for key in api_keys
-            if key not in exclude_keys and self._key_status.get(key, True)
-        ]
-
-        if not available_keys:
-            return api_keys[0] if api_keys else None

        async with self._lock:
+            for key in api_keys:
+                if key not in self._key_stats:
+                    self._key_stats[key] = KeyStats()
+
+            available_keys = [
+                key
+                for key in api_keys
+                if key not in exclude_keys and self._key_stats[key].is_available
+            ]
+
+            if not available_keys:
+                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

-            self._provider_key_index[provider_name] = (current_index + 1) % len(
-                available_keys
+            total_usage = (
+                self._key_stats[selected_key].success_count
+                + self._key_stats[selected_key].failure_count
            )
-
-            import time
-
-            self._key_usage_count[selected_key] = (
-                self._key_usage_count.get(selected_key, 0) + 1
-            )
-            self._key_last_used[selected_key] = time.time()
-
            logger.debug(
                f"轮询选择API密钥: {self._get_key_id(selected_key)} "
-                f"(使用次数: {self._key_usage_count[selected_key]})"
+                f"(使用次数: {total_usage})"
            )
-
            return selected_key

-    async def record_success(self, api_key: str):
-        """记录成功使用"""
+    async def record_success(self, api_key: str, latency: float):
+        """记录成功使用，并持久化"""
        async with self._lock:
-            self._key_status[api_key] = True
-        logger.debug(f"记录API密钥成功使用: {self._get_key_id(api_key)}")
+            stats = self._key_stats.setdefault(api_key, KeyStats())
+            stats.cooldown_until = 0.0
+            stats.success_count += 1
+            stats.total_latency += latency
+            stats.last_error_info = None
+            await self._save_to_file_internal()
+        logger.debug(
+            f"记录API密钥成功使用: {self._get_key_id(api_key)}, 延迟: {latency:.2f}ms"
+        )

-    async def record_failure(self, api_key: str, status_code: int | None):
+    async def record_failure(
+        self, api_key: str, status_code: int | None, error_message: str
+    ):
        """
-        记录失败使用
+        记录失败使用，并设置冷却时间

        参数:
            api_key: API密钥。
            status_code: HTTP状态码。
+            error_message: 错误信息。
        """
        key_id = self._get_key_id(api_key)
+        now = time.time()
+        cooldown_duration = 300
+
+        if status_code in [401, 403, 404]:
+            cooldown_duration = 31536000
+            log_level = "error"
+            log_message = f"API密钥认证/权限/路径错误，将永久禁用: {key_id}"
+        elif status_code == 429:
+            cooldown_duration = 60
+            log_level = "warning"
+            log_message = f"API密钥被限流，冷却60秒: {key_id}"
+        else:
+            log_level = "warning"
+            log_message = f"API密钥遇到临时性错误，冷却{cooldown_duration}秒: {key_id}"
+
        async with self._lock:
-            if status_code in [401, 403]:
-                self._key_status[api_key] = False
-                logger.warning(
-                    f"API密钥认证失败，标记为不可用: {key_id} (状态码: {status_code})"
-                )
-            else:
-                logger.debug(f"记录API密钥失败使用: {key_id} (状态码: {status_code})")
+            stats = self._key_stats.setdefault(api_key, KeyStats())
+            stats.cooldown_until = now + cooldown_duration
+            stats.failure_count += 1
+            stats.last_error_info = error_message[:256]
+            await self._save_to_file_internal()
+
+        getattr(logger, log_level)(log_message)

    async def reset_key_status(self, api_key: str):
-        """重置密钥状态（用于恢复机制）"""
+        """重置密钥状态，并持久化"""
        async with self._lock:
-            self._key_status[api_key] = True
+            stats = self._key_stats.setdefault(api_key, KeyStats())
+            stats.cooldown_until = 0.0
+            stats.last_error_info = None
+            await self._save_to_file_internal()
        logger.info(f"重置API密钥状态: {self._get_key_id(api_key)}")

    async def get_key_stats(self, api_keys: list[str]) -> dict[str, dict]:
        """
-        获取密钥使用统计
+        获取密钥使用统计，并计算出用于展示的派生数据。

        参数:
            api_keys: API密钥列表。

        返回:
-            dict[str, dict]: 密钥统计信息字典。
+            dict[str, dict]: 包含丰富状态和统计信息的密钥字典。
        """
-        stats = {}
+        stats_dict = {}
+        now = time.time()
        async with self._lock:
            for key in api_keys:
                key_id = self._get_key_id(key)
-                stats[key_id] = {
-                    "available": self._key_status.get(key, True),
-                    "usage_count": self._key_usage_count.get(key, 0),
-                    "last_used": self._key_last_used.get(key, 0),
+                stats = self._key_stats.get(key, KeyStats())
+
+                stats_dict[key_id] = {
+                    "status_enum": stats.status,
+                    "cooldown_seconds_left": max(0, stats.cooldown_until - now),
+                    "total_calls": stats.success_count + stats.failure_count,
+                    "success_count": stats.success_count,
+                    "failure_count": stats.failure_count,
+                    "success_rate": stats.success_rate,
+                    "avg_latency": stats.avg_latency,
+                    "last_error": stats.last_error_info,
+                    "suggested_action": stats.suggested_action,
                }
-        return stats
+        return stats_dict

    def _get_key_id(self, api_key: str) -> str:
        """获取API密钥的标识符（用于日志）"""
@@ -446,3 +630,8 @@ class KeyStatusStore:


 key_store = KeyStatusStore()
+
+
+@driver.on_shutdown
+async def _shutdown_key_store():
+    await key_store.shutdown()