#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
测试DashScope流式API响应格式
"""

import os
import json
import requests
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

def test_dashscope_stream():
    """测试DashScope的流式响应格式"""

    api_key = os.getenv('ONEAPI_KEY')
    if not api_key:
        print("请设置ONEAPI_KEY环境变量")
        return

    # API配置
    url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json',
        'X-DashScope-SSE': 'enable',  # 启用SSE
        'Accept': 'text/event-stream'
    }

    # 测试请求
    payload = {
        "model": "qwen-turbo",
        "input": {
            "messages": [
                {"role": "user", "content": "你好，请简单介绍一下RAG系统"}
            ]
        },
        "parameters": {
            "max_tokens": 100,
            "temperature": 0.7
        }
    }

    print("发送请求到DashScope API...")
    print(f"URL: {url}")
    print(f"Model: {payload['model']}")
    print("=" * 80)

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

        if response.status_code != 200:
            print(f"错误: HTTP {response.status_code}")
            print(response.text[:500])
            return

        print("开始接收流式响应...")
        print("=" * 80)

        event_count = 0
        full_text = ""

        for line in response.iter_lines():
            if not line:
                continue

            line_str = line.decode('utf-8')

            if line_str.startswith('data:'):
                data_str = line_str[5:].strip()

                if data_str == '[DONE]':
                    print("\n流式响应结束")
                    break

                try:
                    data = json.loads(data_str)
                    event_count += 1

                    # 打印前3个完整响应
                    if event_count <= 3:
                        print(f"\n事件 #{event_count}:")
                        print(json.dumps(data, ensure_ascii=False, indent=2))
                        print("-" * 40)

                    # 提取文本
                    if 'output' in data:
                        output = data['output']

                        # 尝试不同的字段
                        text = None
                        if 'text' in output:
                            text = output['text']
                            print(f"[text字段] 当前累积长度: {len(text)} 字符")

                        if 'choices' in output:
                            choices = output['choices']
                            if choices and len(choices) > 0:
                                choice = choices[0]
                                if 'delta' in choice:
                                    delta = choice['delta'].get('content', '')
                                    print(f"[delta字段] 增量: {repr(delta)}")
                                elif 'message' in choice:
                                    message = choice['message'].get('content', '')
                                    print(f"[message字段] 消息长度: {len(message)} 字符")

                        # 检查finish_reason
                        if output.get('finish_reason'):
                            print(f"[结束原因] {output['finish_reason']}")

                except json.JSONDecodeError as e:
                    print(f"JSON解析错误: {e}")
                    print(f"原始数据: {data_str[:100]}")

        print("\n" + "=" * 80)
        print(f"总共接收了 {event_count} 个事件")

    except Exception as e:
        print(f"请求失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    test_dashscope_stream()