lagent/hooks/logger.py

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', '%(asctime)s %(levelname)8s %(name)8s - %(message)s')
        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))