# Please install OpenAI SDK first: `pip3 install openai`
import os
import json
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()


class DeepSeek:
    system_prompt = (
        "You are a professional QA answer assistant.\n"
        "The user will provide a JSON object containing a question.\n"
        "Your job:\n"
        "1. Read the 'title' and the 'chooies'.\n"
        "2. Determine the correct answer(s).\n"
        "3. Output ONLY the 'num' of the correct choices.\n"
        "4. If multiple answers exist, join them with English commas, e.g., 'A,C,D'.\n"
        "5. If you cannot determine the answer confidently, return an empty string.\n"
        "6. Do not output explanations or any extra text."
    )

    def __init__(self) -> None:
        self.key = os.environ.get("DEEPSEEK_API_KEY")
        if self.key:
            raise Exception("找不到 DEEPSEEK_API_KEY 请设置环境变量在运行代码")
        self.client = OpenAI(
            api_key=self.key,
            base_url="https://api.deepseek.com",
        )

    def chat(self, question_json: dict) -> str:
        response = self.client.chat.completions.create(
            model="deepseek-chat",
            # model="deepseek-reasoner",
            messages=[
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": json.dumps(question_json)},
            ],
            stream=False,
        )

        # print(response.choices[0].message.content)
        return (
            response.choices[0].message.content
            if response.choices[0].message.content
            else ""
        )


if __name__ == "__main__":
    question = {
        "id": "25594265",
        "index": "1",
        "type": "单选",
        "title": "下列哪个属于AI视觉识别应用？",
        "chooies": [
            {"num": "A", "txt": "音乐合成"},
            {"num": "B", "txt": "植物识别"},
            {"num": "C", "txt": "文案创作"},
            {"num": "D", "txt": "机器翻译"},
        ],
    }
    DeepSeek().chat(question)