Crawl/backend/app/routers/group_keys_test.py

"""
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),
        )