change log framework and fix llm_utils.py's logs

XianBW · XianBW · commit cce68ad81f7d · 2024-07-09T04:04:42.000Z
diff --git a/rdagent/core/conf.py b/rdagent/core/conf.py
@@ -16,6 +16,10 @@
 
 class RDAgentSettings(BaseSettings):
     # TODO: (xiao) I think most of the config should be in oai.config
+    # Log configs
+    log_trace_path: str | None = None
+    log_llm_chat_content: bool = True
+
     use_azure: bool = True
     use_azure_token_provider: bool = False
     managed_identity_client_id: str | None = None
@@ -28,7 +32,6 @@ class RDAgentSettings(BaseSettings):
     prompt_cache_path: str = str(Path.cwd() / "prompt_cache.db")
     session_cache_folder_location: str = str(Path.cwd() / "session_cache_folder/")
     max_past_message_include: int = 10
-    log_llm_chat_content: bool = True
 
     # Chat configs
     chat_openai_api_key: str = ""
diff --git a/rdagent/core/log.py b/rdagent/core/log.py
@@ -1,11 +1,138 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+import re
+import sys
+import pickle
+import json
+import inspect
+from typing import TYPE_CHECKING, Sequence, Literal
+
+if TYPE_CHECKING:
+    from loguru import Logger, Message, Record
 
 from loguru import logger
+from abc import abstractmethod
+from datetime import datetime, timezone
+from pathlib import Path
+from functools import partial
 
-if TYPE_CHECKING:
-    from loguru import Logger
+from rdagent.core.conf import RD_AGENT_SETTINGS
+from rdagent.core.utils import SingletonBaseClass
+
+
+def get_caller_info():
+    # Get the current stack information
+    stack = inspect.stack()
+    # The second element is usually the caller's information
+    caller_info = stack[2]
+    frame = caller_info[0]
+    info = {
+        'line': caller_info.lineno,
+        'name': frame.f_globals['__name__'],  # Get the module name from the frame's globals
+    }
+    return info
+
+
+class Storage:
+    """
+    Basic storage to support saving objects;
+
+    # Usage:
+
+    The storage has mainly two kind of users:
+    - The logging end: you can choose any of the following method to use the object
+        - We can use it directly with the native logging storage
+        - We can use it with other logging tools; For example, serve as a handler for loggers
+    - The view end:
+        - Mainly for the subclass of `logging.base.View`
+        - It should provide two kind of ways to provide content
+            - offline content provision.
+            - online content preovision.
+    """
+
+    @abstractmethod
+    def log(self, obj: object, name: str = "", **kwargs: dict) -> str | Path | None:
+        """
+
+        Parameters
+        ----------
+        obj : object
+            The object for logging.
+        name : str
+            The name of the object.  For example "a.b.c"
+            We may log a lot of objects to a same name
+        """
+        ...
+
+
+class View:
+    """
+    Motivation:
+
+    Display the content in the storage
+    """
+
+
+class FileStorage(Storage):
+    """
+    The info are logginged to the file systems
+
+    TODO: describe the storage format
+    """
+
+    def __init__(self, path: str = "./log/") -> None:
+        self.path = Path(path)
+        self.path.mkdir(parents=True, exist_ok=True)
+
+    def log(self,
+            obj: object,
+            name: str = "",
+            save_type: Literal["json", "text", "pkl", "short-text"] = "short-text",
+            timestamp: datetime | None = None,
+            split_name: bool = True,
+            ) -> Path:
+        if timestamp is None:
+            timestamp = datetime.now(timezone.utc)
+        else:
+            timestamp = timestamp.astimezone(timezone.utc)
+        
+        cur_p = self.path
+        if split_name:
+            uri_l = name.split(".")
+            for u in uri_l:
+                cur_p = cur_p / u
+        else:
+            cur_p = cur_p / name
+        cur_p.mkdir(parents=True, exist_ok=True)
+
+        path = cur_p / f"{timestamp.strftime('%Y-%m-%d_%H-%M-%S-%f')}.log"
+
+        if save_type == "json":
+            path = path.with_suffix(".json")
+            with path.open("w") as f:
+                try:
+                    json.dump(obj, f)
+                except TypeError:
+                    json.dump(json.loads(str(obj)), f)
+            return path
+        elif save_type == "pkl":
+            path = path.with_suffix(".pkl")
+            with path.open("wb") as f:
+                pickle.dump(obj, f)
+            return path
+        elif save_type == "text":
+            obj = str(obj)
+            with path.open("w") as f:
+                f.write(obj)
+            return path
+        else:
+            obj = str(obj).strip()
+            if obj == "":
+                return
+            path = cur_p / "common_logs.log"
+            with path.open("a") as f:
+                f.write(f"{timestamp.isoformat()}: {obj}\n\n") # add a new line to separate logs
+            return path
 
 
 class LogColors:
@@ -54,40 +181,67 @@ def render(self, text: str, color: str = "", style: str = "") -> str:
         return f"{style}{text}{self.END}"
 
 
-class RDAgentLog:
-    # logger.add(loguru_handler, level="INFO")   # you can add use storage as a loguru handler
+def loguru2storage_handler(storage: Storage, record: Message) -> None:
+    msg = f"{record.record['level']} | {record.record['name']}:{record.record['line']} - {RDAgentLog.remove_ansi_codes(record.record['message'])}"
+    storage.log(msg, timestamp=record.record["time"], save_type="short-text")
+
+
+class RDAgentLog(SingletonBaseClass):
+
+    def __init__(self, log_trace_path: str | None = RD_AGENT_SETTINGS.log_trace_path) -> None:
+        if log_trace_path is None:
+            timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d_%H-%M-%S-%f")
+            log_trace_path: Path = Path.cwd() / "log" / timestamp
+            log_trace_path.mkdir(parents=True, exist_ok=True)
+        
+        self.storage = FileStorage(log_trace_path)
+
+        # add handler to save log to storage
+        logger.add(partial(loguru2storage_handler, self.storage))
+
+        self.log_stream = self.LogStreamContextManager(self.storage)
+
+    @staticmethod
+    def remove_ansi_codes(s: str) -> str:
+        ansi_escape = re.compile(r'\x1B\[[0-?]*[ -/]*[@-~]')
+        return ansi_escape.sub('', s)
+
+    class LogStreamContextManager:
+        def __init__(self, storage: Storage) -> None:
+            self.captured_logs = []
+            self.storage = storage
+
+        def capture(self, message: Message) -> None:
+            self.captured_logs.append(message.record["message"])
+
+        def __enter__(self):
+            logger.remove()
+            logger.add(sys.stderr, format=lambda x: x["message"])
+            logger.add(self.capture)
 
-    def __init__(self) -> None:
-        self.logger: Logger = logger
+        def __exit__(self, exc_type, exc_value, traceback) -> None:
+            logger.info('\n')
+            logger.remove()
+            logger.add(partial(loguru2storage_handler, self.storage))
+            logger.info("[stream log] " + "".join(self.captured_logs))
+            logger.add(sys.stderr)
 
-    def info(self, *args: Sequence, plain: bool = False) -> None:
-        if plain:
-            return self.plain_info(*args)
-        for arg in args:
-            info = f"{LogColors.WHITE}{arg}{LogColors.END}"
-            self.logger.info(info)
-        return None
+    def log_objects(self, *objs: Sequence[object]) -> None:
+        caller_info = get_caller_info()
+        for obj in objs:
+            logp = self.storage.log(obj, name=f"{type(obj).__module__}.{type(obj).__name__}", save_type="pkl", split_name=False)
 
-    def __getstate__(self) -> dict:
-        return {}
+            logger.patch(lambda r: r.update(caller_info)).info(f"Logging object in {logp.absolute()}")
 
-    def __setstate__(self, _: str) -> None:
-        self.logger = logger
+    def info(self, msg: str) -> None:
+        caller_info = get_caller_info()
+        logger.patch(lambda r: r.update(caller_info)).info(msg)
 
-    def plain_info(self, *args: Sequence) -> None:
-        for arg in args:
-            info = f"""
-                {LogColors.YELLOW}{LogColors.BOLD}
-                Info:{LogColors.END}{LogColors.WHITE}{arg}{LogColors.END}
-            """
-            self.logger.info(info)
+    def warning(self, msg: str) -> None:
+        caller_info = get_caller_info()
+        logger.patch(lambda r: r.update(caller_info)).warning(msg)
 
-    def warning(self, *args: Sequence) -> None:
-        for arg in args:
-            info = f"{LogColors.BLUE}{LogColors.BOLD}Warning:{LogColors.END}{arg}"
-            self.logger.warning(info)
+    def error(self, msg: str) -> None:
+        caller_info = get_caller_info()
+        logger.patch(lambda r: r.update(caller_info)).error(msg)
 
-    def error(self, *args: Sequence) -> None:
-        for arg in args:
-            info = f"{LogColors.RED}{LogColors.BOLD}Error:{LogColors.END}{arg}"
-            self.logger.error(info)
diff --git a/rdagent/oai/llm_utils.py b/rdagent/oai/llm_utils.py
