s1

pr_agent/algo/utils.py

@@ -336,11 +336,14 @@ 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 Labels(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]}"
+        variables["custom_labels_class"] += f"\n    {k.lower().replace(' ','_')} = '{k}' # {v['description']}"
 
 
 def get_user_labels(current_labels: List[str] = None):

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(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")
+    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[Labels] = Field(description="one or more custom labels that describe the PR")
+{%- 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 @@ class PRDescription:
             "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'
         }
 
         self.user_description = self.git_provider.get_user_description()
@@ -172,16 +171,16 @@ class PRDescription:
         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 @@ class PRDescription:
         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 @@ class PRDescription:
         # 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 @@ class PRDescription:
                     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"