from zhipuai import ZhipuAI

class ZhipuService:
    def __init__(self):
        self.model_name = "glm-4"
        self.app_secret_key = "d54f764a1d67c17d857bd3983b772016.GRjowY0fyiMNurLc"

    def talk_to_zhipu(self, message):
        client = ZhipuAI(api_key=self.app_secret_key) # 请填写您自己的APIKey
        response = client.chat.completions.create(
            model=self.model_name,  # 填写需要调用的模型名称
            messages=[
                {"role": "user", "content": message},
            ],
            stream=True, # 流式输出
            temperature= 0.01,  #随机度，越大越发散，0.01则是一个比较确定和稳定的输出
            top_p= 0.1, #选择前 10% 概率的 tokens 作为候选，也是影响随机程度
        )
        accum_resp = ""
        for chunk in response:
            print(chunk.choices[0].delta.content)
            accum_resp = accum_resp + chunk.choices[0].delta.content
        print(accum_resp)
        return accum_resp
    
    def talk_to_zhipu_sse(self, message):
        client = ZhipuAI(api_key=self.app_secret_key) # 请填写您自己的APIKey
        response = client.chat.completions.create(
            model=self.model_name,  # 填写需要调用的模型名称
            messages=[
                {"role": "user", "content": message},
            ],
            stream=True, # 流式输出
            temperature= 0.01,  #随机度，越大越发散，0.01则是一个比较确定和稳定的输出
            top_p= 0.1, #选择前 10% 概率的 tokens 作为候选，也是影响随机程度
        )

        for chunk in response:
            print(chunk.choices[0].delta.content)
            yield chunk.choices[0].delta.content