Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Do more to prevent logs from being dropped when a Live() display is a…
Browse files Browse the repository at this point in the history
…ctive
cpsievert committed Dec 10, 2024
1 parent a3d09e5 commit 590e92a
Showing 2 changed files with 49 additions and 10 deletions.
14 changes: 10 additions & 4 deletions chatlas/_display.py
Original file line number Diff line number Diff line change
@@ -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):
45 changes: 39 additions & 6 deletions chatlas/_logging.py
Original file line number Diff line number Diff line change
@@ -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:

0 comments on commit 590e92a

Please sign in to comment.