feat: BingImageCreator output four images

.env.example

@@ -11,5 +11,6 @@ BARD_TOKEN="xxxxxxxxxxxxxxxxxxxx", # Optional, for !bard command
 JAILBREAKENABLED="true" # Optional
 BING_AUTH_COOKIE="xxxxxxxxxxxxxxxxxxx" # _U cookie, Optional, for Bing Image Creator
 MARKDOWN_FORMATTED="true" # Optional
+OUTPUT_FOUR_IMAGES="true" # Optional
 IMPORT_KEYS_PATH="element-keys.txt" # Optional
 IMPORT_KEYS_PASSWORD="xxxxxxx" # Optional

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 @@ FORWARDED_IP = (
     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 @@ class ImageGenAsync:
     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 @@ class ImageGenAsync:
             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 @@ class ImageGenAsync:
         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 @@ class ImageGenAsync:
             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"
+        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
 
-        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
+        return image_path_list

README.md

@@ -26,6 +26,7 @@ sudo docker compose up -d
 
 <hr>
 Normal Method:<br>
+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:<br>
+- `!help` help message
 
 - `!gpt` To generate a one time response:
 
@@ -115,6 +117,8 @@ https://github.com/hibobmaster/matrix_chatgpt_bot/wiki/ <br>
 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/)
+
 <a href="https://jb.gg/OpenSourceSupport" target="_blank">
 <img src="https://resources.jetbrains.com/storage/products/company/brand/logos/jb_beam.png" alt="JetBrains Logo (Main) logo." width="200" height="200">
 </a>

bot.py

@@ -42,6 +42,7 @@ class Bot:
         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 @@ class Bot:
         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 @@ class Bot:
                                         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 @@ class Bot:
         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 @@ class Bot:
     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 @@ class Bot:
                                             )
                     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 @@ class Bot:
                         )
                     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 @@ class Bot:
                         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 @@ class Bot:
             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 @@ class Bot:
             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 @@ class Bot:
         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 @@ class Bot:
         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 @@ class Bot:
         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 @@ class Bot:
         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 @@ class Bot:
             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 @@ class Bot:
                         "!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 @@ class Bot:
             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:

config.json.sample

@@ -11,6 +11,7 @@
   "bard_token": "xxxxxxx",
   "bing_auth_cookie": "xxxxxxxxxxx",
   "markdown_formatted": true,
+  "output_four_images": true,
   "import_keys_path": "element-keys.txt",
   "import_keys_password": "xxxxxxxxx"
 }

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)