Merge pull request #154 from n3d1117/feature/allow-disable-transcription-and-image-generation

Allow disabling image generation and transcriptions

Author: n3d1117

README.md

@@ -52,6 +52,8 @@ Customize the configuration by copying `.env.example` and renaming it to `.env`,
 ### Optional configuration
 | Parameter                          | Description                                                                                                                                                                                                                      | Default value                  |
 |------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------|
+| `ENABLE_IMAGE_GENERATION`          | Whether to enable image generation via the `/image` command                                                                                                                                                                      | true                           |
+| `ENABLE_TRANSCRIPTION`             | Whether to enable transcriptions of audio and video messages                                                                                                                                                                     | true                           |
 | `MONTHLY_USER_BUDGETS`             | A comma-separated list of $-amounts per user from list `ALLOWED_TELEGRAM_USER_IDS` to set custom usage limit of OpenAI API costs for each.  **Note**: by default, *no limits* for anyone (`*`)                                   | `*`                            |
 | `MONTHLY_GUEST_BUDGET`             | $-amount as usage limit for all guest users. Guest users are users in group chats that are not in the `ALLOWED_TELEGRAM_USER_IDS` list. Value is ignored if no usage limits are set in user budgets (`MONTHLY_USER_BUDGETS`="*") | `100.0`                        |
 | `PROXY`                            | Proxy to be used for OpenAI and Telegram bot (e.g. `http://localhost:8080`)                                                                                                                                                      | -                              |

bot/main.py

@@ -48,6 +48,8 @@ def main():
         'token': os.environ['TELEGRAM_BOT_TOKEN'],
         'admin_user_ids': os.environ.get('ADMIN_USER_IDS', '-'),
         'allowed_user_ids': os.environ.get('ALLOWED_TELEGRAM_USER_IDS', '*'),
+        'enable_image_generation': os.environ.get('ENABLE_IMAGE_GENERATION', 'true').lower() == 'true',
+        'enable_transcription': os.environ.get('ENABLE_TRANSCRIPTION', 'true').lower() == 'true',
         'monthly_user_budgets': os.environ.get('MONTHLY_USER_BUDGETS', '*'),
         'monthly_guest_budget': float(os.environ.get('MONTHLY_GUEST_BUDGET', '100.0')),
         'stream': os.environ.get('STREAM', 'true').lower() == 'true',

bot/telegram_bot.py

@@ -173,7 +173,7 @@ async def image(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
         """
         Generates an image for the given prompt using DALL·E APIs
         """
-        if not await self.check_allowed_and_within_budget(update, context):
+        if not self.config['enable_image_generation'] or not await self.check_allowed_and_within_budget(update, context):
             return
 
         chat_id = update.effective_chat.id
@@ -215,7 +215,7 @@ async def transcribe(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
         """
         Transcribe audio messages.
         """
-        if not await self.check_allowed_and_within_budget(update, context):
+        if not self.config['enable_transcription'] or not await self.check_allowed_and_within_budget(update, context):
             return
 
         if self.is_group_chat(update) and self.config['ignore_group_transcriptions']: