From 7e0f85d538bb46b9be72ebeec6e647441e31aae9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=9D=92=E8=BD=A9?= <fuqingxu.fqx@alibaba-inc.com>
Date: Mon, 12 Aug 2024 19:23:24 +0800
Subject: [PATCH] improve logging

---
 memoryscope/core/operation/base_workflow.py |  4 ++-
 memoryscope/core/utils/logger.py            | 35 +++++++++++++++++++--
 memoryscope/core/worker/memory_manager.py   |  7 +++--
 3 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/memoryscope/core/operation/base_workflow.py b/memoryscope/core/operation/base_workflow.py
index de3c41fc..b8670395 100644
--- a/memoryscope/core/operation/base_workflow.py
+++ b/memoryscope/core/operation/base_workflow.py
@@ -132,6 +132,7 @@ class BaseWorkflow(object):
             if name not in self.memoryscope_context.worker_conf_dict:
                 raise RuntimeError(f"worker={name} is not exists in worker config!")
 
+            # note: shared context object in all workers
             self.worker_dict[name] = init_instance_by_config(
                 config=self.memoryscope_context.worker_conf_dict[name],
                 name=name,
@@ -165,13 +166,14 @@ class BaseWorkflow(object):
             **kwargs: Additional keyword arguments to be passed to context.
         """
         with Timer(f"workflow.{self.name}", time_log_type="wrap"):
-            self.logger.info(f"\n\n\n++++++++++++++++++++++++ [{self.name}] ++++++++++++++++++++++++")
+            self.logger.info(f"\n\n\n++++++++++++++++++++++++ [Operation: {self.name}] ++++++++++++++++++++++++")
 
             self.context.clear()
             self.context.update({WORKFLOW_NAME: self.name, **kwargs})
             n_stage = len(self.workflow_worker_list)
             # Iterate over each part of the workflow
             for index, workflow_part in enumerate(self.workflow_worker_list):
+                self.logger.info(self.logger.format_current_context(self.context))
                 # Sequential execution for single-item parts
                 if len(workflow_part) == 1:
                     self.logger.info(f"\n-----------------------------------------------------")
diff --git a/memoryscope/core/utils/logger.py b/memoryscope/core/utils/logger.py
index 3caa5c7b..6a533f26 100644
--- a/memoryscope/core/utils/logger.py
+++ b/memoryscope/core/utils/logger.py
@@ -2,6 +2,7 @@ import logging
 from datetime import datetime
 from logging.handlers import RotatingFileHandler
 from pathlib import Path
+from rich.console import Console
 
 LOG_FORMAT = "%(asctime)s %(levelname)s %(threadName)s %(module)s:%(lineno)d] %(message)s"
 DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
@@ -64,6 +65,36 @@ class Logger(logging.Logger):
 
         self.info(f"logger={name} is inited.")  # Logs an initialization message
 
+    def format_current_context(self, context):
+        from rich.panel import Panel
+        from rich.text import Text
+        import pprint
+        pp = pprint.PrettyPrinter()
+        pretty_string = pp.pformat(context)
+
+        def rich2text(rich_table):
+            console = Console(width=150)
+            with console.capture() as capture:
+                console.print(rich_table)
+            return '\n' + str(Text.from_ansi(capture.get()))
+
+        return rich2text(Panel(pretty_string, width=128))
+    
+    def format_current_memory(self, context):
+        from rich.panel import Panel
+        from rich.text import Text
+        import pprint
+        pp = pprint.PrettyPrinter()
+        pretty_string = pp.pformat(context)
+
+        def rich2text(rich_table):
+            console = Console(width=150)
+            with console.capture() as capture:
+                console.print(rich_table)
+            return '\n' + str(Text.from_ansi(capture.get()))
+
+        return rich2text(Panel(pretty_string, width=128))
+    
     def format_chat_message(self, message):
         buf = '\n'
         buf += f"++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\n"
@@ -111,7 +142,7 @@ class Logger(logging.Logger):
         file_path = Path().joinpath(self.dir_path, f"{self.name}.{self.file_type}")
         file_path.parent.mkdir(exist_ok=True)  # Ensure the directory exists
         file_name = file_path.as_posix()  # Get the absolute path as a string
-        print(f"[{self.name}] Registering Logger to file at: ", file_name)
+        Console().print(f"[{self.name}] Registering Logger to file at: " + file_name, style="bold blue")
 
         # Instantiate a rotating file handler with specified parameters
         file_handler = RotatingFileHandler(
@@ -221,4 +252,4 @@ class Logger(logging.Logger):
 
     @staticmethod
     def append_timestamp(name: str) -> str:
-        return f"{name}_{datetime.now().strftime(r'%Y%m%d_%H%M%S')}"
\ No newline at end of file
+        return f"{name}_{datetime.now().strftime(r'%Y%m%d_%H%M')}"
\ No newline at end of file
diff --git a/memoryscope/core/worker/memory_manager.py b/memoryscope/core/worker/memory_manager.py
index 936912a3..7d577993 100644
--- a/memoryscope/core/worker/memory_manager.py
+++ b/memoryscope/core/worker/memory_manager.py
@@ -24,7 +24,7 @@ class MemoryManager(object):
         # dict: key -> memory_id
         self._key_id_dict: Dict[str, List[str]] = {}
 
-        self.logger = Logger.get_logger()
+        self.logger = Logger.get_logger(Logger.append_timestamp("memory_manager"))
 
     @property
     def memory_store(self) -> BaseMemoryStore:
@@ -95,7 +95,10 @@ class MemoryManager(object):
             self.logger.info(f"add to memory context memory id={node.memory_id} content={node.content} "
                              f"store_status={node.store_status} action_status={node.action_status}")
 
-        self._key_id_dict[key] = [n.memory_id for n in nodes]
+        self.logger.info(self.logger.format_current_memory(
+            [f"{node.memory_type} | {node.content}" for node in nodes]
+        ))
+        
 
     def get_memories(self, keys: str | List[str]) -> List[MemoryNode]:
         """