check.py

#!/usr/bin/env python3
import requests
import time
import json
import concurrent.futures
import subprocess
import platform
import os

API_KEY = os.getenv("ANTHROPIC_AUTH_TOKEN")

print(f"🔑 使用的API Key: {API_KEY if API_KEY else ''}")
if not API_KEY:
    print("⚠️ 请设置环境变量 ANTHROPIC_AUTH_TOKEN")
    exit(1)

def test_ping_latency(host, count=15):
    """测试ping延迟"""
    try:
        # 检测操作系统
        if platform.system().lower() == "darwin":  # macOS
            cmd = ["ping", "-c", str(count), "-t", "3", host]
        else:  # Linux
            cmd = ["ping", "-c", str(count), "-W", "3", host]

        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)

        if result.returncode != 0:
            return {
                "host": host,
                "avg_latency": float("inf"),
                "packet_loss": 100.0,
                "success": False,
                "error": "连接失败",
            }

        # 解析ping结果
        lines = result.stdout.split("\n")
        latencies = []
        transmitted = 0
        received = 0

        for line in lines:
            if "time=" in line:
                try:
                    time_part = line.split("time=")[1].split()[0]
                    latency = float(time_part.replace("ms", ""))
                    latencies.append(latency)
                    received += 1
                except:
                    pass
            elif "packets transmitted" in line:
                try:
                    parts = line.split()
                    transmitted = int(parts[0])
                    received = int(parts[3])
                except:
                    transmitted = count

        if not latencies:
            return {
                "host": host,
                "avg_latency": float("inf"),
                "packet_loss": 100.0,
                "success": False,
                "error": "无有效ping响应",
            }

        avg_latency = sum(latencies) / len(latencies)
        packet_loss = (
            ((transmitted - received) / transmitted) * 100 if transmitted > 0 else 100
        )

        return {
            "host": host,
            "avg_latency": avg_latency,
            "packet_loss": packet_loss,
            "success": True,
        }

    except subprocess.TimeoutExpired:
        return {
            "host": host,
            "avg_latency": float("inf"),
            "packet_loss": 100.0,
            "success": False,
            "error": "ping超时",
        }
    except Exception as e:
        return {
            "host": host,
            "avg_latency": float("inf"),
            "packet_loss": 100.0,
            "success": False,
            "error": str(e),
        }


def test_claude_api_stream_first_token(base_url, timeout=30):
    """测试Claude API流模式的首字响应时间"""
    url = f"{base_url}/v1/messages"

    headers = {
        "Content-Type": "application/json",
        "x-api-key": API_KEY,
        "anthropic-version": "2023-06-01",
    }

    payload = {
        "model": "claude-sonnet-4-20250514",  # 使用haiku模型更快
        # "model": "claude-3-5-haiku-20241022",  # 使用haiku模型更快
        "max_tokens": 1,  # 只需要1个token测试首字时间
        "messages": [{"role": "user", "content": "1"}],  # 最短的输入
        "stream": True,
    }

    try:
        start_time = time.time()
        response = requests.post(
            url, headers=headers, json=payload, timeout=timeout, stream=True
        )

        # 检查响应状态
        if response.status_code != 200:
            return {
                "url": base_url,
                "first_token_time": float("inf"),
                "status_code": response.status_code,
                "success": False,
                "error": f"HTTP {response.status_code}",
            }

        # 读取流响应，寻找第一个内容token
        first_token_time = None
        for line in response.iter_lines():
            if line:
                line = line.decode("utf-8")
                if line.startswith("data: "):
                    data_str = line[6:]  # 移除 "data: " 前缀
                    if data_str.strip() == "[DONE]":
                        break

                    try:
                        data = json.loads(data_str)
                        # 检查是否是内容块的delta
                        if data.get("type") == "content_block_delta":
                            delta = data.get("delta", {})
                            if delta.get("type") == "text_delta" and delta.get("text"):
                                if first_token_time is None:
                                    first_token_time = time.time()
                                    break
                    except json.JSONDecodeError:
                        continue

        if first_token_time is not None:
            response_time = (first_token_time - start_time) * 1000
            return {
                "url": base_url,
                "first_token_time": response_time,
                "status_code": response.status_code,
                "success": True,
            }
        else:
            return {
                "url": base_url,
                "first_token_time": float("inf"),
                "status_code": response.status_code,
                "success": False,
                "error": "未收到有效的流响应",
            }

    except requests.exceptions.Timeout:
        return {
            "url": base_url,
            "first_token_time": float("inf"),
            "status_code": None,
            "success": False,
            "error": "请求超时",
        }
    except requests.exceptions.ConnectionError:
        return {
            "url": base_url,
            "first_token_time": float("inf"),
            "status_code": None,
            "success": False,
            "error": "连接失败",
        }
    except Exception as e:
        return {
            "url": base_url,
            "first_token_time": float("inf"),
            "status_code": None,
            "success": False,
            "error": str(e),
        }


def test_comprehensive(base_url, count=1):
    """综合测试ping和API响应时间"""
    # 从URL提取主机名
    host = base_url.replace("https://", "").replace("http://", "")

    print(f"📡 正在测试 {host}...")

    # 并发执行ping测试和API测试
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        ping_future = executor.submit(test_ping_latency, host)
        api_future = executor.submit(test_multiple_times, base_url, count)

        ping_result = ping_future.result()
        api_result = api_future.result()

    # 计算综合得分 (ping延迟 + API响应时间 + 丢包率权重)
    if ping_result["success"] and api_result["success_count"] > 0:
        comprehensive_score = (
            ping_result["avg_latency"]
            + api_result["avg_time"]
            + (ping_result["packet_loss"] * 10)
        )
    else:
        comprehensive_score = float("inf")

    return {
        "url": base_url,
        "host": host,
        "ping_latency": (
            ping_result["avg_latency"] if ping_result["success"] else float("inf")
        ),
        "packet_loss": ping_result["packet_loss"] if ping_result["success"] else 100.0,
        "api_response_time": (
            api_result["avg_time"] if api_result["success_count"] > 0 else float("inf")
        ),
        "comprehensive_score": comprehensive_score,
        "ping_success": ping_result["success"],
        "api_success": api_result["success_count"] > 0,
        "ping_error": ping_result.get("error"),
        "api_error": api_result.get("last_error"),
    }


def test_multiple_times(base_url, count=1):
    """单次测试避免重复消耗"""
    result = test_claude_api_stream_first_token(base_url)

    if result["success"]:
        return {
            "url": base_url,
            "avg_time": result["first_token_time"],
            "min_time": result["first_token_time"],
            "max_time": result["first_token_time"],
            "success_count": 1,
            "total_tests": 1,
            "last_error": None,
        }
    else:
        return {
            "url": base_url,
            "avg_time": float("inf"),
            "min_time": float("inf"),
            "max_time": float("inf"),
            "success_count": 0,
            "total_tests": 1,
            "last_error": result.get("error", "未知错误"),
        }


def main():
    # 测试的线路列表
    endpoints = [
        "https://sg.instcopilot-api.com",
        "https://hk.instcopilot-api.com",
        "https://cn.instcopilot-api.com",
        "https://jp.instcopilot-api.com",
    ]

    region_names = {
        "sg.instcopilot-api.com": "新加坡",
        "hk.instcopilot-api.com": "中国香港",
        "cn.instcopilot-api.com": "中国大陆",
        "jp.instcopilot-api.com": "日本",
    }

    print("🌐 开始综合测试ping延迟和Claude API响应时间...")
    print("=" * 70)

    # 使用线程池并发执行测试
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(endpoints)) as executor:
        # 提交所有测试任务
        future_to_endpoint = {
            executor.submit(test_comprehensive, endpoint, 1): endpoint
            for endpoint in endpoints
        }

        results = []

        # 收集结果
        for future in concurrent.futures.as_completed(future_to_endpoint):
            endpoint = future_to_endpoint[future]
            try:
                result = future.result()
                results.append(result)

                host = result["host"]
                region_name = region_names.get(host, host)

                if result["ping_success"] and result["api_success"]:
                    print(f"✅ {region_name} ({host})")
                    print(
                        f"   Ping延迟: {result['ping_latency']:.2f}ms, 丢包率: {result['packet_loss']:.1f}%"
                    )
                    print(f"   API响应: {result['api_response_time']:.2f}ms")
                    print(f"   综合得分: {result['comprehensive_score']:.2f}")
                else:
                    print(f"❌ {region_name} ({host}) - 测试失败")
                    if not result["ping_success"]:
                        print(f"   Ping错误: {result['ping_error']}")
                    if not result["api_success"]:
                        print(f"   API错误: {result['api_error']}")
                print()

            except Exception as exc:
                print(f"📡 {endpoint} 测试异常: {exc}")
                print()

    # 排序结果 - 按综合得分排序
    successful_results = [r for r in results if r["ping_success"] and r["api_success"]]
    successful_results.sort(key=lambda x: x["comprehensive_score"])

    print("🏆 综合性能排名:")
    print("=" * 70)

    if successful_results:
        print(
            f"{'排名':<4} {'区域':<8} {'Ping延迟':<12} {'丢包率':<10} {'API响应':<12} {'综合得分':<12}"
        )
        print("-" * 70)

        for i, result in enumerate(successful_results, 1):
            host = result["host"]
            region_name = region_names.get(host, host)
            print(
                f"{i:<4} {region_name:<8} {result['ping_latency']:.2f}ms{'':<6} {result['packet_loss']:.1f}%{'':<6} {result['api_response_time']:.2f}ms{'':<6} {result['comprehensive_score']:.2f}"
            )

        print()
        best = successful_results[0]
        best_host = best["host"]
        best_region = region_names.get(best_host, best_host)

        print(f"🚀 推荐线路: {best_region} ({best['url']})")
        print("📊 详细指标:")
        print(f"   • Ping延迟: {best['ping_latency']:.2f}ms")
        print(f"   • 丢包率: {best['packet_loss']:.1f}%")
        print(f"   • API首字响应: {best['api_response_time']:.2f}ms")
        print(f"   • 综合得分: {best['comprehensive_score']:.2f}")

        print("\n🔧 设置命令:")
        print(f"   export ANTHROPIC_BASE_URL=\"{best['url']}\"")

    else:
        print("❌ 所有线路都无法正常工作")
        print("\n常见问题:")
        print("1. 检查API key是否有效")
        print("2. 检查网络连接")
        print("3. 检查防火墙设置")


if __name__ == "__main__":
    main()