[Feat] Support message logging (#267)

* add a message logging hook

* update

* change the default logger name

---------

Co-authored-by: wangzy <[email protected]>

Author: braisedpork1964

lagent/agents/agent.py

@@ -168,6 +168,9 @@ def reset(self, session_id=0):
         if self.memory:
             self.memory.reset(session_id=session_id)
 
+    def __repr__(self):
+        return f"{self.__class__.__name__}(name='{self.name}', description='{self.description or ''}')"
+
 
 class AsyncAgent(Agent):
 
@@ -273,15 +276,15 @@ def _backup(d):
                 setattr(cls, method, wrap_api(getattr(cls, method)))
 
 
-class AgentList(Agent, UserList, AgentContainerMixin):
+class AgentList(UserList, Agent, AgentContainerMixin):
 
     def __init__(self,
                  agents: Optional[Iterable[Union[Agent, AsyncAgent]]] = None):
         Agent.__init__(self, memory=None)
         UserList.__init__(self, agents)
 
 
-class AgentDict(Agent, UserDict, AgentContainerMixin):
+class AgentDict(UserDict, Agent, AgentContainerMixin):
 
     def __init__(self,
                  agents: Optional[Mapping[str, Union[Agent,

lagent/hooks/__init__.py

@@ -1,6 +1,8 @@
 from .action_preprocessor import ActionPreprocessor, InternLMActionProcessor
 from .hook import Hook, RemovableHandle
+from .logger import MessageLogger
 
 __all__ = [
-    'Hook', 'RemovableHandle', 'ActionPreprocessor', 'InternLMActionProcessor'
+    'Hook', 'RemovableHandle', 'ActionPreprocessor', 'InternLMActionProcessor',
+    'MessageLogger'
 ]

lagent/hooks/hook.py

@@ -1,37 +1,40 @@
 from itertools import count
+from typing import Tuple
+
+from lagent.schema import AgentMessage
 
 
 class Hook:
 
     def before_agent(
         self,
         agent,
-        message,
-        session_id,
+        message: Tuple[AgentMessage],
+        session_id: int,
     ):
         pass
 
     def after_agent(
         self,
         agent,
-        message,
-        session_id,
+        message: AgentMessage,
+        session_id: int,
     ):
         pass
 
     def before_action(
         self,
         executor,
-        message,
-        session_id,
+        message: AgentMessage,
+        session_id: int,
     ):
         pass
 
     def after_action(
         self,
         executor,
-        message,
-        session_id,
+        message: AgentMessage,
+        session_id: int,
     ):
         pass
 

lagent/hooks/logger.py

@@ -0,0 +1,36 @@
+import random
+from typing import Optional
+
+from termcolor import COLORS, colored
+
+from lagent.utils import get_logger
+from .hook import Hook
+
+
+class MessageLogger(Hook):
+
+    def __init__(self, name: str = 'lagent'):
+        self.logger = get_logger(name, 'info')
+        self.sender2color = {}
+
+    def before_agent(self, agent, messages, session_id):
+        for message in messages:
+            self._process_message(message, session_id)
+
+    def after_agent(self, agent, message, session_id):
+        self._process_message(message, session_id)
+
+    def before_action(self, executor, message, session_id):
+        self._process_message(message, session_id)
+
+    def after_action(self, executor, message, session_id):
+        self._process_message(message, session_id)
+
+    def _process_message(self, message, session_id):
+        sender = message.sender
+        color = self.sender2color.setdefault(sender,
+                                             random.choice(list(COLORS)))
+        self.logger.info(
+            colored(
+                f'session id: {session_id}, message sender: {sender}\n'
+                f'{message.content}', color))

lagent/utils/__init__.py

@@ -1,7 +1,14 @@
 from .package import is_module_exist
-from .util import GeneratorWithReturn, async_as_completed, create_object, filter_suffix, load_class_from_string
+from .util import (
+    GeneratorWithReturn,
+    async_as_completed,
+    create_object,
+    filter_suffix,
+    get_logger,
+    load_class_from_string,
+)
 
 __all__ = [
-    'is_module_exist', 'filter_suffix', 'create_object',
+    'is_module_exist', 'filter_suffix', 'create_object', 'get_logger',
     'load_class_from_string', 'async_as_completed', 'GeneratorWithReturn'
 ]

lagent/utils/util.py

@@ -1,8 +1,13 @@
 import asyncio
 import importlib
 import inspect
+import logging
+import os
+import os.path as osp
 import sys
+import time
 from functools import partial
+from logging.handlers import RotatingFileHandler
 from typing import Any, Dict, Generator, Iterable, List, Optional, Union
 
 
@@ -87,6 +92,41 @@ def filter_suffix(response: Union[str, List[str]],
     return processed
 
 
+def get_logger(
+    name: str = 'lagent',
+    level: str = 'debug',
+    fmt:
+    str = '%(asctime)s %(levelname)8s %(filename)20s %(lineno)4s - %(message)s',
+    add_file_handler: bool = False,
+    log_dir: str = 'log',
+    log_file: str = time.strftime('%Y-%m-%d.log', time.localtime()),
+    max_bytes: int = 5 * 1024 * 1024,
+    backup_count: int = 3,
+):
+    logger = logging.getLogger(name)
+    logger.propagate = False
+    logger.setLevel(getattr(logging, level.upper(), logging.DEBUG))
+
+    formatter = logging.Formatter(fmt)
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+
+    if add_file_handler:
+        if not osp.exists(log_dir):
+            os.makedirs(log_dir)
+        log_file_path = osp.join(log_dir, log_file)
+        file_handler = RotatingFileHandler(
+            log_file_path,
+            maxBytes=max_bytes,
+            backupCount=backup_count,
+            encoding='utf-8')
+        file_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)
+
+    return logger
+
+
 class GeneratorWithReturn:
     """Generator wrapper to capture the return value."""