Merge branch 'experimental'

README.md

@@ -19,9 +19,7 @@ A [Telegram bot](https://core.telegram.org/bots/api) that integrates with OpenAI
 - [x] (NEW!) Customizable model parameters (see [configuration](#configuration) section)
 - [x] (NEW!) See token usage after each answer
 - [x] (NEW!) Multi-chat support
-
-## Coming soon
-- [ ] Image generation using DALL·E APIs
+- [x] (NEW!) Image generation using DALL·E APIs via the `/image` command
 
 ## Additional Features - help needed!
 - [ ] Group chat support
@@ -64,7 +62,7 @@ Additional optional model parameters can be configured from the `main.py` file:
     # How many answers to generate for each input message. Defaults to 1
     'n_choices': 1,
     
-    # The maximum number of tokens allowed for the generated answer. Defaults to 4096 minus prompt tokens
+    # The maximum number of tokens allowed for the generated answer. Defaults to 1200
     'max_tokens': 1200,
     
     # Number between -2.0 and 2.0. Positive values penalize new tokens based on whether
@@ -73,7 +71,10 @@ Additional optional model parameters can be configured from the `main.py` file:
     
     # Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing
     # frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim. Defaults to 0
-    'frequency_penalty': 0
+    'frequency_penalty': 0,
+    
+    # The DALL·E generated image size
+    'image_size': '512x512'
 }
 ```
 See the [official API reference](https://platform.openai.com/docs/api-reference/chat) for more details.

main.py

@@ -3,7 +3,7 @@ import os
 
 from dotenv import load_dotenv
 
-from gpt_helper import GPTHelper
+from openai_helper import OpenAIHelper
 from telegram_bot import ChatGPT3TelegramBot
 
 
@@ -25,7 +25,7 @@ def main():
         exit(1)
 
     # Setup configurations
-    gpt_config = {
+    openai_config = {
         'api_key': os.environ['OPENAI_API_KEY'],
         'show_usage': os.environ.get('SHOW_USAGE', 'true').lower() == 'true',
 
@@ -51,7 +51,10 @@ def main():
 
         # Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing
         # frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim.
-        'frequency_penalty': 0
+        'frequency_penalty': 0,
+
+        # The DALL·E generated image size
+        'image_size': '512x512'
     }
 
     telegram_config = {
@@ -60,8 +63,8 @@ def main():
     }
 
     # Setup and run ChatGPT and Telegram bot
-    gpt = GPTHelper(config=gpt_config)
-    telegram_bot = ChatGPT3TelegramBot(config=telegram_config, gpt=gpt)
+    openai_helper = OpenAIHelper(config=openai_config)
+    telegram_bot = ChatGPT3TelegramBot(config=telegram_config, openai=openai_helper)
     telegram_bot.run()
 
 

openai_helper.py

@@ -2,14 +2,14 @@ import logging
 import openai
 
 
-class GPTHelper:
+class OpenAIHelper:
     """
     ChatGPT helper class.
     """
 
     def __init__(self, config: dict):
         """
-        Initializes the GPT helper class with the given configuration.
+        Initializes the OpenAI helper class with the given configuration.
         :param config: A dictionary containing the GPT configuration
         """
         openai.api_key = config['api_key']
@@ -17,7 +17,7 @@ class GPTHelper:
         self.sessions: dict[int: list] = dict() # {chat_id: history}
 
 
-    def get_response(self, chat_id: int, query: str) -> str:
+    def get_chat_response(self, chat_id: int, query: str) -> str:
         """
         Gets a response from the GPT-3 model.
         :param chat_id: The chat ID
@@ -26,7 +26,7 @@ class GPTHelper:
         """
         try:
             if chat_id not in self.sessions:
-                self.reset_history(chat_id)
+                self.reset_chat_history(chat_id)
 
             self.__add_to_history(chat_id, role="user", content=query)
 
@@ -78,7 +78,25 @@ class GPTHelper:
             return f"⚠️ _An error has occurred_ ⚠️\n{str(e)}"
 
 
-    def reset_history(self, chat_id):
+    def generate_image(self, prompt: str) -> str:
+        """
+        Generates an image from the given prompt using DALL·E model.
+        :param prompt: The prompt to send to the model
+        :return: The image URL
+        """
+        try:
+            response = openai.Image.create(
+                prompt=prompt,
+                n=1,
+                size=self.config['image_size']
+            )
+            return response['data'][0]['url']
+
+        except Exception as e:
+            logging.exception(e)
+            raise e
+
+    def reset_chat_history(self, chat_id):
         """
         Resets the conversation history.
         """

telegram_bot.py

@@ -4,7 +4,7 @@ import telegram.constants as constants
 from telegram import Update
 from telegram.ext import ApplicationBuilder, ContextTypes, CommandHandler, MessageHandler, filters
 
-from gpt_helper import GPTHelper
+from openai_helper import OpenAIHelper
 
 
 class ChatGPT3TelegramBot:
@@ -12,14 +12,14 @@ class ChatGPT3TelegramBot:
     Class representing a Chat-GPT3 Telegram Bot.
     """
 
-    def __init__(self, config: dict, gpt: GPTHelper):
+    def __init__(self, config: dict, openai: OpenAIHelper):
         """
         Initializes the bot with the given configuration and GPT-3 bot object.
         :param config: A dictionary containing the bot configuration
-        :param gpt: GPTHelper object
+        :param openai: OpenAIHelper object
         """
         self.config = config
-        self.gpt = gpt
+        self.openai = openai
         self.disallowed_message = "Sorry, you are not allowed to use this bot. You can check out the source code at " \
                                   "https://github.com/n3d1117/chatgpt-telegram-bot"
 
@@ -28,6 +28,7 @@ class ChatGPT3TelegramBot:
         Shows the help menu.
         """
         await update.message.reply_text("/reset - Reset conversation\n"
+                                        "/image <prompt> - Generate image\n"
                                         "/help - Help menu\n\n"
                                         "Open source at https://github.com/n3d1117/chatgpt-telegram-bot",
                                         disable_web_page_preview=True)
@@ -42,8 +43,39 @@ class ChatGPT3TelegramBot:
             return
 
         logging.info(f'Resetting the conversation for user {update.message.from_user.name}...')
-        self.gpt.reset_history(chat_id=update.effective_chat.id)
-        await context.bot.send_message(chat_id=update.effective_chat.id, text='Done!')
+        chat_id = update.effective_chat.id
+        self.openai.reset_chat_history(chat_id=chat_id)
+        await context.bot.send_message(chat_id=chat_id, text='Done!')
+
+    async def image(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
+        """
+        Generates an image for the given prompt using DALL·E APIs
+        """
+        if not self.is_allowed(update):
+            logging.warning(f'User {update.message.from_user.name} is not allowed to generate images')
+            await self.send_disallowed_message(update, context)
+            return
+
+        chat_id = update.effective_chat.id
+        image_query = update.message.text.replace('/image', '').strip()
+        if image_query == '':
+            await context.bot.send_message(chat_id=chat_id, text='Please provide a prompt!')
+            return
+
+        await context.bot.send_chat_action(chat_id=chat_id, action=constants.ChatAction.UPLOAD_PHOTO)
+        try:
+            image_url = self.openai.generate_image(prompt=image_query)
+            await context.bot.send_photo(
+                chat_id=chat_id,
+                reply_to_message_id=update.message.message_id,
+                photo=image_url
+            )
+        except:
+            await context.bot.send_message(
+                chat_id=chat_id,
+                reply_to_message_id=update.message.message_id,
+                text='Failed to generate image'
+            )
 
     async def prompt(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
         """
@@ -55,11 +87,12 @@ class ChatGPT3TelegramBot:
             return
 
         logging.info(f'New message received from user {update.message.from_user.name}')
+        chat_id = update.effective_chat.id
 
-        await context.bot.send_chat_action(chat_id=update.effective_chat.id, action=constants.ChatAction.TYPING)
-        response = self.gpt.get_response(chat_id=update.effective_chat.id, query=update.message.text)
+        await context.bot.send_chat_action(chat_id=chat_id, action=constants.ChatAction.TYPING)
+        response = self.openai.get_chat_response(chat_id=chat_id, query=update.message.text)
         await context.bot.send_message(
-            chat_id=update.effective_chat.id,
+            chat_id=chat_id,
             reply_to_message_id=update.message.message_id,
             text=response,
             parse_mode=constants.ParseMode.MARKDOWN
@@ -97,6 +130,7 @@ class ChatGPT3TelegramBot:
 
         application.add_handler(CommandHandler('reset', self.reset))
         application.add_handler(CommandHandler('help', self.help))
+        application.add_handler(CommandHandler('image', self.image))
         application.add_handler(CommandHandler('start', self.help))
         application.add_handler(MessageHandler(filters.TEXT & (~filters.COMMAND), self.prompt))