#!/usr/bin/env python3
"""Test streaming with thinking enabled."""

import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from core.zai_client import ZAiClient


api_key = os.environ.get("ZAI_API_KEY")
if not api_key:
    print("ERROR: ZAI_API_KEY not set")
    sys.exit(1)

print("Testing streaming with thinking...\n")

with ZAiClient(api_key, timeout=30.0) as client:
    reasoning = ""
    content = ""

    print("Response: ", end="", flush=True)

    for r_chunk, c_chunk in client.chat_stream(
        model="glm-4.7-flash",
        messages=[{"role": "user", "content": "What is 2+2?"}],
        thinking={"type": "enabled"},
        max_tokens=100,
    ):
        if r_chunk:
            reasoning += r_chunk
        if c_chunk:
            content += c_chunk
            print(c_chunk, end="", flush=True)

    print()
    print(f"\nReasoning ({len(reasoning)} chars): {reasoning[:200]}...")
    print(f"Content ({len(content)} chars): {content}")