Enhance logging configuration to support optional Rich logging

src/litserve/utils.py

@@ -98,23 +98,38 @@ class WorkerSetupStatus:
     FINISHED: str = "finished"
 
 
+def _get_default_handler(stream, format):
+    handler = logging.StreamHandler(stream)
+    formatter = logging.Formatter(format)
+    handler.setFormatter(formatter)
+    return handler
+
+
 def configure_logging(
-    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", stream=sys.stdout
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", stream=sys.stdout, use_rich=False
 ):
     """Configure logging for the entire library with sensible defaults.
 
     Args:
         level (int): Logging level (default: logging.INFO)
         format (str): Log message format string
         stream (file-like): Output stream for logs
+        use_rich (bool): Whether to use rich for logging
 
     """
-    # Create a library-wide handler
-    handler = logging.StreamHandler(stream)
+    if use_rich:
+        try:
+            from rich.logging import RichHandler
+            from rich.traceback import install
 
-    # Set formatter with user-configurable format
-    formatter = logging.Formatter(format)
-    handler.setFormatter(formatter)
+            install(show_locals=True)
+            handler = RichHandler(rich_tracebacks=True, show_time=True, show_path=True)
+        except ImportError:
+            logger.warning("Rich is not installed, using default logging")
+            handler = _get_default_handler(stream, format)
+
+    else:
+        handler = _get_default_handler(stream, format)
 
     # Configure root library logger
     library_logger = logging.getLogger("litserve")

tests/test_utils.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import pickle
 import sys
@@ -7,7 +8,13 @@
 import pytest
 from fastapi import HTTPException
 
-from litserve.utils import call_after_stream, dump_exception, generate_random_zmq_address, set_trace_if_debug
+from litserve.utils import (
+    call_after_stream,
+    configure_logging,
+    dump_exception,
+    generate_random_zmq_address,
+    set_trace_if_debug,
+)
 
 
 def test_dump_exception():
@@ -53,6 +60,18 @@ def test_generate_random_zmq_address_non_windows(tmpdir):
     assert os.path.commonpath([temp_dir, address2[6:]]) == temp_dir
 
 
+def test_configure_logging():
+    configure_logging(use_rich=False)
+    assert logging.getLogger("litserve").handlers[0].__class__.__name__ == "StreamHandler"
+
+
+def test_configure_logging_rich_not_installed():
+    # patch builtins.__import__ to raise ImportError
+    with mock.patch("builtins.__import__", side_effect=ImportError):
+        configure_logging(use_rich=True)
+        assert logging.getLogger("litserve").handlers[0].__class__.__name__ == "StreamHandler"
+
+
 @mock.patch("litserve.utils.set_trace")
 def test_set_trace_if_debug(mock_set_trace):
     # mock environ