Handle message streaming

README.md

@@ -6,6 +6,7 @@ Your self-hosted telegram bot for chatting with ChatGPT.
 
 - GPT-3.5 chatbot.
 - Handle voice messages
+- Message streaming
 - Toggleable answer on mentions, for a better integration in group chats.
 
 <div style="text-align:center">

src/bot.py

@@ -44,9 +44,7 @@ def run_bot() -> None:
  )
  )
  application.add_handler(
- MessageHandler(
- filters.VOICE & user_filter, handlers.voice_handler
- )
+ MessageHandler(filters.VOICE & user_filter, handlers.voice_handler)
  )
  application.add_handler(
  CommandHandler("reset", handlers.reset_handler, filters=user_filter)

src/handlers.py

@@ -4,6 +4,7 @@
 import pydub
 from telegram import Update, User
 from telegram.constants import ParseMode
+from telegram.error import BadRequest
 from telegram.ext import CallbackContext
 
 import config
@@ -25,8 +26,42 @@ async def message_handler(
  await update.message.reply_text("⚠ Conversation reset due to timeout.")
 
  await update.message.chat.send_action(action="typing")
- answer: str = await openai_instance.complete(message)
- await update.message.reply_text(answer)
+ gen = openai_instance.complete(message)
+ prev_answer = ""
+ message_sent = False
+ async for gen_item in gen:
+ status, answer = gen_item
+ if not message_sent:
+ message_sent = True
+ try:
+ sent_message = await update.message.reply_text(answer)
+ except BadRequest as e:
+ if str(e).startswith(
+ "Message must be non-empty"
+ ): # first answer chunk from openai was empty
+ message_sent = False
+ continue
+ else:
+ sent_message = await update.message.reply_text(answer)
+ continue
+ if abs(len(answer) - len(prev_answer)) < 10 and status != "done":
+ continue
+
+ try:
+ await context.bot.edit_message_text(
+ answer, chat_id=sent_message.chat_id, message_id=sent_message.message_id
+ )
+ except BadRequest as e:
+ if str(e).startswith("Message is not modified"):
+ continue
+ else:
+ await context.bot.edit_message_text(
+ answer,
+ chat_id=sent_message.chat_id,
+ message_id=sent_message.message_id,
+ )
+
+ prev_answer = answer
 
 
 async def voice_handler(update: Update, context: CallbackContext) -> None:

src/openai_utils.py

@@ -17,16 +17,22 @@ def __init__(self, openai_api_key: str) -> None:
  self.model = "gpt-3.5-turbo"
  self.discussions = Discussion()
 
- async def complete(self, message: str) -> str:
+ async def complete(self, message: str):
  self.discussions.add_message(Author.USER, message)
  response = await openai.ChatCompletion.acreate(
  model=self.model,
  messages=self.discussions.get_messages(),
+ stream=True,
  **OPENAI_COMPLETION_OPTIONS,
  )
- answer: str = response.choices[0].message["content"]
+ answer = ""
+ async for response_part in response:
+ delta = response_part.choices[0].delta
+ if "content" in delta:
+ answer += delta.content
+ yield "waiting", answer
  self.discussions.add_message(Author.ASSISTANT, answer)
- return answer
+ yield "done", answer
 
  def reset(self) -> None:
  self.discussions.reset_discussion()