#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
简单测试DashScope流式输出
"""

import os
import sys
from dotenv import load_dotenv

# 添加项目路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# 加载环境变量
load_dotenv()

from retriver.langgraph.dashscope_llm import DashScopeLLM

def test_stream():
    """测试流式输出"""

    # 创建LLM实例
    llm = DashScopeLLM()

    # 测试prompt
    prompt = "请用30字简单介绍RAG系统"

    print("开始流式生成...")
    print("=" * 50)

    # 收集所有token
    tokens = []
    full_text = ""

    try:
        for token in llm.stream(prompt, max_tokens=100, temperature=0.7):
            # 打印每个token
            print(f"Token: {repr(token)}")
            tokens.append(token)
            full_text += token

    except Exception as e:
        print(f"错误: {e}")
        import traceback
        traceback.print_exc()

    print("=" * 50)
    print(f"总共 {len(tokens)} 个token")
    print(f"完整文本({len(full_text)}字): {full_text}")

if __name__ == "__main__":
    test_stream()