# 先安装: pip install openai
# 导包
from dotenv import load_dotenv
load_dotenv()
import os
from openai import OpenAI

api_key = os.getenv("")
print("api_key-->", api_key)
base_url = os.getenv("BASE_URL")
print("base_url-->", base_url)


SYSTEM_PROMPT = """
你是一个专业的中文新闻分类器，你的任务是将输入的中文新闻标题分类到以下10个类别之一：

0 - finance (金融财经)
1 - realty (房地产)
2 - stocks (股市)
3 - education (教育)
4 - science (科技科学)
5 - society (社会民生)
6 - politics (政治)
7 - sports (体育)
8 - 悦儿 (yueer)
9 - game (游戏)

分类规则：
1. 仔细分析新闻标题的内容和主题
2. 只从上述10个类别中选择最匹配的一个
3. 返回格式只需返回对应的英文类别名称
4. 如果确实无法确定类别，选择最接近的类别

示例输入输出：
输入："体验2D巅峰 倚天屠龙记十大创新概览"
输出："悦儿"

输入："同步A股首秀：港股缩量回调"
输出："stocks"

输入："金科西府 名墅天成"
输出："realty"

现在请对以下新闻标题进行分类：
"""


def get_deepseek_res(prompt):
    # 创建客户端
    client = OpenAI(api_key="sk-0797e0871d044ca5a9baf9d30216cd1a", base_url=os.getenv("BASE_URL"))
    # 创建会话
    response = client.chat.completions.create(
        # 模型deepseek-chat是deepseek-v3
        # 模型deepseek-reasoner是deepseek-r1
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ],
        stream=False
    )
    # 返回结果
    # return response.choices[0].message.content
    return response


if __name__ == '__main__':
    prompt = "倚天屠龙记"
    res = get_deepseek_res(prompt)
    print(res.choices[0].message.content)