From 5496168d07ed18e15f94cc4861aa10cf64faa2fc Mon Sep 17 00:00:00 2001
From: you-n-g <you-n-g@users.noreply.github.com>
Date: Tue, 15 Jul 2025 22:26:06 +0800
Subject: [PATCH] feat: add enable_cache toggle for UI data caching (#1075)

---
 rdagent/log/ui/conf.py     | 2 ++
 rdagent/log/ui/ds_trace.py | 6 +++++-
 rdagent/log/ui/llm_st.py   | 6 +++++-
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/rdagent/log/ui/conf.py b/rdagent/log/ui/conf.py
index ee32412c..d689d856 100644
--- a/rdagent/log/ui/conf.py
+++ b/rdagent/log/ui/conf.py
@@ -18,5 +18,7 @@ class UIBasePropSetting(ExtendedBaseSettings):
 
     trace_folder: str = "./traces"
 
+    enable_cache: bool = True
+
 
 UI_SETTING = UIBasePropSetting()
diff --git a/rdagent/log/ui/ds_trace.py b/rdagent/log/ui/ds_trace.py
index 071daf27..f10a1522 100644
--- a/rdagent/log/ui/ds_trace.py
+++ b/rdagent/log/ui/ds_trace.py
@@ -14,6 +14,7 @@ from streamlit import session_state as state
 
 from rdagent.app.data_science.loop import DataScienceRDLoop
 from rdagent.log.storage import FileStorage
+from rdagent.log.ui.conf import UI_SETTING
 from rdagent.log.ui.utils import curve_figure, load_times, trace_figure
 from rdagent.log.utils import (
     LogColors,
@@ -45,7 +46,6 @@ def convert_defaultdict_to_dict(d):
     return d
 
 
-@st.cache_data(persist=True)
 def load_data(log_path: Path):
     data = defaultdict(lambda: defaultdict(dict))
     llm_data = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
@@ -132,6 +132,10 @@ def load_data(log_path: Path):
     )
 
 
+if UI_SETTING.enable_cache:
+    load_data = st.cache_data(persist=True)(load_data)
+
+
 def load_stdout(stdout_path: Path):
     if stdout_path.exists():
         stdout = stdout_path.read_text()
diff --git a/rdagent/log/ui/llm_st.py b/rdagent/log/ui/llm_st.py
index 7438acb1..e3a963a8 100644
--- a/rdagent/log/ui/llm_st.py
+++ b/rdagent/log/ui/llm_st.py
@@ -8,6 +8,7 @@ from pathlib import Path
 import streamlit as st
 from streamlit import session_state
 
+from rdagent.log.ui.conf import UI_SETTING
 from rdagent.log.utils import extract_evoid, extract_loopid_func_name
 
 st.set_page_config(layout="wide", page_title="debug_llm", page_icon="🎓", initial_sidebar_state="expanded")
@@ -18,7 +19,6 @@ parser.add_argument("--log_dir", type=str, help="Path to the log directory")
 args = parser.parse_args()
 
 
-@st.cache_data
 def get_folders_sorted(log_path):
     """缓存并返回排序后的文件夹列表，并加入进度打印"""
     with st.spinner("正在加载文件夹列表..."):
@@ -31,6 +31,10 @@ def get_folders_sorted(log_path):
     return [folder.name for folder in folders]
 
 
+if UI_SETTING.enable_cache:
+    get_folders_sorted = st.cache_data(get_folders_sorted)
+
+
 # 设置主日志路径
 main_log_path = Path(args.log_dir) if args.log_dir else Path("./log")
 if not main_log_path.exists():