Chore: close session when exit
This commit is contained in:
parent
bce867acd2
commit
ff30d06eda
6 changed files with 53 additions and 26 deletions
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
8
bing.py
8
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:
|
||||
|
|
47
bot.py
47
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:
|
||||
|
|
3
main.py
3
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())
|
||||
|
||||
|
|
2
v3.py
2
v3.py
|
@ -1,7 +1,7 @@
|
|||
import json
|
||||
import os
|
||||
from typing import AsyncGenerator
|
||||
|
||||
import asyncio
|
||||
import httpx
|
||||
import requests
|
||||
import tiktoken
|
||||
|
|
Loading…
Reference in a new issue