#!/usr/bin/env python3
"""
测试LLM API调用
"""

import requests
import json
import time

# API配置
BASE_URL = "https://oneapi.gongjulian.cn/v1"
API_KEY = "sk-lB2Fc9ssY5UuwmiV5dD441F997364d29Be547e008dF5Cf41"
MODEL = "deepseek-ai/deepseek-v3.2"

def test_api():
    """测试API连接"""
    url = f"{BASE_URL}/models"
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    print("=" * 60)
    print("测试1: 获取模型列表")
    print("=" * 60)

    try:
        response = requests.get(url, headers=headers, timeout=30)
        print(f"状态码: {response.status_code}")
        print(f"响应: {response.text[:500]}")
        return response.status_code == 200
    except requests.exceptions.Timeout:
        print("超时: 请求超过30秒")
        return False
    except requests.exceptions.ConnectionError as e:
        print(f"连接错误: {e}")
        return False
    except Exception as e:
        print(f"错误: {e}")
        return False

def test_chat_completion():
    """测试聊天完成接口"""
    url = f"{BASE_URL}/chat/completions"
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    payload = {
        "model": MODEL,
        "messages": [
            {
                "role": "user",
                "content": "请简单介绍一下你自己。"
            }
        ],
        "max_tokens": 100,
        "temperature": 0.7
    }

    print("\n" + "=" * 60)
    print("测试2: 聊天完成接口")
    print("=" * 60)
    print(f"URL: {url}")
    print(f"Model: {MODEL}")

    try:
        start_time = time.time()
        response = requests.post(url, headers=headers, json=payload, timeout=60)
        elapsed = time.time() - start_time

        print(f"状态码: {response.status_code}")
        print(f"耗时: {elapsed:.2f}秒")

        if response.status_code == 200:
            result = response.json()
            content = result.get("choices", [{}])[0].get("message", {}).get("content", "")
            print(f"响应内容:\n{content}")
            return True
        else:
            print(f"错误响应: {response.text[:500]}")
            return False

    except requests.exceptions.Timeout:
        print(f"超时: 请求超过60秒")
        return False
    except requests.exceptions.ConnectionError as e:
        print(f"连接错误: {e}")
        return False
    except Exception as e:
        print(f"错误: {e}")
        return False

def test_simple_prompt():
    """测试简单提示词"""
    url = f"{BASE_URL}/chat/completions"
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    payload = {
        "model": MODEL,
        "messages": [
            {
                "role": "user",
                "content": "1+1等于几？"
            }
        ],
        "max_tokens": 50,
        "temperature": 0.1
    }

    print("\n" + "=" * 60)
    print("测试3: 简单数学问题")
    print("=" * 60)

    try:
        start_time = time.time()
        response = requests.post(url, headers=headers, json=payload, timeout=30)
        elapsed = time.time() - start_time

        print(f"状态码: {response.status_code}")
        print(f"耗时: {elapsed:.2f}秒")

        if response.status_code == 200:
            result = response.json()
            content = result.get("choices", [{}])[0].get("message", {}).get("content", "")
            print(f"响应: {content}")
            return True
        else:
            print(f"错误: {response.text[:300]}")
            return False

    except requests.exceptions.Timeout:
        print("超时")
        return False
    except Exception as e:
        print(f"错误: {e}")
        return False

if __name__ == "__main__":
    print("LLM API 连接测试")
    print(f"API地址: {BASE_URL}")
    print(f"模型: {MODEL}")
    print()

    results = []
    results.append(("模型列表", test_api()))
    results.append(("聊天完成", test_chat_completion()))
    results.append(("简单问题", test_simple_prompt()))

    print("\n" + "=" * 60)
    print("测试结果汇总")
    print("=" * 60)
    for name, success in results:
        status = "通过" if success else "失败"
        print(f"{name}: {status}")

    all_passed = all(r[1] for r in results)
    print()
    if all_passed:
        print("所有测试通过！")
    else:
        print("部分测试失败，请检查API配置和网络连接")