support live answers, updating as the bot types

README.md

@@ -14,12 +14,14 @@ A [Telegram bot](https://core.telegram.org/bots/api) that integrates with OpenAI
 - [x] Can reset conversation thread with the `/reset` command
 - [x] Typing indicator while generating a response
 - [x] Access can be restricted by specifying a list of allowed users
-- [x] Docker support
+- [x] (NEW!) Docker support
+- [x] (NEW!) Live answer updating as the bot types
+
 
 ## Additional Features - help needed!
-- [ ] Multi-chat support
+- [ ] Multi-chat support (ongoing, see [pull/22](https://github.com/n3d1117/chatgpt-telegram-bot/pull/22))
   - Idea: cache different instances of `ChatGPT3Bot`, one for every chat id (maybe even persist them), so that every user has their own private conversation
-- [ ] Support group chats
+- [ ] Support group chats (ongoing, see [pull/17](https://github.com/n3d1117/chatgpt-telegram-bot/pull/17))
   - Allow the bot to be used in group chats with specific commands
 - [ ] Advanced commands
   - With premade ad-hoc prompts
@@ -49,10 +51,12 @@ Additional optional (but recommended) configuration values:
 ```bash
 ALLOWED_TELEGRAM_USER_IDS="<USER_ID_1>,<USER_ID_2>,..." # Defaults to "*"
 PROXY="<HTTP/HTTPS_PROXY>" # E.g. "http://localhost:8080", defaults to none
+USE_STREAM=false # Defaults to true
 DEBUG=false # Defaults to true
 ```
 * `ALLOWED_TELEGRAM_USER_IDS`: A comma-separated list of Telegram user IDs that are allowed to interact with the bot (use [getidsbot](https://t.me/getidsbot) to find your user ID). **Important**: by default, *everyone* is allowed (`*`)
 * `PROXY`: Proxy to be used when authenticating with OpenAI
+* `USE_STREAM`: Streams the response as the bot types. Set to `false` to only answer once the response is fully generated
 * `DEBUG`: Enable debug logging for the [revChatGpt](https://github.com/acheong08/ChatGPT) package
 
 ### Installing

main.py

@@ -31,7 +31,8 @@ def main():
     }
     telegram_config = {
         'token': os.environ['TELEGRAM_BOT_TOKEN'],
-        'allowed_user_ids': os.environ.get('ALLOWED_TELEGRAM_USER_IDS', '*')
+        'allowed_user_ids': os.environ.get('ALLOWED_TELEGRAM_USER_IDS', '*'),
+        'use_stream': os.environ.get('USE_STREAM', 'true').lower() == 'true'
     }
 
     if os.environ.get('PROXY', None) is not None:

telegram_bot.py

@@ -3,7 +3,8 @@ import logging
 
 import telegram.constants as constants
 from revChatGPT.revChatGPT import AsyncChatbot as ChatGPT3Bot
-from telegram import Update
+from telegram import Update, Message
+from telegram.error import RetryAfter, BadRequest
 from telegram.ext import ApplicationBuilder, ContextTypes, CommandHandler, MessageHandler, filters
 
 
@@ -58,7 +59,7 @@ class ChatGPT3TelegramBot:
         self.gpt3_bot.reset_chat()
         await context.bot.send_message(chat_id=update.effective_chat.id, text="Done!")
 
-    async def send_typing_periodically(self, update: Update, context: ContextTypes.DEFAULT_TYPE, every_seconds):
+    async def send_typing_periodically(self, update: Update, context: ContextTypes.DEFAULT_TYPE, every_seconds: float):
         """
         Sends the typing action periodically to the chat
         """
@@ -78,18 +79,64 @@ class ChatGPT3TelegramBot:
         logging.info(f'New message received from user {update.message.from_user.name}')
 
         # Send "Typing..." action periodically every 4 seconds until the response is received
-        typing_task = asyncio.get_event_loop().create_task(
+        typing_task = context.application.create_task(
             self.send_typing_periodically(update, context, every_seconds=4)
         )
-        response = await self.get_chatgpt_response(update.message.text)
-        typing_task.cancel()
 
-        await context.bot.send_message(
-            chat_id=update.effective_chat.id,
-            reply_to_message_id=update.message.message_id,
-            text=response['message'],
-            parse_mode=constants.ParseMode.MARKDOWN
-        )
+        if self.config['use_stream']:
+            initial_message: Message or None = None
+            chunk_index, chunk_text = (0, '')
+
+            async def message_update(every_seconds: float):
+                """
+                Edits the `initial_message` periodically with the updated text from the latest chunk
+                """
+                while True:
+                    try:
+                        if initial_message is not None and chunk_text != initial_message.text:
+                            await initial_message.edit_text(chunk_text)
+                    except RetryAfter as e:
+                        logging.info(f'Rate limit exceeded, retrying in {e.retry_after} seconds')
+                        await asyncio.sleep(e.retry_after)
+                    except BadRequest:
+                        # Failed to edit message (new=current), ignoring...
+                        pass
+                    except Exception as e:
+                        logging.info(f'Error while editing the message: {str(e)}')
+
+                    await asyncio.sleep(every_seconds)
+
+            # Start task to update the initial message periodically every 0.5 seconds
+            # If you're frequently hitting rate limits, increase this interval
+            message_update_task = context.application.create_task(message_update(every_seconds=0.5))
+
+            # Stream the response
+            async for chunk in await self.gpt3_bot.get_chat_response(update.message.text, output='stream'):
+                if chunk_index == 0 and initial_message is None:
+                    # Sends the initial message, to be edited later with updated text
+                    initial_message = await context.bot.send_message(
+                        chat_id=update.effective_chat.id,
+                        reply_to_message_id=update.message.message_id,
+                        text=chunk['message']
+                    )
+                    typing_task.cancel()
+                chunk_index, chunk_text = (chunk_index + 1, chunk['message'])
+
+            message_update_task.cancel()
+
+            # Final edit, including Markdown formatting
+            await initial_message.edit_text(chunk_text, parse_mode=constants.ParseMode.MARKDOWN)
+
+        else:
+            response = await self.get_chatgpt_response(update.message.text)
+            typing_task.cancel()
+
+            await context.bot.send_message(
+                chat_id=update.effective_chat.id,
+                reply_to_message_id=update.message.message_id,
+                text=response['message'],
+                parse_mode=constants.ParseMode.MARKDOWN
+            )
 
     async def get_chatgpt_response(self, message) -> dict:
         """