修复部分情况下上游返回正常的空消息导致程序错误 (#1101)

* 修复部分情况下上游返回正常的空消息导致程序错误 * 修复程序出错后仍然计时的问题
lss233 · Jul 31, 2023 · 5a376b8 · 5a376b8
1 parent 8288978
commit 5a376b8
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/adapter/chatgpt/api.py b/adapter/chatgpt/api.py
@@ -52,7 +52,11 @@ async def rollback(self, session_id: str = "default", n: int = 1) -> None:
  raise
 
  def add_to_conversation(self, message: str, role: str, session_id: str = "default") -> None:
- self.conversation[session_id].append({"role": role, "content": message})
+ if role and message is not None:
+ self.conversation[session_id].append({"role": role, "content": message})
+ else:
+ logger.warning("出现错误！返回消息为空，不添加到会话。")
+ raise ValueError("出现错误！返回消息为空，不添加到会话。")
 
  # https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
  def count_tokens(self, session_id: str = "default", model: str = DEFAULT_ENGINE):

diff --git a/middlewares/timeout.py b/middlewares/timeout.py
@@ -32,6 +32,12 @@ async def handle_request(self, session_id: str, prompt: str, respond: Callable,
  del self.timeout_task[session_id]
  except asyncio.TimeoutError:
  await respond(config.response.cancel_wait_too_long)
+ except Exception as e:
+ logger.error(f"发生错误: {e}")
+ if session_id in self.timeout_task:
+ self.timeout_task[session_id].cancel()
+ del self.timeout_task[session_id]
+ raise e
 
  async def on_respond(self, session_id: str, prompt: str, rendered: str):
  if rendered and session_id in self.timeout_task: