""" Key健康度检测路由 手动检测key是否可用,更新状态 """ from __future__ import annotations import asyncio import logging import time from datetime import datetime import httpx from fastapi import APIRouter, HTTPException from pydantic import BaseModel from app.db import get_pool router = APIRouter(tags=["group_keys_test"]) logger = logging.getLogger(__name__) # 检测配置 TEST_TIMEOUT = 10 # 超时时间(秒) MAX_CONSECUTIVE_FAILURES = 3 # 连续失败多少次后禁用 class TestResult(BaseModel): success: bool status: str # active / rate_limited / failed / banned error: str | None = None latency_ms: int = 0 class TestAllResult(BaseModel): total: int active: int failed: int banned: int health_rate: float async def _test_key(api_key: str, base_url: str, provider: str) -> tuple[bool, str]: """测试key是否可用""" try: async with httpx.AsyncClient(timeout=TEST_TIMEOUT) as client: # 根据provider选择不同的测试方式 if provider == "openai": url = f"{base_url.rstrip('/')}/v1/chat/completions" headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} data = {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hi"}], "max_tokens": 5} response = await client.post(url, headers=headers, json=data) elif provider == "dashscope": url = f"{base_url.rstrip('/')}/compatible-mode/v1/chat/completions" headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} data = {"model": "qwen-turbo", "messages": [{"role": "user", "content": "Hi"}], "max_tokens": 5} response = await client.post(url, headers=headers, json=data) else: # 默认OpenAI兼容 url = f"{base_url.rstrip('/')}/v1/chat/completions" headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} data = {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hi"}], "max_tokens": 5} response = await client.post(url, headers=headers, json=data) status_code = response.status_code if status_code in [200, 201]: return True, "" elif status_code == 429: return False, "429_rate_limited" elif status_code in [401, 403]: return False, f"{status_code}_invalid_key" elif status_code == 400: try: error_data = response.json() error_msg = error_data.get('error', {}).get('message', '') if 'model' in error_msg.lower(): return True, "" # 模型名错误但key可能是好的 except: pass return False, "400_bad_request" elif status_code in [500, 502, 503, 504]: return False, f"{status_code}_server_error" else: return False, f"{status_code}_unknown" except httpx.TimeoutException: return False, "timeout" except httpx.ConnectError as e: return False, f"network: {str(e)[:50]}" except Exception as e: return False, f"unknown: {str(e)[:50]}" @router.post("/group-keys/{key_id}/test") async def test_key(key_id: int): """检测单个key的健康状态""" pool = get_pool() async with pool.acquire() as conn: # 获取key信息 row = await conn.fetchrow(""" SELECT gk.id, gk.api_key_id, gk.group_id, gk.status, gk.is_active, ak.key_value, ak.name as key_name, mg.name as group_name FROM group_keys gk JOIN api_keys ak ON ak.id = gk.api_key_id JOIN model_groups mg ON mg.id = gk.group_id WHERE gk.id = $1 """, key_id) if not row: raise HTTPException(status_code=404, detail="Key不存在") # 获取分组的provider(从模型推断) provider = "dashscope" # 默认 base_url = "https://dashscope.aliyuncs.com" # DashScope默认URL # 测试key start_time = time.time() is_healthy, error_type = await _test_key(row['key_value'], base_url, provider) end_time = time.time() latency_ms = int((end_time - start_time) * 1000) # 更新状态 if is_healthy: new_status = "active" await conn.execute(""" UPDATE group_keys SET status = 'active', status_msg = NULL, updated_at = NOW() WHERE id = $1 """, key_id) else: if error_type == "429_rate_limited": new_status = "rate_limited" else: new_status = "failed" await conn.execute(""" UPDATE group_keys SET status = $1, status_msg = $2, updated_at = NOW() WHERE id = $3 """, new_status, error_type, key_id) return TestResult( success=is_healthy, status=new_status, error=error_type if not is_healthy else None, latency_ms=latency_ms, ) @router.post("/groups/{group_id}/test-all") async def test_all_keys(group_id: int): """批量检测分组下所有key""" pool = get_pool() async with pool.acquire() as conn: # 获取分组下所有key rows = await conn.fetch(""" SELECT gk.id, gk.api_key_id, gk.status, gk.is_active, ak.key_value FROM group_keys gk JOIN api_keys ak ON ak.id = gk.api_key_id WHERE gk.group_id = $1 AND gk.is_active = TRUE """, group_id) if not rows: return TestAllResult(total=0, active=0, failed=0, banned=0, health_rate=0) total = len(rows) active = 0 failed = 0 banned = 0 provider = "dashscope" base_url = "https://dashscope.aliyuncs.com" for row in rows: try: start_time = time.time() is_healthy, error_type = await _test_key(row['key_value'], base_url, provider) end_time = time.time() latency_ms = int((end_time - start_time) * 1000) if is_healthy: active += 1 await conn.execute(""" UPDATE group_keys SET status = 'active', status_msg = NULL, updated_at = NOW() WHERE id = $1 """, row['id']) else: failed += 1 status = "rate_limited" if error_type == "429_rate_limited" else "failed" await conn.execute(""" UPDATE group_keys SET status = $1, status_msg = $2, updated_at = NOW() WHERE id = $3 """, status, error_type, row['id']) except Exception as e: failed += 1 await conn.execute(""" UPDATE group_keys SET status = 'failed', status_msg = $1, updated_at = NOW() WHERE id = $2 """, str(e)[:100], row['id']) health_rate = (active / total * 100) if total > 0 else 0 return TestAllResult( total=total, active=active, failed=failed, banned=banned, health_rate=round(health_rate, 2), )