From da682cb9fdf3cf7023a2085b3a11fc5327aa9267 Mon Sep 17 00:00:00 2001 From: hibobmaster Date: Tue, 11 Apr 2023 13:42:43 +0800 Subject: [PATCH] format code --- BingImageGen.py | 9 ++++----- log.py | 6 ++++-- main.py | 53 +++++++++++++++++++++++++------------------------ send_image.py | 8 +++++--- v3.py | 28 +++++++++++++++----------- 5 files changed, 56 insertions(+), 48 deletions(-) diff --git a/BingImageGen.py b/BingImageGen.py index 812010d..4226ae0 100644 --- a/BingImageGen.py +++ b/BingImageGen.py @@ -49,7 +49,6 @@ wait_message = "Waiting for results..." download_message = "\nDownloading images..." - def debug(debug_file, text_var): """helper function for debug""" with open(f"{debug_file}", "a") as f: @@ -74,7 +73,6 @@ class ImageGen: if self.debug_file: self.debug = partial(debug, self.debug_file) - def get_images(self, prompt: str) -> list: """ Fetches image links from Bing @@ -106,7 +104,8 @@ class ImageGen: if response.status_code != 302: # if rt4 fails, try rt3 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 self.debug_file: self.debug(f"ERROR: {error_redirect}") @@ -165,7 +164,7 @@ class ImageGen: """ Saves images to output directory """ - + # image name image_name = str(uuid4()) # 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): output_file.write(chunk) return f"{output_dir}/{image_name}.jpeg" - + except aiohttp.client_exceptions.InvalidURL as url_exception: raise Exception( "Inappropriate contents found in the generated images. Please try again or try another prompt.", diff --git a/log.py b/log.py index f56a19a..ae671e3 100644 --- a/log.py +++ b/log.py @@ -14,8 +14,10 @@ def getlogger(): info_handler.setLevel(logging.INFO) # create formatters - warn_format = logging.Formatter('%(name)s - %(funcName)s - %(levelname)s - %(message)s') - error_format = logging.Formatter('%(asctime)s - %(name)s - %(funcName)s - %(levelname)s - %(message)s') + warn_format = logging.Formatter( + '%(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') # set formatter diff --git a/main.py b/main.py index 59b6076..cc7d15b 100644 --- a/main.py +++ b/main.py @@ -6,38 +6,41 @@ from log import getlogger logger = getlogger() + async def main(): - + if os.path.exists('config.json'): fp = open('config.json', 'r', encoding="utf8") config = json.load(fp) - + matrix_bot = Bot(homeserver=config.get('homeserver'), - user_id=config.get('user_id') , - password=config.get('password'), - device_id=config.get('device_id'), - room_id=config.get('room_id'), - api_key=config.get('api_key'), - bing_api_endpoint=config.get('bing_api_endpoint'), - access_token=config.get('access_token'), - jailbreakEnabled=config.get('jailbreakEnabled'), - bing_auth_cookie=config.get('bing_auth_cookie'), - markdown_formatted=config.get('markdown_formatted'), - ) + user_id=config.get('user_id'), + password=config.get('password'), + device_id=config.get('device_id'), + room_id=config.get('room_id'), + api_key=config.get('api_key'), + bing_api_endpoint=config.get('bing_api_endpoint'), + access_token=config.get('access_token'), + jailbreakEnabled=config.get('jailbreakEnabled'), + bing_auth_cookie=config.get('bing_auth_cookie'), + markdown_formatted=config.get('markdown_formatted'), + ) else: matrix_bot = Bot(homeserver=os.environ.get('HOMESERVER'), - user_id=os.environ.get('USER_ID') , - password=os.environ.get('PASSWORD'), - device_id=os.environ.get("DEVICE_ID"), - room_id=os.environ.get("ROOM_ID"), - api_key=os.environ.get("OPENAI_API_KEY"), - bing_api_endpoint=os.environ.get("BING_API_ENDPOINT"), - access_token=os.environ.get("ACCESS_TOKEN"), - jailbreakEnabled=os.environ.get("JAILBREAKENABLED", "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'), - ) + user_id=os.environ.get('USER_ID'), + password=os.environ.get('PASSWORD'), + device_id=os.environ.get("DEVICE_ID"), + room_id=os.environ.get("ROOM_ID"), + api_key=os.environ.get("OPENAI_API_KEY"), + bing_api_endpoint=os.environ.get("BING_API_ENDPOINT"), + access_token=os.environ.get("ACCESS_TOKEN"), + jailbreakEnabled=os.environ.get( + "JAILBREAKENABLED", "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.sync_encryption_key() @@ -50,5 +53,3 @@ async def main(): if __name__ == "__main__": logger.info("matrix chatgpt bot start.....") asyncio.run(main()) - - diff --git a/send_image.py b/send_image.py index 70f4818..3378d44 100644 --- a/send_image.py +++ b/send_image.py @@ -13,7 +13,7 @@ logger = getlogger() async def send_room_image(client: AsyncClient, - room_id: str, image: str): + room_id: str, image: str): """ image: image path """ @@ -36,7 +36,8 @@ async def send_room_image(client: AsyncClient, await client.room_send( room_id, 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, ) return @@ -56,4 +57,5 @@ async def send_room_image(client: AsyncClient, try: await client.room_send(room_id, message_type="m.room.message", content=content) 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) diff --git a/v3.py b/v3.py index 2c1e6ee..b91dca3 100644 --- a/v3.py +++ b/v3.py @@ -53,7 +53,8 @@ class Chatbot: }, ) 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: @@ -79,7 +80,6 @@ class Chatbot: ], } - def add_to_conversation( self, message: str, @@ -105,7 +105,6 @@ class Chatbot: else: break - def get_token_count(self, convo_id: str = "default") -> int: """ Get token count @@ -158,8 +157,10 @@ class Chatbot: self.__truncate_conversation(convo_id=convo_id) # Get response response = self.session.post( - os.environ.get("API_URL") or "https://api.openai.com/v1/chat/completions", - headers={"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"}, + os.environ.get( + "API_URL") or "https://api.openai.com/v1/chat/completions", + headers={ + "Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"}, json={ "model": self.engine, "messages": self.conversation[convo_id], @@ -182,7 +183,7 @@ class Chatbot: timeout=kwargs.get("timeout", self.timeout), stream=True, ) - + response_role: str = None full_response: str = "" for line in response.iter_lines(): @@ -205,7 +206,8 @@ class Chatbot: content = delta["content"] full_response += 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( self, @@ -225,8 +227,10 @@ class Chatbot: # Get response async with self.aclient.stream( "post", - os.environ.get("API_URL") or "https://api.openai.com/v1/chat/completions", - headers={"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"}, + os.environ.get( + "API_URL") or "https://api.openai.com/v1/chat/completions", + headers={ + "Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"}, json={ "model": self.engine, "messages": self.conversation[convo_id], @@ -274,7 +278,8 @@ class Chatbot: content: str = delta["content"] full_response += 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( self, @@ -314,11 +319,10 @@ class Chatbot: full_response: str = "".join(response) return full_response - def reset(self, convo_id: str = "default", system_prompt: str = None) -> None: """ Reset the conversation """ self.conversation[convo_id] = [ {"role": "system", "content": system_prompt or self.system_prompt}, - ] \ No newline at end of file + ]