From 1229fba346664d5d347a15571a0066f8e55ea9e6 Mon Sep 17 00:00:00 2001
From: idavidov <idavidov@gmail.com>
Date: Tue, 25 Jul 2023 16:37:13 +0300
Subject: [PATCH] +  settings.github.ratelimit_retries setup in 
 configuration.toml

---
 pr_agent/algo/pr_processing.py            |  7 ++--
 pr_agent/git_providers/git_provider.py    |  5 +++
 pr_agent/git_providers/github_provider.py | 44 +++++++++++++----------
 3 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/pr_agent/algo/pr_processing.py b/pr_agent/algo/pr_processing.py
index 59b342a0..94cc1a00 100644
--- a/pr_agent/algo/pr_processing.py
+++ b/pr_agent/algo/pr_processing.py
@@ -10,8 +10,7 @@ from pr_agent.algo.token_handler import TokenHandler
 from pr_agent.algo.utils import load_large_diff
 from pr_agent.config_loader import settings
 from pr_agent.git_providers.git_provider import GitProvider
-from github import GithubException
-
+from git_provider import RateLimitExceeded
 
 DELETED_FILES_ = "Deleted files:\n"
 
@@ -43,8 +42,8 @@ def get_pr_diff(git_provider: GitProvider, token_handler: TokenHandler, model: s
 
     try:
         diff_files = list(git_provider.get_diff_files())
-    except GithubException.RateLimitExceededException as e:
-        logging.error(f"Rate limit exceeded for GitHub API. original message {e}")
+    except RateLimitExceededException as e:
+        logging.error(f"Rate limit exceeded for git provider API. original message {e}")
         raise
 
     # get pr languages
diff --git a/pr_agent/git_providers/git_provider.py b/pr_agent/git_providers/git_provider.py
index 3f7c1ef2..97dbe8ab 100644
--- a/pr_agent/git_providers/git_provider.py
+++ b/pr_agent/git_providers/git_provider.py
@@ -136,3 +136,8 @@ class IncrementalPR:
         self.commits_range = None
         self.first_new_commit_sha = None
         self.last_seen_commit_sha = None
+
+
+class RateLimitExceeded(Exception):
+    """Raised when the git provider API rate limit has been exceeded."""
+    pass
\ No newline at end of file
diff --git a/pr_agent/git_providers/github_provider.py b/pr_agent/git_providers/github_provider.py
index 64e4ab46..a9979a78 100644
--- a/pr_agent/git_providers/github_provider.py
+++ b/pr_agent/git_providers/github_provider.py
@@ -4,6 +4,7 @@ from typing import Optional, Tuple
 from urllib.parse import urlparse
 
 from github import AppAuthentication, Github, Auth, GithubException
+from git_provider import RateLimitExceeded
 
 from pr_agent.config_loader import settings
 
@@ -81,26 +82,31 @@ class GithubProvider(GitProvider):
     @retry(exceptions=(GithubException.RateLimitExceededException),
            tries=settings.github.ratelimit_retries, delay=2, backoff=2, jitter=(1, 3))
     def get_diff_files(self) -> list[FilePatchInfo]:
-        files = self.get_files()
-        diff_files = []
-        for file in files:
-            if is_valid_file(file.filename):
-                new_file_content_str = self._get_pr_file_content(file, self.pr.head.sha)
-                patch = file.patch
-                if self.incremental.is_incremental and self.file_set:
-                    original_file_content_str = self._get_pr_file_content(file, self.incremental.last_seen_commit_sha)
-                    patch = load_large_diff(file,
-                                            new_file_content_str,
-                                            original_file_content_str,
-                                             None)
-                    self.file_set[file.filename] = patch
-                else:
-                    original_file_content_str = self._get_pr_file_content(file, self.pr.base.sha)
+        try:
+            files = self.get_files()
+            diff_files = []
+            for file in files:
+                if is_valid_file(file.filename):
+                    new_file_content_str = self._get_pr_file_content(file, self.pr.head.sha)
+                    patch = file.patch
+                    if self.incremental.is_incremental and self.file_set:
+                        original_file_content_str = self._get_pr_file_content(file,
+                                                                              self.incremental.last_seen_commit_sha)
+                        patch = load_large_diff(file,
+                                                new_file_content_str,
+                                                original_file_content_str,
+                                                None)
+                        self.file_set[file.filename] = patch
+                    else:
+                        original_file_content_str = self._get_pr_file_content(file, self.pr.base.sha)
 
-                diff_files.append(
-                    FilePatchInfo(original_file_content_str, new_file_content_str, patch, file.filename))
-        self.diff_files = diff_files
-        return diff_files
+                    diff_files.append(
+                        FilePatchInfo(original_file_content_str, new_file_content_str, patch, file.filename))
+            self.diff_files = diff_files
+            return diff_files
+        except GithubException.RateLimitExceededException as e:
+            logging.error(f"Rate limit exceeded for GitHub API. Original message: {e}")
+            raise RateLimitExceeded("Rate limit exceeded for GitHub API.") from e
 
     def publish_description(self, pr_title: str, pr_body: str):
         self.pr.edit(title=pr_title, body=pr_body)