From 8c564fcb6017c8597011af8f73a3b5c4bd915dfb Mon Sep 17 00:00:00 2001
From: yudong <1721464347@qq.com>
Date: Mon, 19 Feb 2024 22:23:33 +0800
Subject: [PATCH] =?UTF-8?q?=E9=80=82=E9=85=8D=E6=99=BA=E8=B0=B1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 base/func_zhipu.py   | 48 ++++++++++++++++++++++++++++++++++++++++++++
 config.yaml.template |  4 ++++
 configuration.py     |  1 +
 constants.py         |  1 +
 robot.py             |  5 +++++
 5 files changed, 59 insertions(+)
 create mode 100644 base/func_zhipu.py

diff --git a/base/func_zhipu.py b/base/func_zhipu.py
new file mode 100644
index 0000000..6efb5a0
--- /dev/null
+++ b/base/func_zhipu.py
@@ -0,0 +1,48 @@
+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)
\ No newline at end of file
diff --git a/config.yaml.template b/config.yaml.template
index 653b232..ddaf102 100644
--- a/config.yaml.template
+++ b/config.yaml.template
@@ -78,3 +78,7 @@ bard: # -----bard配置这行不填-----
   # 提示词尽可能用英文，bard对中文提示词的效果不是很理想，下方提示词为英语老师的示例，请按实际需要修改,默认设置的提示词为谷歌创造的AI大语言模型
   # I want you to act as a spoken English teacher and improver. I will speak to you in English and you will reply to me in English to practice my spoken English. I want you to keep your reply neat, limiting the reply to 100 words. I want you to strictly correct my grammar mistakes, typos, and factual errors. I want you to ask me a question in your reply. Now let's start practicing, you could ask me a question first. Remember, I want you to strictly correct my grammar mistakes, typos, and factual errors.
   prompt: You am a large language model, trained by Google.
+
+zhipu:  # -----zhipu配置这行不填-----
+  api_key:  #api key
+  model: glm-4  # 模型类型
\ No newline at end of file
diff --git a/configuration.py b/configuration.py
index 0aa5b45..611a526 100644
--- a/configuration.py
+++ b/configuration.py
@@ -36,3 +36,4 @@ class Config(object):
         self.XINGHUO_WEB = yconfig.get("xinghuo_web", {})
         self.CHATGLM = yconfig.get("chatglm", {})
         self.BardAssistant = yconfig.get("bard", {})
+        self.ZhiPu = yconfig.get("zhipu", {})
diff --git a/constants.py b/constants.py
index 87c6c50..a1f87f1 100644
--- a/constants.py
+++ b/constants.py
@@ -9,6 +9,7 @@ class ChatType(IntEnum):
     XINGHUO_WEB = 3  # 讯飞星火
     CHATGLM = 4  # ChatGLM
     BardAssistant = 5  # Google Bard
+    ZhiPu = 6  # ZhiPu
 
     @staticmethod
     def is_in_chat_types(chat_type: int) -> bool:
diff --git a/robot.py b/robot.py
index bc14f1b..f9c8712 100644
--- a/robot.py
+++ b/robot.py
@@ -6,6 +6,7 @@ import time
 import xml.etree.ElementTree as ET
 from queue import Empty
 from threading import Thread
+from base.func_zhipu import ZhiPu
 
 from wcferry import Wcf, WxMsg
 
@@ -45,6 +46,8 @@ class Robot(Job):
                 self.chat = ChatGLM(self.config.CHATGLM)
             elif chat_type == ChatType.BardAssistant.value and BardAssistant.value_check(self.config.BardAssistant):
                 self.chat = BardAssistant(self.config.BardAssistant)
+            elif chat_type == ChatType.ZhiPu.value and ZhiPu.value_check(self.config.ZHIPU):
+                self.chat = ZhiPu(self.config.ZHIPU)
             else:
                 self.LOG.warning("未配置模型")
                 self.chat = None
@@ -59,6 +62,8 @@ class Robot(Job):
                 self.chat = ChatGLM(self.config.CHATGLM)
             elif BardAssistant.value_check(self.config.BardAssistant):
                 self.chat = BardAssistant(self.config.BardAssistant)
+            elif ZhiPu.value_check(self.config.ZhiPu):
+                self.chat = ZhiPu(self.config.ZhiPu)
             else:
                 self.LOG.warning("未配置模型")
                 self.chat = None