From 0ac5e5b35e6f63795cdb98a76e689c4e7da70fa9 Mon Sep 17 00:00:00 2001
From: Rock Chin <1010553892@qq.com>
Date: Fri, 3 Mar 2023 21:13:31 +0800
Subject: [PATCH] =?UTF-8?q?fix(session.py):=20=E9=94=99=E8=AF=AF=E7=9A=84u?=
 =?UTF-8?q?ndo()=E6=96=B9=E6=B3=95=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/openai/session.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/pkg/openai/session.py b/pkg/openai/session.py
index 9bd9e72d..6932a852 100644
--- a/pkg/openai/session.py
+++ b/pkg/openai/session.py
@@ -228,16 +228,15 @@ class Session:
     def undo(self) -> str:
         self.last_interact_timestamp = int(time.time())
 
-        # 删除上一回合
-        if self.prompt[-1]['role'] != 'user':
-            res = self.prompt[-1]['content']
-            self.prompt.remove(self.prompt[-2])
-        else:
-            res = self.prompt[-2]['content']
-        self.prompt.remove(self.prompt[-1])
-
+        # 删除最后两个消息
+        if len(self.prompt) < 2:
+            raise Exception('之前无对话，无法撤销')
+    
+        question = self.prompt[-2]['content']
+        self.prompt = self.prompt[:-2]
+        
         # 返回上一回合的问题
-        return res
+        return question
 
     # 构建对话体
     def cut_out(self, msg: str, max_tokens: int) -> list: