Demo (#101)

New demo windows structure

rdagent/log/storage.py

@@ -100,8 +100,11 @@ class FileStorage(Storage):
                     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)
+                    try:
+                        with pkl_path.open("rb") as f:
+                            m.content = pickle.load(f)
+                    except:
+                        continue
 
                 msg_l.append(m)
 

rdagent/log/ui/app.py

@@ -1,4 +1,4 @@
-from rdagent.log.ui.web import WebView, QlibTraceWindow, TraceObjWindow, mock_msg
+from rdagent.log.ui.web import WebView, SimpleTraceWindow, TraceObjWindow, mock_msg, TraceWindow
 from rdagent.log.storage import FileStorage, Message
 from rdagent.core.proposal import Trace
 from pathlib import Path
@@ -6,13 +6,15 @@ import pickle
 
 
 # show logs folder
-# WebView(QlibTraceWindow(show_common_logs=False, show_llm=False)).display(FileStorage("./log/2024-07-22_03-01-12-021659"))
+WebView(TraceWindow()).display(FileStorage("/data/home/bowen/workspace/RD-Agent/log/yuante/2024-07-24_04-03-33-691119"))
+# WebView(TraceWindow()).display(FileStorage("/data/home/bowen/workspace/RD-Agent/log/2024-07-22_03-01-12-021659"))
+# WebView(TraceWindow()).display(FileStorage("./log/2024-07-18_08-37-00-477228"))
 
 
 # load Trace obj
-with Path('./log/step_trace.pkl').open('rb') as f:
-    obj = pickle.load(f)
-    trace: Trace = obj[-1]
+# with Path('./log/step_trace.pkl').open('rb') as f:
+#     obj = pickle.load(f)
+#     trace: Trace = obj[-1]
 
 # show Trace obj
 # TraceObjWindow().consume_msg(mock_msg(trace))

rdagent/log/ui/web.py

@@ -1,12 +1,14 @@
 import pandas as pd
 import streamlit as st
 import plotly.express as px
+import time
 
 from rdagent.log.base import Storage, View
 from rdagent.log.base import Message
 from datetime import timezone, datetime
 from collections import defaultdict
 from copy import deepcopy
+from rdagent.core.proposal import Trace
 
 from typing import Callable, Type
 from streamlit.delta_generator import DeltaGenerator
@@ -21,49 +23,11 @@ from rdagent.components.coder.model_coder.CoSTEER.evaluators import ModelCoderFe
 from rdagent.components.coder.model_coder.model import ModelTask, ModelFBWorkspace
 
 
-
 st.set_page_config(layout="wide")
 
+TIME_DELAY = 0.001
+
 class WebView(View):
-    r"""
-
-    We have tree structure for sequence
-
-    session
-    |      \
-    ... defined by user ...
-    |                              \
-    info1 -> info2 -> ... -> info3 -> ...  overtime.
-
-    <message dispature>
-          |  | -  dispatch according to uri(e.g. `a.b.c. ...`)
-    Frontend is composed of windows.
-    Each window can individually display the message flow.
-
-    Some design principles:
-        session1.module(e.g. implement).
-        `s.log(a.b.1.c) s.log(a.b.2.c)` should not handed over to users.
-
-    An display example:
-
-        W1 write factor
-        W2 evaluate factor
-        W3 backtest
-
-        W123
-            R
-            RX
-            RXX
-            RX
-
-        W4
-            trace r1 r2 r3 r4
-
-    What to do next?
-    1. Data structure
-    2. Map path like `a.b.c` to frontend components
-    3. Display logic
-    """
 
     def __init__(self, ui: 'StWindow'):
         self.ui = ui
@@ -98,12 +62,6 @@ class LLMWindow(StWindow):
         self.container.chat_message('user').markdown(f"{msg.content}")
 
 
-class CodeWindow(StWindow):
-
-    def consume_msg(self, msg: Message):
-        self.container.code(msg.content, language="python")
-
-
 class ProgressTabsWindow(StWindow):
     '''
     For windows with stream messages, will refresh when a new tab is created.
@@ -182,11 +140,37 @@ class ObjectsTabsWindow(StWindow):
             self.inner_class(tabs[id]).consume_msg(splited_msg)
 
 
+class RoundTabsWindow(StWindow):
+
+    def __init__(self,
+                 container: 'DeltaGenerator',
+                 new_tab_func: Callable[[Message], bool],
+                 inner_class: Type[StWindow] = StWindow,
+                 title: str = 'Round tabs'):
+
+        container.markdown(f"### **{title}**")
+        self.inner_class = inner_class
+        self.new_tab_func = new_tab_func
+        self.round = 0
+
+        self.current_win = StWindow(container)
+        self.tabs_c = container.empty()
+
+
+    def consume_msg(self, msg: Message):
+        if self.new_tab_func(msg):
+            self.round += 1
+            self.current_win = self.inner_class(self.tabs_c.tabs([str(i) for i in range(1, self.round+1)])[-1])
+
+        self.current_win.consume_msg(msg)
+
+
 class HypothesisWindow(StWindow):
     
-    def consume_msg(self, msg: Message):
-        h: Hypothesis = msg.content
-        self.container.subheader('Hypothesis')
+    def consume_msg(self, msg: Message | Hypothesis):
+        h: Hypothesis = msg.content if isinstance(msg, Message) else msg
+
+        self.container.markdown('#### **Hypothesis💡**')
         self.container.markdown(f"""
 - **Hypothesis**: {h.hypothesis}
 - **Reason**: {h.reason}""")
@@ -194,9 +178,10 @@ class HypothesisWindow(StWindow):
 
 class HypothesisFeedbackWindow(StWindow):
 
-    def consume_msg(self, msg: Message):
-        h: HypothesisFeedback = msg.content
-        self.container.subheader('Hypothesis Feedback')
+    def consume_msg(self, msg: Message | HypothesisFeedback):
+        h: HypothesisFeedback = msg.content if isinstance(msg, Message) else msg
+
+        self.container.markdown('#### **Hypothesis Feedback🔍**')
         self.container.markdown(f"""
 - **Observations**: {h.observations}
 - **Hypothesis Evaluation**: {h.hypothesis_evaluation}
@@ -207,8 +192,8 @@ class HypothesisFeedbackWindow(StWindow):
 
 class FactorTaskWindow(StWindow):
 
-    def consume_msg(self, msg: Message):
-        ft: FactorTask = msg.content
+    def consume_msg(self, msg: Message | FactorTask):
+        ft: FactorTask = msg.content if isinstance(msg, Message) else msg
 
         self.container.markdown(f"**Factor Name**: {ft.factor_name}")
         self.container.markdown(f"**Description**: {ft.factor_description}")
@@ -222,8 +207,8 @@ class FactorTaskWindow(StWindow):
 
 class ModelTaskWindow(StWindow):
 
-    def consume_msg(self, msg: Message):
-        mt: ModelTask = msg.content
+    def consume_msg(self, msg: Message | ModelTask):
+        mt: ModelTask = msg.content if isinstance(msg, Message) else msg
 
         self.container.markdown(f"**Model Name**: {mt.name}")
         self.container.markdown(f"**Model Type**: {mt.model_type}")
@@ -237,8 +222,9 @@ class ModelTaskWindow(StWindow):
 
 class FactorFeedbackWindow(StWindow):
 
-    def consume_msg(self, msg: Message):
-        fb: FactorSingleFeedback = msg.content
+    def consume_msg(self, msg: Message | FactorSingleFeedback):
+        fb: FactorSingleFeedback = msg.content if isinstance(msg, Message) else msg
+
         self.container.markdown(f"""### :blue[Factor Execution Feedback]
 {fb.execution_feedback}
 ### :blue[Factor Code Feedback]
@@ -254,8 +240,9 @@ This implementation is {'SUCCESS' if fb.final_decision else 'FAIL'}.
 
 class ModelFeedbackWindow(StWindow):
 
-    def consume_msg(self, msg: Message):
-        mb: ModelCoderFeedback = msg.content
+    def consume_msg(self, msg: Message | ModelCoderFeedback):
+        mb: ModelCoderFeedback = msg.content if isinstance(msg, Message) else msg
+
         self.container.markdown(f"""### :blue[Model Execution Feedback]
 {mb.execution_feedback}
 ### :blue[Model Shape Feedback]
@@ -272,50 +259,57 @@ This implementation is {'SUCCESS' if mb.final_decision else 'FAIL'}.
 
 
 class WorkspaceWindow(StWindow):
+    def __init__(self, container: 'DeltaGenerator', show_task_info: bool = False):
+        self.container = container
+        self.show_task_info = show_task_info
 
-    def consume_msg(self, msg: Message):
-        ws: FactorFBWorkspace | ModelFBWorkspace = msg.content
+    def consume_msg(self, msg: Message | FactorFBWorkspace | ModelFBWorkspace):
+        ws: FactorFBWorkspace | ModelFBWorkspace = msg.content if isinstance(msg, Message) else msg
 
         # no workspace
         if ws is None: return
 
         # task info
-        task_msg = deepcopy(msg)
-        task_msg.content = ws.target_task
-        if isinstance(ws, FactorFBWorkspace):
-            self.container.subheader('Factor Info')
-            FactorTaskWindow(self.container.container()).consume_msg(task_msg)
-        else:
-            self.container.subheader('Model Info')
-            ModelTaskWindow(self.container.container()).consume_msg(task_msg)
+        if self.show_task_info:
+            task_msg = deepcopy(msg)
+            task_msg.content = ws.target_task
+            if isinstance(ws, FactorFBWorkspace):
+                self.container.subheader('Factor Info')
+                FactorTaskWindow(self.container.container()).consume_msg(task_msg)
+            else:
+                self.container.subheader('Model Info')
+                ModelTaskWindow(self.container.container()).consume_msg(task_msg)
 
         # task codes
-        self.container.subheader('Codes')
         for k,v in ws.code_dict.items():
             self.container.markdown(f"`{k}`")
             self.container.code(v, language="python")
 
         # executed_factor_value_dataframe
-        if isinstance(ws, FactorFBWorkspace):
-            self.container.subheader('Executed Factor Value Dataframe')
-            self.container.dataframe(ws.executed_factor_value_dataframe)
+        # if isinstance(ws, FactorFBWorkspace):
+        #     self.container.dataframe(ws.executed_factor_value_dataframe)
 
 
 class QlibFactorExpWindow(StWindow):
+    def __init__(self, container: DeltaGenerator, show_task_info: bool = False):
+        self.container = container
+        self.show_task_info = show_task_info
 
-    def consume_msg(self, msg: Message):
-        exp: QlibFactorExperiment = msg.content
+    def consume_msg(self, msg: Message | QlibFactorExperiment):
+        exp: QlibFactorExperiment = msg.content if isinstance(msg, Message) else msg
 
         # factor tasks
-        ftm_msg = deepcopy(msg)
-        ftm_msg.content = [ws for ws in exp.sub_workspace_list if ws]
-        ObjectsTabsWindow(self.container.expander('Factor Tasks'),
-                          inner_class=WorkspaceWindow,
-                          mapper=lambda x: x.target_task.factor_name,
-                          ).consume_msg(ftm_msg)
+        if self.show_task_info:
+            ftm_msg = deepcopy(msg)
+            ftm_msg.content = [ws for ws in exp.sub_workspace_list if ws]
+            self.container.markdown('**Factor Tasks**')
+            ObjectsTabsWindow(self.container.container(),
+                            inner_class=WorkspaceWindow,
+                            mapper=lambda x: x.target_task.factor_name,
+                            ).consume_msg(ftm_msg)
 
         # result
-        self.container.subheader('Results', divider=True)
+        self.container.markdown('**Results**', divider=True)
         results = pd.DataFrame({f'base_exp_{id}':e.result for id, e in enumerate(exp.based_experiments)})
         results['now'] = exp.result
 
@@ -329,17 +323,22 @@ class QlibFactorExpWindow(StWindow):
 
 
 class QlibModelExpWindow(StWindow):
+    def __init__(self, container: DeltaGenerator, show_task_info: bool = False):
+        self.container = container
+        self.show_task_info = show_task_info
 
-    def consume_msg(self, msg: Message):
-        exp: QlibModelExperiment = msg.content
+    def consume_msg(self, msg: Message | QlibModelExperiment):
+        exp: QlibModelExperiment = msg.content if isinstance(msg, Message) else msg
 
         # model tasks
-        _msg = deepcopy(msg)
-        _msg.content = [ws for ws in exp.sub_workspace_list if ws]
-        ObjectsTabsWindow(self.container.expander('Model Tasks'),
-                          inner_class=WorkspaceWindow,
-                          mapper=lambda x: x.target_task.name,
-                          ).consume_msg(_msg)
+        if self.show_task_info:
+            _msg = deepcopy(msg)
+            _msg.content = [ws for ws in exp.sub_workspace_list if ws]
+            self.container.markdown('**Model Tasks**')
+            ObjectsTabsWindow(self.container.container(),
+                            inner_class=WorkspaceWindow,
+                            mapper=lambda x: x.target_task.name,
+                            ).consume_msg(_msg)
 
         # result
         self.container.subheader('Results', divider=True)
@@ -349,9 +348,9 @@ class QlibModelExpWindow(StWindow):
         self.container.expander('results table').table(results)
 
 
-class QlibTraceWindow(StWindow):
+class SimpleTraceWindow(StWindow):
 
-    def __init__(self, container: 'DeltaGenerator' = st.container(), show_llm: bool = False, show_common_logs: bool = True):
+    def __init__(self, container: 'DeltaGenerator' = st.container(), show_llm: bool = False, show_common_logs: bool = False):
         super().__init__(container)
         self.show_llm = show_llm
         self.show_common_logs = show_common_logs
@@ -429,14 +428,16 @@ def mock_msg(obj) -> Message:
     return Message(tag='mock', level='INFO', timestamp=datetime.now(), pid_trace='000', caller='mock',content=obj)
 
 
-from rdagent.core.proposal import Trace
 class TraceObjWindow(StWindow):
 
     def __init__(self, container: 'DeltaGenerator' = st.container()):
         self.container = container
 
-    def consume_msg(self, msg: Message):
-        trace:Trace = msg.content
+    def consume_msg(self, msg: Message | Trace):
+        if isinstance(msg, Message):
+            trace: Trace = msg.content
+        else:
+            trace = msg
 
         for id, (h, e, hf) in enumerate(trace.hist):
             self.container.header(f'Trace History {id}', divider=True)
@@ -447,3 +448,163 @@ class TraceObjWindow(StWindow):
                 QlibModelExpWindow(self.container).consume_msg(mock_msg(e))
             HypothesisFeedbackWindow(self.container).consume_msg(mock_msg(hf))
 
+
+class ResearchWindow(StWindow):
+
+    def consume_msg(self, msg: Message):
+        if msg.tag.endswith('hypothesis generation'):
+            HypothesisWindow(self.container.container()).consume_msg(msg)
+        elif msg.tag.endswith('experiment generation'):
+            if isinstance(msg, list):
+                if isinstance(msg.content[0], FactorTask):
+                    self.container.markdown('**Factor Tasks**')
+                    ObjectsTabsWindow(self.container.container(), FactorTaskWindow, lambda x: x.factor_name).consume_msg(msg)
+                elif isinstance(msg.content[0], ModelTask):
+                    self.container.markdown('**Model Tasks**')
+                    ObjectsTabsWindow(self.container.container(), ModelTaskWindow, lambda x: x.name).consume_msg(msg)
+
+
+class EvolvingWindow(StWindow):
+    def __init__(self, container: 'DeltaGenerator'):
+        self.container = container
+        self.evolving_tasks: list[str] = []
+
+    def consume_msg(self, msg: Message):
+        if msg.tag.endswith('evolving code'):
+            if isinstance(msg.content, list):
+                msg.content = [m for m in msg.content if m]
+                if len(msg.content) == 0:
+                    return
+                if isinstance(msg.content[0], FactorFBWorkspace):
+                    self.container.markdown('**Factor Codes**')
+                    ObjectsTabsWindow(self.container.container(),
+                                        inner_class=WorkspaceWindow,
+                                        mapper=lambda x: x.target_task.factor_name).consume_msg(msg)
+                    self.evolving_tasks = [m.target_task.factor_name for m in msg.content]
+                elif isinstance(msg.content[0], ModelFBWorkspace):
+                    self.container.markdown('**Model Codes**')
+                    ObjectsTabsWindow(self.container.container(),
+                                        inner_class=WorkspaceWindow,
+                                        mapper=lambda x: x.target_task.name).consume_msg(msg)
+                    self.evolving_tasks = [m.target_task.name for m in msg.content]
+        elif msg.tag.endswith('evolving feedback'):
+            if isinstance(msg.content, list):
+                msg.content = [m for m in msg.content if m]
+                if len(msg.content) == 0:
+                    return
+                if isinstance(msg.content[0], FactorSingleFeedback):
+                    self.container.markdown('**Factor Feedbacks🔍**')
+                    ObjectsTabsWindow(self.container.container(),
+                                        inner_class=FactorFeedbackWindow,
+                                        tab_names=self.evolving_tasks).consume_msg(msg)
+                elif isinstance(msg.content[0], ModelCoderFeedback):
+                    self.container.markdown('**Model Feedbacks🔍**')
+                    ObjectsTabsWindow(self.container.container(),
+                                        inner_class=ModelFeedbackWindow,
+                                        tab_names=self.evolving_tasks).consume_msg(msg)
+
+
+class DevelopmentWindow(StWindow):
+
+    def __init__(self, container: 'DeltaGenerator'):
+        self.E_win = RoundTabsWindow(container.container(),
+                                     new_tab_func=lambda x: x.tag.endswith('evolving code'),
+                                     inner_class=EvolvingWindow,
+                                     title='Evolving Loops🔧')
+
+    def consume_msg(self, msg: Message):
+        if 'evolving' in msg.tag:
+            self.E_win.consume_msg(msg)
+        # elif msg.tag.endswith('result'):
+        #     self.container.subheader('Results')
+        #     if isinstance(msg.content[0], FactorFBWorkspace):
+        #         ObjectsTabsWindow(self.container.expander('Factor Workspaces'),
+        #                             inner_class=WorkspaceWindow,
+        #                             mapper=lambda x: x.target_task.factor_name).consume_msg(msg)
+        #     elif isinstance(msg.content[0], ModelFBWorkspace):
+        #         ObjectsTabsWindow(self.container.expander('Model Workspaces'),
+        #                             inner_class=WorkspaceWindow,
+        #                             mapper=lambda x: x.target_task.name).consume_msg(msg)
+
+
+class FeedbackWindow(StWindow):
+    
+    def __init__(self, container: 'DeltaGenerator'):
+        self.container = container
+
+    def consume_msg(self, msg: Message):
+        if isinstance(msg.content, HypothesisFeedback):
+            HypothesisFeedbackWindow(self.container.container(border=True)).consume_msg(msg)
+        elif isinstance(msg.content, QlibModelExperiment):
+            QlibModelExpWindow(self.container.container(border=True)).consume_msg(msg)
+        elif isinstance(msg.content, QlibFactorExperiment):
+            QlibFactorExpWindow(self.container.container(border=True)).consume_msg(msg)
+
+
+class SingleRDLoopWindow(StWindow):
+    
+    def __init__(self, container: 'DeltaGenerator'):
+        self.container = container
+        col1, col2 = self.container.columns([2, 3])
+        self.R_win = ResearchWindow(col1.container(border=True))
+        self.F_win = FeedbackWindow(col1.container(border=True))
+        self.D_win = DevelopmentWindow(col2.container(border=True))
+
+    def consume_msg(self, msg: Message):
+        tags = msg.tag.split('.')
+        if 'r' in tags:
+            self.R_win.consume_msg(msg)
+        elif 'd' in tags:
+            self.D_win.consume_msg(msg)
+        elif 'ef' in tags:
+            self.F_win.consume_msg(msg)
+
+
+class TraceWindow(StWindow):
+
+    def __init__(self, container: 'DeltaGenerator' = st.container(), show_llm: bool = False, show_common_logs: bool = False):
+        self.show_llm = show_llm
+        self.show_common_logs = show_common_logs
+
+        top_container = container.container()
+        col1, col2 = top_container.columns([2,3])
+        chart_c = col2.container(border=True, height=300)
+        chart_c.markdown('**Metrics📈**')
+        self.chart_c = chart_c.empty()
+        hypothesis_status_c = col1.container(border=True, height=300)
+        hypothesis_status_c.markdown('**Hypotheses🏅**')
+        self.summary_c = hypothesis_status_c.empty()
+
+        self.RDL_win = RoundTabsWindow(container.container(),
+                                       new_tab_func=lambda x: x.tag.endswith('hypothesis generation'),
+                                       inner_class=SingleRDLoopWindow,
+                                       title='R&D Loops♾️')
+        
+        self.hypothesis_decisions = defaultdict(bool)
+        self.current_hypothesis = None
+
+        self.results = []
+
+    def consume_msg(self, msg: Message):
+        if not self.show_llm and 'llm_messages' in msg.tag:
+            return
+        if not self.show_common_logs and isinstance(msg.content, str):
+            return
+        if isinstance(msg.content, dict):
+            return
+        if msg.tag.endswith('hypothesis generation'):
+            self.current_hypothesis = msg.content.hypothesis
+        elif msg.tag.endswith('ef.feedback'):
+            self.hypothesis_decisions[self.current_hypothesis] = msg.content.decision
+            self.summary_c.markdown('\n'.join(f"{id+1}. :green[{h}]\n" if d else f"{id+1}. {h}\n" for id,(h,d) in enumerate(self.hypothesis_decisions.items())))
+        elif msg.tag.endswith('ef.model runner result') or msg.tag.endswith('ef.factor runner result'):
+            self.results.append(msg.content.result)
+            if len(self.results) == 1:
+                self.chart_c.table(self.results[0])
+            else:
+                df = pd.DataFrame(self.results, index=range(1, len(self.results)+1))
+                fig = px.line(df, x=df.index, y=df.columns, markers=True)
+                self.chart_c.plotly_chart(fig)
+
+        self.RDL_win.consume_msg(msg)
+        # time.sleep(TIME_DELAY)