diff --git a/chatlas/_display.py b/chatlas/_display.py index 6492cc7..33bc5bb 100644 --- a/chatlas/_display.py +++ b/chatlas/_display.py @@ -1,3 +1,4 @@ +import logging from abc import ABC, abstractmethod from typing import Any from uuid import uuid4 @@ -66,10 +67,15 @@ def update(self, content: str): def __enter__(self): self.live.__enter__() - # When logging is enabled, the RichHandler should be the second handler. - # That handler needs to know about the live console so it can add logs to it. - if len(logger.handlers) > 1 and isinstance(logger.handlers[1], RichHandler): - logger.handlers[1].console = self.live.console + # Live() isn't smart enough to know to automatically display logs when + # when they get handled while it Live() is active. + # However, if the logging handler is a RichHandler, it can be told + # about the live console so it can add logs to the top of the Live console. + handlers = [*logging.getLogger().handlers, *logger.handlers] + for h in handlers: + if isinstance(h, RichHandler): + h.console = self.live.console + return self def __exit__(self, exc_type, exc_value, traceback): diff --git a/chatlas/_logging.py b/chatlas/_logging.py index f22fc4c..214c9dc 100644 --- a/chatlas/_logging.py +++ b/chatlas/_logging.py @@ -1,19 +1,52 @@ import logging import os +import warnings from rich.logging import RichHandler -logger = logging.getLogger("chatlas") -if len(logger.handlers) == 0: - logger.addHandler(logging.NullHandler()) -if os.getenv("CHATLAS_LOG", "").lower() == "info": +def _rich_handler() -> RichHandler: formatter = logging.Formatter("%(name)s - %(message)s") handler = RichHandler() handler.setFormatter(formatter) - handler.setLevel(logging.INFO) - logger.addHandler(handler) + return handler + + +logger = logging.getLogger("chatlas") + +if os.environ.get("CHATLAS_LOG") == "info": + # By adding a RichHandler to chatlas' logger, we can guarantee that they + # never get dropped, even if the root logger's handlers are not + # RichHandlers. logger.setLevel(logging.INFO) + logger.addHandler(_rich_handler()) + logger.propagate = False + + # Add a RichHandler to the root logger if there are no handlers. Note that + # if chatlas is imported before other libraries that set up logging, (like + # openai, anthropic, or httpx), this will ensure that logs from those + # libraries are also displayed in the rich console. + root = logging.getLogger() + if not root.handlers: + root.addHandler(_rich_handler()) + + # Warn if there are non-RichHandler handlers on the root logger. + # TODO: we could consider something a bit more abusive here, like removing + # non-RichHandler handlers from the root logger, but that could be + # surprising to users. + bad_handlers = [ + h.get_name() for h in root.handlers if not isinstance(h, RichHandler) + ] + if len(bad_handlers) > 0: + warnings.warn( + "When setting up logging handlers for CHATLAS_LOG, chatlas detected " + f"non-rich handler(s) on the root logger named {bad_handlers}. " + "As a result, logs handled those handlers may be dropped when the " + "`echo` argument of `.chat()`, `.stream()`, etc., is something " + "other than 'none'. This problem can likely be fixed by importing " + "`chatlas` before other libraries that set up logging, or adding a " + "RichHandler to the root logger before loading other libraries.", + ) def log_model_default(model: str) -> str: