"""
LLM 流式调用测试
"""

from src.Module.llm.llm import StreamingLLM, chat


def test_streaming_chat():
    """测试流式对话"""
    print("=" * 60)
    print("     LLM 流式对话测试")
    print("=" * 60)

    llm = StreamingLLM(
        model='deepseek-v3.2',
        enable_thinking=False,  # 不开启思考模式
    )

    message = "用一句话介绍你自己"
    print(f"\n用户: {message}\n")
    print("助手: ", end="", flush=True)

    full_response = ""
    for chunk in llm.chat(message):
        if chunk.error:
            print(f"\n[错误] {chunk.error}")
            return False
        if chunk.content:
            print(chunk.content, end="", flush=True)
            full_response += chunk.content

    print("\n")
    print(f"完整响应长度: {len(full_response)} 字符")

    return len(full_response) > 0


def test_thinking_mode():
    """测试思考模式"""
    print("\n" + "=" * 60)
    print("思考模式测试")
    print("=" * 60)

    llm = StreamingLLM(
        model='deepseek-v3.2',
        enable_thinking=True,  # 开启思考模式
    )

    message = "1+1等于几？简单回答"
    print(f"\n用户: {message}\n")

    reasoning = ""
    content = ""
    is_answering = False

    for chunk in llm.chat(message):
        if chunk.error:
            print(f"\n[错误] {chunk.error}")
            return False

        if chunk.reasoning_content:
            if not is_answering:
                if not reasoning:
                    print("思考过程: ", end="", flush=True)
                print(chunk.reasoning_content, end="", flush=True)
            reasoning += chunk.reasoning_content

        if chunk.content:
            if not is_answering:
                print(f"\n\n回复: ", end="", flush=True)
                is_answering = True
            print(chunk.content, end="", flush=True)
            content += chunk.content

    print("\n")
    print(f"思考过程: {len(reasoning)} 字符")
    print(f"回复内容: {len(content)} 字符")

    return len(content) > 0


if __name__ == '__main__':
    results = []

    success1 = test_streaming_chat()
    results.append(("流式对话", success1))

    success2 = test_thinking_mode()
    results.append(("思考模式", success2))

    print("\n" + "=" * 60)
    print("测试结果:")
    for name, success in results:
        status = "✓ 通过" if success else "✗ 失败"
        print(f"  {name}: {status}")