From 2a8a223473e459201fa493d952bf60a2e1aaef56 Mon Sep 17 00:00:00 2001
From: Divided by Zer0 <mail@dbzer0.com>
Date: Sun, 11 Sep 2022 23:55:48 +0200
Subject: [PATCH] added init logs

---
 aiserver.py | 12 +++---------
 logger.py   | 17 +++++++++++++----
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/aiserver.py b/aiserver.py
index 4400c724..8663c52f 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -18,12 +18,6 @@ from eventlet import tpool
 
 import logging
 from logger import logger
-logger.generation("This is a generation message")
-logger.debug("That's it, beautiful and simple logging!")
-logger.info("That's it, beautiful and simple logging!")
-logger.error("That's it, beautiful and simple logging!")
-logger.critical("That's it, beautiful and simple logging!")
-
 
 logging.getLogger("urllib3").setLevel(logging.ERROR)
 
@@ -78,7 +72,7 @@ global tpu_mtj_backend
 
 
 if lupa.LUA_VERSION[:2] != (5, 4):
-    print(f"Please install lupa==1.10. You have lupa {lupa.__version__}.", file=sys.stderr)
+    logging.error(f"Please install lupa==1.10. You have lupa {lupa.__version__}.")
 
 patch_causallm_patched = False
 
@@ -409,7 +403,7 @@ log = logging.getLogger('werkzeug')
 log.setLevel(logging.ERROR)
 
 # Start flask & SocketIO
-print("{0}Initializing Flask... {1}".format(colors.PURPLE, colors.END), end="")
+logger.init("Flask Initialization", status="Starting")
 from flask import Flask, render_template, Response, request, copy_current_request_context, send_from_directory, session, jsonify, abort, redirect
 from flask_socketio import SocketIO
 from flask_socketio import emit as _emit
@@ -422,7 +416,7 @@ app.config['SESSION_TYPE'] = 'filesystem'
 app.config['TEMPLATES_AUTO_RELOAD'] = True
 Session(app)
 socketio = SocketIO(app, async_method="eventlet")
-print("{0}OK!{1}".format(colors.GREEN, colors.END))
+logger.init("Flask Initialization", status="OK")
 
 old_socketio_on = socketio.on
 def new_socketio_on(*a, **k):
diff --git a/logger.py b/logger.py
index 46a785b1..2549920b 100644
--- a/logger.py
+++ b/logger.py
@@ -3,31 +3,40 @@ from functools import partialmethod
 from loguru import logger
 
 STDOUT_LEVELS = ["GENERATION", "PROMPT"]
+INIT_LEVELS = ["INIT"]
 
 def is_stdout_log(record):
     if record["level"].name in STDOUT_LEVELS:
         return(True)
     return(False)
 
+def is_init_log(record):
+    if record["level"].name in INIT_LEVELS:
+        return(True)
+    return(False)
+
 def is_stderr_log(record):
-    if record["level"].name not in STDOUT_LEVELS:
+    if record["level"].name not in STDOUT_LEVELS + INIT_LEVELS:
         return(True)
     return(False)
 
 logfmt = "<level>{level: <10}</level> | <green>{name}</green>:<green>{function}</green>:<green>{line}</green> - <level>{message}</level>"
 genfmt = "<level>{level: <10}</level> @ <green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{message}</level>"
-promptfmt = "<level>{level: <10}</level> @ <green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{message}</level>"
+initfmt = "<level>{level: <10}</level> | <green>{extra[status]: <8}</green> | <level>{message}</level>"
 
-new_level = logger.level("GENERATION", no=22, color="<cyan>")
-new_level = logger.level("PROMPT", no=21, color="<yellow>")
+logger.level("GENERATION", no=24, color="<cyan>")
+logger.level("PROMPT", no=23, color="<yellow>")
+logger.level("INIT", no=21, color="<magenta>")
 
 logger.__class__.generation = partialmethod(logger.__class__.log, "GENERATION")
 logger.__class__.prompt = partialmethod(logger.__class__.log, "PROMPT")
+logger.__class__.init = partialmethod(logger.__class__.log, "INIT")
 
 config = {
     "handlers": [
         {"sink": sys.stderr, "format": logfmt, "colorize":True, "filter": is_stderr_log},
         {"sink": sys.stdout, "format": genfmt, "level": "PROMPT", "colorize":True, "filter": is_stdout_log},
+        {"sink": sys.stdout, "format": initfmt, "level": "INIT", "colorize":True, "filter": is_init_log}
     ],
 }
 logger.configure(**config)