Skip to content

Commit 72b5527

Browse files
n3d1117n3d1117
authored
Merge pull request #341 from n3d1117/feature/rate-limit-retry
Auto retry while rate limited using tenacity
2 parents fdd76ee + 5fd70f4 commit 72b5527

File tree

4 files changed

+81
-66
lines changed

4 files changed

+81
-66
lines changed

‎bot/main.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ def main():
1616
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
1717
level=logging.INFO
1818
)
19+
logging.getLogger("httpx").setLevel(logging.WARNING)
1920

2021
# Check if the required environment variables are set
2122
required_values = ['TELEGRAM_BOT_TOKEN', 'OPENAI_API_KEY']

‎bot/openai_helper.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
from datetime import date
1313
from calendar import monthrange
1414

15+
from tenacity import retry, stop_after_attempt, wait_fixed, retry_if_exception_type
16+
1517
# Models can be found here: https://platform.openai.com/docs/models/overview
1618
GPT_3_MODELS = ("gpt-3.5-turbo", "gpt-3.5-turbo-0301", "gpt-3.5-turbo-0613")
1719
GPT_3_16K_MODELS = ("gpt-3.5-turbo-16k", "gpt-3.5-turbo-16k-0613")
@@ -145,6 +147,12 @@ async def get_chat_response_stream(self, chat_id: int, query: str):
145147

146148
yield answer, tokens_used
147149

150+
@retry(
151+
reraise=True,
152+
retry=retry_if_exception_type(openai.error.RateLimitError),
153+
wait=wait_fixed(20),
154+
stop=stop_after_attempt(3)
155+
)
148156
async def __common_get_chat_response(self, chat_id: int, query: str, stream=False):
149157
"""
150158
Request a response from the GPT model.
@@ -190,7 +198,7 @@ async def __common_get_chat_response(self, chat_id: int, query: str, stream=Fals
190198
)
191199

192200
except openai.error.RateLimitError as e:
193-
raise Exception(f"⚠️ _{localized_text('openai_rate_limit', bot_language)}._ ⚠️\n{str(e)}") from e
201+
raise e
194202

195203
except openai.error.InvalidRequestError as e:
196204
raise Exception(f"⚠️ _{localized_text('openai_invalid', bot_language)}._ ⚠️\n{str(e)}") from e

‎bot/telegram_bot.py‎

Lines changed: 68 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -388,84 +388,88 @@ async def prompt(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
388388
total_tokens = 0
389389

390390
if self.config['stream']:
391-
await update.effective_message.reply_chat_action(
392-
action=constants.ChatAction.TYPING,
393-
message_thread_id=get_thread_id(update)
394-
)
391+
async def _reply():
392+
nonlocal total_tokens
393+
await update.effective_message.reply_chat_action(
394+
action=constants.ChatAction.TYPING,
395+
message_thread_id=get_thread_id(update)
396+
)
395397

396-
stream_response = self.openai.get_chat_response_stream(chat_id=chat_id, query=prompt)
397-
i = 0
398-
prev = ''
399-
sent_message = None
400-
backoff = 0
401-
stream_chunk = 0
402-
403-
async for content, tokens in stream_response:
404-
if len(content.strip()) == 0:
405-
continue
406-
407-
stream_chunks = split_into_chunks(content)
408-
if len(stream_chunks) > 1:
409-
content = stream_chunks[-1]
410-
if stream_chunk != len(stream_chunks) - 1:
411-
stream_chunk += 1
412-
try:
413-
await edit_message_with_retry(context, chat_id, str(sent_message.message_id),
414-
stream_chunks[-2])
415-
except:
416-
pass
398+
stream_response = self.openai.get_chat_response_stream(chat_id=chat_id, query=prompt)
399+
i = 0
400+
prev = ''
401+
sent_message = None
402+
backoff = 0
403+
stream_chunk = 0
404+
405+
async for content, tokens in stream_response:
406+
if len(content.strip()) == 0:
407+
continue
408+
409+
stream_chunks = split_into_chunks(content)
410+
if len(stream_chunks) > 1:
411+
content = stream_chunks[-1]
412+
if stream_chunk != len(stream_chunks) - 1:
413+
stream_chunk += 1
414+
try:
415+
await edit_message_with_retry(context, chat_id, str(sent_message.message_id),
416+
stream_chunks[-2])
417+
except:
418+
pass
419+
try:
420+
sent_message = await update.effective_message.reply_text(
421+
message_thread_id=get_thread_id(update),
422+
text=content if len(content) > 0 else "..."
423+
)
424+
except:
425+
pass
426+
continue
427+
428+
cutoff = get_stream_cutoff_values(update, content)
429+
cutoff += backoff
430+
431+
if i == 0:
417432
try:
433+
if sent_message is not None:
434+
await context.bot.delete_message(chat_id=sent_message.chat_id,
435+
message_id=sent_message.message_id)
418436
sent_message = await update.effective_message.reply_text(
419437
message_thread_id=get_thread_id(update),
420-
text=content if len(content) > 0 else "..."
438+
reply_to_message_id=get_reply_to_message_id(self.config, update),
439+
text=content
421440
)
422441
except:
423-
pass
424-
continue
425-
426-
cutoff = get_stream_cutoff_values(update, content)
427-
cutoff += backoff
442+
continue
428443

429-
if i == 0:
430-
try:
431-
if sent_message is not None:
432-
await context.bot.delete_message(chat_id=sent_message.chat_id,
433-
message_id=sent_message.message_id)
434-
sent_message = await update.effective_message.reply_text(
435-
message_thread_id=get_thread_id(update),
436-
reply_to_message_id=get_reply_to_message_id(self.config, update),
437-
text=content
438-
)
439-
except:
440-
continue
444+
elif abs(len(content) - len(prev)) > cutoff or tokens != 'not_finished':
445+
prev = content
441446

442-
elif abs(len(content) - len(prev)) > cutoff or tokens != 'not_finished':
443-
prev = content
447+
try:
448+
use_markdown = tokens != 'not_finished'
449+
await edit_message_with_retry(context, chat_id, str(sent_message.message_id),
450+
text=content, markdown=use_markdown)
444451

445-
try:
446-
use_markdown = tokens != 'not_finished'
447-
await edit_message_with_retry(context, chat_id, str(sent_message.message_id),
448-
text=content, markdown=use_markdown)
452+
except RetryAfter as e:
453+
backoff += 5
454+
await asyncio.sleep(e.retry_after)
455+
continue
449456

450-
except RetryAfter as e:
451-
backoff += 5
452-
await asyncio.sleep(e.retry_after)
453-
continue
457+
except TimedOut:
458+
backoff += 5
459+
await asyncio.sleep(0.5)
460+
continue
454461

455-
except TimedOut:
456-
backoff += 5
457-
await asyncio.sleep(0.5)
458-
continue
462+
except Exception:
463+
backoff += 5
464+
continue
459465

460-
except Exception:
461-
backoff += 5
462-
continue
466+
await asyncio.sleep(0.01)
463467

464-
await asyncio.sleep(0.01)
468+
i += 1
469+
if tokens != 'not_finished':
470+
total_tokens = int(tokens)
465471

466-
i += 1
467-
if tokens != 'not_finished':
468-
total_tokens = int(tokens)
472+
await wrap_with_indicator(update, context, _reply, constants.ChatAction.TYPING)
469473

470474
else:
471475
async def _reply():

‎requirements.txt‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,6 @@ python-dotenv~=1.0.0
22
pydub~=0.25.1
33
tiktoken==0.4.0
44
openai==0.27.8
5-
python-telegram-bot==20.2
5+
python-telegram-bot==20.3
6+
requests~=2.31.0
7+
tenacity==8.2.2

0 commit comments

Comments
 (0)