import os
import json
from openai import OpenAI


def deepseek(user_query, need_extra=False):
    """
    调用 deepseek 接口生成回答。

    参数：
        user_query (str): 用户输入的查询内容。
        need_extra (bool): 是否额外返回 usage 信息（completion_tokens），默认为 False。

    返回：
        如果 need_extra 为 True，则返回一个元组 (full_response, completion_tokens)；
        否则仅返回 full_response。
    """
    client = OpenAI(
        # 如果没有配置环境变量，请直接将 api_key 替换为实际的 API Key，例如：api_key="sk-xxx"
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )

    # 调用接口，开启流式输出并包含 usage 信息
    completion = client.chat.completions.create(
        model="deepseek-r1-distill-qwen-32b",
        temperature=0.5,
        messages=[
            {'role': 'user', 'content': user_query}
        ],
        stream=True,
        stream_options={"include_usage": True}
    )

    full_response = ""
    completion_tokens = 0

    for chunk in completion:
        # 解析 chunk 数据：优先使用 to_dict 方法，否则采用 model_dump_json 进行转换
        if hasattr(chunk, 'to_dict'):
            chunk_data = chunk.to_dict()
        else:
            chunk_data = json.loads(chunk.model_dump_json())

        # 处理 usage 信息
        usage = chunk_data.get('usage')
        if usage is not None:
            completion_tokens = usage.get('completion_tokens', 0)

        # 处理 choices 信息
        choices = chunk_data.get('choices', [])
        if choices:
            choice = choices[0]
            delta = choice.get('delta', {})
            content = delta.get('content', '')
            if content:
                full_response += content
                # 如需实时输出可取消注释下面一行
                print(content, end='', flush=True)
            # 若需要处理完成原因，可在此处增加相应逻辑
            if choice.get('finish_reason'):
                pass

    # 根据 need_extra 决定返回值
    if need_extra:
        return full_response, completion_tokens
    else:
        return full_response


# 示例调用
if __name__ == "__main__":
    query = "1+1等于几？请用json格式回答，键名为'答案',键值为你的回答"
    result = deepseek(query, need_extra=True)
    print("\n完整内容为：", result)