#!/usr/bin/env python3
"""
简单的 Amazon Q API 测试脚本
"""

import requests
import json

def test_api():
    """测试 API 基本功能"""
    url = "http://localhost:8000/v1/chat/completions"
    
    payload = {
        "model": "claude-sonnet-4",
        "messages": [
            {"role": "user", "content": "你好，请简单介绍一下你自己"}
        ],
        "stream": False
    }
    
    try:
        response = requests.post(url, json=payload, timeout=30)
        
        if response.status_code == 200:
            result = response.json()
            content = result.get("choices", [{}])[0].get("message", {}).get("content", "")
            print("✅ API 测试成功!")
            print(f"响应内容: {content}")
            return True
        else:
            print(f"❌ API 测试失败: {response.status_code}")
            print(f"错误信息: {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ 请求异常: {e}")
        return False

def test_streaming():
    """测试流式响应"""
    url = "http://localhost:8000/v1/chat/completions"
    
    payload = {
        "model": "claude-sonnet-4",
        "messages": [
            {"role": "user", "content": "请用一句话介绍人工智能"}
        ],
        "stream": True
    }
    
    try:
        response = requests.post(url, json=payload, stream=True, timeout=30)
        
        if response.status_code == 200:
            print("✅ 流式测试成功!")
            print("流式响应:")
            for line in response.iter_lines():
                if line:
                    line_str = line.decode('utf-8')
                    if line_str.startswith('data: ') and line_str != 'data: [DONE]':
                        data = json.loads(line_str[6:])
                        content = data.get("choices", [{}])[0].get("delta", {}).get("content", "")
                        if content:
                            print(content, end='', flush=True)
            print()  # 换行
            return True
        else:
            print(f"❌ 流式测试失败: {response.status_code}")
            return False
            
    except Exception as e:
        print(f"❌ 流式请求异常: {e}")
        return False

if __name__ == "__main__":
    print("🧪 Amazon Q API 快速测试")
    print("=" * 40)
    
    # 测试基本功能
    print("1. 测试基本对话功能...")
    test_api()
    
    print("\n2. 测试流式响应...")
    test_streaming()
    
    print("\n✨ 测试完成!")