mirror of
https://github.com/bentoml/OpenLLM.git
synced 2026-02-07 14:22:08 -05:00
83 lines
2.4 KiB
Python
83 lines
2.4 KiB
Python
from __future__ import annotations
|
|
import enum, functools, logging, os, typing as t
|
|
import click, inflection, orjson
|
|
from openllm_core._typing_compat import DictStrAny, TypedDict
|
|
from openllm_core.utils import get_debug_mode
|
|
|
|
logger = logging.getLogger('openllm')
|
|
|
|
|
|
class Level(enum.IntEnum):
|
|
NOTSET = logging.DEBUG
|
|
DEBUG = logging.DEBUG
|
|
INFO = logging.INFO
|
|
WARNING = logging.WARNING
|
|
ERROR = logging.ERROR
|
|
CRITICAL = logging.CRITICAL
|
|
|
|
@property
|
|
def color(self) -> str | None:
|
|
return {
|
|
Level.NOTSET: None,
|
|
Level.DEBUG: 'cyan',
|
|
Level.INFO: 'green',
|
|
Level.WARNING: 'yellow',
|
|
Level.ERROR: 'red',
|
|
Level.CRITICAL: 'red',
|
|
}[self]
|
|
|
|
@classmethod
|
|
def from_logging_level(cls, level: int) -> Level:
|
|
return {
|
|
logging.DEBUG: Level.DEBUG,
|
|
logging.INFO: Level.INFO,
|
|
logging.WARNING: Level.WARNING,
|
|
logging.ERROR: Level.ERROR,
|
|
logging.CRITICAL: Level.CRITICAL,
|
|
}[level]
|
|
|
|
|
|
class JsonLog(TypedDict):
|
|
log_level: Level
|
|
content: str
|
|
|
|
|
|
def log(content: str, level: Level = Level.INFO, fg: str | None = None) -> None:
|
|
if get_debug_mode():
|
|
echo(content, fg=fg)
|
|
else:
|
|
echo(orjson.dumps(JsonLog(log_level=level, content=content)).decode(), fg=fg, json=True)
|
|
|
|
|
|
warning = functools.partial(log, level=Level.WARNING)
|
|
error = functools.partial(log, level=Level.ERROR)
|
|
critical = functools.partial(log, level=Level.CRITICAL)
|
|
debug = functools.partial(log, level=Level.DEBUG)
|
|
info = functools.partial(log, level=Level.INFO)
|
|
notset = functools.partial(log, level=Level.NOTSET)
|
|
|
|
|
|
def echo(text: t.Any, fg: str | None = None, *, _with_style: bool = True, json: bool = False, **attrs: t.Any) -> None:
|
|
if json:
|
|
text = orjson.loads(text)
|
|
if 'content' in text and 'log_level' in text:
|
|
content = text['content']
|
|
fg = Level.from_logging_level(text['log_level']).color
|
|
else:
|
|
content = orjson.dumps(text).decode()
|
|
fg = Level.INFO.color if not get_debug_mode() else Level.DEBUG.color
|
|
else:
|
|
content = t.cast(str, text)
|
|
attrs['fg'] = fg
|
|
|
|
(click.echo if not _with_style else click.secho)(content, **attrs)
|
|
|
|
|
|
COLUMNS: int = int(os.environ.get('COLUMNS', str(120)))
|
|
CONTEXT_SETTINGS: DictStrAny = {
|
|
'help_option_names': ['-h', '--help'],
|
|
'max_content_width': COLUMNS,
|
|
'token_normalize_func': inflection.underscore,
|
|
}
|
|
__all__ = ['COLUMNS', 'CONTEXT_SETTINGS', 'Level', 'critical', 'debug', 'echo', 'error', 'info', 'log', 'warning']
|