import os
from dashscope import Generation
import dashscope
dashscope.base_http_api_url = "https://dashscope.aliyuncs.com/api/v1/"

messages = [{"role": "user", "content": "你是谁？"}]

completion = Generation.call(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    model="kimi-k2-thinking",
    messages=messages,
    result_format="message", 
    stream=True,
    incremental_output=True, 
)

is_answering = False

print("=" * 20 + "思考过程" + "=" * 20)

for chunk in completion:
    if (
        chunk.output.choices[0].message.content == ""
        and chunk.output.choices[0].message.reasoning_content == ""
    ):
        pass
    else:
        if (
            chunk.output.choices[0].message.reasoning_content != ""
            and chunk.output.choices[0].message.content == ""
        ):
            print(chunk.output.choices[0].message.reasoning_content, end="", flush=True)
        elif chunk.output.choices[0].message.content != "":
            if not is_answering:
                print("\n" + "=" * 20 + "完整回复" + "=" * 20)
                is_answering = True
            print(chunk.output.choices[0].message.content, end="", flush=True)