Merge pull request #100 from Codium-ai/tr/describe_labels

Enhancement of Code Review Functionality

pr_agent/git_providers/git_provider.py

@@ -53,8 +53,11 @@ class GitProvider(ABC):
         pass
 
     @abstractmethod
-    def publish_code_suggestion(self, body: str, relevant_file: str,
+    def publish_code_suggestions(self, code_suggestions: list):
-                                relevant_lines_start: int, relevant_lines_end: int):
+        pass
+
+    @abstractmethod
+    def publish_labels(self, labels):
         pass
 
     @abstractmethod

pr_agent/git_providers/github_provider.py

@@ -13,7 +13,7 @@ from ..algo.utils import load_large_diff
 
 
 class GithubProvider(GitProvider):
-    def __init__(self, pr_url: Optional[str] = None, incremental: Optional[IncrementalPR] = False):
+    def __init__(self, pr_url: Optional[str] = None, incremental=IncrementalPR(False)):
         self.repo_obj = None
         self.installation_id = settings.get("GITHUB.INSTALLATION_ID")
         self.github_client = self._get_github_client()
@@ -146,24 +146,32 @@ class GithubProvider(GitProvider):
     def publish_inline_comments(self, comments: list[dict]):
         self.pr.create_review(commit=self.last_commit_id, comments=comments)
 
-    def publish_code_suggestion(self, body: str,
+    def publish_code_suggestions(self, code_suggestions: list):
-                                relevant_file: str,
+        """
-                                relevant_lines_start: int,
+        Publishes code suggestions as comments on the PR.
-                                relevant_lines_end: int):
+        In practice current APU enables to send only one code suggestion per comment. Might change in the future.
+        """
+        post_parameters_list = []
+        import github.PullRequestComment
+        for suggestion in code_suggestions:
+            body = suggestion['body']
+            relevant_file = suggestion['relevant_file']
+            relevant_lines_start = suggestion['relevant_lines_start']
+            relevant_lines_end = suggestion['relevant_lines_end']
+
             if not relevant_lines_start or relevant_lines_start == -1:
                 if settings.config.verbosity_level >= 2:
-                logging.exception(f"Failed to publish code suggestion, relevant_lines_start is {relevant_lines_start}")
+                    logging.exception(
-            return False
+                        f"Failed to publish code suggestion, relevant_lines_start is {relevant_lines_start}")
+                continue
 
             if relevant_lines_end < relevant_lines_start:
                 if settings.config.verbosity_level >= 2:
                     logging.exception(f"Failed to publish code suggestion, "
                                       f"relevant_lines_end is {relevant_lines_end} and "
                                       f"relevant_lines_start is {relevant_lines_start}")
-            return False
+                continue
 
-        try:
-            import github.PullRequestComment
             if relevant_lines_end > relevant_lines_start:
                 post_parameters = {
                     "body": body,
@@ -181,6 +189,8 @@ class GithubProvider(GitProvider):
                     "line": relevant_lines_start,
                     "side": "RIGHT",
                 }
+
+            try:
                 headers, data = self.pr._requester.requestJsonAndCheck(
                     "POST", f"{self.pr.url}/comments", input=post_parameters
                 )
@@ -306,3 +316,16 @@ class GithubProvider(GitProvider):
         except Exception:
             file_content_str = ""
         return file_content_str
+
+    def publish_labels(self, pr_types):
+        try:
+            label_color_map = {"Bug fix": "1d76db", "Tests": "e99695", "Bug fix with tests": "c5def5", "Refactoring": "bfdadc", "Enhancement": "bfd4f2", "Documentation": "d4c5f9", "Other": "d1bcf9"}
+            post_parameters = []
+            for p in pr_types:
+                color = label_color_map.get(p, "d1bcf9")  # default to "Other" color
+                post_parameters.append({"name": p, "color": color})
+            headers, data = self.pr._requester.requestJsonAndCheck(
+                "PUT", f"{self.pr.issue_url}/labels", input=post_parameters
+            )
+        except:
+            logging.exception("Failed to publish labels")

pr_agent/git_providers/gitlab_provider.py

@@ -135,10 +135,13 @@ class GitLabProvider(GitProvider):
             self.mr.discussions.create({'body': body,
                                         'position': pos_obj})
 
-    def publish_code_suggestion(self, body: str,
+    def publish_code_suggestions(self, code_suggestions: list):
-                                relevant_file: str,
+        for suggestion in code_suggestions:
-                                relevant_lines_start: int,
+            body = suggestion['body']
-                                relevant_lines_end: int):
+            relevant_file = suggestion['relevant_file']
+            relevant_lines_start = suggestion['relevant_lines_start']
+            relevant_lines_end = suggestion['relevant_lines_end']
+
             self.diff_files = self.diff_files if self.diff_files else self.get_diff_files()
             target_file = None
             for file in self.diff_files:
@@ -254,3 +257,6 @@ class GitLabProvider(GitProvider):
 
     def get_user_id(self):
         return None
+
+    def publish_labels(self, labels):
+        pass

pr_agent/servers/help.py

@@ -1,9 +1,9 @@
-commands_text = "> /review [-i]: Request a review of your Pull Request. For an incremental review, which only " \
+commands_text = "> **/review [-i]**: Request a review of your Pull Request. For an incremental review, which only " \
                 "considers changes since the last review, include the '-i' option.\n" \
-                "> /describe: Modify the PR title and description based on the contents of the PR.\n" \
+                "> **/describe**: Modify the PR title and description based on the contents of the PR.\n" \
-                "> /improve: Suggest improvements to the code in the PR. " \
+                "> **/improve**: Suggest improvements to the code in the PR. " \
                 "These will be provided as pull request comments, ready to commit.\n" \
-                "> /ask \\<QUESTION\\>: Pose a question about the PR.\n"
+                "> **/ask \\<QUESTION\\>**: Pose a question about the PR.\n"
 
 
 def bot_help_text(user: str):

pr_agent/settings/pr_description_prompts.toml

@@ -10,9 +10,9 @@ You must use the following JSON schema to format your answer:
       "type": "string",
       "description": "an informative title for the PR, describing its main theme"
   },
-  "Type of PR": {
+  "PR Type": {
       "type": "string",
-      "enum": ["Bug fix", "Tests", "Bug fix with tests", "Refactoring", "Enhancement", "Documentation", "Other"]
+      "description": possible values are: ["Bug fix", "Tests", "Bug fix with tests", "Refactoring", "Enhancement", "Documentation", "Other"]
     },
   "PR Description": {
       "type": "string",

pr_agent/tools/pr_code_suggestions.py

@@ -79,7 +79,6 @@ class PRCodeSuggestions:
 
     def _prepare_pr_code_suggestions(self) -> str:
         review = self.prediction.strip()
-        data = None
         try:
             data = json.loads(review)
         except json.decoder.JSONDecodeError:
@@ -89,6 +88,7 @@ class PRCodeSuggestions:
         return data
 
     def push_inline_code_suggestions(self, data):
+        code_suggestions = []
         for d in data['Code suggestions']:
             if settings.config.verbosity_level >= 2:
                 logging.info(f"suggestion: {d}")
@@ -100,6 +100,16 @@ class PRCodeSuggestions:
             new_code_snippet = d['improved code']
 
             if new_code_snippet:
+                new_code_snippet = self.dedent_code(relevant_file, relevant_lines_start, new_code_snippet)
+
+            body = f"**Suggestion:** {content}\n```suggestion\n" + new_code_snippet + "\n```"
+            code_suggestions.append({'body': body,'relevant_file': relevant_file,
+                                     'relevant_lines_start': relevant_lines_start,
+                                     'relevant_lines_end': relevant_lines_end})
+
+        self.git_provider.publish_code_suggestions(code_suggestions)
+
+    def dedent_code(self, relevant_file, relevant_lines_start, new_code_snippet):
         try:  # dedent code snippet
             self.diff_files = self.git_provider.diff_files if self.git_provider.diff_files \
                 else self.git_provider.get_diff_files()
@@ -119,8 +129,4 @@ class PRCodeSuggestions:
             if settings.config.verbosity_level >= 2:
                 logging.info(f"Could not dedent code snippet for file {relevant_file}, error: {e}")
 
-            body = f"**Suggestion:** {content}\n```suggestion\n" + new_code_snippet + "\n```"
+        return new_code_snippet
-            self.git_provider.publish_code_suggestion(body=body,
-                                                      relevant_file=relevant_file,
-                                                      relevant_lines_start=relevant_lines_start,
-                                                      relevant_lines_end=relevant_lines_end)

pr_agent/tools/pr_description.py

@@ -42,13 +42,14 @@ class PRDescription:
         logging.info('Getting AI prediction...')
         self.prediction = await self._get_prediction()
         logging.info('Preparing answer...')
-        pr_title, pr_body, markdown_text = self._prepare_pr_answer()
+        pr_title, pr_body, pr_types, markdown_text = self._prepare_pr_answer()
         if settings.config.publish_output:
             logging.info('Pushing answer...')
             if settings.pr_description.publish_description_as_comment:
                 self.git_provider.publish_comment(markdown_text)
             else:
                 self.git_provider.publish_description(pr_title, pr_body)
+                self.git_provider.publish_labels(pr_types)
             self.git_provider.remove_initial_comment()
         return ""
 
@@ -73,6 +74,9 @@ class PRDescription:
             markdown_text += f"## {key}\n\n"
             markdown_text += f"{value}\n\n"
         pr_body = ""
+        pr_types = []
+        if 'PR Type' in data:
+            pr_types = data['PR Type'].split(',')
         title = data['PR Title']
         del data['PR Title']
         for key, value in data.items():
@@ -83,4 +87,4 @@ class PRDescription:
                 pr_body += f"**{value}**\n\n___\n"
         if settings.config.verbosity_level >= 2:
             logging.info(f"title:\n{title}\n{pr_body}")
-        return title, pr_body, markdown_text
+        return title, pr_body, pr_types, markdown_text