import zhipuai
from .base import register_llm


def ask_chatglm(message: str, api_key: str):
    zhipuai.api_key = api_key

    response = zhipuai.model_api.invoke(
        model="chatglm_turbo",
        prompt=[{
            "role": "user",
            "content": message
        }],
        top_p=0.7,
        temperature=0.9,
    )

    response_msg = response['data']['choices'][0]['content']
    # strip the front and end '"'
    if len(response_msg) >= 2:
        response_msg = response_msg[1:-1]

    return response_msg


register_llm('chatglm', ask_chatglm)