Spaces:
Sleeping
Sleeping
File size: 1,216 Bytes
e679d69 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
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))
|