diff --git a/pr_agent/tools/pr_reviewer.py b/pr_agent/tools/pr_reviewer.py
index 6210b188..bc9d503d 100644
--- a/pr_agent/tools/pr_reviewer.py
+++ b/pr_agent/tools/pr_reviewer.py
@@ -252,27 +252,8 @@ class PRReviewer:
             else:
                 markdown_text += actions_help_text
 
-        if (get_settings().pr_reviewer.enable_review_labels_security or
-                get_settings().pr_reviewer.enable_review_labels_effort):
-            try:
-                review_labels = []
-                if get_settings().pr_reviewer.enable_review_labels_effort:
-                    estimated_effort = data['PR Analysis']['Estimated effort to review [1-5]']
-                    estimated_effort_number = int(estimated_effort.split(',')[0])
-                    if 1 <= estimated_effort_number <= 5:
-                        review_labels.append(f'Review effort [1-5]: {estimated_effort_number}')
-                if get_settings().pr_reviewer.enable_review_labels_security:
-                    security_concerns = data['PR Analysis']['Security concerns']
-                    security_concerns_bool = 'yes' in security_concerns.lower() or 'true' in security_concerns.lower()
-                    if security_concerns_bool:
-                        review_labels.append('Possible security concern')
-
-                if review_labels:
-                    current_labels = self.git_provider.get_labels()
-                    current_labels_filtered = [label for label in current_labels if not label.startswith('Review effort [1-5]:') and not label.startswith('Possible security concern')]
-                    self.git_provider.publish_labels(review_labels + current_labels_filtered)
-            except:
-                pass
+        # Add custom labels from the review prediction (effort, security)
+        self.set_review_labels(data)
 
         # Log markdown response if verbosity level is high
         if get_settings().config.verbosity_level >= 2:
@@ -394,3 +375,28 @@ class PRReviewer:
             )
             return False
         return True
+
+    def set_review_labels(self, data):
+        if (get_settings().pr_reviewer.enable_review_labels_security or
+                get_settings().pr_reviewer.enable_review_labels_effort):
+            try:
+                review_labels = []
+                if get_settings().pr_reviewer.enable_review_labels_effort:
+                    estimated_effort = data['PR Analysis']['Estimated effort to review [1-5]']
+                    estimated_effort_number = int(estimated_effort.split(',')[0])
+                    if 1 <= estimated_effort_number <= 5:
+                        review_labels.append(f'Review effort [1-5]: {estimated_effort_number}')
+                if get_settings().pr_reviewer.enable_review_labels_security:
+                    security_concerns = data['PR Analysis']['Security concerns']
+                    security_concerns_bool = 'yes' in security_concerns.lower() or 'true' in security_concerns.lower()
+                    if security_concerns_bool:
+                        review_labels.append('Possible security concern')
+
+                if review_labels:
+                    current_labels = self.git_provider.get_labels()
+                    current_labels_filtered = [label for label in current_labels if
+                                               not label.lower().startswith('review effort [1-5]:') and not label.lower().startswith(
+                                                   'possible security concern')]
+                    self.git_provider.publish_labels(review_labels + current_labels_filtered)
+            except:
+                pass