Chore: close session when exit

This commit is contained in:
hibobmaster 2023-04-13 21:31:11 +08:00
parent bce867acd2
commit ff30d06eda
No known key found for this signature in database
GPG Key ID: 316B77D7914D713C
6 changed files with 53 additions and 26 deletions

View File

@ -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

View File

@ -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

View File

@ -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:

47
bot.py
View File

@ -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:

View File

@ -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())

2
v3.py
View File

@ -1,7 +1,7 @@
import json
import os
from typing import AsyncGenerator
import asyncio
import httpx
import requests
import tiktoken