From 1c06d0758f32c2425adcd5ab51b6eb74584c1523 Mon Sep 17 00:00:00 2001
From: Changhua <lichanghua0821@gmail.com>
Date: Sat, 4 Mar 2023 23:21:19 +0800
Subject: [PATCH] Add ChatGPT

---
 config.yaml.template |  3 +++
 configuration.py     |  1 +
 func_chatgpt.py      | 40 ++++++++++++++++++++++++++++++++++++++++
 main.py              |  2 +-
 robot.py             | 40 ++++++++++++++++++++++++----------------
 5 files changed, 69 insertions(+), 17 deletions(-)
 create mode 100644 func_chatgpt.py

diff --git a/config.yaml.template b/config.yaml.template
index 1b2a946..e80393a 100644
--- a/config.yaml.template
+++ b/config.yaml.template
@@ -40,3 +40,6 @@ logging:
 
 groups:
   enable: [] # 允许响应的群 roomId，可以通过 WxGetContacts 获取
+
+chatgpt:
+  key: your key
diff --git a/configuration.py b/configuration.py
index 29dbcd3..ed24f11 100644
--- a/configuration.py
+++ b/configuration.py
@@ -27,3 +27,4 @@ class Config(object):
         yconfig = self._load_config()
         logging.config.dictConfig(yconfig["logging"])
         self.GROUPS = yconfig["groups"]["enable"]
+        self.CHAT_KEY = yconfig["chatgpt"]["key"]
diff --git a/func_chatgpt.py b/func_chatgpt.py
new file mode 100644
index 0000000..649f506
--- /dev/null
+++ b/func_chatgpt.py
@@ -0,0 +1,40 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import openai
+
+
+class ChatGPT():
+    def __init__(self, key) -> None:
+        openai.api_key = key
+
+    def get_answer(self, question: str) -> str:
+        try:
+            ret = openai.ChatCompletion.create(
+                model="gpt-3.5-turbo",
+                messages=[
+                    {"role": "user", "content": question},
+                ],
+                temperature=0.5,
+            )
+            rsp = ret["choices"][0]["message"]["content"]
+            rsp = rsp[2:] if rsp.startswith("\n\n") else rsp
+            rsp = rsp.replace("\n\n", "\n")
+        except Exception as e:
+            ret = ""
+
+        return rsp
+
+
+if __name__ == "__main__":
+    from datetime import datetime
+    chat = ChatGPT("your key")
+    while True:
+        q = input(">>> ")
+        try:
+            t1 = datetime.now()
+            print(chat.get_answer(q))
+            t2 = datetime.now()
+            print(f"{t2-t1}s")
+        except Exception as e:
+            print(e)
diff --git a/main.py b/main.py
index 5ad872b..a29cede 100644
--- a/main.py
+++ b/main.py
@@ -23,7 +23,7 @@ def weather_report(robot: Robot) -> None:
 
 
 def main():
-    wcf = Wcf()
+    wcf = Wcf(debug=True)
 
     def handler(sig, frame):
         wcf.cleanup()  # 退出前清理环境
diff --git a/robot.py b/robot.py
index a5243ef..75b1d06 100644
--- a/robot.py
+++ b/robot.py
@@ -9,6 +9,7 @@ from wcferry import Wcf
 
 from configuration import Config
 from func_chengyu import cy
+from func_chatgpt import ChatGPT
 from job_mgmt import Job
 
 
@@ -22,6 +23,7 @@ class Robot(Job):
         self.LOG = logging.getLogger("Robot")
         self.wxid = self.wcf.get_self_wxid()
         self.allContacts = self.getAllContacts()
+        self.chat = ChatGPT(self.config.CHAT_KEY)
 
     def toAt(self, msg: Wcf.WxMsg) -> bool:
         """
@@ -29,11 +31,7 @@ class Robot(Job):
         :param msg: 微信消息结构
         :return: 处理状态，`True` 成功，`False` 失败
         """
-        status = True
-        rsp = "你@我干嘛？"
-        self.sendTextMsg(rsp, msg.roomid, msg.sender)
-
-        return status
+        return self.toChitchat(msg)
 
     def toChengyu(self, msg: Wcf.WxMsg) -> bool:
         """
@@ -62,10 +60,21 @@ class Robot(Job):
 
         return status
 
-    def toChitchat(self, msg: Wcf.WxMsg) -> None:
-        """闲聊，目前未实现
+    def toChitchat(self, msg: Wcf.WxMsg) -> bool:
+        """闲聊，接入 ChatGPT
         """
-        pass
+        q = re.sub(r"@.*?[\u2005|\s]", "", msg.content).replace(" ", "")
+        rsp = self.chat.get_answer(q)
+        if rsp:
+            if msg.from_group():
+                self.sendTextMsg(rsp, msg.roomid, msg.sender)
+            else:
+                self.sendTextMsg(rsp, msg.sender)
+
+            return True
+        else:
+            self.LOG.error(f"无法从ChatGPT获得答案")
+            return False
 
     def processMsg(self, msg: Wcf.WxMsg) -> None:
         """当接收到消息的时候，会调用本方法。如果不实现本方法，则打印原始消息。
@@ -88,8 +97,10 @@ class Robot(Job):
             else:                # 其他消息
                 self.toChengyu(msg)
 
-        # 非群聊信息
-        elif msg.type == 37:     # 好友请求
+            return  # 处理完群聊信息，后面就不需要处理了
+
+        # 非群聊信息，按消息类型进行处理
+        if msg.type == 37:     # 好友请求
             self.autoAcceptFriendRequest(msg)
 
         elif msg.type == 10000:  # 系统信息
@@ -100,15 +111,12 @@ class Robot(Job):
             if msg.from_self() and msg.content == "^更新$":
                 self.config.reload()
                 self.LOG.info("已更新")
-                return
-
-            # 闲聊
-            self.toChitchat(msg)
+            else:
+                self.toChitchat(msg)  # 闲聊
 
     def onMsg(self, msg: Wcf.WxMsg) -> int:
-        self.LOG.info(msg)  # 打印信息
-
         try:
+            self.LOG.info(msg)  # 打印信息
             self.processMsg(msg)
         except Exception as e:
             self.LOG.error(e)