demo (#80)

save logs and load it.

---------

Co-authored-by: Young <afe.young@gmail.com>

rdagent/log/base.py

@@ -1,7 +1,22 @@
 from __future__ import annotations
 
 from abc import abstractmethod
+from collections.abc import Generator
+from datetime import datetime
 from pathlib import Path
+from typing import Literal, Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class Message:
+    """The info unit of the storage"""
+    tag: str  # namespace like like a.b.c
+    level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]  # The level of the logging
+    timestamp: datetime  # The time when the message is generated
+    caller: Optional[str] # The caller of the logging like `rdagent.oai.llm_utils:_create_chat_completion_inner_function:55`(file:func:line)
+    pid_trace: Optional[str]  # The process id trace;  A-B-C represents A create B, B create C
+    content: object  # The content
 
 
 class Storage:
@@ -40,6 +55,16 @@ class Storage:
         """
         ...
 
+    @abstractmethod
+    def iter_msg(self, watch: bool = False) -> Generator[Message, None, None]:
+        """
+        Parameters
+        ----------
+        watch : bool
+            should we watch the new content and display them
+        """
+        ...
+
 
 class View:
     """

rdagent/log/storage.py

@@ -1,10 +1,11 @@
+import re
 import json
 import pickle
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Literal
+from typing import Literal, Generator
 
-from .base import Storage
+from .base import Message, Storage
 
 
 class FileStorage(Storage):
@@ -54,3 +55,55 @@ class FileStorage(Storage):
             with path.open("w") as f:
                 f.write(obj)
             return path
+
+    def iter_msg(self, watch: bool = False) -> Generator[Message, None, None]:
+        log_pattern = re.compile(
+            r"(?P<timestamp>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3}) \| "
+            r"(?P<level>DEBUG|INFO|WARNING|ERROR|CRITICAL) *\| "
+            r"(?P<caller>.+:.+:\d+) - "
+        )
+        msg_l = []
+        for file in self.path.glob("**/*.log"):
+            tag = '.'.join(str(file.relative_to(self.path)).replace("/", ".").split(".")[:-3])
+            pid = file.parent.name
+
+            with file.open("r") as f:
+                content = f.read()
+
+            matches, next_matches = log_pattern.finditer(content), 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 = 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()
+
+                m = Message(
+                    tag=tag,
+                    level=level,
+                    timestamp=timestamp,
+                    caller=caller,
+                    pid_trace=pid,
+                    content=message_content
+                )
+
+                if "Logging object in" in m.content:
+                    absolute_p = m.content.split("Logging object in ")[1]
+                    relative_p = "." + absolute_p.split(self.path.name)[1]
+                    pkl_path = self.path / relative_p
+                    with pkl_path.open("rb") as f:
+                        m.content = pickle.load(f)
+
+                msg_l.append(m)
+
+        msg_l.sort(key=lambda x: x.timestamp)
+        for m in msg_l:
+            yield m

rdagent/log/ui/web.py

@@ -1,7 +1,19 @@
+import pickle
+import streamlit as st
 from pathlib import Path
-
 from rdagent.log.base import Storage, View
+from rdagent.log.storage import FileStorage
+from rdagent.log.base import Message
+from datetime import timezone, datetime
+from collections import defaultdict
 
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from streamlit.delta_generator import DeltaGenerator
+    from rdagent.core.proposal import Hypothesis, HypothesisFeedback
+    from rdagent.scenarios.qlib.experiment.factor_experiment import QlibFactorExperiment
+    from rdagent.components.coder.factor_coder.factor import FactorTask, FactorFBWorkspace
+    from rdagent.components.coder.factor_coder.CoSTEER.evaluators import FactorSingleFeedback
 
 class ProcessView(View):
     def __init__(self, trace_path: Path):
@@ -54,38 +66,100 @@ class WebView(View):
     3. Display logic
     """
 
-    def __init__(self, trace_path: Path):
-        pass
+    def __init__(self, ui: 'StWindow'):
+        self.ui = ui
         # Save logs to your desired data structure
         # ...
 
-    def display(s: Storage, watch: bool = False):
-        ui = STLUI()
+    def display(self, s: Storage, watch: bool = False):
+
         for msg in s.iter_msg():  # iterate overtime
             # NOTE:  iter_msg will correctly seperate the information.
             # TODO: msg may support streaming mode.
-            ui.dispatch(msg)
-        pass
+            self.ui.consume_msg(msg)
+
 
 
 # TODO: Implement the following classes
-class STLWindow:
-    ...
+class StWindow:
 
-    def consume_msg(self, msg):
-        ...  # update it's view
+    def __init__(self, container: 'DeltaGenerator'):
+        self.container = container
+
+    def consume_msg(self, msg: Message):
+        msg_str = f"{msg.timestamp.astimezone(timezone.utc).isoformat()} | {msg.level} | {msg.caller} - {msg.content}"
+        self.container.write(msg_str)
 
 
-class STLUI:
-    wd_l: list[STLWindow]
+class LLMWindow(StWindow):
+    def __init__(self, container: 'DeltaGenerator', session_name: str="common"):
+        self.container = container
+        self.container.subheader(f"{session_name} Messages")
 
-    def __init__(self):
-        self.build_ui()
+    def consume_msg(self, msg: Message):
+        self.container.chat_message('User').write(f"{msg.content}")
 
-    def build_ui(self):
-        # control the dispaly of windows
-        ...
 
-    def dispatch(self, msg):
-        # map the message to a specific window
-        ...
+class CodeWindow(StWindow):
+    def __init__(self, container: 'DeltaGenerator'):
+        self.container = container.empty()
+    
+    def consume_msg(self, msg: Message):
+        self.container.code(msg.content, language="python")
+
+
+class MultiProcessWindow(StWindow):
+    def __init__(self, container: 'DeltaGenerator', inner_class: str = "STLWindow"):
+        '''
+        inner_class: STLWindow 子类名称, 用来实例化多进程窗口的内部窗口实例
+        '''
+        self.container = container.empty()
+        self.tabs_cache = defaultdict(list)
+        self.inner_class = eval(inner_class)
+        
+    
+    def consume_msg(self, msg: Message):
+        name = msg.pid_trace.split("-")[-1]
+        self.tabs_cache[name].append(msg)
+
+        tabs = self.container.tabs(list(self.tabs_cache.keys()))
+        for i, name in enumerate(self.tabs_cache):
+            inner_win: StWindow = self.inner_class(tabs[i])
+            for m in self.tabs_cache[name]:
+                inner_win.consume_msg(m)
+
+
+class HypothesisRelatedWindow(StWindow):
+    def __init__(self, container: 'DeltaGenerator'):
+        self.container = container
+    
+    def consume_msg(self, msg: Message):
+        h: Hypothesis | HypothesisFeedback = msg.content
+        self.container.text(str(h))
+
+
+class QlibFactorUI(StWindow):
+
+    def __init__(self, container: 'DeltaGenerator' = st.container()):
+        super().__init__(container)
+        self.pid_level = 0
+        self.tag_level = 0
+        self.current_win = container.container()
+
+    def consume_msg(self, msg: Message):
+        pid_level = msg.pid_trace.count("-")
+        tag_level = msg.tag.count(".")
+
+        if pid_level > self.pid_level:
+            self.pid_level = pid_level
+            self.current_win = MultiProcessWindow(st.container(), "STLWindow")
+        
+        if tag_level > self.tag_level:
+            self.tag_level = tag_level
+
+
+        self.current_win.consume_msg(msg)
+
+
+if __name__ == "__main__":
+    WebView(QlibFactorUI()).display(FileStorage("./log/test_trace"))