From e96f3800e8b7ff092cc81fed96b41d9d2dfefd6a Mon Sep 17 00:00:00 2001
From: Changhua <lichanghua0821@gmail.com>
Date: Sun, 2 Apr 2023 17:32:55 +0800
Subject: [PATCH] Add proxy and optimize ChatGPT

---
 config.yaml.template | 10 ++++++----
 configuration.py     | 14 +++++++-------
 func_chatgpt.py      | 15 +++++++++++++--
 robot.py             | 10 ++++++++--
 4 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/config.yaml.template b/config.yaml.template
index 537591a..508c111 100644
--- a/config.yaml.template
+++ b/config.yaml.template
@@ -39,8 +39,10 @@ logging:
     handlers: [console, info_file_handler, error_file_handler]
 
 groups:
-  enable: [] # 允许响应的群 roomId，可以通过 WxGetContacts 获取
+  enable: [] # 允许响应的群 roomId，大概长这样：2xxxxxxxxx3@chatroom
 
-chatgpt:
-  key: your key
-  api: https://api.openai.com/v1
+# 如果要使用 ChatGPT，取消下面的注释并填写相关内容
+# chatgpt:
+#   key: 填写你 ChatGPT 的 key
+#   api: https://api.openai.com/v1 # 如果你不知道这是干嘛的，就不要改
+#   proxy: # 如果你在国内，你可能需要魔法，大概长这样：http://域名或者IP地址:端口号
diff --git a/configuration.py b/configuration.py
index da5bd79..8eacc21 100644
--- a/configuration.py
+++ b/configuration.py
@@ -1,9 +1,11 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
-import os
-import yaml
 import logging.config
+import os
+import shutil
+
+import yaml
 
 
 class Config(object):
@@ -16,10 +18,9 @@ class Config(object):
             with open(f"{pwd}/config.yaml", "rb") as fp:
                 yconfig = yaml.safe_load(fp)
         except FileNotFoundError:
-            with open(f"{pwd}/config.yaml.template", "rb") as fp:
+            shutil.copyfile(f"{pwd}/config.yaml.template", f"{pwd}/config.yaml")
+            with open(f"{pwd}/config.yaml", "rb") as fp:
                 yconfig = yaml.safe_load(fp)
-                with open(f"{pwd}/config.yaml", "w+") as yf:
-                    yaml.dump(yconfig, yf, default_flow_style=False)
 
         return yconfig
 
@@ -27,5 +28,4 @@ class Config(object):
         yconfig = self._load_config()
         logging.config.dictConfig(yconfig["logging"])
         self.GROUPS = yconfig["groups"]["enable"]
-        self.CHAT_KEY = yconfig["chatgpt"]["key"]
-        self.CHAT_API = yconfig["chatgpt"]["api"]
+        self.CHATGPT = yconfig.get("chatgpt")
diff --git a/func_chatgpt.py b/func_chatgpt.py
index 77d44e8..b74b029 100644
--- a/func_chatgpt.py
+++ b/func_chatgpt.py
@@ -1,5 +1,6 @@
 #! /usr/bin/env python3
 # -*- coding: utf-8 -*-
+
 from datetime import datetime
 
 import openai
@@ -7,10 +8,12 @@ import openai
 
 class ChatGPT():
 
-    def __init__(self, key: str, api: str = openai.api_base) -> None:
+    def __init__(self, key: str, api: str, proxy: str) -> None:
         openai.api_key = key
         # 自己搭建或第三方代理的接口
         openai.api_base = api
+        if proxy:
+            openai.proxy = {"http": proxy, "https": proxy}
         self.conversation_list = {}
         self.system_content_msg = {"role": "system",
                                    "content": "你是智能聊天机器人,你叫小小，调皮可爱喜欢二次元的小萝莉，"
@@ -77,7 +80,15 @@ class ChatGPT():
 
 
 if __name__ == "__main__":
-    chat = ChatGPT("your key", "https://api.openai.com/v1")
+    from configuration import Config
+    config = Config().CHATGPT
+    if not config:
+        exit(0)
+
+    key = config.get("key")
+    api = config.get("api")
+    proxy = config.get("proxy")
+    chat = ChatGPT(key, api, proxy)
     while True:
         q = input(">>> ")
         try:
diff --git a/robot.py b/robot.py
index 27e96a5..b287577 100644
--- a/robot.py
+++ b/robot.py
@@ -23,7 +23,10 @@ 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, self.config.CHAT_API)
+        self.chat = None
+        chatgpt = self.config.CHATGPT
+        if chatgpt:
+            self.chat = ChatGPT(chatgpt.get("key"), chatgpt.get("api"), chatgpt.get("proxy"))
 
     def toAt(self, msg: Wcf.WxMsg) -> bool:
         """
@@ -63,8 +66,11 @@ class Robot(Job):
     def toChitchat(self, msg: Wcf.WxMsg) -> bool:
         """闲聊，接入 ChatGPT
         """
+        if not self.chat:
+            return False
+
         q = re.sub(r"@.*?[\u2005|\s]", "", msg.content).replace(" ", "")
-        rsp = self.chat.get_answer(q,(msg.roomid if msg.from_group() else msg.sender))
+        rsp = self.chat.get_answer(q, (msg.roomid if msg.from_group() else msg.sender))
         if rsp:
             if msg.from_group():
                 self.sendTextMsg(rsp, msg.roomid, msg.sender)