From ff30d06eda060cad6707ee5ca219562d1d72023f Mon Sep 17 00:00:00 2001 From: hibobmaster Date: Thu, 13 Apr 2023 21:31:11 +0800 Subject: [PATCH] Chore: close session when exit --- BingImageGen.py | 11 +++++++++++ askgpt.py | 8 ++++---- bing.py | 8 ++++---- bot.py | 47 ++++++++++++++++++++++++++++++----------------- main.py | 3 +++ v3.py | 2 +- 6 files changed, 53 insertions(+), 26 deletions(-) diff --git a/BingImageGen.py b/BingImageGen.py index 4226ae0..bcdf629 100644 --- a/BingImageGen.py +++ b/BingImageGen.py @@ -211,6 +211,17 @@ class ImageGenAsync: async def __aexit__(self, *excinfo) -> None: await self.session.close() + def __del__(self): + try: + loop = asyncio.get_event_loop() + except RuntimeError as e: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(self._close()) + + async def _close(self): + await self.session.close() + async def get_images(self, prompt: str) -> list: """ Fetches image links from Bing diff --git a/askgpt.py b/askgpt.py index 1b1c083..e7b2d86 100644 --- a/askgpt.py +++ b/askgpt.py @@ -6,8 +6,8 @@ logger = getlogger() class askGPT: - def __init__(self): - self.session = aiohttp.ClientSession() + def __init__(self, session: aiohttp.ClientSession): + self.session = session async def oneTimeAsk(self, prompt: str, api_endpoint: str, headers: dict) -> str: jsons = { @@ -19,11 +19,11 @@ class askGPT: }, ], } - max_try = 3 + max_try = 2 while max_try > 0: try: async with self.session.post(url=api_endpoint, - json=jsons, headers=headers, timeout=60) as response: + json=jsons, headers=headers, timeout=120) as response: status_code = response.status if not status_code == 200: # print failed reason diff --git a/bing.py b/bing.py index d14b3a7..a06c837 100644 --- a/bing.py +++ b/bing.py @@ -7,13 +7,13 @@ logger = getlogger() class BingBot: - def __init__(self, bing_api_endpoint: str, jailbreakEnabled: bool = False): + def __init__(self, session: aiohttp.ClientSession, bing_api_endpoint: str, jailbreakEnabled: bool = False): self.data = { 'clientOptions.clientToUse': 'bing', } self.bing_api_endpoint = bing_api_endpoint - self.session = aiohttp.ClientSession() + self.session = session self.jailbreakEnabled = jailbreakEnabled @@ -22,10 +22,10 @@ class BingBot: async def ask_bing(self, prompt) -> str: self.data['message'] = prompt - max_try = 3 + max_try = 2 while max_try > 0: try: - resp = await self.session.post(url=self.bing_api_endpoint, json=self.data, timeout=60) + resp = await self.session.post(url=self.bing_api_endpoint, json=self.data, timeout=120) status_code = resp.status body = await resp.read() if not status_code == 200: diff --git a/bot.py b/bot.py index 9764e68..98fd43b 100644 --- a/bot.py +++ b/bot.py @@ -1,5 +1,6 @@ import sys import asyncio +import aiohttp import re import os from functools import partial @@ -69,6 +70,8 @@ class Bot: self.api_key = api_key self.chatgpt_api_endpoint = chatgpt_api_endpoint + self.session = aiohttp.ClientSession() + if bing_api_endpoint is None: self.bing_api_endpoint = '' else: @@ -121,7 +124,7 @@ class Bot: # initialize chatbot and chatgpt_api_endpoint if self.api_key != '': - self.chatbot = Chatbot(api_key=self.api_key, timeout=60) + self.chatbot = Chatbot(api_key=self.api_key, timeout=120) self.chatgpt_api_endpoint = self.chatgpt_api_endpoint # request header for !gpt command @@ -136,19 +139,29 @@ class Bot: } # initialize askGPT class - self.askgpt = askGPT() + self.askgpt = askGPT(self.session) # initialize bingbot if self.bing_api_endpoint != '': self.bingbot = BingBot( - bing_api_endpoint, jailbreakEnabled=self.jailbreakEnabled) + self.session, bing_api_endpoint, jailbreakEnabled=self.jailbreakEnabled) # initialize BingImageGenAsync if self.bing_auth_cookie != '': self.imageGen = ImageGenAsync(self.bing_auth_cookie, quiet=True) - # get current event loop - self.loop = asyncio.get_running_loop() + + def __del__(self): + try: + loop = asyncio.get_event_loop() + except RuntimeError as e: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(self._close()) + + + async def _close(self): + await self.session.close() # message_callback RoomMessageText event async def message_callback(self, room: MatrixRoom, event: RoomMessageText) -> None: @@ -194,7 +207,7 @@ class Bot: ) except Exception as e: - logger.error(e) + logger.error(e, exc_info=True) await send_room_message(self.client, room_id, reply_message=str(e)) else: logger.warning("No API_KEY provided") @@ -211,7 +224,7 @@ class Bot: raw_user_message ) except Exception as e: - logger.error(e) + logger.error(e, exc_info=True) await send_room_message(self.client, room_id, reply_message=str(e)) # bing ai @@ -230,6 +243,7 @@ class Bot: ) except Exception as e: + logger.error(e, exc_info=True) await send_room_message(self.client, room_id, reply_message=str(e)) # Image Generation by Microsoft Bing @@ -240,6 +254,7 @@ class Bot: try: await self.pic(room_id, prompt) except Exception as e: + logger.error(e, exc_info=True) await send_room_message(self.client, room_id, reply_message=str(e)) # help command @@ -491,12 +506,9 @@ class Bot: # !chat command async def chat(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message): - await self.client.room_typing(room_id, timeout=180000) + await self.client.room_typing(room_id, timeout=120000) try: - text = await asyncio.wait_for(self.chatbot.ask_async(prompt), timeout=180) - except TimeoutError as e: - logger.error(f"TimeoutException: {e}", exc_info=True) - raise Exception("Timeout error") + text = await self.chatbot.ask_async(prompt) except Exception as e: raise Exception(e) @@ -509,12 +521,12 @@ class Bot: logger.error(f"Error: {e}", exc_info=True) # !gpt command - async def gpt(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message): + async def gpt(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message) -> None: try: # sending typing state - await self.client.room_typing(room_id, timeout=180000) - # timeout 120s - text = await asyncio.wait_for(self.askgpt.oneTimeAsk(prompt, self.chatgpt_api_endpoint, self.headers), timeout=180) + await self.client.room_typing(room_id, timeout=240000) + # timeout 240s + text = await asyncio.wait_for(self.askgpt.oneTimeAsk(prompt, self.chatgpt_api_endpoint, self.headers), timeout=240) except TimeoutError: logger.error("TimeoutException", exc_info=True) raise Exception("Timeout error") @@ -535,7 +547,7 @@ class Bot: # sending typing state await self.client.room_typing(room_id, timeout=180000) # timeout 120s - text = await asyncio.wait_for(self.bingbot.ask_bing(prompt), timeout=180) + text = await asyncio.wait_for(self.bingbot.ask_bing(prompt), timeout=240) except TimeoutError: logger.error("timeoutException", exc_info=True) raise Exception("Timeout error") @@ -561,6 +573,7 @@ class Bot: image_path = await self.imageGen.save_images(links, "images") except Exception as e: logger.error(f"Image Generation error: {e}", exc_info=True) + raise Exception(e) # send image try: diff --git a/main.py b/main.py index cc7d15b..b139cbd 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,6 @@ import json import os +import sys import asyncio from bot import Bot from log import getlogger @@ -52,4 +53,6 @@ async def main(): if __name__ == "__main__": logger.info("matrix chatgpt bot start.....") + print("matrix chatgpt bot start.....") asyncio.run(main()) + diff --git a/v3.py b/v3.py index b91dca3..e109b52 100644 --- a/v3.py +++ b/v3.py @@ -1,7 +1,7 @@ import json import os from typing import AsyncGenerator - +import asyncio import httpx import requests import tiktoken