From 62fd63d70bdb37f3152886bf047511ef12cd67c8 Mon Sep 17 00:00:00 2001
From: LanQian <5499636+LanQian528@users.noreply.github.com>
Date: Tue, 30 Apr 2024 05:42:33 +0800
Subject: [PATCH] fix bugs

---
 app.py                  | 13 +++++++++----
 chatgpt/ChatService.py  |  5 +++--
 chatgpt/refreshToken.py |  2 ++
 utils/Client.py         |  4 ++++
 4 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/app.py b/app.py
index 714338f..5a0c11c 100644
--- a/app.py
+++ b/app.py
@@ -3,6 +3,7 @@ import types
 from fastapi import FastAPI, Request, Depends, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse, JSONResponse
+from starlette.background import BackgroundTask
 
 from chatgpt.ChatService import ChatService
 from chatgpt.reverseProxy import chatgpt_reverse_proxy
@@ -41,16 +42,20 @@ async def send_conversation(request: Request, token=Depends(verify_token)):
         raise HTTPException(status_code=400, detail={"error": "Invalid JSON body"})
 
     chat_service = await async_retry(to_send_conversation, request_data, access_token)
-    try:
-        await chat_service.prepare_send_conversation()
 
+    await chat_service.prepare_send_conversation()
+    try:
         res = await chat_service.send_conversation()
         if isinstance(res, types.AsyncGeneratorType):
-            return StreamingResponse(res, media_type="text/event-stream")
+            background = BackgroundTask(await chat_service.close_client)
+            return StreamingResponse(res, media_type="text/event-stream", background=background)
         else:
             return JSONResponse(res, media_type="application/json")
+    except Exception:
+        raise HTTPException(status_code=500, detail="Server error")
     finally:
-        await chat_service.close_client()
+        if res and not isinstance(res, types.AsyncGeneratorType):
+            await chat_service.close_client()
 
 
 @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "PATCH", "TRACE"])
diff --git a/chatgpt/ChatService.py b/chatgpt/ChatService.py
index 22198bc..b262922 100644
--- a/chatgpt/ChatService.py
+++ b/chatgpt/ChatService.py
@@ -89,6 +89,8 @@ class ChatService:
                         self.arkose_token = r2esp.get('token')
                     except Exception:
                         raise HTTPException(status_code=403, detail="Failed to get Arkose token")
+                    finally:
+                        await arkose_client.close()
 
                 proofofwork_required = proofofwork.get('required')
                 if proofofwork_required:
@@ -310,5 +312,4 @@ class ChatService:
             return False
 
     async def close_client(self):
-        await self.s.session.close()
-        self.s.session = None
\ No newline at end of file
+        await self.s.close()
\ No newline at end of file
diff --git a/chatgpt/refreshToken.py b/chatgpt/refreshToken.py
index c52d661..ed87858 100644
--- a/chatgpt/refreshToken.py
+++ b/chatgpt/refreshToken.py
@@ -44,3 +44,5 @@ async def chat_refresh(refresh_token):
             raise Exception("Unknown or invalid refresh token.")
     except Exception as e:
         raise HTTPException(status_code=401, detail=str(e))
+    finally:
+        await client.close()
diff --git a/utils/Client.py b/utils/Client.py
index f8d5cde..b552125 100644
--- a/utils/Client.py
+++ b/utils/Client.py
@@ -39,3 +39,7 @@ class Client:
     async def put(self, *args, headers=None, cookies=None, **kwargs):
         r = await self.session.put(*args, impersonate=self.impersonate, **kwargs)
         return r
+
+    async def close(self):
+        await self.session.close()
+        self.session = None
\ No newline at end of file