#!/usr/bin/env python3
"""测试多模型并行调用"""
import asyncio
import httpx
import os

AI_SCORER_MODEL = os.getenv("AI_SCORER_MODEL", "qwen2.5:7b")
AI_SUMMARIZER_MODEL = os.getenv("AI_SUMMARIZER_MODEL", "qwen2.5:7b")
AI_SUMMARY_MODEL = os.getenv("AI_SUMMARY_MODEL", "qwen3.6:35b")

SAMPLE_NEWS = "美联储宣布加息25个基点，将联邦基金利率目标区间上调至5.25%-5.5%，为2007年以来的最高水平。"

async def test_model(model: str, prompt: str, task_name: str):
    """测试单个模型"""
    print(f"  [{task_name}] 模型={model} 开始...")
    start = asyncio.get_event_loop().time()
    
    try:
        async with httpx.AsyncClient(timeout=30) as client:
            resp = await client.post(
                "http://localhost:11434/api/generate",
                json={"model": model, "prompt": prompt, "stream": False}
            )
            elapsed = asyncio.get_event_loop().time() - start
            result = resp.json().get("response", "").strip()
            print(f"  [{task_name}] ✓ {elapsed:.1f}s: {result[:80]}...")
            return {"model": model, "task": task_name, "elapsed": elapsed, "result": result}
    except Exception as e:
        elapsed = asyncio.get_event_loop().time() - start
        print(f"  [{task_name}] ✗ {elapsed:.1f}s: {e}")
        return {"model": model, "task": task_name, "elapsed": elapsed, "error": str(e)}

async def main():
    print("=" * 60)
    print("多模型并行测试")
    print("=" * 60)
    print(f"评分模型: {AI_SCORER_MODEL}")
    print(f"摘要模型: {AI_SUMMARIZER_MODEL}")
    print(f"总结模型: {AI_SUMMARY_MODEL}")
    print()
    
    # 并行调用三个模型
    tasks = [
        test_model(AI_SCORER_MODEL, f"评分（0-5分）: {SAMPLE_NEWS}", "评分"),
        test_model(AI_SUMMARIZER_MODEL, f"一句话概括（≤30字）: {SAMPLE_NEWS}", "摘要"),
        test_model(AI_SUMMARY_MODEL, f"一句话概括（≤30字）: {SAMPLE_NEWS}", "总结"),
    ]
    
    print("开始并行调用...")
    overall_start = asyncio.get_event_loop().time()
    results = await asyncio.gather(*tasks)
    overall_elapsed = asyncio.get_event_loop().time() - overall_start
    
    print()
    print("=" * 60)
    print(f"总耗时: {overall_elapsed:.1f}s")
    print("=" * 60)
    for r in results:
        if "error" in r:
            print(f"  {r['task']}: 失败 ({r['elapsed']:.1f}s)")
        else:
            print(f"  {r['task']}: {r['elapsed']:.1f}s - {r['model']}")
    
    print()
    print("并行优势: 如果串行需要 ~" + f"{sum(r['elapsed'] for r in results):.1f}s，并行只需 {overall_elapsed:.1f}s")

if __name__ == "__main__":
    asyncio.run(main())