security-platform/test_llm_timeout.py

#!/usr/bin/env python3
"""
测试不同LLM提供商的超时设置
"""
import requests
import time
import sys

def test_llm_generation(model_name):
    """测试指定模型的生成速度和超时"""
    print(f"\n🧪 测试模型: {model_name}")
    print("=" * 50)
    
    # 切换模型
    print(f"1. 切换到 {model_name}...")
    try:
        import subprocess
        result = subprocess.run([
            sys.executable, "switch_model.py", model_name
        ], capture_output=True, text=True)
        
        if result.returncode == 0:
            print("✅ 模型切换成功")
        else:
            print(f"❌ 模型切换失败: {result.stderr}")
            return
    except Exception as e:
        print(f"❌ 切换失败: {e}")
        return
    
    # 等待服务重启
    print("2. 等待服务重启...")
    time.sleep(5)
    
    # 检查LLM状态
    print("3. 检查LLM状态...")
    try:
        response = requests.get("http://localhost:8000/api/debug/llm-status", timeout=10)
        if response.status_code == 200:
            data = response.json()
            current_model = data.get("current_model", "unknown")
            llm_available = data.get("llm_available", False)
            print(f"   当前模型: {current_model}")
            print(f"   LLM可用: {llm_available}")
        else:
            print("❌ 无法获取LLM状态")
            return
    except requests.RequestException as e:
        print(f"❌ 连接失败: {e}")
        return
    
    # 测试生成速度
    print("4. 测试报告生成...")
    start_time = time.time()
    try:
        response = requests.post(
            "http://localhost:8000/api/reports/generate?use_llm=true", 
            timeout=120  # 2分钟测试超时
        )
        end_time = time.time()
        
        if response.status_code == 200:
            data = response.json()
            generated = data.get("generated", 0)
            duration = end_time - start_time
            print(f"✅ 生成成功!")
            print(f"   生成报告数: {generated}")
            print(f"   耗时: {duration:.1f} 秒")
            
            if generated > 0 and "report_ids" in data:
                # 获取报告内容样本
                report_id = data["report_ids"][0]
                report_response = requests.get(f"http://localhost:8000/api/reports/{report_id}")
                if report_response.status_code == 200:
                    report_data = report_response.json()
                    summary_length = len(report_data.get("summary", ""))
                    print(f"   报告长度: {summary_length} 字符")
        else:
            print(f"❌ 生成失败: {response.status_code}")
            print(f"   错误信息: {response.text}")
            
    except requests.Timeout:
        end_time = time.time()
        duration = end_time - start_time
        print(f"⏰ 请求超时 ({duration:.1f}秒)")
    except requests.RequestException as e:
        print(f"❌ 请求失败: {e}")

def main():
    print("🚀 LLM超时测试工具")
    print("测试不同模型的生成速度和超时配置")
    
    models_to_test = [
        ("local", "qwen3:latest - 8B模型"),
        ("local_powerful", "qwen3:32b - 32B模型"),
    ]
    
    # 可选添加大模型测试 (需要更多资源)
    add_large_model = input("\n是否测试70B大模型? (y/N): ").lower().strip()
    if add_large_model == 'y':
        models_to_test.append(("local_deepseek", "deepseek-r1:70b - 70B模型"))
    
    for model_key, description in models_to_test:
        test_llm_generation(model_key)
        
        # 询问是否继续
        if model_key != models_to_test[-1][0]:  # 不是最后一个
            continue_test = input(f"\n继续测试下一个模型? (Y/n): ").lower().strip()
            if continue_test == 'n':
                break
    
    print(f"\n🏁 测试完成!")
    print("\n建议的超时设置:")
    print("- 8B模型 (qwen3:latest): 120秒")
    print("- 32B模型 (qwen3:32b): 300秒") 
    print("- 70B模型 (deepseek-r1:70b): 600秒")

if __name__ == "__main__":
    main()