import os
import sys
import json
from openai import OpenAI

sys.stdout.reconfigure(encoding="utf-8")

client = OpenAI(
base_url="https://token-plan-cn.xiaomimimo.com/v1", // API 地址
    api_key=os.environ.get("OPENAI_API_KEY", "你的key")
)
model = "mimo-v2.5-pro"



def dump_request(label, **kwargs):
    print(f"\n{'=' * 50}")
    print(f"[请求] {label}")
    print(f"{'=' * 50}")
    print(json.dumps(kwargs, ensure_ascii=False, indent=2, default=str))


def dump_response(label, response):
    print(f"\n{'=' * 50}")
    print(f"[响应] {label}")
    print(f"{'=' * 50}")
    print(f"Model: {response.model}")
    print(f"Finish reason: {response.choices[0].finish_reason}")
    message = response.choices[0].message
    print(f"Reasoning: {message.reasoning_content}")
    if message.content:
        print(f"Content: {message.content}")


# ========================================
# 测试 1：单轮纯文本对话
# ========================================
print("\n" + "#" * 60)
print("# 测试 1：单轮纯文本对话")
print("#" * 60)

messages1 = [
    {"role": "user", "content": "用一句话介绍量子计算"}
]

request1 = dict(model=model, messages=messages1)
dump_request("单轮对话", **request1)

response1 = client.chat.completions.create(**request1)
dump_response("单轮对话", response1)


# ========================================
# 测试 2：多轮对话（带角色设定）
# ========================================
print("\n" + "#" * 60)
print("# 测试 2：多轮对话（带角色设定）")
print("#" * 60)

messages2 = [
    {"role": "system", "content": "你是一个厨师"},
]

# 第一轮
messages2.append({"role": "user", "content": "红烧肉怎么做？"})

request2 = dict(model=model, messages=messages2)
dump_request("第一轮", **request2)

response2 = client.chat.completions.create(**request2)
dump_response("第一轮", response2)

# 第二轮：验证角色是否保持
messages2.append(response2.choices[0].message.model_dump(exclude_none=True))
messages2.append({"role": "user", "content": "你是谁？"})

request2b = dict(model=model, messages=messages2)
dump_request("第二轮：验证角色", **request2b)

response2b = client.chat.completions.create(**request2b)
dump_response("第二轮：验证角色", response2b)


# ========================================
# 测试 3：系统提示词对比
# ========================================
print("\n" + "#" * 60)
print("# 测试 3：系统提示词对比")
print("#" * 60)

question = "介绍一下你自己"

# 无 system
messages3a = [
    {"role": "user", "content": question}
]
request3a = dict(model=model, messages=messages3a)
dump_request("无 system", **request3a)
response3a = client.chat.completions.create(**request3a)
dump_response("无 system", response3a)

# 有 system
messages3b = [
    {"role": "system", "content": "你是一个海盗船长，说话要带海盗口吻。"},
    {"role": "user", "content": question}
]
request3b = dict(model=model, messages=messages3b)
dump_request("有 system（海盗船长）", **request3b)
response3b = client.chat.completions.create(**request3b)
dump_response("有 system（海盗船长）", response3b)

import os
import sys
import json
from openai import OpenAI


client = OpenAI(
base_url="https://token-plan-cn.xiaomimimo.com/v1",
    api_key=os.environ.get("OPENAI_API_KEY", "你的key")
)
model="mimo-v2.5-pro"

sys.stdout.reconfigure(encoding="utf-8")


def dump_request(label, **kwargs):
    print(f"\n{'=' * 50}")
    print(f"[请求] {label}")
    print(f"{'=' * 50}")
    print(json.dumps(kwargs, ensure_ascii=False, indent=2, default=str))


def dump_response(label, response):
    print(f"\n{'=' * 50}")
    print(f"[响应] {label}")
    print(f"{'=' * 50}")
    message = response.choices[0].message
    print(f"Finish reason: {response.choices[0].finish_reason}")
    if hasattr(message, 'reasoning_content') and message.reasoning_content:
        print(f"Reasoning: {message.reasoning_content}")
    if message.content:
        print(f"Content: {message.content}")
    if message.tool_calls:
        print(f"Tool calls: {json.dumps([tc.model_dump() for tc in message.tool_calls], ensure_ascii=False, indent=2)}")
    # 打印 token 消耗
    if response.usage:
        u = response.usage
        print(f"Usage: prompt={u.prompt_tokens}, completion={u.completion_tokens}, total={u.total_tokens}")
        if hasattr(u, 'prompt_tokens_details') and u.prompt_tokens_details:
            d = u.prompt_tokens_details
            print(f"  cached_tokens={getattr(d, 'cached_tokens', 'N/A')}")




# --- 模拟工具函数 ---
def get_weather(city: str) -> dict:
    return {"city": city, "temperature": "22°C", "condition": "晴", "humidity": "45%"}


def get_fx_rate(**kwargs) -> dict:
    return {"from": kwargs["from"], "to": kwargs["to"], "rate": 7.24}


def get_air_quality(city: str) -> dict:
    return {"city": city, "aqi": 35, "level": "优", "pm25": 12}


TOOL_MAP = {
    "get_weather": lambda **kw: get_weather(kw["city"]),
    "get_fx_rate": get_fx_rate,
    "get_air_quality": lambda **kw: get_air_quality(kw["city"]),
}

# --- 工具定义 ---
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "查询城市当前天气",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string"}
                },
                "required": ["city"],
                "additionalProperties": False,
            },
            "strict": True,
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_fx_rate",
            "description": "查询两种货币的汇率",
            "parameters": {
                "type": "object",
                "properties": {
                    "from": {"type": "string"},
                    "to": {"type": "string"},
                },
                "required": ["from", "to"],
                "additionalProperties": False,
            },
            "strict": True,
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_air_quality",
            "description": "查询城市空气质量",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string"},
                },
                "required": ["city"],
                "additionalProperties": False,
            },
            "strict": True,
        },
    },
]

# --- 第一轮：发送用户消息 ---
messages = [
    {"role": "user", "content": "我想去北京旅游，先帮我查一下天气，如果天气不错再查一下空气质量。"}
]

request1 = dict(model=model, tools=tools, messages=messages, tool_choice="auto")
dump_request("第一轮：发送用户消息", **request1)

response = client.chat.completions.create(**request1)
dump_response("第一轮：模型请求调用工具", response)

message = response.choices[0].message

# --- 第二轮：将工具调用结果返回给模型 ---
if message.tool_calls:
    # 把 assistant 的 tool_calls 消息加入对话历史（不回传 reasoning_content）
    assistant_msg = message.model_dump(exclude_none=True)
    assistant_msg.pop("reasoning_content", None)
    messages.append(assistant_msg)

    for tc in message.tool_calls:
        func_name = tc.function.name
        func_args = json.loads(tc.function.arguments)

        # 执行对应的模拟函数
        result = TOOL_MAP[func_name](**func_args)

        # 把工具执行结果加入对话历史
        messages.append({
            "role": "tool",
            "tool_call_id": tc.id,
            "content": json.dumps(result, ensure_ascii=False),
        })

    request2 = dict(model=model, tools=tools, messages=messages, tool_choice="auto")
    dump_request("第二轮：将工具结果返回给模型", **request2)

    response2 = client.chat.completions.create(**request2)
    dump_response("第二轮：模型基于工具结果回答", response2)

    message2 = response2.choices[0].message

    # --- 第三轮：如果模型再次请求调用工具 ---
    if message2.tool_calls:
        assistant_msg2 = message2.model_dump(exclude_none=True)
        assistant_msg2.pop("reasoning_content", None)  # 不回传 reasoning_content
        messages.append(assistant_msg2)

        for tc in message2.tool_calls:
            func_name = tc.function.name
            func_args = json.loads(tc.function.arguments)
            result = TOOL_MAP[func_name](**func_args)

            messages.append({
                "role": "tool",
                "tool_call_id": tc.id,
                "content": json.dumps(result, ensure_ascii=False),
            })

        request3 = dict(model=model, tools=tools, messages=messages, tool_choice="auto")
        dump_request("第三轮：将工具结果返回给模型", **request3)

        response3 = client.chat.completions.create(**request3)
        dump_response("第三轮：模型最终回答", response3)