From 7ed3cd1240641f20cb2239e8e5d987ca2872b7d5 Mon Sep 17 00:00:00 2001
From: Tim <illking@foxmail.com>
Date: Tue, 22 Jul 2025 17:02:56 +0800
Subject: [PATCH] chore: show merge loop in figure (#1105)

---
 rdagent/log/ui/ds_trace.py | 10 +++++++++-
 rdagent/log/ui/utils.py    |  7 ++++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/rdagent/log/ui/ds_trace.py b/rdagent/log/ui/ds_trace.py
index c8000768..68d6e655 100644
--- a/rdagent/log/ui/ds_trace.py
+++ b/rdagent/log/ui/ds_trace.py
@@ -588,7 +588,15 @@ def summarize_win():
             final_trace_loop_id = max_id
             while "record" not in state.data[final_trace_loop_id]:
                 final_trace_loop_id -= 1
-            st.pyplot(trace_figure(state.data[final_trace_loop_id]["record"]["trace"]))
+            merge_loops = []
+            for loop_id in state.llm_data.keys():
+                if "direct_exp_gen" not in state.llm_data[loop_id]:
+                    continue
+                if "scenarios.data_science.proposal.exp_gen.merge:trace" in [
+                    i["obj"]["uri"] for i in state.llm_data[loop_id]["direct_exp_gen"]["no_tag"] if "uri" in i["obj"]
+                ]:
+                    merge_loops.append(loop_id)
+            st.pyplot(trace_figure(state.data[final_trace_loop_id]["record"]["trace"], merge_loops))
         df = pd.DataFrame(
             columns=[
                 "Component",
diff --git a/rdagent/log/ui/utils.py b/rdagent/log/ui/utils.py
index bdd5b053..c3000404 100644
--- a/rdagent/log/ui/utils.py
+++ b/rdagent/log/ui/utils.py
@@ -598,7 +598,7 @@ def curve_figure(scores: pd.DataFrame) -> go.Figure:
     return fig
 
 
-def trace_figure(trace: Trace):
+def trace_figure(trace: Trace, merge_loops: list = []):
     G = nx.DiGraph()
 
     # Calculate the number of ancestors for each node (root node is 0, more ancestors means lower level)
@@ -610,7 +610,7 @@ def trace_figure(trace: Trace):
         """
         Convert to index in the queue (enque id) to loop_idx for easier understanding.
         """
-        if hasattr(trace, "idx2loop_id"):
+        if hasattr(trace, "idx2loop_id") and idx in trace.idx2loop_id:
             # FIXME: only keep me after it is stable. Just for compatibility.
             return f"L{trace.idx2loop_id[idx]} ({idx})"
         return f"L{idx}"
@@ -684,7 +684,8 @@ def trace_figure(trace: Trace):
                 pos[node] = (x, y)
 
     fig, ax = plt.subplots(figsize=(8, 6))
-    nx.draw(G, pos, with_labels=True, arrows=True, node_color="skyblue", node_size=100, font_size=5, ax=ax)
+    color_map = ["tomato" if node in [get_display_name(idx) for idx in merge_loops] else "skyblue" for node in G]
+    nx.draw(G, pos, with_labels=True, arrows=True, node_color=color_map, node_size=100, font_size=5, ax=ax)
     return fig