[Bug] Fix root logger handler

clarifai/utils/logging.py

@@ -1,6 +1,6 @@
 import logging
 from collections import defaultdict
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from rich import print as rprint
 from rich.logging import RichHandler
@@ -73,18 +73,29 @@ def _get_library_name() -> str:
   return __name__.split(".")[0]
 
 
-def _configure_logger(logger_level: str = "ERROR") -> None:
-  logging.basicConfig(
-      level=logger_level,
-      datefmt='%Y-%m-%d %H:%M:%S',
-      handlers=[RichHandler(rich_tracebacks=True)])
+def _configure_logger(name: str, logger_level: Union[int, str] = logging.NOTSET) -> None:
+  """Configure the logger with the specified name."""
 
+  logger = logging.getLogger(name)
+  logger.setLevel(logger_level)
 
-def get_logger(logger_level: str = "ERROR", name: Optional[str] = None) -> logging.Logger:
+  # Remove existing handlers
+  for handler in logger.handlers[:]:
+    logger.removeHandler(handler)
+
+  # Add the new rich handler and formatter
+  handler = RichHandler(rich_tracebacks=True, log_time_format="%Y-%m-%d %H:%M:%S")
+  formatter = logging.Formatter('%(name)s:  %(message)s')
+  handler.setFormatter(formatter)
+  logger.addHandler(handler)
+
+
+def get_logger(logger_level: Union[int, str] = logging.NOTSET,
+               name: Optional[str] = None) -> logging.Logger:
   """Return a logger with the specified name."""
 
   if name is None:
     name = _get_library_name()
 
-  _configure_logger(logger_level)
+  _configure_logger(name, logger_level)
   return logging.getLogger(name)

tests/test_misc.py

@@ -0,0 +1,19 @@
+import logging
+
+from rich.logging import RichHandler
+
+from clarifai.utils.logging import _get_library_name, get_logger
+
+
+def test_get_logger():
+  logger = get_logger("DEBUG", "test_logger")
+  assert logger.level == logging.DEBUG
+  assert logger.name == "test_logger"
+  assert isinstance(logger.handlers[0], RichHandler)
+
+
+def test_get_logger_defaults():
+  logger = get_logger()
+  assert logger.level == logging.NOTSET
+  assert logger.name == _get_library_name()
+  assert isinstance(logger.handlers[0], RichHandler)