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