48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
from pyexpat import model
|
|
from zhipuai import ZhipuAI
|
|
|
|
class ZhiPu():
|
|
def __init__(self, conf: dict) -> None:
|
|
api_key = conf.get("api_key")
|
|
model = conf.get("model", "glm-4") # 默认使用 glm-4 模型
|
|
self.api_key = api_key
|
|
self.model = model
|
|
self.client = ZhipuAI(api_key=api_key)
|
|
self.converstion_list = {}
|
|
|
|
@staticmethod
|
|
def value_check(conf: dict) -> bool:
|
|
if conf:
|
|
if conf.get("api_key") :
|
|
return True
|
|
return False
|
|
|
|
def __repr__(self):
|
|
return 'ZhiPu'
|
|
|
|
def get_answer(self, msg: str, wxid: str, **args) -> str:
|
|
self._update_message(wxid, str(msg), "user")
|
|
response = self.client.chat.completions.create(
|
|
model=self.model,
|
|
messages=self.converstion_list[wxid]
|
|
)
|
|
resp_msg = response.choices[0].message
|
|
answer = resp_msg.content
|
|
self._update_message(wxid, answer, "assistant")
|
|
return answer
|
|
|
|
def _update_message(self, wxid: str, msg: str, role: str) -> None:
|
|
if wxid not in self.converstion_list.keys():
|
|
self.converstion_list[wxid] = []
|
|
content = {"role": role, "content": str(msg)}
|
|
self.converstion_list[wxid].append(content)
|
|
|
|
if __name__ == "__main__":
|
|
from configuration import Config
|
|
config = Config().ZHIPU
|
|
if not config:
|
|
exit(0)
|
|
|
|
zhipu = ZhiPu(config)
|
|
rsp = zhipu.get_answer("你好")
|
|
print(rsp) |