diff --git a/aider/coder.py b/aider/coder.py index 10e356df1..57d364c77 100755 --- a/aider/coder.py +++ b/aider/coder.py @@ -2,10 +2,10 @@ import os import sys -import time import traceback from pathlib import Path +import backoff import git import openai import requests @@ -393,23 +393,19 @@ class Coder: return prompts.added_files.format(fnames=", ".join(mentioned_rel_fnames)) + @backoff.on_exception( + backoff.expo, + (RateLimitError, requests.exceptions.ConnectionError), + max_tries=5, + on_backoff=lambda details: print(f"Retry in {details['wait']} seconds."), + ) def send_with_retries(self, model, messages): - while True: - try: - return openai.ChatCompletion.create( - model=model, - messages=messages, - temperature=0, - stream=True, - ) - except RateLimitError as err: - self.io.tool_error(f"RateLimitError: {err}") - except requests.exceptions.ConnectionError as err: - self.io.tool_error(f"ConnectionError: {err}") - - retry_after = 1 - self.io.tool_error(f"Retry in {retry_after} seconds.") - time.sleep(retry_after) + return openai.ChatCompletion.create( + model=model, + messages=messages, + temperature=0, + stream=True, + ) def send(self, messages, model=None, silent=False): if not model: diff --git a/tests/test_coder.py b/tests/test_coder.py index eee43897e..cfc45db20 100644 --- a/tests/test_coder.py +++ b/tests/test_coder.py @@ -121,8 +121,8 @@ class TestCoder(unittest.TestCase): self.assertEqual(result, 'a good "commit message"') @patch("aider.coder.openai.ChatCompletion.create") - @patch("aider.coder.time.sleep") - def test_send_with_retries_rate_limit_error(self, mock_sleep, mock_chat_completion_create): + @patch("builtins.print") + def test_send_with_retries_rate_limit_error(self, mock_print, mock_chat_completion_create): # Mock the IO object mock_io = MagicMock() @@ -139,12 +139,12 @@ class TestCoder(unittest.TestCase): # Call the send_with_retries method coder.send_with_retries("model", ["message"]) - # Assert that time.sleep was called once - mock_sleep.assert_called_once() + # Assert that print was called once + mock_print.assert_called_once() @patch("aider.coder.openai.ChatCompletion.create") - @patch("aider.coder.time.sleep") - def test_send_with_retries_connection_error(self, mock_sleep, mock_chat_completion_create): + @patch("builtins.print") + def test_send_with_retries_connection_error(self, mock_print, mock_chat_completion_create): # Mock the IO object mock_io = MagicMock() @@ -161,8 +161,8 @@ class TestCoder(unittest.TestCase): # Call the send_with_retries method coder.send_with_retries("model", ["message"]) - # Assert that time.sleep was called once - mock_sleep.assert_called_once() + # Assert that print was called once + mock_print.assert_called_once() if __name__ == "__main__":