浏览代码

print a warning if stream task ever times out

Alex Cheema 10 月之前
父节点
当前提交
d8c40bb4f1
共有 1 个文件被更改,包括 4 次插入1 次删除
  1. 4 1
      exo/api/chatgpt_api.py

+ 4 - 1
exo/api/chatgpt_api.py

@@ -200,7 +200,10 @@ class ChatGPTAPI:
                 _, tokens, _ = await callback.wait(on_result, timeout=self.response_timeout_secs)
                 _, tokens, _ = await callback.wait(on_result, timeout=self.response_timeout_secs)
                 if stream_task: # in case there is still a stream task running, wait for it to complete
                 if stream_task: # in case there is still a stream task running, wait for it to complete
                     if DEBUG >= 2: print(f"Pending stream task. Waiting for stream task to complete.")
                     if DEBUG >= 2: print(f"Pending stream task. Waiting for stream task to complete.")
-                    await stream_task
+                    try:
+                        await asyncio.wait_for(stream_task, timeout=30)
+                    except asyncio.TimeoutError:
+                        print("WARNING: Stream task timed out. This should not happen.")
                 await response.write_eof()
                 await response.write_eof()
                 return response
                 return response
             else:
             else: