format code

This commit is contained in:
hibobmaster 2023-04-11 13:42:43 +08:00
parent f5ccc67258
commit da682cb9fd
Signed by: bobmaster
GPG key ID: 316B77D7914D713C
5 changed files with 56 additions and 48 deletions

View file

@ -49,7 +49,6 @@ wait_message = "Waiting for results..."
download_message = "\nDownloading images..." download_message = "\nDownloading images..."
def debug(debug_file, text_var): def debug(debug_file, text_var):
"""helper function for debug""" """helper function for debug"""
with open(f"{debug_file}", "a") as f: with open(f"{debug_file}", "a") as f:
@ -74,7 +73,6 @@ class ImageGen:
if self.debug_file: if self.debug_file:
self.debug = partial(debug, self.debug_file) self.debug = partial(debug, self.debug_file)
def get_images(self, prompt: str) -> list: def get_images(self, prompt: str) -> list:
""" """
Fetches image links from Bing Fetches image links from Bing
@ -106,7 +104,8 @@ class ImageGen:
if response.status_code != 302: if response.status_code != 302:
# if rt4 fails, try rt3 # if rt4 fails, try rt3
url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=3&FORM=GENCRE" url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=3&FORM=GENCRE"
response3 = self.session.post(url, allow_redirects=False, timeout=200) response3 = self.session.post(
url, allow_redirects=False, timeout=200)
if response3.status_code != 302: if response3.status_code != 302:
if self.debug_file: if self.debug_file:
self.debug(f"ERROR: {error_redirect}") self.debug(f"ERROR: {error_redirect}")
@ -165,7 +164,7 @@ class ImageGen:
""" """
Saves images to output directory Saves images to output directory
""" """
# image name # image name
image_name = str(uuid4()) image_name = str(uuid4())
# since matrix only support one media attachment per message, we just need one link # since matrix only support one media attachment per message, we just need one link
@ -308,7 +307,7 @@ class ImageGenAsync:
async for chunk in response.content.iter_chunked(8192): async for chunk in response.content.iter_chunked(8192):
output_file.write(chunk) output_file.write(chunk)
return f"{output_dir}/{image_name}.jpeg" return f"{output_dir}/{image_name}.jpeg"
except aiohttp.client_exceptions.InvalidURL as url_exception: except aiohttp.client_exceptions.InvalidURL as url_exception:
raise Exception( raise Exception(
"Inappropriate contents found in the generated images. Please try again or try another prompt.", "Inappropriate contents found in the generated images. Please try again or try another prompt.",

6
log.py
View file

@ -14,8 +14,10 @@ def getlogger():
info_handler.setLevel(logging.INFO) info_handler.setLevel(logging.INFO)
# create formatters # create formatters
warn_format = logging.Formatter('%(name)s - %(funcName)s - %(levelname)s - %(message)s') warn_format = logging.Formatter(
error_format = logging.Formatter('%(asctime)s - %(name)s - %(funcName)s - %(levelname)s - %(message)s') '%(name)s - %(funcName)s - %(levelname)s - %(message)s')
error_format = logging.Formatter(
'%(asctime)s - %(name)s - %(funcName)s - %(levelname)s - %(message)s')
info_format = logging.Formatter('%(message)s') info_format = logging.Formatter('%(message)s')
# set formatter # set formatter

53
main.py
View file

@ -6,38 +6,41 @@ from log import getlogger
logger = getlogger() logger = getlogger()
async def main(): async def main():
if os.path.exists('config.json'): if os.path.exists('config.json'):
fp = open('config.json', 'r', encoding="utf8") fp = open('config.json', 'r', encoding="utf8")
config = json.load(fp) config = json.load(fp)
matrix_bot = Bot(homeserver=config.get('homeserver'), matrix_bot = Bot(homeserver=config.get('homeserver'),
user_id=config.get('user_id') , user_id=config.get('user_id'),
password=config.get('password'), password=config.get('password'),
device_id=config.get('device_id'), device_id=config.get('device_id'),
room_id=config.get('room_id'), room_id=config.get('room_id'),
api_key=config.get('api_key'), api_key=config.get('api_key'),
bing_api_endpoint=config.get('bing_api_endpoint'), bing_api_endpoint=config.get('bing_api_endpoint'),
access_token=config.get('access_token'), access_token=config.get('access_token'),
jailbreakEnabled=config.get('jailbreakEnabled'), jailbreakEnabled=config.get('jailbreakEnabled'),
bing_auth_cookie=config.get('bing_auth_cookie'), bing_auth_cookie=config.get('bing_auth_cookie'),
markdown_formatted=config.get('markdown_formatted'), markdown_formatted=config.get('markdown_formatted'),
) )
else: else:
matrix_bot = Bot(homeserver=os.environ.get('HOMESERVER'), matrix_bot = Bot(homeserver=os.environ.get('HOMESERVER'),
user_id=os.environ.get('USER_ID') , user_id=os.environ.get('USER_ID'),
password=os.environ.get('PASSWORD'), password=os.environ.get('PASSWORD'),
device_id=os.environ.get("DEVICE_ID"), device_id=os.environ.get("DEVICE_ID"),
room_id=os.environ.get("ROOM_ID"), room_id=os.environ.get("ROOM_ID"),
api_key=os.environ.get("OPENAI_API_KEY"), api_key=os.environ.get("OPENAI_API_KEY"),
bing_api_endpoint=os.environ.get("BING_API_ENDPOINT"), bing_api_endpoint=os.environ.get("BING_API_ENDPOINT"),
access_token=os.environ.get("ACCESS_TOKEN"), access_token=os.environ.get("ACCESS_TOKEN"),
jailbreakEnabled=os.environ.get("JAILBREAKENABLED", "false").lower() in ('true', '1', 't'), jailbreakEnabled=os.environ.get(
bing_auth_cookie=os.environ.get("BING_AUTH_COOKIE"), "JAILBREAKENABLED", "false").lower() in ('true', '1', 't'),
markdown_formatted=os.environ.get("MARKDOWN_FORMATTED", "false").lower() in ('true', '1', 't'), bing_auth_cookie=os.environ.get("BING_AUTH_COOKIE"),
) markdown_formatted=os.environ.get(
"MARKDOWN_FORMATTED", "false").lower() in ('true', '1', 't'),
)
await matrix_bot.login() await matrix_bot.login()
await matrix_bot.sync_encryption_key() await matrix_bot.sync_encryption_key()
@ -50,5 +53,3 @@ async def main():
if __name__ == "__main__": if __name__ == "__main__":
logger.info("matrix chatgpt bot start.....") logger.info("matrix chatgpt bot start.....")
asyncio.run(main()) asyncio.run(main())

View file

@ -13,7 +13,7 @@ logger = getlogger()
async def send_room_image(client: AsyncClient, async def send_room_image(client: AsyncClient,
room_id: str, image: str): room_id: str, image: str):
""" """
image: image path image: image path
""" """
@ -36,7 +36,8 @@ async def send_room_image(client: AsyncClient,
await client.room_send( await client.room_send(
room_id, room_id,
message_type="m.room.message", message_type="m.room.message",
content={"msgtype": "m.text", "body": f"Failed to generate image. Failure response: {resp}", }, content={"msgtype": "m.text",
"body": f"Failed to generate image. Failure response: {resp}", },
ignore_unverified_devices=True, ignore_unverified_devices=True,
) )
return return
@ -56,4 +57,5 @@ async def send_room_image(client: AsyncClient,
try: try:
await client.room_send(room_id, message_type="m.room.message", content=content) await client.room_send(room_id, message_type="m.room.message", content=content)
except Exception as e: except Exception as e:
logger.error(f"Image send of file {image} failed.\n Error: {e}", exc_info=True) logger.error(
f"Image send of file {image} failed.\n Error: {e}", exc_info=True)

28
v3.py
View file

@ -53,7 +53,8 @@ class Chatbot:
}, },
) )
proxy = ( proxy = (
proxy or os.environ.get("all_proxy") or os.environ.get("ALL_PROXY") or None proxy or os.environ.get(
"all_proxy") or os.environ.get("ALL_PROXY") or None
) )
if proxy: if proxy:
@ -79,7 +80,6 @@ class Chatbot:
], ],
} }
def add_to_conversation( def add_to_conversation(
self, self,
message: str, message: str,
@ -105,7 +105,6 @@ class Chatbot:
else: else:
break break
def get_token_count(self, convo_id: str = "default") -> int: def get_token_count(self, convo_id: str = "default") -> int:
""" """
Get token count Get token count
@ -158,8 +157,10 @@ class Chatbot:
self.__truncate_conversation(convo_id=convo_id) self.__truncate_conversation(convo_id=convo_id)
# Get response # Get response
response = self.session.post( response = self.session.post(
os.environ.get("API_URL") or "https://api.openai.com/v1/chat/completions", os.environ.get(
headers={"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"}, "API_URL") or "https://api.openai.com/v1/chat/completions",
headers={
"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"},
json={ json={
"model": self.engine, "model": self.engine,
"messages": self.conversation[convo_id], "messages": self.conversation[convo_id],
@ -182,7 +183,7 @@ class Chatbot:
timeout=kwargs.get("timeout", self.timeout), timeout=kwargs.get("timeout", self.timeout),
stream=True, stream=True,
) )
response_role: str = None response_role: str = None
full_response: str = "" full_response: str = ""
for line in response.iter_lines(): for line in response.iter_lines():
@ -205,7 +206,8 @@ class Chatbot:
content = delta["content"] content = delta["content"]
full_response += content full_response += content
yield content yield content
self.add_to_conversation(full_response, response_role, convo_id=convo_id) self.add_to_conversation(
full_response, response_role, convo_id=convo_id)
async def ask_stream_async( async def ask_stream_async(
self, self,
@ -225,8 +227,10 @@ class Chatbot:
# Get response # Get response
async with self.aclient.stream( async with self.aclient.stream(
"post", "post",
os.environ.get("API_URL") or "https://api.openai.com/v1/chat/completions", os.environ.get(
headers={"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"}, "API_URL") or "https://api.openai.com/v1/chat/completions",
headers={
"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"},
json={ json={
"model": self.engine, "model": self.engine,
"messages": self.conversation[convo_id], "messages": self.conversation[convo_id],
@ -274,7 +278,8 @@ class Chatbot:
content: str = delta["content"] content: str = delta["content"]
full_response += content full_response += content
yield content yield content
self.add_to_conversation(full_response, response_role, convo_id=convo_id) self.add_to_conversation(
full_response, response_role, convo_id=convo_id)
async def ask_async( async def ask_async(
self, self,
@ -314,11 +319,10 @@ class Chatbot:
full_response: str = "".join(response) full_response: str = "".join(response)
return full_response return full_response
def reset(self, convo_id: str = "default", system_prompt: str = None) -> None: def reset(self, convo_id: str = "default", system_prompt: str = None) -> None:
""" """
Reset the conversation Reset the conversation
""" """
self.conversation[convo_id] = [ self.conversation[convo_id] = [
{"role": "system", "content": system_prompt or self.system_prompt}, {"role": "system", "content": system_prompt or self.system_prompt},
] ]