From c641eb4ad6cfbaa266cdb74fbb8c815531fed67f Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Wed, 15 Feb 2023 01:10:46 -0600
Subject: [PATCH] Improve rendering log and error stacktraces using the Rich
 package

- Use Rich to render uvicorn, fastAPI logs as well
  The previous CustomFormatter only worked on khoj logs
- Improve rendering stacktrace on errors using Rich
---
 setup.py                  |  1 +
 src/khoj/main.py          | 19 +++++++++----------
 src/khoj/utils/helpers.py | 26 +-------------------------
 3 files changed, 11 insertions(+), 35 deletions(-)

diff --git a/setup.py b/setup.py
index 2ea55970..8d7b6920 100644
--- a/setup.py
+++ b/setup.py
@@ -37,6 +37,7 @@ setup(
         "pyqt6 == 6.3.1",
         "defusedxml == 0.7.1",
         'schedule == 1.1.0',
+        "rich >= 13.3.1",
     ],
     include_package_data=True,
     entry_points={"console_scripts": ["khoj = khoj.main:run"]},
diff --git a/src/khoj/main.py b/src/khoj/main.py
index 8461f1b2..8c21aee3 100644
--- a/src/khoj/main.py
+++ b/src/khoj/main.py
@@ -17,6 +17,7 @@ from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from PyQt6 import QtWidgets
 from PyQt6.QtCore import QThread, QTimer
+from rich.logging import RichHandler
 import schedule
 
 # Internal Packages
@@ -26,7 +27,6 @@ from khoj.routers.api_beta import api_beta
 from khoj.routers.web_client import web_client
 from khoj.utils import constants, state
 from khoj.utils.cli import cli
-from khoj.utils.helpers import CustomFormatter
 from khoj.interface.desktop.main_window import MainWindow
 from khoj.interface.desktop.system_tray import create_system_tray
 
@@ -38,8 +38,12 @@ app.include_router(api, prefix="/api")
 app.include_router(api_beta, prefix="/api/beta")
 app.include_router(web_client)
 
-logger = logging.getLogger('khoj')
+# Setup Logger
+rich_handler = RichHandler(rich_tracebacks=True)
+rich_handler.setFormatter(fmt=logging.Formatter(fmt="%(message)s", datefmt="[%X]"))
+logging.basicConfig(handlers=[rich_handler])
 
+logger = logging.getLogger('khoj')
 
 def run():
     # Turn Tokenizers Parallelism Off. App does not support it.
@@ -53,7 +57,7 @@ def run():
     # Create app directory, if it doesn't exist
     state.config_file.parent.mkdir(parents=True, exist_ok=True)
 
-    # Setup Logger
+    # Set Logging Level
     if args.verbose == 0:
         logger.setLevel(logging.WARN)
     elif args.verbose == 1:
@@ -61,11 +65,6 @@ def run():
     elif args.verbose >= 2:
         logger.setLevel(logging.DEBUG)
 
-    # Set Log Format
-    ch = logging.StreamHandler()
-    ch.setFormatter(CustomFormatter())
-    logger.addHandler(ch)
-
     # Set Log File
     fh = logging.FileHandler(state.config_file.parent / 'khoj.log')
     fh.setLevel(logging.DEBUG)
@@ -141,9 +140,9 @@ def set_state(args):
 
 def start_server(app, host=None, port=None, socket=None):
     if socket:
-        uvicorn.run(app, proxy_headers=True, uds=socket)
+        uvicorn.run(app, proxy_headers=True, uds=socket, log_level="debug", use_colors=True, log_config=None)
     else:
-        uvicorn.run(app, host=host, port=port)
+        uvicorn.run(app, host=host, port=port, log_level="debug", use_colors=True, log_config=None)
 
 
 def poll_task_scheduler():
diff --git a/src/khoj/utils/helpers.py b/src/khoj/utils/helpers.py
index b361efd2..8e3e2bd5 100644
--- a/src/khoj/utils/helpers.py
+++ b/src/khoj/utils/helpers.py
@@ -116,28 +116,4 @@ class LRU(OrderedDict):
         super().__setitem__(key, value)
         if len(self) > self.capacity:
             oldest = next(iter(self))
-            del self[oldest]
-
-
-class CustomFormatter(logging.Formatter):
-    blue = "\x1b[1;34m"
-    green = "\x1b[1;32m"
-    grey = "\x1b[38;20m"
-    yellow = "\x1b[33;20m"
-    red = "\x1b[31;20m"
-    bold_red = "\x1b[31;1m"
-    reset = "\x1b[0m"
-    format_str = "%(levelname)s: %(asctime)s: %(name)s | %(message)s"
-
-    FORMATS = {
-        logging.DEBUG: blue + format_str + reset,
-        logging.INFO: green + format_str + reset,
-        logging.WARNING: yellow + format_str + reset,
-        logging.ERROR: red + format_str + reset,
-        logging.CRITICAL: bold_red + format_str + reset
-    }
-
-    def format(self, record):
-        log_fmt = self.FORMATS.get(record.levelno)
-        formatter = logging.Formatter(log_fmt)
-        return formatter.format(record)
\ No newline at end of file
+            del self[oldest]
\ No newline at end of file