From 590e92a6fc371c65907094b147437a679aa778e7 Mon Sep 17 00:00:00 2001
From: Carson <cpsievert1@gmail.com>
Date: Tue, 10 Dec 2024 13:39:43 -0600
Subject: [PATCH] Do more to prevent logs from being dropped when a Live()
 display is active

---
 chatlas/_display.py | 14 ++++++++++----
 chatlas/_logging.py | 45 +++++++++++++++++++++++++++++++++++++++------
 2 files changed, 49 insertions(+), 10 deletions(-)

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: