From aeb1bd8dbcde2a4e9c9b5a0fb758dd6afacb713b Mon Sep 17 00:00:00 2001
From: Abhinav Kumar <abhinavkumar130503@gmail.com>
Date: Wed, 2 Jul 2025 18:11:19 +0530
Subject: [PATCH] fix suggestions

---
 pr_agent/tools/pr_description.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pr_agent/tools/pr_description.py b/pr_agent/tools/pr_description.py
index 41cc4006..3ce633dd 100644
--- a/pr_agent/tools/pr_description.py
+++ b/pr_agent/tools/pr_description.py
@@ -538,9 +538,10 @@ class PRDescription:
                 get_logger().error(f"Failing to process walkthrough {self.pr_id}: {e}")
                 body = body.replace('pr_agent:walkthrough', "")
 
+        # Add support for pr_agent:diagram marker (plain and HTML comment formats)
         ai_diagram = self.data.get('changes_diagram')
-        if ai_diagram and not re.search(r'<!--\s*pr_agent:diagram\s*-->', body):
-            body = body.replace('pr_agent:diagram', ai_diagram)
+        if ai_diagram:
+            body = re.sub(r'<!--\s*pr_agent:diagram\s*-->|pr_agent:diagram', ai_diagram, body)
 
         return title, body, walkthrough_gfm, pr_file_changes