from zhipuai import ZhipuAI
from model.util import InputData, OutputData, getModelAPIKey

# client = ZhipuAI(api_key="") # 填写您自己的APIKey
# response = client.chat.completions.create(
#     model="glm-4-0520",  # 填写需要调用的模型编码
#     messages=[
#         {"role": "user", "content": "作为一名营销专家，请为我的产品创作一个吸引人的slogan"},
#         {"role": "assistant", "content": "当然，为了创作一个吸引人的slogan，请告诉我一些关于您产品的信息"},
#         {"role": "user", "content": "智谱AI开放平台"},
#         {"role": "assistant", "content": "智启未来，谱绘无限一智谱AI，让创新触手可及!"},
#         {"role": "user", "content": "创造一个更精准、吸引人的slogan"}
#     ],
# )
# print(response.choices[0].message)



def predict(input_data:InputData):
    client = ZhipuAI(api_key=getModelAPIKey("glm-4-flash"))
    response = client.chat.completions.create(
        model="glm-4-flash",  # 填写需要调用的模型编码
        messages=[
            {"role": "user", "content": input_data.message}],
    )
    if  response.choices[0].finish_reason == "stop":
        return OutputData(response.choices[0].message.content,200,response.usage.total_tokens)
    elif response.choices[0].finish_reason == "length":
        return OutputData(response.choices[0].message.content,201,response.usage.total_tokens)
    elif response.choices[0].finish_reason == "network_error":
        return OutputData("Server Network Error",500,0)
    else: return OutputData("Unknown Error",500,0)