From fb7e52b4dcf29c280ea05a4f2fda209d793af53d Mon Sep 17 00:00:00 2001
From: Kenxu2022 <kenxu2022@gmail.com>
Date: Mon, 18 Nov 2024 22:19:19 +0800
Subject: [PATCH] add AI functions

---
 WeiBanHelper.py  | 57 ++++++++++++++++++++++++++++++++++++++++++++++--
 ai.conf          |  4 ++++
 requirements.txt |  1 +
 3 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 ai.conf

diff --git a/WeiBanHelper.py b/WeiBanHelper.py
index 38a169e..ce56eee 100644
--- a/WeiBanHelper.py
+++ b/WeiBanHelper.py
@@ -17,6 +17,8 @@
 
 import encrypted
 
+from openai import OpenAI
+import configparser
 
 class WeibanHelper:
     tenantCode = 0
@@ -388,6 +390,46 @@ def get_verify_code():
             now = time.time()
             content = retry_request_2("GET", get_verify_code_url + str(now), headers=self.headers).content
             return self.ocr.classification(content), now
+        
+        def ai_response(input, type):
+            client = OpenAI(base_url = config['AI']['API_ENDPOINT'],api_key = config['AI']['API_KEY'])
+
+            if type == 1:
+                completion = client.chat.completions.create(
+                    model = config['AI']['MODEL'],
+                    messages=[
+                        {
+                            "role": "system", 
+                            "content": "本题为单选题，请根据题目和选项回答问题，以json格式输出正确的选项对应的id（即正确选项'id'键对应的值）和内容（即正确选项'content'键对应的值），除此之外不要输出任何多余的内容。"
+                        },
+                        {
+                            "role": "user",
+                            "content": input
+                        }
+                    ]
+                )
+            if type == 2:
+                completion = client.chat.completions.create(
+                    model = config['AI']['MODEL'],
+                    messages=[
+                        {
+                            "role": "system", 
+                            "content": "本题为多选题，请根据题目和选项回答问题，以json格式输出正确的选项对应的id（即正确选项'id'键对应的值）和内容（即正确选项'content'键对应的值），并使用逗号连接多个值，除此之外不要输出任何多余的内容。"
+                        },
+                        {
+                            "role": "user",
+                            "content": input
+                        }
+                    ]
+                )
+
+            response = completion.choices[0].message.content
+            data = json.loads(response)
+
+            id_value = data['id']
+            content_value = data['content']
+
+            return id_value, content_value
 
         # 获取当前系统时间
         current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
@@ -464,9 +506,12 @@ def get_verify_code():
             # 提取题目列表
             question_list = paper_data['questionList']
             match_count = 0
+            answerIds = None
 
             for question in question_list:
                 question_title = question['title']
+                question_type = question['type'] # 1是单选，2是多选
+                question_type_name = question['typeLabel']
                 option_list = question['optionList']
                 submit_answer_id_list = []
 
@@ -475,6 +520,8 @@ def get_verify_code():
 
                 print(f"题目: {question_title}")
 
+                config = configparser.ConfigParser()
+                config.read('ai.conf')
                 # 检查题目标题是否匹配
                 if answer_list:
                     found_match = False
@@ -490,12 +537,18 @@ def get_verify_code():
                         print("<===答案匹配成功===>\n")
                     else:
                         print("<——————————!!!题目匹配但选项未找到匹配项!!!——————————>\n")
+                elif not config['AI'].get('API_ENDPOINT') or not config['AI'].get('API_KEY') or not config['AI'].get('MODEL'):
+                    print("<——————————!!!未匹配到答案，可配置ai.conf文件通过大模型答题!!!——————————>\n")
                 else:
-                    print("<——————————!!!未匹配到答案，题库暂未收录此题!!!——————————>\n")
+                    print("<——————————未匹配到答案，将使用AI获取答案——————————>\n")
+                    problemInput = f"{question_title}\n{option_list}"
+                    answerIds, content = ai_response(problemInput, question_type)
+                    print(f"{question_type_name}，AI获取的答案: {content}")
+                    match_count += 1
 
                 # 记录答案
                 record_data = {
-                    "answerIds": ",".join(submit_answer_id_list),
+                    "answerIds": answerIds if answerIds is not None else ",".join(submit_answer_id_list),
                     "questionId": question['id'],
                     "tenantCode": self.tenantCode,
                     "userId": self.userId,
diff --git a/ai.conf b/ai.conf
new file mode 100644
index 0000000..a0cff47
--- /dev/null
+++ b/ai.conf
@@ -0,0 +1,4 @@
+[AI]
+API_ENDPOINT=
+API_KEY=
+MODEL=
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index f42a34f..65d63a1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,3 +5,4 @@ pycryptodomex
 fastapi
 uvicorn
 ddddocr
+openai
\ No newline at end of file