Unify logging

Signed-off-by: András Schmelczer <andras@schmelczer.dev>

good_ai/src/good_ai/good_ai/deploy/serve.py

@@ -1,7 +1,7 @@
 from typing import Any, Callable
 
 import uvicorn
-from fastapi import status
+from fastapi import FastAPI, status
 from fastapi.middleware.wsgi import WSGIMiddleware
 from fastapi.responses import RedirectResponse
 from typing_extensions import Never
@@ -18,6 +18,7 @@ def serve(
     function: Callable[..., Any],
     disable_docs: bool = False,
     disable_metrics: bool = False,
+    configure: Callable[[FastAPI], None]=lambda _:None
 ) -> Never:
     app = create_fastapi_app(function.__name__, disable_docs=disable_docs)
 
@@ -36,4 +37,36 @@ def serve(
             output = t.log_output(result)
         return output
 
-    uvicorn.run(app, host="0.0.0.0", port=5050)
+    configure(app)
+
+    uvicorn.run(app, host="0.0.0.0", port=5050, log_config={
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "default": {
+                "()": "good_ai.logger.CustomFormatter",
+                "fmt": "%(asctime)s | %(levelname)8s | %(message)s",
+            },
+            "access": {
+                "()": "good_ai.logger.CustomFormatter",
+                "fmt": '%(asctime)s | %(levelname)8s | %(message)s',  # noqa: E501
+            },
+        },
+        "handlers": {
+            "default": {
+                "formatter": "default",
+                "class": "logging.StreamHandler",
+                "stream": "ext://sys.stderr",
+            },
+            "access": {
+                "formatter": "access",
+                "class": "logging.StreamHandler",
+                "stream": "ext://sys.stdout",
+            },
+        },
+        "loggers": {
+            "uvicorn": {"handlers": ["default"], "level": "INFO"},
+            "uvicorn.error": {"level": "INFO"},
+            "uvicorn.access": {"handlers": ["access"], "level": "INFO", "propagate": False},
+        },
+    })

good_ai/src/good_ai/utilities/logger/__init__.py

@@ -1 +1,2 @@
 from .create_logger import create_logger
+from .custom_formatter import CustomFormatter

good_ai/src/good_ai/utilities/logger/custom_formatter.py

@@ -6,16 +6,16 @@ from .colors import BLUE, BOLD_RED, GREY, RED, RESET, YELLOW
 class CustomFormatter(logging.Formatter):
     def __init__(self, fmt: str):
         super().__init__()
-        self.fmt = fmt
+        self._fmt = fmt
-        self.FORMATS = {
+        self._color_mapping = {
-            logging.DEBUG: GREY + self.fmt + RESET,
+            logging.DEBUG: GREY + self._fmt + RESET,
-            logging.INFO: BLUE + self.fmt + RESET,
+            logging.INFO: BLUE + self._fmt + RESET,
-            logging.WARNING: YELLOW + self.fmt + RESET,
+            logging.WARNING: YELLOW + self._fmt + RESET,
-            logging.ERROR: RED + self.fmt + RESET,
+            logging.ERROR: RED + self._fmt + RESET,
-            logging.CRITICAL: BOLD_RED + self.fmt + RESET,
+            logging.CRITICAL: BOLD_RED + self._fmt + RESET,
         }
 
     def format(self, record: logging.LogRecord) -> str:
-        log_fmt = self.FORMATS.get(record.levelno)
+        log_fmt = self._color_mapping.get(record.levelno)
         formatter = logging.Formatter(log_fmt)
         return formatter.format(record)