From bc950c308f46dcaec87d9608703b7fe145d086a5 Mon Sep 17 00:00:00 2001 From: hibobmaster Date: Sat, 20 May 2023 09:46:16 +0800 Subject: [PATCH] feat: BingImageCreator output four images --- BingImageGen.py | 63 +++++++++++++++++++++++++------------------ README.md | 4 +++ bot.py | 71 ++++++++++++++++++++++++++++++++++--------------- main.py | 20 +++++++++----- 4 files changed, 105 insertions(+), 53 deletions(-) diff --git a/BingImageGen.py b/BingImageGen.py index b3c3a4f..7fb32d6 100644 --- a/BingImageGen.py +++ b/BingImageGen.py @@ -2,6 +2,7 @@ Code derived from: https://github.com/acheong08/EdgeGPT/blob/f940cecd24a4818015a8b42a2443dd97c3c2a8f4/src/ImageGen.py """ + from log import getlogger from uuid import uuid4 import os @@ -20,16 +21,17 @@ f"13.{random.randint(104, 107)}.{random.randint(0, 255)}.{random.randint(0, 255)}" ) HEADERS = { - "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7", + "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7", # noqa: E501 "accept-language": "en-US,en;q=0.9", "cache-control": "max-age=0", "content-type": "application/x-www-form-urlencoded", "referrer": "https://www.bing.com/images/create/", "origin": "https://www.bing.com", - "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36 Edg/110.0.1587.63", + "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36 Edg/110.0.1587.63", # noqa: E501 "x-forwarded-for": FORWARDED_IP, } + class ImageGenAsync: """ Image generation by Microsoft Bing @@ -53,7 +55,7 @@ async def __aexit__(self, *excinfo) -> None: def __del__(self): try: loop = asyncio.get_running_loop() - except RuntimeError as e: + except RuntimeError: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) loop.run_until_complete(self._close()) @@ -76,7 +78,7 @@ async def get_images(self, prompt: str) -> list: content = await response.text() if "this prompt has been blocked" in content.lower(): raise Exception( - "Your prompt has been blocked by Bing. Try to change any bad words and try again.", + "Your prompt has been blocked by Bing. Try to change any bad words and try again.", # noqa: E501 ) if response.status != 302: # if rt4 fails, try rt3 @@ -97,7 +99,7 @@ async def get_images(self, prompt: str) -> list: request_id = redirect_url.split("id=")[-1] await self.session.get(f"{BING_URL}{redirect_url}") # https://www.bing.com/images/create/async/results/{ID}?q={PROMPT} - polling_url = f"{BING_URL}/images/create/async/results/{request_id}?q={url_encoded_prompt}" + polling_url = f"{BING_URL}/images/create/async/results/{request_id}?q={url_encoded_prompt}" # noqa: E501 # Poll for results if not self.quiet: print("Waiting for results...") @@ -134,31 +136,40 @@ async def get_images(self, prompt: str) -> list: raise Exception("No images") return normal_image_links - async def save_images(self, links: list, output_dir: str) -> str: + async def save_images(self, links: list, output_dir: str, + output_four_images: bool) -> list: """ Saves images to output directory """ - if not self.quiet: - print("\nDownloading images...") with contextlib.suppress(FileExistsError): os.mkdir(output_dir) - # image name - image_name = str(uuid4()) - # since matrix only support one media attachment per message, we just need one link - if links: - link = links.pop() + image_path_list = [] - image_path = os.path.join(output_dir, f"{image_name}.jpeg") - try: - async with self.session.get(link, raise_for_status=True) as response: - # save response to file - with open(image_path, "wb") as output_file: - 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.", - ) from url_exception + if output_four_images: + for link in links: + image_name = str(uuid4()) + image_path = os.path.join(output_dir, f"{image_name}.jpeg") + try: + async with self.session.get(link, raise_for_status=True) as response: + with open(image_path, "wb") as output_file: + async for chunk in response.content.iter_chunked(8192): + output_file.write(chunk) + image_path_list.append(image_path) + except aiohttp.client_exceptions.InvalidURL as url_exception: + raise Exception("Inappropriate contents found in the generated images. Please try again or try another prompt.") from url_exception # noqa: E501 + else: + image_name = str(uuid4()) + if links: + link = links.pop() + try: + async with self.session.get(link, raise_for_status=True) as response: + image_path = os.path.join(output_dir, f"{image_name}.jpeg") + with open(image_path, "wb") as output_file: + async for chunk in response.content.iter_chunked(8192): + output_file.write(chunk) + image_path_list.append(image_path) + except aiohttp.client_exceptions.InvalidURL as url_exception: + raise Exception("Inappropriate contents found in the generated images. Please try again or try another prompt.") from url_exception # noqa: E501 + + return image_path_list diff --git a/README.md b/README.md index 8305d5e..561ac43 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ sudo docker compose up -d
Normal Method:
+system dependece: `libolm-dev` 1. Clone the repository and create virtual environment: @@ -74,6 +75,7 @@ python main.py ## Usage To interact with the bot, simply send a message to the bot in the Matrix room with one of the two prompts:
+- `!help` help message - `!gpt` To generate a one time response: @@ -115,6 +117,8 @@ https://github.com/hibobmaster/matrix_chatgpt_bot/wiki/
1. [matrix-nio](https://github.com/poljar/matrix-nio) 2. [acheong08](https://github.com/acheong08) 3. [node-chatgpt-api](https://github.com/waylaidwanderer/node-chatgpt-api) +4. [8go](https://github.com/8go/) + JetBrains Logo (Main) logo. diff --git a/bot.py b/bot.py index 176097b..a8ce610 100644 --- a/bot.py +++ b/bot.py @@ -42,6 +42,7 @@ def __init__( jailbreakEnabled: Union[bool, None] = True, bing_auth_cookie: Union[str, None] = '', markdown_formatted: Union[bool, None] = False, + output_four_images: Union[bool, None] = False, import_keys_path: Optional[str] = None, import_keys_password: Optional[str] = None, ): @@ -88,6 +89,11 @@ def __init__( else: self.markdown_formatted = markdown_formatted + if output_four_images is None: + self.output_four_images = False + else: + self.output_four_images = output_four_images + # initialize AsyncClient object self.store_path = os.getcwd() self.config = AsyncClientConfig(store=SqliteStore, @@ -95,7 +101,8 @@ def __init__( store_sync_tokens=True, encryption_enabled=True, ) - self.client = AsyncClient(homeserver=self.homeserver, user=self.user_id, device_id=self.device_id, + self.client = AsyncClient(homeserver=self.homeserver, user=self.user_id, + device_id=self.device_id, config=self.config, store_path=self.store_path,) if self.access_token is not None: @@ -111,7 +118,7 @@ def __init__( self.client.add_to_device_callback( self.to_device_callback, (KeyVerificationEvent, )) - # regular expression to match keyword [!gpt {prompt}] [!chat {prompt}] [!bing {prompt}] [!pic {prompt}] [!bard {prompt}] + # regular expression to match keyword commands self.gpt_prog = re.compile(r"^\s*!gpt\s*(.+)$") self.chat_prog = re.compile(r"^\s*!chat\s*(.+)$") self.bing_prog = re.compile(r"^\s*!bing\s*(.+)$") @@ -149,7 +156,7 @@ def __init__( def __del__(self): try: loop = asyncio.get_running_loop() - except RuntimeError as e: + except RuntimeError: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) loop.run_until_complete(self._close()) @@ -202,10 +209,12 @@ async def message_callback(self, room: MatrixRoom, event: RoomMessageText) -> No ) except Exception as e: logger.error(e, exc_info=True) - await send_room_message(self.client, room_id, reply_message=str(e)) + await send_room_message(self.client, room_id, + reply_message=str(e)) else: logger.warning("No API_KEY provided") - await send_room_message(self.client, room_id, reply_message="API_KEY not provided") + await send_room_message(self.client, room_id, + reply_message="API_KEY not provided") m = self.gpt_prog.match(content_body) if m: @@ -239,7 +248,8 @@ async def message_callback(self, room: MatrixRoom, event: RoomMessageText) -> No ) except Exception as e: logger.error(e, exc_info=True) - await send_room_message(self.client, room_id, reply_message=str(e)) + await send_room_message(self.client, room_id, + reply_message=str(e)) # Image Generation by Microsoft Bing if self.bing_auth_cookie != '': @@ -250,7 +260,8 @@ async def message_callback(self, room: MatrixRoom, event: RoomMessageText) -> No asyncio.create_task(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)) + await send_room_message(self.client, room_id, + reply_message=str(e)) # Google's Bard if self.bard_token is not None: @@ -281,7 +292,8 @@ async def decryption_failure(self, room: MatrixRoom, event: MegolmEvent) -> None return logger.error( - f"Failed to decrypt message: {event.event_id} from {event.sender} in {room.room_id}\n" + + f"Failed to decrypt message: {event.event_id} \ + from {event.sender} in {room.room_id}\n" + "Please make sure the bot current session is verified" ) @@ -501,7 +513,8 @@ async def to_device_callback(self, event: KeyVerificationEvent) -> None: logger.info(estr) # !chat command - async def chat(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message): + async def chat(self, room_id, reply_to_event_id, prompt, sender_id, + raw_user_message): await self.client.room_typing(room_id, timeout=120000) try: text = await self.chatbot.ask_async(prompt) @@ -511,17 +524,23 @@ async def chat(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_mes try: text = text.strip() await send_room_message(self.client, room_id, reply_message=text, - reply_to_event_id="", sender_id=sender_id, user_message=raw_user_message, markdown_formatted=self.markdown_formatted) + reply_to_event_id="", sender_id=sender_id, + user_message=raw_user_message, + markdown_formatted=self.markdown_formatted) except Exception as e: 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) -> None: + 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=240000) # timeout 240s - text = await asyncio.wait_for(self.askgpt.oneTimeAsk(prompt, self.chatgpt_api_endpoint, self.headers), timeout=240) + 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") @@ -531,12 +550,15 @@ async def gpt(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_mess try: text = text.strip() await send_room_message(self.client, room_id, reply_message=text, - reply_to_event_id="", sender_id=sender_id, user_message=raw_user_message, markdown_formatted=self.markdown_formatted) + reply_to_event_id="", sender_id=sender_id, + user_message=raw_user_message, + markdown_formatted=self.markdown_formatted) except Exception as e: logger.error(f"Error: {e}", exc_info=True) # !bing command - async def bing(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message) -> None: + async def bing(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) @@ -551,12 +573,15 @@ async def bing(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_mes try: text = text.strip() await send_room_message(self.client, room_id, reply_message=text, - reply_to_event_id="", sender_id=sender_id, user_message=raw_user_message, markdown_formatted=self.markdown_formatted) + reply_to_event_id="", sender_id=sender_id, + user_message=raw_user_message, + markdown_formatted=self.markdown_formatted) except Exception as e: logger.error(e, exc_info=True) # !bard command - async def bard(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message) -> None: + async def bard(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) @@ -567,7 +592,9 @@ async def bard(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_mes try: content = str(response['content']).strip() await send_room_message(self.client, room_id, reply_message=content, - reply_to_event_id="", sender_id=sender_id, user_message=raw_user_message, markdown_formatted=self.markdown_formatted) + reply_to_event_id="", sender_id=sender_id, + user_message=raw_user_message, + markdown_formatted=self.markdown_formatted) except Exception as e: logger.error(e, exc_info=True) @@ -580,14 +607,16 @@ async def pic(self, room_id, prompt): try: links = await self.imageGen.get_images(prompt) - image_path = await self.imageGen.save_images(links, "images") + image_path_list = await self.imageGen.save_images(links, "images", + self.output_four_images) except Exception as e: logger.error(f"Image Generation error: {e}", exc_info=True) raise Exception(e) # send image try: - await send_room_image(self.client, room_id, image_path) + for image_path in image_path_list: + await send_room_image(self.client, room_id, image_path) await self.client.room_typing(room_id, typing_state=False) except Exception as e: logger.error(e, exc_info=True) @@ -605,7 +634,7 @@ async def help(self, room_id): "!bing [content], chat with context conversation powered by Bing AI\n" + \ "!bard [content], chat with Google's Bard\n" + \ "!pic [prompt], Image generation by Microsoft Bing\n" + \ - "!help, help message" + "!help, help message" # noqa: E501 await send_room_message(self.client, room_id, reply_message=help_info) except Exception as e: @@ -635,7 +664,7 @@ async def import_keys(self): logger.error(f"import_keys failed with {resp}") else: logger.info( - f"import_keys success, please remove import_keys configuration!!!") + "import_keys success, please remove import_keys configuration!!!") # sync messages in the room async def sync_forever(self, timeout=30000, full_state=True) -> None: diff --git a/main.py b/main.py index 8338bf1..b467a9e 100644 --- a/main.py +++ b/main.py @@ -25,11 +25,13 @@ async def main(): jailbreakEnabled=config.get('jailbreakEnabled'), bing_auth_cookie=config.get('bing_auth_cookie'), markdown_formatted=config.get('markdown_formatted'), + output_four_images=config.get('output_four_images'), import_keys_path=config.get('import_keys_path'), import_keys_password=config.get( 'import_keys_password'), ) - if config.get('import_keys_path') and config.get('import_keys_password') is not None: + if config.get('import_keys_path') and \ + config.get('import_keys_password') is not None: need_import_keys = True else: @@ -43,21 +45,27 @@ async def main(): access_token=os.environ.get("ACCESS_TOKEN"), bard_token=os.environ.get("BARD_TOKEN"), jailbreakEnabled=os.environ.get( - "JAILBREAKENABLED", "false").lower() in ('true', '1', 't'), + "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'), + "MARKDOWN_FORMATTED", "false").lower() \ + in ('true', '1', 't'), + output_four_images=os.environ.get( + "OUTPUT_FOUR_IMAGES", "false").lower() \ + in ('true', '1', 't'), import_keys_path=os.environ.get("IMPORT_KEYS_PATH"), import_keys_password=os.environ.get( "IMPORT_KEYS_PASSWORD"), ) - if os.environ.get("IMPORT_KEYS_PATH") and os.environ.get("IMPORT_KEYS_PASSWORD") is not None: + if os.environ.get("IMPORT_KEYS_PATH") \ + and os.environ.get("IMPORT_KEYS_PASSWORD") is not None: need_import_keys = True await matrix_bot.login() if need_import_keys: - logger.info("start import_keys process, this may take a while...") - await matrix_bot.import_keys() + logger.info("start import_keys process, this may take a while...") + await matrix_bot.import_keys() await matrix_bot.sync_forever(timeout=30000, full_state=True)