Skip to content

Commit

Permalink
fix: added retry_anthropic decorator for rate limits
Browse files Browse the repository at this point in the history
  • Loading branch information
ErikBjare committed Aug 14, 2024
1 parent 123344f commit d37d7fb
Showing 1 changed file with 22 additions and 0 deletions.
22 changes: 22 additions & 0 deletions gptme/llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,10 @@
import shutil
import sys
from collections.abc import Generator, Iterator
from functools import wraps
from time import sleep

import anthropic
from anthropic import Anthropic
from openai import AzureOpenAI, OpenAI
from rich import print
Expand Down Expand Up @@ -88,6 +91,24 @@ def _chat_complete_openai(messages: list[Message], model: str) -> str:
return content


def retry_anthropic(func):
@wraps(func)
def wrapper(*args, _retries=0, **kwargs):
try:
return func(*args, **kwargs)
except anthropic.RateLimitError:
if _retries >= 5:
raise
# exponential backoff
backoff = 5 + 2**_retries
print(f"Rate limit exceeded. Retrying in {backoff} seconds...")
sleep(backoff)
return wrapper(*args, **kwargs, _retries=_retries + 1)

return wrapper


@retry_anthropic
def _chat_complete_anthropic(messages: list[Message], model: str) -> str:
assert anthropic_client, "LLM not initialized"
messages, system_message = _transform_system_messages_anthropic(messages)
Expand Down Expand Up @@ -179,6 +200,7 @@ def _stream_openai(messages: list[Message], model: str) -> Generator[str, None,
logger.debug(f"Stop reason: {stop_reason}")


@retry_anthropic
def _stream_anthropic(
messages: list[Message], model: str
) -> Generator[str, None, None]:
Expand Down

0 comments on commit d37d7fb

Please sign in to comment.