Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
KoboldAI
GitHub Repository: KoboldAI/KoboldAI-Client
Path: blob/main/logger.py
471 views
1
import sys
2
from functools import partialmethod
3
from loguru import logger
4
5
STDOUT_LEVELS = ["GENERATION", "PROMPT"]
6
INIT_LEVELS = ["INIT", "INIT_OK", "INIT_WARN", "INIT_ERR"]
7
MESSAGE_LEVELS = ["MESSAGE"]
8
# By default we're at error level or higher
9
verbosity = 20
10
quiet = 0
11
12
def set_logger_verbosity(count):
13
global verbosity
14
# The count comes reversed. So count = 0 means minimum verbosity
15
# While count 5 means maximum verbosity
16
# So the more count we have, the lowe we drop the versbosity maximum
17
verbosity = 20 - (count * 10)
18
19
def quiesce_logger(count):
20
global quiet
21
# The bigger the count, the more silent we want our logger
22
quiet = count * 10
23
24
def is_stdout_log(record):
25
if record["level"].name not in STDOUT_LEVELS:
26
return(False)
27
if record["level"].no < verbosity + quiet:
28
return(False)
29
return(True)
30
31
def is_init_log(record):
32
if record["level"].name not in INIT_LEVELS:
33
return(False)
34
if record["level"].no < verbosity + quiet:
35
return(False)
36
return(True)
37
38
def is_msg_log(record):
39
if record["level"].name not in MESSAGE_LEVELS:
40
return(False)
41
if record["level"].no < verbosity + quiet:
42
return(False)
43
return(True)
44
45
def is_stderr_log(record):
46
if record["level"].name in STDOUT_LEVELS + INIT_LEVELS + MESSAGE_LEVELS:
47
return(False)
48
if record["level"].no < verbosity + quiet:
49
return(False)
50
return(True)
51
52
def test_logger():
53
logger.generation("This is a generation message\nIt is typically multiline\nThee Lines".encode("unicode_escape").decode("utf-8"))
54
logger.prompt("This is a prompt message")
55
logger.debug("Debug Message")
56
logger.info("Info Message")
57
logger.warning("Info Warning")
58
logger.error("Error Message")
59
logger.critical("Critical Message")
60
logger.init("This is an init message", status="Starting")
61
logger.init_ok("This is an init message", status="OK")
62
logger.init_warn("This is an init message", status="Warning")
63
logger.init_err("This is an init message", status="Error")
64
logger.message("This is user message")
65
sys.exit()
66
67
68
logfmt = "<level>{level: <10}</level> | <green>{name}</green>:<green>{function}</green>:<green>{line}</green> - <level>{message}</level>"
69
genfmt = "<level>{level: <10}</level> @ <green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{message}</level>"
70
initfmt = "<magenta>INIT </magenta> | <level>{extra[status]: <10}</level> | <magenta>{message}</magenta>"
71
msgfmt = "<level>{level: <10}</level> | <level>{message}</level>"
72
73
logger.level("GENERATION", no=24, color="<cyan>")
74
logger.level("PROMPT", no=23, color="<yellow>")
75
logger.level("INIT", no=31, color="<white>")
76
logger.level("INIT_OK", no=31, color="<green>")
77
logger.level("INIT_WARN", no=31, color="<yellow>")
78
logger.level("INIT_ERR", no=31, color="<red>")
79
# Messages contain important information without which this application might not be able to be used
80
# As such, they have the highest priority
81
logger.level("MESSAGE", no=61, color="<green>")
82
83
logger.__class__.generation = partialmethod(logger.__class__.log, "GENERATION")
84
logger.__class__.prompt = partialmethod(logger.__class__.log, "PROMPT")
85
logger.__class__.init = partialmethod(logger.__class__.log, "INIT")
86
logger.__class__.init_ok = partialmethod(logger.__class__.log, "INIT_OK")
87
logger.__class__.init_warn = partialmethod(logger.__class__.log, "INIT_WARN")
88
logger.__class__.init_err = partialmethod(logger.__class__.log, "INIT_ERR")
89
logger.__class__.message = partialmethod(logger.__class__.log, "MESSAGE")
90
91
config = {
92
"handlers": [
93
{"sink": sys.stderr, "format": logfmt, "colorize":True, "filter": is_stderr_log},
94
{"sink": sys.stdout, "format": genfmt, "level": "PROMPT", "colorize":True, "filter": is_stdout_log},
95
{"sink": sys.stdout, "format": initfmt, "level": "INIT", "colorize":True, "filter": is_init_log},
96
{"sink": sys.stdout, "format": msgfmt, "level": "MESSAGE", "colorize":True, "filter": is_msg_log}
97
],
98
}
99
logger.configure(**config)
100
101