Skip to content

Commit

Permalink
Improve logger
Browse files Browse the repository at this point in the history
  • Loading branch information
@yxjiang
yxjiang committed Jul 23, 2024
1 parent ace4ca6 commit e821f29
Show file tree
Hide file tree
Showing 3 changed files with 107 additions and 185 deletions.
43 changes: 32 additions & 11 deletions polymind/core/logger.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ def __init__(
load_dotenv(override=True)

if display_level is None:
env_level = os.getenv("LOGGING_LEVEL", "INFO")
env_level = os.getenv("LOGGING_LEVEL", "DEBUG") # Change default to DEBUG
self.logging_level = self.LoggingLevel.from_string(env_level)
elif isinstance(display_level, str):
self.logging_level = self.LoggingLevel.from_string(display_level)
Expand All @@ -61,7 +61,12 @@ def __init__(
self.console_handler.setFormatter(self.formatter)
self.logger.addHandler(self.console_handler)

def log(self, message: str, level: LoggingLevel, color: str = ansi.Fore.GREEN) -> None:
# Add custom log levels
logging.addLevelName(self.LoggingLevel.TOOL.value, "TOOL")
logging.addLevelName(self.LoggingLevel.TASK.value, "TASK")
logging.addLevelName(self.LoggingLevel.THOUGHT_PROCESS.value, "THOUGHT_PROCESS")

def _log(self, message: str, level: LoggingLevel, color: str) -> None:
if level.value >= self.logging_level.value:
if len(inspect.stack()) >= 4:
caller_frame = inspect.stack()[3]
Expand All @@ -71,28 +76,44 @@ def log(self, message: str, level: LoggingLevel, color: str = ansi.Fore.GREEN) -
caller_line = caller_frame.lineno
message = f"{caller_name}({caller_line}): {message}"
log_message = color + message + Fore.RESET
self.logger.log(level.value, log_message)

if level == self.LoggingLevel.DEBUG:
self.logger.debug(log_message)
elif level == self.LoggingLevel.INFO:
self.logger.info(log_message)
elif level == self.LoggingLevel.TOOL:
self.logger.log(level.value, log_message)
elif level == self.LoggingLevel.TASK:
self.logger.log(level.value, log_message)
elif level == self.LoggingLevel.THOUGHT_PROCESS:
self.logger.log(level.value, log_message)
elif level == self.LoggingLevel.WARNING:
self.logger.warning(log_message)
elif level == self.LoggingLevel.ERROR:
self.logger.error(log_message)
elif level == self.LoggingLevel.CRITICAL:
self.logger.critical(log_message)

def debug(self, message: str) -> None:
self.log(message, Logger.LoggingLevel.DEBUG, Fore.BLACK)
self._log(message, self.LoggingLevel.DEBUG, Fore.BLACK)

def info(self, message: str) -> None:
self.log(message, Logger.LoggingLevel.INFO, Fore.WHITE)
self._log(message, self.LoggingLevel.INFO, Fore.WHITE)

def tool_log(self, message: str) -> None:
self.log(message, Logger.LoggingLevel.TOOL, Fore.YELLOW)
self._log(message, self.LoggingLevel.TOOL, Fore.YELLOW)

def task_log(self, message: str) -> None:
self.log(message, Logger.LoggingLevel.TASK, Fore.BLUE)
self._log(message, self.LoggingLevel.TASK, Fore.BLUE)

def thought_process_log(self, message: str) -> None:
self.log(message, Logger.LoggingLevel.THOUGHT_PROCESS, Fore.GREEN)
self._log(message, self.LoggingLevel.THOUGHT_PROCESS, Fore.GREEN)

def warning(self, message: str) -> None:
self.log(message, Logger.LoggingLevel.WARNING, Fore.YELLOW)
self._log(message, self.LoggingLevel.WARNING, Fore.YELLOW)

def error(self, message: str) -> None:
self.log(message, Logger.LoggingLevel.ERROR, Fore.RED)
self._log(message, self.LoggingLevel.ERROR, Fore.RED)

def critical(self, message: str) -> None:
self.log(message, Logger.LoggingLevel.CRITICAL, Fore.MAGENTA)
self._log(message, self.LoggingLevel.CRITICAL, Fore.MAGENTA)
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "polymind"
version = "0.0.57" # Update this version before publishing to PyPI
version = "0.0.58" # Update this version before publishing to PyPI
description = "PolyMind is a customizable collaborative multi-agent framework for collective intelligence and distributed problem solving."
authors = ["TechTao"]
license = "MIT License"
Expand Down
247 changes: 74 additions & 173 deletions tests/polymind/core/test_logger.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,185 +17,86 @@ class TestLogger:
def setup_logger(self, tmp_path):
log_folder = tmp_path / "logs"
log_folder.mkdir()
logger = Logger(logger_name="test_logger", display_level="INFO")
logger = Logger(logger_name="test_logger", display_level="DEBUG") # Change to DEBUG
return logger

def test_initialization(self, setup_logger):
logger = setup_logger
assert (
logger.logging_level == Logger.LoggingLevel.INFO
), f"Logger level is incorrect: {logger.logging_level} != INFO"
assert logger.logger.level == logging.INFO, f"Logger level is not set correctly: {logger.logger.level} != INFO"
assert len(logger.logger.handlers) == 1, "Logger handlers are not initialized properly"

def test_logging_levels(self, setup_logger):
logger = setup_logger
logger.info("Info message")
logger.warning("Warning message")
logger.logging_level == Logger.LoggingLevel.DEBUG
), f"Logger level is incorrect: {logger.logging_level} != DEBUG"
assert (
logger.logging_level == Logger.LoggingLevel.INFO
), f"Logger level is incorrect: {logger.logging_level} != INFO"
assert logger.logger.level == logging.INFO, f"Logger level is not set correctly: {logger.logger.level} != INFO"
logger.logger.level == Logger.LoggingLevel.DEBUG.value
), f"Logger level is not set correctly: {logger.logger.level} != {Logger.LoggingLevel.DEBUG.value}"
assert len(logger.logger.handlers) == 1, "Logger handlers are not initialized properly"

@pytest.mark.parametrize(
"log_method, log_level, color, expected_log_method",
[
("debug", Logger.LoggingLevel.DEBUG, Fore.BLACK, "debug"),
("info", Logger.LoggingLevel.INFO, Fore.WHITE, "info"),
("tool_log", Logger.LoggingLevel.TOOL, Fore.YELLOW, "log"),
("task_log", Logger.LoggingLevel.TASK, Fore.BLUE, "log"),
("thought_process_log", Logger.LoggingLevel.THOUGHT_PROCESS, Fore.GREEN, "log"),
("warning", Logger.LoggingLevel.WARNING, Fore.YELLOW, "warning"),
("error", Logger.LoggingLevel.ERROR, Fore.RED, "error"),
("critical", Logger.LoggingLevel.CRITICAL, Fore.MAGENTA, "critical"),
],
)
@patch("inspect.stack", return_value=[None, None, None, MagicMock(function="test_func", lineno=42)])
def test_log_method(self, mock_stack, setup_logger):
logger = setup_logger
logger.log("Test log message", Logger.LoggingLevel.INFO, color=Fore.WHITE)
log_message = "test_func(42): Test log message"
log_output = f"{Fore.WHITE}{log_message}{Fore.RESET}"
formatted_message = logger.formatter.format(
logging.LogRecord(
name="test_logger",
level=logging.INFO,
pathname="",
lineno=42,
msg=log_output,
args=None,
exc_info=None,
)
)
assert log_message in formatted_message, "Log message is incorrect"

def test_debug_log(self, setup_logger):
logger = setup_logger
logger.debug("Debug message")
log_message = "pytest_pyfunc_call(1): Debug message"
log_output = f"{Fore.BLACK}{log_message}{Fore.RESET}"
formatted_message = logger.formatter.format(
logging.LogRecord(
name="test_logger",
level=logging.DEBUG,
pathname="",
lineno=1,
msg=log_output,
args=None,
exc_info=None,
)
)
assert log_message in formatted_message, "Debug log message is incorrect"

def test_info_log(self, setup_logger):
logger = setup_logger
logger.info("Info message")
log_message = "pytest_pyfunc_call(1): Info message"
log_output = f"{Fore.WHITE}{log_message}{Fore.RESET}"
formatted_message = logger.formatter.format(
logging.LogRecord(
name="test_logger",
level=logging.INFO,
pathname="",
lineno=1,
msg=log_output,
args=None,
exc_info=None,
)
)
assert log_message in formatted_message, "Info log message is incorrect"

def test_tool_log(self, setup_logger):
logger = setup_logger
logger.tool_log("Tool log message")
log_message = "pytest_pyfunc_call(1): Tool log message"
log_output = f"{Fore.YELLOW}{log_message}{Fore.RESET}"
formatted_message = logger.formatter.format(
logging.LogRecord(
name="test_logger",
level=logging.INFO,
pathname="",
lineno=1,
msg=log_output,
args=None,
exc_info=None,
)
)
assert log_message in formatted_message, "Tool log message is incorrect"

def test_task_log(self, setup_logger):
logger = setup_logger
logger.task_log("Task log message")
log_message = "pytest_pyfunc_call(1): Task log message"
log_output = f"{Fore.BLUE}{log_message}{Fore.RESET}"
formatted_message = logger.formatter.format(
logging.LogRecord(
name="test_logger",
level=logging.INFO,
pathname="",
lineno=1,
msg=log_output,
args=None,
exc_info=None,
)
)
assert log_message in formatted_message, "Task log message is incorrect"

def test_thought_process_log(self, setup_logger):
def test_log_methods(self, mock_stack, setup_logger, log_method, log_level, color, expected_log_method):
logger = setup_logger
logger.thought_process_log("Thought process log message")
log_message = "pytest_pyfunc_call(1): Thought process log message"
log_output = f"{Fore.GREEN}{log_message}{Fore.RESET}"
formatted_message = logger.formatter.format(
logging.LogRecord(
name="test_logger",
level=logging.INFO,
pathname="",
lineno=1,
msg=log_output,
args=None,
exc_info=None,
)
)
assert log_message in formatted_message, "Thought process log message is incorrect"

def test_warning_log(self, setup_logger):
logger = setup_logger
logger.warning("Warning message")
log_message = "pytest_pyfunc_call(1): Warning message"
log_output = f"{Fore.YELLOW}{log_message}{Fore.RESET}"
formatted_message = logger.formatter.format(
logging.LogRecord(
name="test_logger",
level=logging.WARNING,
pathname="",
lineno=1,
msg=log_output,
args=None,
exc_info=None,
)
)
assert log_message in formatted_message, "Warning log message is incorrect"

def test_error_log(self, setup_logger):
logger = setup_logger
logger.error("Error message")
log_message = "pytest_pyfunc_call(1): Error message"
log_output = f"{Fore.RED}{log_message}{Fore.RESET}"
formatted_message = logger.formatter.format(
logging.LogRecord(
name="test_logger",
level=logging.ERROR,
pathname="",
lineno=1,
msg=log_output,
args=None,
exc_info=None,
)
)
assert log_message in formatted_message, "Error log message is incorrect"

def test_critical_log(self, setup_logger):
logger = setup_logger
logger.critical("Critical message")
log_message = "pytest_pyfunc_call(1): Critical message"
log_output = f"{Fore.MAGENTA}{log_message}{Fore.RESET}"
formatted_message = logger.formatter.format(
logging.LogRecord(
name="test_logger",
level=logging.CRITICAL,
pathname="",
lineno=1,
msg=log_output,
args=None,
exc_info=None,
)
)
assert log_message in formatted_message, "Critical log message is incorrect"
log_func = getattr(logger, log_method)

with patch.object(logger.logger, expected_log_method) as mock_log_method:
log_func("Test message")

log_message = "test_func(42): Test message"
expected_log_message = f"{color}{log_message}{Fore.RESET}"

if expected_log_method == "log":
mock_log_method.assert_called_once_with(log_level.value, expected_log_message)
else:
mock_log_method.assert_called_once_with(expected_log_message)

def test_logging_levels_order(self):
assert (
Logger.LoggingLevel.DEBUG.value
< Logger.LoggingLevel.INFO.value
< Logger.LoggingLevel.TOOL.value
< Logger.LoggingLevel.TASK.value
< Logger.LoggingLevel.THOUGHT_PROCESS.value
< Logger.LoggingLevel.WARNING.value
< Logger.LoggingLevel.ERROR.value
< Logger.LoggingLevel.CRITICAL.value
), "Logging levels are not in the correct order"

def test_custom_logging_levels(self):
assert Logger.LoggingLevel.TOOL.value == 25, "TOOL logging level is incorrect"
assert Logger.LoggingLevel.TASK.value == 26, "TASK logging level is incorrect"
assert Logger.LoggingLevel.THOUGHT_PROCESS.value == 27, "THOUGHT_PROCESS logging level is incorrect"

def test_from_string_method(self):
assert Logger.LoggingLevel.from_string("DEBUG") == Logger.LoggingLevel.DEBUG
assert Logger.LoggingLevel.from_string("info") == Logger.LoggingLevel.INFO
assert Logger.LoggingLevel.from_string("CRITICAL") == Logger.LoggingLevel.CRITICAL

with pytest.raises(ValueError):
Logger.LoggingLevel.from_string("INVALID_LEVEL")

def test_logger_singleton(self):
logger1 = Logger(logger_name="test_logger1")
logger2 = Logger(logger_name="test_logger2")
assert logger1 is logger2, "Logger is not implementing the singleton pattern correctly"

@patch("logging.getLogger")
def test_custom_log_levels_added(self, mock_get_logger):
mock_logger = MagicMock()
mock_get_logger.return_value = mock_logger

Logger(logger_name="test_logger")

mock_logger.addHandler.assert_called()
assert logging.getLevelName(25) == "TOOL", "TOOL log level not added"
assert logging.getLevelName(26) == "TASK", "TASK log level not added"
assert logging.getLevelName(27) == "THOUGHT_PROCESS", "THOUGHT_PROCESS log level not added"

0 comments on commit e821f29

Please sign in to comment.