diff --git a/rdagent/log/base.py b/rdagent/log/base.py index 8cec23cc..fc62ed39 100644 --- a/rdagent/log/base.py +++ b/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: """ diff --git a/rdagent/log/storage.py b/rdagent/log/storage.py index 92470eab..d07baa58 100644 --- a/rdagent/log/storage.py +++ b/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\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3}) \| " + r"(?PDEBUG|INFO|WARNING|ERROR|CRITICAL) *\| " + r"(?P.+:.+:\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 diff --git a/rdagent/log/ui/web.py b/rdagent/log/ui/web.py index 6fae6658..9d75870b 100644 --- a/rdagent/log/ui/web.py +++ b/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")) \ No newline at end of file