Add the use of the /pic command to generate images, and add support for the dall-e-3 model

Author: yym68686

bot.py

@@ -193,6 +193,47 @@ async def search(update, context, title, robot):
             result = re.sub(r",", '，', result)
         await context.bot.edit_message_text(chat_id=update.message.chat_id, message_id=messageid, text=escape(result), parse_mode='MarkdownV2', disable_web_page_preview=True)
 
+async def image(update, context):
+    print("\033[32m", update.effective_user.username, update.effective_user.id, update.message.text, "\033[0m")
+    if (len(context.args) == 0):
+        message = (
+            f"格式错误哦~，示例：\n\n"
+            f"`/pic 一只可爱长毛金渐层在趴在路由器上`\n\n"
+            f"👆点击上方命令复制格式\n\n"
+        )
+        await context.bot.send_message(chat_id=update.effective_chat.id, text=escape(message), parse_mode='MarkdownV2', disable_web_page_preview=True)
+        return
+    message = ' '.join(context.args)
+    result = ""
+    robot = config.dallbot
+    text = message
+    message = await context.bot.send_message(
+        chat_id=update.message.chat_id,
+        text="生成中💭",
+        parse_mode='MarkdownV2',
+        reply_to_message_id=update.message.message_id,
+    )
+    messageid = message.message_id
+
+    try:
+        for data in robot.dall_e_3(text):
+            result = data
+            await context.bot.delete_message(chat_id=update.message.chat_id, message_id=messageid)
+            await context.bot.send_photo(chat_id=update.message.chat_id, photo=result, reply_to_message_id=update.message.message_id)
+    except Exception as e:
+        print('\033[31m')
+        print("response_msg", result)
+        print("error", e)
+        traceback.print_exc()
+        print('\033[0m')
+        if "You exceeded your current quota, please check your plan and billing details." in str(e):
+            print("OpenAI api 已过期！")
+            await context.bot.delete_message(chat_id=update.message.chat_id, message_id=messageid)
+            messageid = ''
+            config.API = ''
+        result += f"`出错啦！{e}`"
+    print(result)
+
 import time
 import threading
 async def delete_message(update, context, messageid, delay=10):
@@ -512,6 +553,7 @@ def setup(token):
 
     run_async(application.bot.set_my_commands([
         BotCommand('info', 'basic information'),
+        BotCommand('pic', 'Generate image'),
         BotCommand('search', 'search Google or duckduckgo'),
         BotCommand('en2zh', 'translate to Chinese'),
         BotCommand('zh2en', 'translate to English'),
@@ -521,6 +563,7 @@ def setup(token):
     ]))
 
     application.add_handler(CommandHandler("start", start))
+    application.add_handler(CommandHandler("pic", image))
     application.add_handler(CommandHandler("search", lambda update, context: search(update, context, title=f"`🤖️ {config.GPT_ENGINE}`\n\n", robot=config.ChatGPTbot)))
     application.add_handler(CallbackQueryHandler(button_press))
     application.add_handler(CommandHandler("reset", reset_chat))

chatgpt2api/chatgpt2api.py

@@ -55,6 +55,51 @@ def get_filtered_keys_from_object(obj: object, *keys: str) -> Set[str]:
     "claude-2",
 ]
 
+class Imagebot:
+    def __init__(
+        self,
+        api_key: str,
+        timeout: float = None,
+    ):
+        self.api_key: str = api_key
+        self.engine: str = "dall-e-3"
+        self.session = requests.Session()
+        self.timeout: float = timeout
+
+    def dall_e_3(
+        self,
+        prompt: str,
+        model: str = None,
+        **kwargs,
+    ):
+        url = (
+            os.environ.get("API_URL").split("chat")[0] + "images/generations"
+            or "https://api.openai.com/v1/images/generations"
+        )
+        headers = {"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"}
+
+        json_post = {
+                "model": os.environ.get("IMAGE_MODEL_NAME") or model or self.engine,
+                "prompt": prompt,
+                "stream": True,
+                "n": 1,
+                "size": "1024x1024",
+        }
+        response = self.session.post(
+            url,
+            headers=headers,
+            json=json_post,
+            timeout=kwargs.get("timeout", self.timeout),
+            stream=True,
+        )
+        if response.status_code != 200:
+            raise t.APIConnectionError(
+                f"{response.status_code} {response.reason} {response.text}",
+            )
+        for line in response.iter_lines():
+            json_data = json.loads(line.decode("utf-8"))
+            url = json_data["data"][0]["url"]
+            yield url
 
 class Chatbot:
     """

config.py

@@ -21,9 +21,11 @@
 systemprompt = f"You are ChatGPT, a large language model trained by OpenAI. Knowledge cutoff: 2021-09. Current date: [ {Current_Date} ]"
 
 from chatgpt2api.chatgpt2api import Chatbot as GPT
+from chatgpt2api.chatgpt2api import Imagebot
 if API:
     ChatGPTbot = GPT(api_key=f"{API}", engine=GPT_ENGINE, system_prompt=systemprompt, temperature=temperature)
     Claude2bot = GPT(api_key=f"{API}", engine="claude-2-web")
+    dallbot = Imagebot(api_key=f"{API}")
 else:
     ChatGPTbot = None