From 0df054295828c8b1395a1b084a89b99688200d0b Mon Sep 17 00:00:00 2001
From: mrT23 <tal.r@codium.ai>
Date: Mon, 13 Nov 2023 15:55:35 +0200
Subject: [PATCH] prompt

---
 pr_agent/algo/utils.py                        | 2 +-
 pr_agent/settings/pr_custom_labels.toml       | 4 ++--
 pr_agent/settings/pr_description_prompts.toml | 6 +++---
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py
index 8d74f48d..bd91477c 100644
--- a/pr_agent/algo/utils.py
+++ b/pr_agent/algo/utils.py
@@ -341,7 +341,7 @@ def set_custom_labels(variables):
     #    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(Enum):"
+    variables["custom_labels_class"] = "class Label(str, Enum):"
     for k, v in labels.items():
         description = v['description'].strip('\n').replace('\n', '\\n')
         variables["custom_labels_class"] += f"\n    {k.lower().replace(' ', '_')} = '{k}' # {description}"
diff --git a/pr_agent/settings/pr_custom_labels.toml b/pr_agent/settings/pr_custom_labels.toml
index 821ec8d2..f61a208c 100644
--- a/pr_agent/settings/pr_custom_labels.toml
+++ b/pr_agent/settings/pr_custom_labels.toml
@@ -20,7 +20,7 @@ The output must be a YAML object equivalent to type $Labels, according to the fo
 {{ custom_labels_class }}
 
 {%- else %}
-class Label(Enum):
+class Label(str, Enum):
     bug_fix = "Bug fix"
     tests = "Tests"
     refactoring = "Refactoring"
@@ -30,7 +30,7 @@ class Label(Enum):
 {%- endif %}
 
 class Labels(BaseModel):
-    labels: List[Label] = Field(min_items=0, description="custom labels that describe the PR")
+    labels: List[Label] = Field(min_items=0, description="custom labels that describe the PR. Return the label value, not the name.")
 '
 
 
diff --git a/pr_agent/settings/pr_description_prompts.toml b/pr_agent/settings/pr_description_prompts.toml
index d59e585d..514a1991 100644
--- a/pr_agent/settings/pr_description_prompts.toml
+++ b/pr_agent/settings/pr_description_prompts.toml
@@ -16,7 +16,7 @@ Extra instructions from the user:
 
 The output must be a YAML object equivalent to type $PRDescription, according to the following Pydantic definitions:
 '
-Class PRType(Enum):
+class PRType(str, Enum):
     bug_fix = "Bug fix"
     tests = "Tests"
     refactoring = "Refactoring"
@@ -35,10 +35,10 @@ class FileWalkthrough(BaseModel):
 
 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")
+    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")
+    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)
 '