Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[HOTFIX] Fix the bug in loguru that cannot display special tags #266

Merged
merged 3 commits into from
May 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 51 additions & 26 deletions src/agentscope/utils/logging_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,17 +24,19 @@
"CRITICAL",
]

LEVEL_CHAT_LOG = "CHAT_LOG"
LEVEL_CHAT_SAVE = "CHAT_SAVE"
LEVEL_SAVE_LOG = "SAVE_LOG"
LEVEL_SAVE_MSG = "SAVE_MSG"
LEVEL_DISPLAY_MSG = "DISPLAY_MSG"

_SPEAKER_COLORS = [
("<blue>", "</blue>"),
("<cyan>", "</cyan>"),
("<green>", "</green>"),
("<magenta>", "</magenta>"),
("<red>", "</red>"),
("<white>", "</white>"),
("<yellow>", "</yellow>"),
("\033[90m", "\033[0m"),
("\033[91m", "\033[0m"),
("\033[92m", "\033[0m"),
("\033[93m", "\033[0m"),
("\033[94m", "\033[0m"),
("\033[95m", "\033[0m"),
("\033[96m", "\033[0m"),
("\033[97m", "\033[0m"),
]

_SPEAKER_TO_COLORS = {}
Expand Down Expand Up @@ -82,9 +84,10 @@ def _chat(
"content" keys, and the message will be logged as "<name/role>:
<content>".
"""
# Save message into file, add default to ignore not serializable objects
# Save message into chat file, add default to ignore not serializable
# objects
logger.log(
LEVEL_CHAT_SAVE,
LEVEL_SAVE_MSG,
json.dumps(message, ensure_ascii=False, default=lambda _: None),
*args,
**kwargs,
Expand All @@ -102,26 +105,43 @@ def _chat(
(m1, m2) = _get_speaker_color(speaker)

print_str = []
print_str_without_markers = []
if contain_content:
print_str.append(
f"{m1}<b>{speaker}</b>{m2}: {message['content']}",
f"{m1}\033[1m{speaker}\033[0m{m2}: {message['content']}",
)
print_str_without_markers.append(
f"{speaker}: {message['content']}",
)

if contain_url:
print_str.append(f"{m1}<b>{speaker}</b>{m2}: {message['url']}")
print_str.append(
f"{m1}\033[1m{speaker}\033[0m{m2}: {message['url']}",
)
print_str_without_markers.append(
f"{speaker}: {message['url']}",
)

if len(print_str) > 0:
print_str = (
"\n".join(print_str).replace("{", "{{").replace("}", "}}")
print_str = "\n".join(print_str)
print_str_without_markers = "\n".join(
print_str_without_markers,
)

logger.log(LEVEL_DISPLAY_MSG, print_str, *args, **kwargs)
logger.log(
LEVEL_SAVE_LOG,
print_str_without_markers,
*args,
**kwargs,
)
logger.log(LEVEL_CHAT_LOG, print_str, *args, **kwargs)

if hasattr(thread_local_data, "uid") and not disable_studio:
log_studio(message, thread_local_data.uid, **kwargs)
return

message = str(message).replace("{", "{{").replace("}", "}}")
logger.log(LEVEL_CHAT_LOG, message, *args, **kwargs)
logger.log(LEVEL_DISPLAY_MSG, message, *args, **kwargs)
logger.log(LEVEL_SAVE_LOG, message, *args, **kwargs)


def log_studio(message: dict, uid: str, **kwargs: Any) -> None:
Expand Down Expand Up @@ -182,8 +202,9 @@ def log_studio(message: dict, uid: str, **kwargs: Any) -> None:

def _level_format(record: dict) -> str:
"""Format the log record."""
if record["level"].name == LEVEL_CHAT_LOG:
return record["message"] + "\n"
# Display the chat message
if record["level"].name in [LEVEL_DISPLAY_MSG, LEVEL_SAVE_LOG]:
return "{message}\n"
DavdGao marked this conversation as resolved.
Show resolved Hide resolved
else:
return (
"<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{"
Expand All @@ -210,16 +231,20 @@ def setup_logger(
# avoid reinit in subprocess
if not hasattr(logger, "chat"):
# add chat function for logger
logger.level(LEVEL_CHAT_LOG, no=21)
logger.level(LEVEL_CHAT_SAVE, no=0)
logger.level(LEVEL_SAVE_LOG, no=51)
logger.level(LEVEL_DISPLAY_MSG, no=52)

# save chat message into file
logger.level(LEVEL_SAVE_MSG, no=53)
logger.chat = _chat

# set logging level
logger.remove()
# standard output for all logging except chat
logger.add(
sys.stdout,
filter=lambda record: record["level"].name != LEVEL_CHAT_SAVE,
filter=lambda record: record["level"].name
not in [LEVEL_SAVE_LOG, LEVEL_SAVE_MSG],
format=_level_format,
enqueue=True,
level=level,
Expand All @@ -236,16 +261,16 @@ def setup_logger(
# save all logging into file
logger.add(
path_log_file,
filter=lambda record: record["level"].name != LEVEL_CHAT_SAVE,
filter=lambda record: record["level"].name
not in [LEVEL_SAVE_MSG, LEVEL_DISPLAY_MSG],
format=_level_format,
enqueue=True,
level=level,
)

logger.add(
path_chat_file,
filter=lambda record: record["level"].name == LEVEL_CHAT_SAVE,
format="{message}",
enqueue=True,
level=LEVEL_CHAT_SAVE,
level=LEVEL_SAVE_MSG,
)
6 changes: 5 additions & 1 deletion tests/logger_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,9 @@ def test_logger_chat(self) -> None:
# dict
logger.chat({"abc": 1})

# html labels
logger.chat({"name": "Bob", "content": "<div>abc</div"})

# To avoid that logging is not finished before the file is read
time.sleep(3)

Expand All @@ -58,13 +61,14 @@ def test_logger_chat(self) -> None:
'"}\n',
'{"name": "Alice", "url": "https://xxx.png"}\n',
'{"abc": 1}\n',
'{"name": "Bob", "content": "<div>abc</div"}\n',
]

self.assertListEqual(lines, ground_truth)

def tearDown(self) -> None:
"""Tear down for LoggerTest."""
logger.stop()
logger.remove()
if os.path.exists(self.run_dir):
shutil.rmtree(self.run_dir)

Expand Down