From d93a24bbf77be3500d25b1a64732464dbac48f6a Mon Sep 17 00:00:00 2001
From: "Hussam.lawen" <hussam789@gmail.com>
Date: Sun, 11 Feb 2024 12:14:25 +0200
Subject: [PATCH] add loading comment to /improve

---
 pr_agent/tools/pr_code_suggestions.py | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/pr_agent/tools/pr_code_suggestions.py b/pr_agent/tools/pr_code_suggestions.py
index b944047ea..83139abca 100644
--- a/pr_agent/tools/pr_code_suggestions.py
+++ b/pr_agent/tools/pr_code_suggestions.py
@@ -65,11 +65,16 @@ def __init__(self, pr_url: str, cli_mode=False, args: list = None,
                                           get_settings().pr_code_suggestions_prompt.system,
                                           get_settings().pr_code_suggestions_prompt.user)
 
+        self.progress = f"## Generating PR code suggestions\n\n"
+        self.progress += f"""\nWork in progress ...<br>\n<img src="https://codium.ai/images/pr_agent/dual_ball_loading-crop.gif" width=48>"""
+        self.progress_response = None
+
     async def run(self):
         try:
             get_logger().info('Generating code suggestions for PR...')
+
             if get_settings().config.publish_output:
-                self.git_provider.publish_comment("Preparing suggestions...", is_temporary=True)
+                self.progress_response = self.git_provider.publish_comment(self.progress)
 
             get_logger().info('Preparing PR code suggestions...')
             if not self.is_extended:
@@ -103,12 +108,20 @@ async def run(self):
                         pr_body += HelpMessage.get_improve_usage_guide()
                         pr_body += "\n</details>\n"
 
-                    self.git_provider.publish_comment(pr_body)
+                    if self.progress_response:
+                        self.git_provider.edit_comment(self.progress_response, body=pr_body)
+                    else:
+                        self.git_provider.publish_comment(pr_body)
+
                 else:
                     get_logger().info('Pushing inline code suggestions...')
                     self.push_inline_code_suggestions(data)
+                    if self.progress_response:
+                        self.progress_response.delete()
         except Exception as e:
             get_logger().error(f"Failed to generate code suggestions for PR, error: {e}")
+            if self.progress_response:
+                self.progress_response.delete()
 
     async def _prepare_prediction(self, model: str):
         get_logger().info('Getting PR diff...')
@@ -162,7 +175,10 @@ def push_inline_code_suggestions(self, data):
 
         if not data['code_suggestions']:
             get_logger().info('No suggestions found to improve this PR.')
-            return self.git_provider.publish_comment('No suggestions found to improve this PR.')
+            if self.progress_response:
+                return self.git_provider.edit_comment(self.progress_response, body='No suggestions found to improve this PR.')
+            else:
+                return self.git_provider.publish_comment('No suggestions found to improve this PR.')
 
         for d in data['code_suggestions']:
             try: