Merge pull request Codium-ai#431 from Codium-ai/hl/type_vs_labels

Refactoring PR Labels Handling and Display
yochail · Nov 6, 2023 · 0d9c042 · 0d9c042
2 parents 7e1bf20 + 6e50b44
commit 0d9c042
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 26 deletions.
diff --git a/pr_agent/settings/configuration.toml b/pr_agent/settings/configuration.toml
@@ -38,6 +38,7 @@ add_original_user_description=false
 keep_original_user_title=false
 use_bullet_points=true
 extra_instructions = ""
+enable_pr_type=true
 
 # markers
 use_description_markers=false

diff --git a/pr_agent/settings/pr_custom_labels.toml b/pr_agent/settings/pr_custom_labels.toml
@@ -39,7 +39,6 @@ PR Type:
 {{ custom_labels_examples }}
 {%- else %}
   - Bug fix
-  - Tests
 {%- endif %}
 ```
 

diff --git a/pr_agent/settings/pr_description_prompts.toml b/pr_agent/settings/pr_description_prompts.toml
@@ -19,22 +19,22 @@ PR Title:
   type: string
   description: an informative title for the PR, describing its main theme
 PR Type:
-  type: array
+  type: string
+  enum:
+    - Bug fix
+    - Tests
+    - Refactoring
+    - Enhancement
+    - Documentation
+    - Other
 {%- if enable_custom_labels %}
-  description: One or more labels that describe the PR type. Don't output the description in the parentheses.
-{%- endif %}
+PR Labels:
+  type: array
+  description: One or more labels that describe the PR labels. Don't output the description in the parentheses.
   items:
     type: string
     enum:
-{%- if enable_custom_labels %}
 {{ custom_labels }}
-{%- else %}
-      - Bug fix
-      - Tests
-      - Refactoring
-      - Enhancement
-      - Documentation
-      - Other
 {%- endif %}
 PR Description:
   type: string
@@ -60,10 +60,10 @@ Example output:
 PR Title: |-
   ...
 PR Type:
+  ...
 {%- if enable_custom_labels %}
+PR Labels:
 {{ custom_labels_examples }}
-{%- else %}
-  - Bug fix
 {%- endif %}
 PR Description: |-
   ...

diff --git a/pr_agent/settings/pr_reviewer_prompts.toml b/pr_agent/settings/pr_reviewer_prompts.toml
@@ -51,22 +51,13 @@ PR Analysis:
     description: summary of the PR in 2-3 sentences.
   Type of PR:
     type: string
-{%- if enable_custom_labels %}
-  description: One or more labels that describe the PR type. Don't output the description in the parentheses.
-{%- endif %}
-  items:
-    type: string
     enum:
-{%- if enable_custom_labels %}
-{{ custom_labels }}
-{%- else %}
       - Bug fix
       - Tests
       - Refactoring
       - Enhancement
       - Documentation
       - Other
-{%- endif %}
 {%- if require_score %}
   Score:
     type: int

diff --git a/pr_agent/tools/pr_description.py b/pr_agent/tools/pr_description.py
@@ -172,12 +172,16 @@ def _prepare_labels(self) -> List[str]:
         pr_types = []
 
         # If the 'PR Type' key is present in the dictionary, split its value by comma and assign it to 'pr_types'
-        if 'PR Type' in self.data:
+        if 'PR Labels' in self.data:
+            if type(self.data['PR Labels']) == list:
+                pr_types = self.data['PR Labels']
+            elif type(self.data['PR Labels']) == str:
+                pr_types = self.data['PR Labels'].split(',')
+        elif 'PR Type' in self.data:
             if type(self.data['PR Type']) == list:
                 pr_types = self.data['PR Type']
             elif type(self.data['PR Type']) == str:
                 pr_types = self.data['PR Type'].split(',')
-
         return pr_types
 
     def _prepare_pr_answer_with_markers(self) -> Tuple[str, str]:
@@ -223,6 +227,11 @@ def _prepare_pr_answer(self) -> Tuple[str, str]:
 
         # Iterate over the dictionary items and append the key and value to 'markdown_text' in a markdown format
         markdown_text = ""
+        # Don't display 'PR Labels'
+        if 'PR Labels' in self.data:
+            self.data.pop('PR Labels')
+        if not get_settings().pr_description.enable_pr_type:
+            self.data.pop('PR Type')
         for key, value in self.data.items():
             markdown_text += f"## {key}\n\n"
             markdown_text += f"{value}\n\n"

diff --git a/pr_agent/tools/pr_reviewer.py b/pr_agent/tools/pr_reviewer.py
@@ -156,7 +156,7 @@ async def _get_prediction(self, model: str) -> str:
         variables["diff"] = self.patches_diff  # update diff
 
         environment = Environment(undefined=StrictUndefined)
-        set_custom_labels(variables)
+        # set_custom_labels(variables)
         system_prompt = environment.from_string(get_settings().pr_review_prompt.system).render(variables)
         user_prompt = environment.from_string(get_settings().pr_review_prompt.user).render(variables)