"""Minimal Moonshot sample that also prints the usage payload."""

from __future__ import annotations

import json
from typing import Dict, List, Tuple

from openai import OpenAI


client = OpenAI(
    api_key="sk-xW0xjfQM6Mp9ZCWMLlnHiRJcpEOIZPTkXcN0dQ15xpZSuw2y",
    base_url="https://api.moonshot.cn/v1",
)

history: List[Dict[str, str]] = [
    {
        "role": "system",
        "content": (
            "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。"
            "你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，"
            "种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。"
        ),
    }
]


def chat(query: str, history_list: List[Dict[str, str]]) -> Tuple[str, Dict[str, int]]:
    history_list.append({"role": "user", "content": query})
    completion = client.chat.completions.create(
        model="kimi-k2-0905-preview",
        messages=history_list,
        temperature=0.6,
    )
    result = completion.choices[0].message.content
    history_list.append({"role": "assistant", "content": result})
    usage_dict: Dict[str, int] = (
        completion.usage.model_dump() if completion.usage else {}
    )
    return result, usage_dict


def main() -> None:
    earth_answer, earth_usage = chat("地球的自转周期是多少？", history)
    moon_answer, moon_usage = chat("月球呢？", history)

    print(earth_answer)
    print(json.dumps(earth_usage, ensure_ascii=False, indent=2))
    print(moon_answer)
    print(json.dumps(moon_usage, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
