From 2c11752a543e1ef02a362e759dcddcb03b5353dd Mon Sep 17 00:00:00 2001
From: XianBW <36835909+XianBW@users.noreply.github.com>
Date: Fri, 16 May 2025 18:29:59 +0800
Subject: [PATCH] chore: log storage refine (#883)

* log storage change

* fix bug

* fix ci
---
 rdagent/log/base.py    |  7 ++--
 rdagent/log/storage.py | 73 ++++++++++++++++++++++++------------------
 2 files changed, 44 insertions(+), 36 deletions(-)

diff --git a/rdagent/log/base.py b/rdagent/log/base.py
index 305fb7d3..7ef274d6 100644
--- a/rdagent/log/base.py
+++ b/rdagent/log/base.py
@@ -66,12 +66,9 @@ class Storage:
         ...
 
     @abstractmethod
-    def iter_msg(self, watch: bool = False) -> Generator[Message, None, None]:
+    def iter_msg(self) -> Generator[Message, None, None]:
         """
-        Parameters
-        ----------
-        watch : bool
-            should we watch the new content and display them
+        Iterate the message in the storage.
         """
         ...
 
diff --git a/rdagent/log/storage.py b/rdagent/log/storage.py
index f92d2ef7..cfc5ff22 100644
--- a/rdagent/log/storage.py
+++ b/rdagent/log/storage.py
@@ -65,44 +65,55 @@ class FileStorage(Storage):
         r"(?P<caller>.+:.+:\d+) - "
     )
 
-    def iter_msg(self, watch: bool = False) -> Generator[Message, None, None]:
+    def iter_msg(self, common: bool = False, tag: str | None = None) -> Generator[Message, None, None]:
         msg_l = []
-        for file in self.path.glob("**/*.log"):
-            tag = ".".join(file.relative_to(self.path).as_posix().replace("/", ".").split(".")[:-3])
-            pid = file.parent.name
+        if common:  # return string logs in common_logs.log
+            for file in self.path.glob("**/*.log"):
+                common_log_tag = ".".join(file.relative_to(self.path).as_posix().replace("/", ".").split(".")[:-3])
 
-            with file.open("r", encoding="utf-8") as f:
-                content = f.read()
-
-            matches, next_matches = self.log_pattern.finditer(content), self.log_pattern.finditer(content)
-            next_match = next(next_matches, None)
-            # NOTE: the content will be the text between `match` and `next_match`
-            for match in matches:
-                next_match = next(next_matches, None)
-
-                timestamp_str = match.group("timestamp")
-                timestamp = datetime.strptime(timestamp_str, "%Y-%m-%d %H:%M:%S.%f").replace(tzinfo=timezone.utc)
-                level: LOG_LEVEL = cast(LOG_LEVEL, match.group("level"))
-                caller = match.group("caller")
-
-                # Extract the message content
-                message_start = match.end()
-                message_end = next_match.start() if next_match else len(content)
-                message_content = content[message_start:message_end].strip()
-
-                if "Logging object in" in message_content:
+                if tag is not None and tag not in common_log_tag:
                     continue
 
-                m = Message(
-                    tag=tag, level=level, timestamp=timestamp, caller=caller, pid_trace=pid, content=message_content
-                )
+                pid = file.parent.name
 
-                msg_l.append(m)
+                with file.open("r", encoding="utf-8") as f:
+                    content = f.read()
 
-        for file in self.path.glob("**/*.pkl"):
+                matches, next_matches = self.log_pattern.finditer(content), self.log_pattern.finditer(content)
+                next_match = next(next_matches, None)
+                # NOTE: the content will be the text between `match` and `next_match`
+                for match in matches:
+                    next_match = next(next_matches, None)
+
+                    timestamp_str = match.group("timestamp")
+                    timestamp = datetime.strptime(timestamp_str, "%Y-%m-%d %H:%M:%S.%f").replace(tzinfo=timezone.utc)
+                    level: LOG_LEVEL = cast(LOG_LEVEL, match.group("level"))
+                    caller = match.group("caller")
+
+                    # Extract the message content
+                    message_start = match.end()
+                    message_end = next_match.start() if next_match else len(content)
+                    message_content = content[message_start:message_end].strip()
+
+                    if "Logging object in" in message_content:
+                        continue
+
+                    m = Message(
+                        tag=common_log_tag,
+                        level=level,
+                        timestamp=timestamp,
+                        caller=caller,
+                        pid_trace=pid,
+                        content=message_content,
+                    )
+
+                    msg_l.append(m)
+
+        pkl_files = "**/*.pkl" if tag is None else f"**/{tag.replace('.','/')}/**/*.pkl"
+        for file in self.path.glob(pkl_files):
             if file.name == "debug_llm.pkl":
                 continue
-            tag = ".".join(file.relative_to(self.path).as_posix().replace("/", ".").split(".")[:-3])
+            pkl_log_tag = ".".join(file.relative_to(self.path).as_posix().replace("/", ".").split(".")[:-3])
             pid = file.parent.name
 
             with file.open("rb") as f:
@@ -110,7 +121,7 @@ class FileStorage(Storage):
 
             timestamp = datetime.strptime(file.stem, "%Y-%m-%d_%H-%M-%S-%f").replace(tzinfo=timezone.utc)
 
-            m = Message(tag=tag, level="INFO", timestamp=timestamp, caller="", pid_trace=pid, content=content)
+            m = Message(tag=pkl_log_tag, level="INFO", timestamp=timestamp, caller="", pid_trace=pid, content=content)
 
             msg_l.append(m)