Refactoring PR label handling, enhancing language detection, and updating CLI commands

pr_agent/algo/utils.py

@@ -336,12 +336,15 @@ def set_custom_labels(variables):
         labels_list = f"      - {labels_list}" if labels_list else ""
         variables["custom_labels"] = labels_list
         return
-    final_labels = ""
+    #final_labels = ""
+    #for k, v in labels.items():
+    #    final_labels += f"      - {k} ({v['description']})\n"
+    #variables["custom_labels"] = final_labels
+    #variables["custom_labels_examples"] = f"      - {list(labels.keys())[0]}"
+    variables["custom_labels_class"] = "class Label(str, Enum):"
     for k, v in labels.items():
-        final_labels += f"      - {k} ({v['description']})\n"
-    variables["custom_labels"] = final_labels
-    variables["custom_labels_examples"] = f"      - {list(labels.keys())[0]}"
-
+        description = v['description'].strip('\n').replace('\n', '\\n')
+        variables["custom_labels_class"] += f"\n    {k.lower().replace(' ', '_')} = '{k}' # {description}"
 
 def get_user_labels(current_labels: List[str] = None):
     """

pr_agent/cli.py

@@ -21,18 +21,22 @@ def run(inargs=None):
 - cli.py --issue_url=... similar_issue
 
 Supported commands:
--review / review_pr - Add a review that includes a summary of the PR and specific suggestions for improvement.
+- review / review_pr - Add a review that includes a summary of the PR and specific suggestions for improvement.
 
--ask / ask_question [question] - Ask a question about the PR.
+- ask / ask_question [question] - Ask a question about the PR.
 
--describe / describe_pr - Modify the PR title and description based on the PR's contents.
+- describe / describe_pr - Modify the PR title and description based on the PR's contents.
 
--improve / improve_code - Suggest improvements to the code in the PR as pull request comments ready to commit.
+- improve / improve_code - Suggest improvements to the code in the PR as pull request comments ready to commit.
 Extended mode ('improve --extended') employs several calls, and provides a more thorough feedback
 
--reflect - Ask the PR author questions about the PR.
+- reflect - Ask the PR author questions about the PR.
 
--update_changelog - Update the changelog based on the PR's contents.
+- update_changelog - Update the changelog based on the PR's contents.
+
+- add_docs
+
+- generate_labels
 
 
 Configuration:

pr_agent/settings/custom_labels.toml

@@ -3,16 +3,16 @@ enable_custom_labels=false
 
 ## template for custom labels
 #[custom_labels."Bug fix"]
-#description = "Fixes a bug in the code"
+#description = """Fixes a bug in the code"""
 #[custom_labels."Tests"]
-#description = "Adds or modifies tests"
+#description = """Adds or modifies tests"""
 #[custom_labels."Bug fix with tests"]
-#description = "Fixes a bug in the code and adds or modifies tests"
+#description = """Fixes a bug in the code and adds or modifies tests"""
 #[custom_labels."Refactoring"]
-#description = "Code refactoring without changing functionality"
+#description = """Code refactoring without changing functionality"""
 #[custom_labels."Enhancement"]
-#description = "Adds new features or functionality"
+#description = """Adds new features or functionality"""
 #[custom_labels."Documentation"]
-#description = "Adds or modifies documentation"
+#description = """Adds or modifies documentation"""
 #[custom_labels."Other"]
-#description = "Other changes that do not fit in any of the above categories"
+#description = """Other changes that do not fit in any of the above categories"""

pr_agent/settings/pr_custom_labels.toml

@@ -1,8 +1,10 @@
 [pr_custom_labels_prompt]
-system="""You are CodiumAI-PR-Reviewer, a language model designed to review git pull requests.
-Your task is to label the type of the PR content.
-- Make sure not to focus the new PR code (the '+' lines).
-- If needed, each YAML output should be in block scalar format ('|-')
+system="""You are CodiumAI-PR-Reviewer, a language model designed to review git Pull Requests (PR).
+Your task is to provide labels that describe the PR content.
+{%- if enable_custom_labels %}
+Thoroughly read the labels name and the provided description, and decide whether the label is relevant to the PR.
+{%- endif %}
+
 {%- if extra_instructions %}
 
 Extra instructions from the user:
@@ -11,38 +13,35 @@ Extra instructions from the user:
 '
 {% endif %}
 
-You must use the following YAML schema to format your answer:
-```yaml
-PR Type:
-  type: array
-{%- 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:
+The output must be a YAML object equivalent to type $Labels, according to the following Pydantic definitions:
+'
 {%- if enable_custom_labels %}
-{{ custom_labels }}
+
+{{ custom_labels_class }}
+
 {%- else %}
-      - Bug fix
-      - Tests
-      - Refactoring
-      - Enhancement
-      - Documentation
-      - Other
+class Label(str, Enum):
+    bug_fix = "Bug fix"
+    tests = "Tests"
+    refactoring = "Refactoring"
+    enhancement = "Enhancement"
+    documentation = "Documentation"
+    other = "Other"
 {%- endif %}
 
+class Labels(BaseModel):
+    labels: List[Label] = Field(min_items=0, description="custom labels that describe the PR. Return the label value, not the name.")
+'
+
+
 Example output:
 ```yaml
-PR Type:
-{%- if enable_custom_labels %}
-{{ custom_labels_examples }}
-{%- else %}
-  - Bug fix
-{%- endif %}
+labels:
+- ...
+- ...
 ```
 
-Make sure to output a valid YAML. Don't repeat the prompt in the answer, and avoid outputting the 'type' and 'description' fields.
+Answer should be a valid YAML, and nothing else.
 """
 
 user="""PR Info:

pr_agent/settings/pr_description_prompts.toml

@@ -13,67 +13,57 @@ Extra instructions from the user:
 '
 {% endif %}
 
-You must use the following YAML schema to format your answer:
-```yaml
-PR Title:
-  type: string
-  description: an informative title for the PR, describing its main theme
-PR Type:
-  type: string
-  enum:
-    - Bug fix
-    - Tests
-    - Refactoring
-    - Enhancement
-    - Documentation
-    - Other
+
+The output must be a YAML object equivalent to type $PRDescription, according to the following Pydantic definitions:
+'
+class PRType(str, Enum):
+    bug_fix = "Bug fix"
+    tests = "Tests"
+    refactoring = "Refactoring"
+    enhancement = "Enhancement"
+    documentation = "Documentation"
+    other = "Other"
+
 {%- if enable_custom_labels %}
-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:
-{{ custom_labels }}
+
+{{ custom_labels_class }}
 {%- endif %}
-PR Description:
-  type: string
-  description: an informative and concise description of the PR.
-  {%- if use_bullet_points %} Use bullet points. {% endif %}
-PR Main Files Walkthrough:
-  type: array
-  maxItems: 10
-  description: |-
-    a walkthrough of the PR changes. Review main files, and shortly describe the changes in each file (up to 10 most important files).
-  items:
-    filename:
-      type: string
-      description: the relevant file full path
-    changes in file:
-      type: string
-      description: minimal and concise description of the changes in the relevant file
-```
+
+class FileWalkthrough(BaseModel):
+    filename: str = Field(description="the relevant file full path")
+    changes_in_file: str = Field(description="minimal and concise description of the changes in the relevant file")
+
+Class PRDescription(BaseModel):
+    title: str = Field(description="an informative title for the PR, describing its main theme")
+    type: List[PRType] = Field(description="one or more types that describe the PR type. . Return the label value, not the name.")
+    description: str = Field(description="an informative and concise description of the PR. {%- if use_bullet_points %} Use bullet points. {% endif %}")
+{%- if enable_custom_labels %}
+    labels: List[Label] = Field(min_items=0, description="custom labels that describe the PR. Return the label value, not the name.")
+{%- endif %}
+    main_files_walkthrough: List[FileWalkthrough] = Field(max_items=10)
+'
 
 
 Example output:
 ```yaml
-PR Title: |-
-  ...
-PR Type:
+title: |-
   ...
+type:
+- ...
+- ...
 {%- if enable_custom_labels %}
-PR Labels:
+labels:
 - ...
 - ...
 {%- endif %}
-PR Description: |-
+description: |-
   ...
-PR Main Files Walkthrough:
-  - ...
-  - ...
+main_files_walkthrough:
+- ...
+- ...
 ```
 
-Make sure to output a valid YAML. Don't repeat the prompt in the answer, and avoid outputting the 'type' and 'description' fields.
+Answer should be a valid YAML, and nothing else. Each YAML output MUST be after a newline, with proper indent, and block scalar indicator ('|-')
 """
 
 user="""PR Info:

pr_agent/tools/pr_description.py

@@ -44,8 +44,7 @@ def __init__(self, pr_url: str, args: list = None):
             "extra_instructions": get_settings().pr_description.extra_instructions,
             "commit_messages_str": self.git_provider.get_commit_messages(),
             "enable_custom_labels": get_settings().config.enable_custom_labels,
-            "custom_labels": "",
-            "custom_labels_examples": "",
+            "custom_labels_class": "",  # will be filled if necessary in 'set_custom_labels' function
         }
 
         self.user_description = self.git_provider.get_user_description()
@@ -172,16 +171,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 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(',')
+        if 'labels' in self.data:
+            if type(self.data['labels']) == list:
+                pr_types = self.data['labels']
+            elif type(self.data['labels']) == str:
+                pr_types = self.data['labels'].split(',')
+        elif 'type' in self.data:
+            if type(self.data['type']) == list:
+                pr_types = self.data['type']
+            elif type(self.data['type']) == str:
+                pr_types = self.data['type'].split(',')
         return pr_types
 
     def _prepare_pr_answer_with_markers(self) -> Tuple[str, str]:
@@ -193,12 +192,12 @@ def _prepare_pr_answer_with_markers(self) -> Tuple[str, str]:
         else:
             ai_header = ""
 
-        ai_type = self.data.get('PR Type')
+        ai_type = self.data.get('type')
         if ai_type and not re.search(r'<!--\s*pr_agent:type\s*-->', body):
             pr_type = f"{ai_header}{ai_type}"
             body = body.replace('pr_agent:type', pr_type)
 
-        ai_summary = self.data.get('PR Description')
+        ai_summary = self.data.get('description')
         if ai_summary and not re.search(r'<!--\s*pr_agent:summary\s*-->', body):
             summary = f"{ai_header}{ai_summary}"
             body = body.replace('pr_agent:summary', summary)
@@ -228,16 +227,16 @@ 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 and self.git_provider.is_supported("get_labels"):
-            self.data.pop('PR Labels')
+        if 'labels' in self.data and self.git_provider.is_supported("get_labels"):
+            self.data.pop('labels')
         if not get_settings().pr_description.enable_pr_type:
-            self.data.pop('PR Type')
+            self.data.pop('type')
         for key, value in self.data.items():
             markdown_text += f"## {key}\n\n"
             markdown_text += f"{value}\n\n"
 
         # Remove the 'PR Title' key from the dictionary
-        ai_title = self.data.pop('PR Title', self.vars["title"])
+        ai_title = self.data.pop('title', self.vars["title"])
         if get_settings().pr_description.keep_original_user_title:
             # Assign the original PR title to the 'title' variable
             title = self.vars["title"]
@@ -256,7 +255,7 @@ def _prepare_pr_answer(self) -> Tuple[str, str]:
                     pr_body += "<details> <summary>files:</summary>\n\n"
                 for file in value:
                     filename = file['filename'].replace("'", "`")
-                    description = file['changes in file']
+                    description = file['changes_in_file']
                     pr_body += f'- `{filename}`: {description}\n'
                 if self.git_provider.is_supported("gfm_markdown"):
                     pr_body +="</details>\n"

pr_agent/tools/pr_generate_labels.py

@@ -43,9 +43,8 @@ def __init__(self, pr_url: str, args: list = None):
             "use_bullet_points": get_settings().pr_description.use_bullet_points,
             "extra_instructions": get_settings().pr_description.extra_instructions,
             "commit_messages_str": self.git_provider.get_commit_messages(),
-            "custom_labels": "",
-            "custom_labels_examples": "",
             "enable_custom_labels": get_settings().config.enable_custom_labels,
+            "custom_labels_class": "",  # will be filled if necessary in 'set_custom_labels' function
         }
 
         # Initialize the token handler
@@ -159,11 +158,11 @@ def _prepare_data(self):
     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 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(',')
+        # If the 'labels' key is present in the dictionary, split its value by comma and assign it to 'pr_types'
+        if 'labels' in self.data:
+            if type(self.data['labels']) == list:
+                pr_types = self.data['labels']
+            elif type(self.data['labels']) == str:
+                pr_types = self.data['labels'].split(',')
 
         return pr_types